Skip to content

Commit eb733ad

Browse files
pj-workspacepj-workspace
authored
GH-1234: Release serialization buffers when compression fails (#1290)
## What's Changed When record batch serialization fails after retaining or compressing buffers, `VectorUnloader` can leave those buffers alive even after the source vectors are closed. Compression allocation failures can also leave an extra reference to the input buffer in `AbstractCompressionCodec`. Release collected batch buffers if traversal or batch construction throws, while preserving the original exception. Use try-with-resources in `AbstractCompressionCodec.compress` so the codec releases the input it owns on both success and failure. Cleanup remains in the component that owns each buffer: the unloader does not release an input already handed to a custom codec. Regression coverage checks allocator exhaustion with zero allocation headroom, runtime exceptions and errors injected on the second compression call, an empty input buffer, and a malformed later vector. The compression-only test stub is explicitly named `SimulatedLz4Codec`; it exercises allocation and ownership, not LZ4 round trips. The tests assert memory usage and reference counts return to their previous values and that the source data remains readable. Additional tests exercise allocation failure with the real LZ4 and ZSTD codecs. The memory-leak report and initial outer-cleanup proposal are from #1234; this change also handles the codec's input ownership. ## Validation - Before the fix: the initial regression run exposed retained references and unreleased allocations; the 5 existing unload/load tests passed. - After the review update: 10 vector cases pass on each of Netty and Unsafe; 53 compression cases pass, including LZ4/ZSTD allocation failures and compressed stream/file round trips (73 test executions, no failures/errors/skips). - JDK 17 compilation and Checkstyle passed. Spotless check passed on JDK 23; `git diff --check` passed. ```sh mvn -B -ntp -pl compression -am test -Dtest=TestVectorUnloadLoad,TestVectorUnloaderFailure,TestCompressionCodec,TestArrowReaderWriterWithCompression,TestCompressionCodecServiceProvider -Dsurefire.failIfNoSpecifiedTests=false ``` - No full repository test run, native dataset/Flight integration run or performance benchmark is claimed. Closes #1234. AI assistance: OpenAI Codex was used for implementation and local validation.
1 parent ce8a19d commit eb733ad

4 files changed

Lines changed: 217 additions & 36 deletions

File tree

‎compression/src/test/java/org/apache/arrow/compression/TestCompressionCodec.java‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@
3737
import java.util.stream.Stream;
3838
import org.apache.arrow.memory.ArrowBuf;
3939
import org.apache.arrow.memory.BufferAllocator;
40+
import org.apache.arrow.memory.OutOfMemoryException;
4041
import org.apache.arrow.memory.RootAllocator;
4142
import org.apache.arrow.util.AutoCloseables;
4243
import org.apache.arrow.vector.IntVector;
@@ -64,6 +65,7 @@
6465
import org.junit.jupiter.params.ParameterizedTest;
6566
import org.junit.jupiter.params.provider.Arguments;
6667
import org.junit.jupiter.params.provider.MethodSource;
68+
import org.junit.jupiter.params.provider.ValueSource;
6769

6870
/** Test cases for {@link CompressionCodec}s. */
6971
class TestCompressionCodec {
@@ -79,6 +81,27 @@ void terminate() {
7981
allocator.close();
8082
}
8183

84+
@ParameterizedTest
85+
@ValueSource(booleans = {false, true})
86+
void testCompressionAllocationFailureReleasesSource(boolean useZstd) {
87+
CompressionCodec codec = useZstd ? new ZstdCompressionCodec() : new Lz4CompressionCodec();
88+
try (IntVector vector = new IntVector("values", allocator);
89+
VectorSchemaRoot root = VectorSchemaRoot.of(vector)) {
90+
vector.allocateNew(1);
91+
vector.set(0, 42);
92+
root.setRowCount(1);
93+
long allocatedBefore = allocator.getAllocatedMemory();
94+
int referencesBefore = vector.getDataBuffer().getReferenceManager().getRefCount();
95+
allocator.setLimit(allocatedBefore);
96+
VectorUnloader unloader = new VectorUnloader(root, true, codec, true);
97+
98+
assertThrows(OutOfMemoryException.class, unloader::getRecordBatch);
99+
assertEquals(allocatedBefore, allocator.getAllocatedMemory());
100+
assertEquals(referencesBefore, vector.getDataBuffer().getReferenceManager().getRefCount());
101+
assertEquals(42, vector.get(0));
102+
}
103+
}
104+
82105
static Collection<Arguments> codecs() {
83106
List<Arguments> params = new ArrayList<>();
84107

‎vector/src/main/java/org/apache/arrow/vector/VectorUnloader.java‎

Lines changed: 17 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
import java.util.ArrayList;
2020
import java.util.List;
2121
import org.apache.arrow.memory.ArrowBuf;
22+
import org.apache.arrow.util.AutoCloseables;
2223
import org.apache.arrow.vector.compression.CompressionCodec;
2324
import org.apache.arrow.vector.compression.CompressionUtil;
2425
import org.apache.arrow.vector.compression.NoCompressionCodec;
@@ -78,18 +79,23 @@ public ArrowRecordBatch getRecordBatch() {
7879
List<ArrowFieldNode> nodes = new ArrayList<>();
7980
List<ArrowBuf> buffers = new ArrayList<>();
8081
List<Long> variadicBufferCounts = new ArrayList<>();
81-
for (FieldVector vector : root.getFieldVectors()) {
82-
appendNodes(vector, nodes, buffers, variadicBufferCounts);
82+
try {
83+
for (FieldVector vector : root.getFieldVectors()) {
84+
appendNodes(vector, nodes, buffers, variadicBufferCounts);
85+
}
86+
// Do NOT retain buffers in ArrowRecordBatch constructor since we have already retained them.
87+
return new ArrowRecordBatch(
88+
root.getRowCount(),
89+
nodes,
90+
buffers,
91+
CompressionUtil.createBodyCompression(codec),
92+
variadicBufferCounts,
93+
alignBuffers, /*retainBuffers*/
94+
false);
95+
} catch (RuntimeException | Error e) {
96+
AutoCloseables.close(e, buffers);
97+
throw e;
8398
}
84-
// Do NOT retain buffers in ArrowRecordBatch constructor since we have already retained them.
85-
return new ArrowRecordBatch(
86-
root.getRowCount(),
87-
nodes,
88-
buffers,
89-
CompressionUtil.createBodyCompression(codec),
90-
variadicBufferCounts,
91-
alignBuffers, /*retainBuffers*/
92-
false);
9399
}
94100

95101
private long getVariadicBufferCount(FieldVector vector) {

‎vector/src/main/java/org/apache/arrow/vector/compression/AbstractCompressionCodec.java‎

Lines changed: 25 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -29,34 +29,34 @@ public abstract class AbstractCompressionCodec implements CompressionCodec {
2929

3030
@Override
3131
public ArrowBuf compress(BufferAllocator allocator, ArrowBuf uncompressedBuffer) {
32-
// GH-1116: capture writerIndex() once so the empty-buffer check, size
33-
// comparison, and uncompressed-length prefix all see the same value.
34-
long uncompressedLength = uncompressedBuffer.writerIndex();
32+
try (uncompressedBuffer) {
33+
// GH-1116: capture writerIndex() once so the empty-buffer check, size
34+
// comparison, and uncompressed-length prefix all see the same value.
35+
long uncompressedLength = uncompressedBuffer.writerIndex();
36+
37+
if (uncompressedLength == 0L) {
38+
// shortcut for empty buffer
39+
ArrowBuf compressedBuffer = allocator.buffer(CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH);
40+
compressedBuffer.setLong(0, 0);
41+
compressedBuffer.writerIndex(CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH);
42+
return compressedBuffer;
43+
}
44+
45+
ArrowBuf compressedBuffer = doCompress(allocator, uncompressedBuffer);
46+
long compressedLength =
47+
compressedBuffer.writerIndex() - CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH;
48+
49+
if (compressedLength > uncompressedLength) {
50+
// compressed buffer is larger, send the raw buffer
51+
compressedBuffer.close();
52+
// XXX: this makes a copy of uncompressedBuffer
53+
compressedBuffer = CompressionUtil.packageRawBuffer(allocator, uncompressedBuffer);
54+
} else {
55+
writeUncompressedLength(compressedBuffer, uncompressedLength);
56+
}
3557

36-
if (uncompressedLength == 0L) {
37-
// shortcut for empty buffer
38-
ArrowBuf compressedBuffer = allocator.buffer(CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH);
39-
compressedBuffer.setLong(0, 0);
40-
compressedBuffer.writerIndex(CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH);
41-
uncompressedBuffer.close();
4258
return compressedBuffer;
4359
}
44-
45-
ArrowBuf compressedBuffer = doCompress(allocator, uncompressedBuffer);
46-
long compressedLength =
47-
compressedBuffer.writerIndex() - CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH;
48-
49-
if (compressedLength > uncompressedLength) {
50-
// compressed buffer is larger, send the raw buffer
51-
compressedBuffer.close();
52-
// XXX: this makes a copy of uncompressedBuffer
53-
compressedBuffer = CompressionUtil.packageRawBuffer(allocator, uncompressedBuffer);
54-
} else {
55-
writeUncompressedLength(compressedBuffer, uncompressedLength);
56-
}
57-
58-
uncompressedBuffer.close();
59-
return compressedBuffer;
6060
}
6161

6262
@Override

‎vector/src/test/java/org/apache/arrow/vector/TestVectorUnloaderFailure.java‎

Lines changed: 152 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,152 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one or more
3+
* contributor license agreements. See the NOTICE file distributed with
4+
* this work for additional information regarding copyright ownership.
5+
* The ASF licenses this file to You under the Apache License, Version 2.0
6+
* (the "License"); you may not use this file except in compliance with
7+
* the License. You may obtain a copy of the License at
8+
*
9+
* http://www.apache.org/licenses/LICENSE-2.0
10+
*
11+
* Unless required by applicable law or agreed to in writing, software
12+
* distributed under the License is distributed on an "AS IS" BASIS,
13+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
* See the License for the specific language governing permissions and
15+
* limitations under the License.
16+
*/
17+
package org.apache.arrow.vector;
18+
19+
import static org.junit.jupiter.api.Assertions.assertEquals;
20+
import static org.junit.jupiter.api.Assertions.assertSame;
21+
import static org.junit.jupiter.api.Assertions.assertThrows;
22+
23+
import java.util.Collections;
24+
import java.util.List;
25+
import org.apache.arrow.memory.ArrowBuf;
26+
import org.apache.arrow.memory.BufferAllocator;
27+
import org.apache.arrow.memory.OutOfMemoryException;
28+
import org.apache.arrow.memory.RootAllocator;
29+
import org.apache.arrow.vector.compression.AbstractCompressionCodec;
30+
import org.apache.arrow.vector.compression.CompressionUtil;
31+
import org.junit.jupiter.api.Test;
32+
import org.junit.jupiter.params.ParameterizedTest;
33+
import org.junit.jupiter.params.provider.ValueSource;
34+
35+
/** Exception-safety tests for record batch serialization. */
36+
class TestVectorUnloaderFailure {
37+
38+
@Test
39+
void compressionAllocationFailureReleasesBuffers() {
40+
try (BufferAllocator allocator = new RootAllocator();
41+
IntVector vector = new IntVector("values", allocator);
42+
VectorSchemaRoot root = VectorSchemaRoot.of(vector)) {
43+
vector.allocateNew(1);
44+
vector.set(0, 42);
45+
root.setRowCount(1);
46+
long allocatedBefore = allocator.getAllocatedMemory();
47+
int referencesBefore = vector.getDataBuffer().getReferenceManager().getRefCount();
48+
// Any compression allocation must fail, regardless of allocator rounding.
49+
allocator.setLimit(allocatedBefore);
50+
VectorUnloader unloader = new VectorUnloader(root, true, new SimulatedLz4Codec(), true);
51+
52+
assertThrows(OutOfMemoryException.class, unloader::getRecordBatch);
53+
assertEquals(allocatedBefore, allocator.getAllocatedMemory());
54+
assertEquals(referencesBefore, vector.getDataBuffer().getReferenceManager().getRefCount());
55+
assertEquals(42, vector.get(0));
56+
}
57+
}
58+
59+
@ParameterizedTest
60+
@ValueSource(booleans = {false, true})
61+
void compressionFailurePreservesExceptionAndSource(boolean throwError) {
62+
Throwable failure =
63+
throwError
64+
? new OutOfMemoryError("compression failed")
65+
: new IllegalStateException("compression failed");
66+
try (BufferAllocator allocator = new RootAllocator();
67+
IntVector vector = new IntVector("values", allocator);
68+
VectorSchemaRoot root = VectorSchemaRoot.of(vector)) {
69+
vector.allocateNew(1);
70+
vector.set(0, 42);
71+
root.setRowCount(1);
72+
long allocatedBefore = allocator.getAllocatedMemory();
73+
int referencesBefore = vector.getDataBuffer().getReferenceManager().getRefCount();
74+
SimulatedLz4Codec codec =
75+
new SimulatedLz4Codec() {
76+
private int calls;
77+
78+
@Override
79+
protected ArrowBuf doCompress(BufferAllocator allocator, ArrowBuf input) {
80+
if (++calls == 2) {
81+
if (failure instanceof Error) {
82+
throw (Error) failure;
83+
}
84+
throw (RuntimeException) failure;
85+
}
86+
return super.doCompress(allocator, input);
87+
}
88+
};
89+
VectorUnloader unloader = new VectorUnloader(root, true, codec, true);
90+
91+
assertSame(failure, assertThrows(failure.getClass(), unloader::getRecordBatch));
92+
assertEquals(allocatedBefore, allocator.getAllocatedMemory());
93+
assertEquals(referencesBefore, vector.getDataBuffer().getReferenceManager().getRefCount());
94+
assertEquals(42, vector.get(0));
95+
}
96+
}
97+
98+
@Test
99+
void emptyBufferAllocationFailureReleasesInput() {
100+
try (BufferAllocator allocator = new RootAllocator()) {
101+
ArrowBuf input = allocator.buffer(8);
102+
allocator.setLimit(allocator.getAllocatedMemory());
103+
104+
assertThrows(
105+
OutOfMemoryException.class, () -> new SimulatedLz4Codec().compress(allocator, input));
106+
assertEquals(0, input.getReferenceManager().getRefCount());
107+
assertEquals(0, allocator.getAllocatedMemory());
108+
}
109+
}
110+
111+
@Test
112+
void invalidLaterVectorReleasesPreviouslyRetainedBuffers() {
113+
try (BufferAllocator allocator = new RootAllocator();
114+
IntVector first = new IntVector("first", allocator)) {
115+
FieldVector invalid =
116+
new NullVector("invalid") {
117+
@Override
118+
public List<ArrowBuf> getFieldBuffers() {
119+
return Collections.singletonList(allocator.getEmpty());
120+
}
121+
};
122+
try (VectorSchemaRoot root = VectorSchemaRoot.of(first, invalid)) {
123+
first.allocateNew(1);
124+
first.set(0, 42);
125+
root.setRowCount(1);
126+
int referencesBefore = first.getDataBuffer().getReferenceManager().getRefCount();
127+
128+
assertThrows(IllegalArgumentException.class, new VectorUnloader(root)::getRecordBatch);
129+
assertEquals(referencesBefore, first.getDataBuffer().getReferenceManager().getRefCount());
130+
assertEquals(42, first.get(0));
131+
}
132+
}
133+
}
134+
135+
// Compression-only stub for allocation and ownership tests, not LZ4 round trips.
136+
private static class SimulatedLz4Codec extends AbstractCompressionCodec {
137+
@Override
138+
protected ArrowBuf doCompress(BufferAllocator allocator, ArrowBuf input) {
139+
return CompressionUtil.packageRawBuffer(allocator, input);
140+
}
141+
142+
@Override
143+
protected ArrowBuf doDecompress(BufferAllocator allocator, ArrowBuf input) {
144+
throw new UnsupportedOperationException();
145+
}
146+
147+
@Override
148+
public CompressionUtil.CodecType getCodecType() {
149+
return CompressionUtil.CodecType.LZ4_FRAME;
150+
}
151+
}
152+
}

0 commit comments

Comments
 (0)