This is an automated email from the ASF dual-hosted git repository.
jbonofre pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-java.git
The following commit(s) were added to refs/heads/main by this push:
new eb733ad64 GH-1234: Release serialization buffers when compression
fails (#1290)
eb733ad64 is described below
commit eb733ad64a158a56cab91df5e60a513536133be2
Author: Sunflower_PJ <[email protected]>
AuthorDate: Mon Oct 5 22:41:51 2026 +0800
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.
---
.../arrow/compression/TestCompressionCodec.java | 23 ++++
.../org/apache/arrow/vector/VectorUnloader.java | 28 ++--
.../compression/AbstractCompressionCodec.java | 50 +++----
.../arrow/vector/TestVectorUnloaderFailure.java | 152 +++++++++++++++++++++
4 files changed, 217 insertions(+), 36 deletions(-)
diff --git
a/compression/src/test/java/org/apache/arrow/compression/TestCompressionCodec.java
b/compression/src/test/java/org/apache/arrow/compression/TestCompressionCodec.java
index d2d292164..4378be96a 100644
---
a/compression/src/test/java/org/apache/arrow/compression/TestCompressionCodec.java
+++
b/compression/src/test/java/org/apache/arrow/compression/TestCompressionCodec.java
@@ -37,6 +37,7 @@ import java.util.function.BiConsumer;
import java.util.stream.Stream;
import org.apache.arrow.memory.ArrowBuf;
import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.memory.OutOfMemoryException;
import org.apache.arrow.memory.RootAllocator;
import org.apache.arrow.util.AutoCloseables;
import org.apache.arrow.vector.IntVector;
@@ -64,6 +65,7 @@ import org.junit.jupiter.api.Test;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.Arguments;
import org.junit.jupiter.params.provider.MethodSource;
+import org.junit.jupiter.params.provider.ValueSource;
/** Test cases for {@link CompressionCodec}s. */
class TestCompressionCodec {
@@ -79,6 +81,27 @@ class TestCompressionCodec {
allocator.close();
}
+ @ParameterizedTest
+ @ValueSource(booleans = {false, true})
+ void testCompressionAllocationFailureReleasesSource(boolean useZstd) {
+ CompressionCodec codec = useZstd ? new ZstdCompressionCodec() : new
Lz4CompressionCodec();
+ try (IntVector vector = new IntVector("values", allocator);
+ VectorSchemaRoot root = VectorSchemaRoot.of(vector)) {
+ vector.allocateNew(1);
+ vector.set(0, 42);
+ root.setRowCount(1);
+ long allocatedBefore = allocator.getAllocatedMemory();
+ int referencesBefore =
vector.getDataBuffer().getReferenceManager().getRefCount();
+ allocator.setLimit(allocatedBefore);
+ VectorUnloader unloader = new VectorUnloader(root, true, codec, true);
+
+ assertThrows(OutOfMemoryException.class, unloader::getRecordBatch);
+ assertEquals(allocatedBefore, allocator.getAllocatedMemory());
+ assertEquals(referencesBefore,
vector.getDataBuffer().getReferenceManager().getRefCount());
+ assertEquals(42, vector.get(0));
+ }
+ }
+
static Collection<Arguments> codecs() {
List<Arguments> params = new ArrayList<>();
diff --git a/vector/src/main/java/org/apache/arrow/vector/VectorUnloader.java
b/vector/src/main/java/org/apache/arrow/vector/VectorUnloader.java
index 4782b3f53..83b985b70 100644
--- a/vector/src/main/java/org/apache/arrow/vector/VectorUnloader.java
+++ b/vector/src/main/java/org/apache/arrow/vector/VectorUnloader.java
@@ -19,6 +19,7 @@ package org.apache.arrow.vector;
import java.util.ArrayList;
import java.util.List;
import org.apache.arrow.memory.ArrowBuf;
+import org.apache.arrow.util.AutoCloseables;
import org.apache.arrow.vector.compression.CompressionCodec;
import org.apache.arrow.vector.compression.CompressionUtil;
import org.apache.arrow.vector.compression.NoCompressionCodec;
@@ -78,18 +79,23 @@ public class VectorUnloader {
List<ArrowFieldNode> nodes = new ArrayList<>();
List<ArrowBuf> buffers = new ArrayList<>();
List<Long> variadicBufferCounts = new ArrayList<>();
- for (FieldVector vector : root.getFieldVectors()) {
- appendNodes(vector, nodes, buffers, variadicBufferCounts);
+ try {
+ for (FieldVector vector : root.getFieldVectors()) {
+ appendNodes(vector, nodes, buffers, variadicBufferCounts);
+ }
+ // Do NOT retain buffers in ArrowRecordBatch constructor since we have
already retained them.
+ return new ArrowRecordBatch(
+ root.getRowCount(),
+ nodes,
+ buffers,
+ CompressionUtil.createBodyCompression(codec),
+ variadicBufferCounts,
+ alignBuffers, /*retainBuffers*/
+ false);
+ } catch (RuntimeException | Error e) {
+ AutoCloseables.close(e, buffers);
+ throw e;
}
- // Do NOT retain buffers in ArrowRecordBatch constructor since we have
already retained them.
- return new ArrowRecordBatch(
- root.getRowCount(),
- nodes,
- buffers,
- CompressionUtil.createBodyCompression(codec),
- variadicBufferCounts,
- alignBuffers, /*retainBuffers*/
- false);
}
private long getVariadicBufferCount(FieldVector vector) {
diff --git
a/vector/src/main/java/org/apache/arrow/vector/compression/AbstractCompressionCodec.java
b/vector/src/main/java/org/apache/arrow/vector/compression/AbstractCompressionCodec.java
index b108173c8..39b5ceae3 100644
---
a/vector/src/main/java/org/apache/arrow/vector/compression/AbstractCompressionCodec.java
+++
b/vector/src/main/java/org/apache/arrow/vector/compression/AbstractCompressionCodec.java
@@ -29,34 +29,34 @@ public abstract class AbstractCompressionCodec implements
CompressionCodec {
@Override
public ArrowBuf compress(BufferAllocator allocator, ArrowBuf
uncompressedBuffer) {
- // GH-1116: capture writerIndex() once so the empty-buffer check, size
- // comparison, and uncompressed-length prefix all see the same value.
- long uncompressedLength = uncompressedBuffer.writerIndex();
+ try (uncompressedBuffer) {
+ // GH-1116: capture writerIndex() once so the empty-buffer check, size
+ // comparison, and uncompressed-length prefix all see the same value.
+ long uncompressedLength = uncompressedBuffer.writerIndex();
+
+ if (uncompressedLength == 0L) {
+ // shortcut for empty buffer
+ ArrowBuf compressedBuffer =
allocator.buffer(CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH);
+ compressedBuffer.setLong(0, 0);
+
compressedBuffer.writerIndex(CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH);
+ return compressedBuffer;
+ }
+
+ ArrowBuf compressedBuffer = doCompress(allocator, uncompressedBuffer);
+ long compressedLength =
+ compressedBuffer.writerIndex() -
CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH;
+
+ if (compressedLength > uncompressedLength) {
+ // compressed buffer is larger, send the raw buffer
+ compressedBuffer.close();
+ // XXX: this makes a copy of uncompressedBuffer
+ compressedBuffer = CompressionUtil.packageRawBuffer(allocator,
uncompressedBuffer);
+ } else {
+ writeUncompressedLength(compressedBuffer, uncompressedLength);
+ }
- if (uncompressedLength == 0L) {
- // shortcut for empty buffer
- ArrowBuf compressedBuffer =
allocator.buffer(CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH);
- compressedBuffer.setLong(0, 0);
-
compressedBuffer.writerIndex(CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH);
- uncompressedBuffer.close();
return compressedBuffer;
}
-
- ArrowBuf compressedBuffer = doCompress(allocator, uncompressedBuffer);
- long compressedLength =
- compressedBuffer.writerIndex() -
CompressionUtil.SIZE_OF_UNCOMPRESSED_LENGTH;
-
- if (compressedLength > uncompressedLength) {
- // compressed buffer is larger, send the raw buffer
- compressedBuffer.close();
- // XXX: this makes a copy of uncompressedBuffer
- compressedBuffer = CompressionUtil.packageRawBuffer(allocator,
uncompressedBuffer);
- } else {
- writeUncompressedLength(compressedBuffer, uncompressedLength);
- }
-
- uncompressedBuffer.close();
- return compressedBuffer;
}
@Override
diff --git
a/vector/src/test/java/org/apache/arrow/vector/TestVectorUnloaderFailure.java
b/vector/src/test/java/org/apache/arrow/vector/TestVectorUnloaderFailure.java
new file mode 100644
index 000000000..828b81fb9
--- /dev/null
+++
b/vector/src/test/java/org/apache/arrow/vector/TestVectorUnloaderFailure.java
@@ -0,0 +1,152 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.arrow.vector;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertSame;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+import java.util.Collections;
+import java.util.List;
+import org.apache.arrow.memory.ArrowBuf;
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.memory.OutOfMemoryException;
+import org.apache.arrow.memory.RootAllocator;
+import org.apache.arrow.vector.compression.AbstractCompressionCodec;
+import org.apache.arrow.vector.compression.CompressionUtil;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
+
+/** Exception-safety tests for record batch serialization. */
+class TestVectorUnloaderFailure {
+
+ @Test
+ void compressionAllocationFailureReleasesBuffers() {
+ try (BufferAllocator allocator = new RootAllocator();
+ IntVector vector = new IntVector("values", allocator);
+ VectorSchemaRoot root = VectorSchemaRoot.of(vector)) {
+ vector.allocateNew(1);
+ vector.set(0, 42);
+ root.setRowCount(1);
+ long allocatedBefore = allocator.getAllocatedMemory();
+ int referencesBefore =
vector.getDataBuffer().getReferenceManager().getRefCount();
+ // Any compression allocation must fail, regardless of allocator
rounding.
+ allocator.setLimit(allocatedBefore);
+ VectorUnloader unloader = new VectorUnloader(root, true, new
SimulatedLz4Codec(), true);
+
+ assertThrows(OutOfMemoryException.class, unloader::getRecordBatch);
+ assertEquals(allocatedBefore, allocator.getAllocatedMemory());
+ assertEquals(referencesBefore,
vector.getDataBuffer().getReferenceManager().getRefCount());
+ assertEquals(42, vector.get(0));
+ }
+ }
+
+ @ParameterizedTest
+ @ValueSource(booleans = {false, true})
+ void compressionFailurePreservesExceptionAndSource(boolean throwError) {
+ Throwable failure =
+ throwError
+ ? new OutOfMemoryError("compression failed")
+ : new IllegalStateException("compression failed");
+ try (BufferAllocator allocator = new RootAllocator();
+ IntVector vector = new IntVector("values", allocator);
+ VectorSchemaRoot root = VectorSchemaRoot.of(vector)) {
+ vector.allocateNew(1);
+ vector.set(0, 42);
+ root.setRowCount(1);
+ long allocatedBefore = allocator.getAllocatedMemory();
+ int referencesBefore =
vector.getDataBuffer().getReferenceManager().getRefCount();
+ SimulatedLz4Codec codec =
+ new SimulatedLz4Codec() {
+ private int calls;
+
+ @Override
+ protected ArrowBuf doCompress(BufferAllocator allocator, ArrowBuf
input) {
+ if (++calls == 2) {
+ if (failure instanceof Error) {
+ throw (Error) failure;
+ }
+ throw (RuntimeException) failure;
+ }
+ return super.doCompress(allocator, input);
+ }
+ };
+ VectorUnloader unloader = new VectorUnloader(root, true, codec, true);
+
+ assertSame(failure, assertThrows(failure.getClass(),
unloader::getRecordBatch));
+ assertEquals(allocatedBefore, allocator.getAllocatedMemory());
+ assertEquals(referencesBefore,
vector.getDataBuffer().getReferenceManager().getRefCount());
+ assertEquals(42, vector.get(0));
+ }
+ }
+
+ @Test
+ void emptyBufferAllocationFailureReleasesInput() {
+ try (BufferAllocator allocator = new RootAllocator()) {
+ ArrowBuf input = allocator.buffer(8);
+ allocator.setLimit(allocator.getAllocatedMemory());
+
+ assertThrows(
+ OutOfMemoryException.class, () -> new
SimulatedLz4Codec().compress(allocator, input));
+ assertEquals(0, input.getReferenceManager().getRefCount());
+ assertEquals(0, allocator.getAllocatedMemory());
+ }
+ }
+
+ @Test
+ void invalidLaterVectorReleasesPreviouslyRetainedBuffers() {
+ try (BufferAllocator allocator = new RootAllocator();
+ IntVector first = new IntVector("first", allocator)) {
+ FieldVector invalid =
+ new NullVector("invalid") {
+ @Override
+ public List<ArrowBuf> getFieldBuffers() {
+ return Collections.singletonList(allocator.getEmpty());
+ }
+ };
+ try (VectorSchemaRoot root = VectorSchemaRoot.of(first, invalid)) {
+ first.allocateNew(1);
+ first.set(0, 42);
+ root.setRowCount(1);
+ int referencesBefore =
first.getDataBuffer().getReferenceManager().getRefCount();
+
+ assertThrows(IllegalArgumentException.class, new
VectorUnloader(root)::getRecordBatch);
+ assertEquals(referencesBefore,
first.getDataBuffer().getReferenceManager().getRefCount());
+ assertEquals(42, first.get(0));
+ }
+ }
+ }
+
+ // Compression-only stub for allocation and ownership tests, not LZ4 round
trips.
+ private static class SimulatedLz4Codec extends AbstractCompressionCodec {
+ @Override
+ protected ArrowBuf doCompress(BufferAllocator allocator, ArrowBuf input) {
+ return CompressionUtil.packageRawBuffer(allocator, input);
+ }
+
+ @Override
+ protected ArrowBuf doDecompress(BufferAllocator allocator, ArrowBuf input)
{
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public CompressionUtil.CodecType getCodecType() {
+ return CompressionUtil.CodecType.LZ4_FRAME;
+ }
+ }
+}
