[dongjoon-hyun]

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

[dongjoon-hyun]

Could you review this PR when you have some time, @viirya ?

[viirya]

LGTM. heldBuffers holds exactly the extra reference at each throw point in Arrow 19's AbstractCompressionCodec.compress, and apache/arrow-java#1290 confirms the 20.0.0 double release that the removal note warns about. One optional suggestion on the test inline.

[viirya]

After the Arrow 20 upgrade, closing the compressed validity buffer a second time ends up as a suppressed exception on e. So this intercept still passes, and the test fails only later in root.close(). Adding assert(e.getSuppressed.isEmpty) makes it fail here, with a clearer message.

[dongjoon-hyun]

Thank you, @viirya. Added in 705be92, with the note about the Arrow upgrade moved next to the assertion. I ran the test with the VectorUnloader and AbstractCompressionCodec changes of apache/arrow-java#1290 on the test classpath. It now fails at this assertion with Array(java.lang.IllegalStateException: RefCnt has gone negative) was not empty, and it passes on Arrow 19.0.0.

[dongjoon-hyun]

Merge Summary:

• merged into master e438d0c
• merged into branch-4.x 2c4ebfe

Posted by merge_spark_pr.py

[dongjoon-hyun]

Thank you so much for your valuable feedback, @viirya !