dongjoon-hyun opened a new pull request, #59991:
URL: https://github.com/apache/spark/pull/59991
### 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 with a
codec that closes its input on failure like the one in Arrow 20.
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 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
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
