GH-1234: Release serialization buffers when compression fails by pj-workspace · Pull Request #1290 · apache/arrow-java

pj-workspace · 2026-09-10T03:38:48Z

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.

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.

Copilot

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

Copilot

Warning

Copilot couldn't run its full agentic review because it didn't start before the timeout. Make sure your repository has a runner available, or add a copilot-code-review.yml file specifying one with the runs-on attribute. See the docs for more details.

Pull request overview

Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.

+  @ParameterizedTest
+  @ValueSource(longs = {0, 16})
+  void compressionAllocationFailureReleasesBuffers(long availableBytes) {
+    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();
+      allocator.setLimit(allocatedBefore + availableBytes);
+      VectorUnloader unloader = new VectorUnloader(root, true, new CopyCodec(), true);
+
+      assertThrows(OutOfMemoryException.class, unloader::getRecordBatch);
+      assertEquals(allocatedBefore, allocator.getAllocatedMemory());
+      assertEquals(referencesBefore, vector.getDataBuffer().getReferenceManager().getRefCount());
+      assertEquals(42, vector.get(0));


+    @Override
+    public CompressionUtil.CodecType getCodecType() {
+      return CompressionUtil.CodecType.LZ4_FRAME;
+    }


jbonofre

Thanks for the fix and good test coverage @pj-workspace!

The separation between VectorUnloader and AbstractCompressionCodec is clean and strictly adheres to the CompressionCodec contract:

using try (uncompressedBuffer) in AbstractCompressionCodec.compress() ensures the input buffer reference on all exit paths (empty check, compression failure, packaging fallback) without double-freeing
the rollback cleanup via AutoCloseable.close(e, buffers) in VectorUnloader.getRecordBatch() reliably frees collected buffers and preserves the root exception
the regression test suite is great.

(nit and non blocker: we could apply similar try-with-resource semantics to AbstractCompressionCodec.decompress() and align StructVectorUnloader in the C module)

Generated-by: OpenAI Codex

…tch fails ### What changes were proposed in this pull request? This PR fixes an Arrow memory leak when compressing a record batch fails, as a workaround until Arrow Java 20.0.0. It adds `ArrowCompressionUtils.createVectorUnloader` and uses it in all nine places that create a `VectorUnloader` with a codec: `ArrowConverters.ArrowBatchIterator`, `ArrowCachedBatchSerializer` (3), `PythonArrowInput` (2), `ColumnarArrowPythonInput` (2) and `CoGroupedArrowPythonRunner`. With a compressing codec, it returns a `ReleasingVectorUnloader`, which extends Arrow's `VectorUnloader` and wraps the codec in a `TrackingCompressionCodec`: - `TrackingCompressionCodec` keeps the buffer that the batch holds for each `compress` call: the input until the call returns, and the output after that. - When `getRecordBatch` throws, `ReleasingVectorUnloader` closes these buffers with `AutoCloseables.close`, which adds a failure of `close()` to the original exception as a suppressed exception. When it returns, the batch owns the buffers, so the list is only cleared. With `NoCompressionCodec`, which cannot fail, `createVectorUnloader` returns Arrow's own `VectorUnloader`, so the default configuration is unchanged. Arrow Java 20.0.0 fixes this in `VectorUnloader` and `AbstractCompressionCodec` (apache/arrow-java#1290). `ReleasingVectorUnloader` must be removed with that upgrade, because it would release the buffers twice. The new test then fails with `RefCnt has gone negative`, which I checked by running it with the `VectorUnloader` and `AbstractCompressionCodec` changes of apache/arrow-java#1290 on the test classpath. Alternatives that were not taken: - Wrap only the codec in `createCompressionCodec`: a codec does not know where a batch ends, so it cannot tell the buffers of the failed batch from those of the batches already returned. - Copy `VectorUnloader.getRecordBatch` and `appendNodes` into Spark: it brings Arrow's traversal and buffer layout checks into Spark. ### Why are the changes needed? Arrow Java up to 19.0.0 does not release buffers when `CompressionCodec.compress` throws inside `VectorUnloader.getRecordBatch` (apache/arrow-java#1234): - `appendNodes` retains each buffer before passing it to `compress`, and `AbstractCompressionCodec.compress` closes its input only on success. So the buffer being compressed keeps the extra reference. - The buffers already compressed for the batch are only in a local list of `getRecordBatch`, which nothing releases. Closing the vectors and the allocator cannot release these buffers, so they stay allocated until the JVM exits. The realistic trigger is direct memory running out, because compression allocates an additional direct buffer per input buffer. Arrow's Netty allocator turns direct memory exhaustion into `org.apache.arrow.memory.OutOfMemoryException`, which is not fatal to the executor. So the executor keeps running with less direct memory for the next tasks. The task still fails with the original exception, because `TaskContext.runTaskWithListeners` only adds a failure of a completion listener as a suppressed exception. Closing the allocator reports the leak, as in this example with a capped allocator. ``` org.apache.arrow.memory.OutOfMemoryException: Unable to allocate buffer of size 256 (rounded from 196) due to memory limit. Current allocation: 65536 at org.apache.arrow.vector.compression.AbstractCompressionCodec.compress(AbstractCompressionCodec.java:41) at org.apache.arrow.vector.VectorUnloader.appendNodes(VectorUnloader.java:131) ... Suppressed: org.apache.spark.util.TaskCompletionListenerException: Memory was leaked by query. Memory leaked: (16384) ``` On the zero-copy paths (`ArrowCachedBatchSerializer.convertArrowBatchZeroCopy` and `ColumnarArrowPythonInput.writeArrowDirect`), the buffers belong to the allocator of the upstream vectors, so the leak is reported by the owner of those vectors rather than by the code that failed. This only happens when `spark.sql.execution.arrow.compression.codec` is `zstd` or `lz4`; the default is `none`. By code reading, the affected paths have existed since Apache Spark 4.1.0 for `ArrowConverters`, `PythonArrowInput` and `CoGroupedArrowPythonRunner` (SPARK-54134, SPARK-54226), since Apache Spark 4.2.0 for `ColumnarArrowPythonInput` (SPARK-56350), and since Apache Spark 4.3.0 for `ArrowCachedBatchSerializer` (SPARK-57268). ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Pass the CIs with the newly added test case in `ArrowCompressionUtilsSuite`. Its codec fails while the second batch compresses its second buffer. The test checks that no buffer is released twice, that the first batch keeps its buffers, and that closing the allocator finds no leak. I also ran the affected paths in tasks with a temporary suite, which is not part of this PR. It caps `ArrowUtils.rootAllocator` when the input rows run out, so that the first allocation inside `compress` fails, and then measures the memory left behind. | Path | Before | After | |---|---|---| | `ArrowBatchIterator` | 16384 bytes | 0 bytes | | `InternalRowToArrowCachedBatchIterator` | 16384 bytes | 0 bytes | | `ColumnarBatchToArrowCachedBatchIterator` zero-copy, in the upstream allocator | 8192 bytes | 0 bytes | ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Claude Opus 5.5 Closes #59991 from dongjoon-hyun/SPARK-60776. Authored-by: Dongjoon Hyun <dongjoon@apache.org> Signed-off-by: Dongjoon Hyun <dongjoon@apache.org>

…tch fails ### What changes were proposed in this pull request? This PR fixes an Arrow memory leak when compressing a record batch fails, as a workaround until Arrow Java 20.0.0. It adds `ArrowCompressionUtils.createVectorUnloader` and uses it in all nine places that create a `VectorUnloader` with a codec: `ArrowConverters.ArrowBatchIterator`, `ArrowCachedBatchSerializer` (3), `PythonArrowInput` (2), `ColumnarArrowPythonInput` (2) and `CoGroupedArrowPythonRunner`. With a compressing codec, it returns a `ReleasingVectorUnloader`, which extends Arrow's `VectorUnloader` and wraps the codec in a `TrackingCompressionCodec`: - `TrackingCompressionCodec` keeps the buffer that the batch holds for each `compress` call: the input until the call returns, and the output after that. - When `getRecordBatch` throws, `ReleasingVectorUnloader` closes these buffers with `AutoCloseables.close`, which adds a failure of `close()` to the original exception as a suppressed exception. When it returns, the batch owns the buffers, so the list is only cleared. With `NoCompressionCodec`, which cannot fail, `createVectorUnloader` returns Arrow's own `VectorUnloader`, so the default configuration is unchanged. Arrow Java 20.0.0 fixes this in `VectorUnloader` and `AbstractCompressionCodec` (apache/arrow-java#1290). `ReleasingVectorUnloader` must be removed with that upgrade, because it would release the buffers twice. The new test then fails with `RefCnt has gone negative`, which I checked by running it with the `VectorUnloader` and `AbstractCompressionCodec` changes of apache/arrow-java#1290 on the test classpath. Alternatives that were not taken: - Wrap only the codec in `createCompressionCodec`: a codec does not know where a batch ends, so it cannot tell the buffers of the failed batch from those of the batches already returned. - Copy `VectorUnloader.getRecordBatch` and `appendNodes` into Spark: it brings Arrow's traversal and buffer layout checks into Spark. ### Why are the changes needed? Arrow Java up to 19.0.0 does not release buffers when `CompressionCodec.compress` throws inside `VectorUnloader.getRecordBatch` (apache/arrow-java#1234): - `appendNodes` retains each buffer before passing it to `compress`, and `AbstractCompressionCodec.compress` closes its input only on success. So the buffer being compressed keeps the extra reference. - The buffers already compressed for the batch are only in a local list of `getRecordBatch`, which nothing releases. Closing the vectors and the allocator cannot release these buffers, so they stay allocated until the JVM exits. The realistic trigger is direct memory running out, because compression allocates an additional direct buffer per input buffer. Arrow's Netty allocator turns direct memory exhaustion into `org.apache.arrow.memory.OutOfMemoryException`, which is not fatal to the executor. So the executor keeps running with less direct memory for the next tasks. The task still fails with the original exception, because `TaskContext.runTaskWithListeners` only adds a failure of a completion listener as a suppressed exception. Closing the allocator reports the leak, as in this example with a capped allocator. ``` org.apache.arrow.memory.OutOfMemoryException: Unable to allocate buffer of size 256 (rounded from 196) due to memory limit. Current allocation: 65536 at org.apache.arrow.vector.compression.AbstractCompressionCodec.compress(AbstractCompressionCodec.java:41) at org.apache.arrow.vector.VectorUnloader.appendNodes(VectorUnloader.java:131) ... Suppressed: org.apache.spark.util.TaskCompletionListenerException: Memory was leaked by query. Memory leaked: (16384) ``` On the zero-copy paths (`ArrowCachedBatchSerializer.convertArrowBatchZeroCopy` and `ColumnarArrowPythonInput.writeArrowDirect`), the buffers belong to the allocator of the upstream vectors, so the leak is reported by the owner of those vectors rather than by the code that failed. This only happens when `spark.sql.execution.arrow.compression.codec` is `zstd` or `lz4`; the default is `none`. By code reading, the affected paths have existed since Apache Spark 4.1.0 for `ArrowConverters`, `PythonArrowInput` and `CoGroupedArrowPythonRunner` (SPARK-54134, SPARK-54226), since Apache Spark 4.2.0 for `ColumnarArrowPythonInput` (SPARK-56350), and since Apache Spark 4.3.0 for `ArrowCachedBatchSerializer` (SPARK-57268). ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Pass the CIs with the newly added test case in `ArrowCompressionUtilsSuite`. Its codec fails while the second batch compresses its second buffer. The test checks that no buffer is released twice, that the first batch keeps its buffers, and that closing the allocator finds no leak. I also ran the affected paths in tasks with a temporary suite, which is not part of this PR. It caps `ArrowUtils.rootAllocator` when the input rows run out, so that the first allocation inside `compress` fails, and then measures the memory left behind. | Path | Before | After | |---|---|---| | `ArrowBatchIterator` | 16384 bytes | 0 bytes | | `InternalRowToArrowCachedBatchIterator` | 16384 bytes | 0 bytes | | `ColumnarBatchToArrowCachedBatchIterator` zero-copy, in the upstream allocator | 8192 bytes | 0 bytes | ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Claude Opus 5.5 Closes #59991 from dongjoon-hyun/SPARK-60776. Authored-by: Dongjoon Hyun <dongjoon@apache.org> Signed-off-by: Dongjoon Hyun <dongjoon@apache.org> (cherry picked from commit e438d0c) Signed-off-by: Dongjoon Hyun <dongjoon@apache.org>

Copilot AI lite review requested due to automatic review settings September 10, 2026 03:38

pj-workspace requested review from jbonofre, laurentgo, lidavidm and wgtmac as code owners September 10, 2026 03:38

Copilot AI reviewed Sep 10, 2026

This comment has been minimized.

Sign in to view

lidavidm added the bug-fix PRs that fix a big. label Sep 10, 2026

lidavidm requested a balanced review from Copilot September 10, 2026 04:05

Copilot AI reviewed Sep 10, 2026

View reviewed changes

Copilot started reviewing on behalf of lidavidm September 10, 2026 04:26

View session

github-actions Bot added this to the 20.0.0 milestone Sep 10, 2026

jbonofre approved these changes Oct 5, 2026

View reviewed changes

pj-workspace added 2 commits October 5, 2026 13:11

apacheGH-1234: Release serialization buffers when compression fails

9e4d373

Generated-by: OpenAI Codex

apacheGH-1234: Make allocation failure test deterministic

a5bf863

Generated-by: OpenAI Codex

jbonofre force-pushed the fix/vector-unloader-exception-cleanup branch from dc61ee7 to a5bf863 Compare October 5, 2026 11:11

jbonofre merged commit eb733ad into apache:main Oct 5, 2026
22 checks passed

wuchong mentioned this pull request Oct 24, 2026

[arrow] Backport Arrow Java serialization buffer cleanup fix apache/fluss-shaded#14

Open

dongjoon-hyun mentioned this pull request Oct 31, 2026

[SPARK-60776][SQL] Fix Arrow memory leak when compressing a record batch fails apache/spark#59991

Closed

Navigation Menu

Repository navigation

GH-1234: Release serialization buffers when compression fails - #1290

GH-1234: Release serialization buffers when compression fails#1290
jbonofre merged 2 commits into
apache:mainfrom
pj-workspace:fix/vector-unloader-exception-cleanup

pj-workspace commented Sep 10, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

This comment has been minimized.

Copilot AI left a comment

Uh oh!

jbonofre left a comment

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Conversation

pj-workspace commented Sep 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What's Changed

Validation

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Uh oh!

This comment has been minimized.

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

jbonofre left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

pj-workspace commented Sep 10, 2026 •

edited

Loading