dongjoon-hyun opened a new pull request, #59998: URL: https://github.com/apache/spark/pull/59998
### What changes were proposed in this pull request? This PR makes `BoundedZstdCompressionCodec` copy a compressed output to a smaller buffer and release the buffer that zstd wrote into, when that at least halves the memory that the output takes. The codec now overrides `compress`. After `super.compress` returns, it checks `allocator.getRoundingPolicy.getRoundedSize(length) <= capacity / 2`. If this holds, it copies the output with its length prefix to a new buffer and closes the big one. - Empty and raw (incompressible) outputs are already in a buffer of their size, so they are not copied. - If the new buffer cannot be allocated (`OutOfMemoryException`), the codec keeps the big buffer. The output is complete, so this only gives up the saving, and a batch that succeeded before does not fail now. - `createVectorUnloader` (SPARK-60776) wraps this codec, so `TrackingCompressionCodec` holds the new buffer. It assumes that a failed `compress` still holds its input, but the shrink runs after `super.compress` released the input. So the shrink must not fail with `OutOfMemoryException`, and a code comment explains this. All zstd write paths get this codec from `ArrowCompressionUtils.createCompressionCodec`: Python UDF input, `toPandas`, Spark Connect results, External UDF, and the Arrow cache. lz4 is not changed, because `Lz4CompressionCodec` already requests a buffer of the output size. The class documentation now limits its removal note to `doCompress`. Arrow 20.0.0 will fix the overrun (apache/arrow-java#1119), but the main branch of Arrow Java does not shrink the output, so `compress` has to stay until it does. Alternatives that were not taken: - Shrink in `doCompress`: the SPARK-60707 test calls `doCompress` directly and checks for a write past the buffer that zstd wrote into, so it would check the copy instead. It would also need its own check for the raw fallback. - Copy at a lower threshold: with 25% saved as the threshold, the 1,000,000-row UUID batch below would hold 22.20 MiB instead of 38.47 MiB, but its peak would be 56.67 MiB instead of 38.47 MiB. - A scratch buffer per batch or per task: it saved no CPU in the measurements and holds the same memory or more. `CompressionCodec` has no `close()`, and some paths create a codec per group or batch, so each path would need new code to release it. - A different `RoundingPolicy`: it applies to all Arrow allocations under the shared `ArrowUtils.rootAllocator`, and the buffer would still be `compressBound`-sized. ### Why are the changes needed? Arrow Java 19.0.0 leaves the zstd output in the buffer that it allocated before compressing, and [its main branch](https://github.com/apache/arrow-java/blob/67ef5abf308a7a158c4fa8eeebb72b87ca5a6f97/compression/src/main/java/org/apache/arrow/compression/ZstdCompressionCodec.java#L41-L55) still does: - `ZstdCompressionCodec.doCompress` requests `8 + compressBound(len)` bytes, and the allocator rounds that up to a power of two below 16 MiB. - `VectorUnloader` keeps the output of every vector buffer until the record batch is closed, that is, until it is serialized. So a batch holds about as much extra direct memory as its uncompressed vectors, whatever the compression ratio. For comparison, [Arrow C++](https://github.com/apache/arrow/blob/926885f1b1ede73badd24c2e6f63c9429441e99a/cpp/src/arrow/ipc/writer.cc#L210-L250) shrinks its output right after compressing. This is not a correctness issue, and the cost depends on the batch size: - A 10,000-row batch, the default `spark.sql.execution.arrow.maxRecordsPerBatch`, costs 0.1 to 2.6 MiB per task with narrow rows. - A 64 MB batch costs 64 to 88 MiB per task, or about 1.0 to 1.4 GiB with 16 concurrent tasks. Batches reach that size with wide rows (the default `spark.sql.execution.arrow.maxBytesPerBatch` is 64MB) or with a larger `maxRecordsPerBatch`. Spark's `MemoryManager` does not track this memory, and when direct memory runs out, the task fails with `OutOfMemoryException`. This only happens when `spark.sql.execution.arrow.compression.codec` is `zstd`; the default is `none`. ### Does this PR introduce _any_ user-facing change? No. The serialized bytes are the same, and only the direct memory in use changes. ### How was this patch tested? Pass the CIs with the newly added test cases in `ArrowCompressionUtilsSuite`: - A much smaller output moves to a buffer of its size and the big buffer is released. This test fails without this change. - An output that would not at least halve its buffer is not copied, so the peak allocation does not grow. - When there is no memory for the new buffer, the output stays in its buffer. I checked these tests by changing the code: each test fails when its branch is removed, and bringing back the SPARK-60707 overrun still fails the SPARK-60707 tests. `ArrowCompressionUtilsSuite`, `ArrowCachedBatchSerializerSuite`, `ArrowConvertersSuite`, and `ExecuteExternalUDFExecSuite` pass locally (151 tests). I also ran `createVectorUnloader` with this codec in a temporary script, which is not part of this PR. A batch that fails after an output was shrunk, and a batch that runs out of memory after a shrink gave up, both release each buffer once: there is no suppressed exception, the allocator goes back to the size of the vectors, and closing it succeeds. I measured the memory with a temporary benchmark, which is not part of this PR. It fills a root with `ArrowWriter` and calls `VectorUnloader.getRecordBatch()`, with an `AllocationListener` on the child allocator of the root to track allocations and releases. The environment is an Apple M3 Max, JDK 25.0.5, Arrow 19.0.0, and zstd level 3. For all measured batches, the serialized IPC bytes are the same as before. | Batch (MiB) | Vectors | Compressed | Held before | Held after | Peak before | Peak after | |---|---:|---:|---:|---:|---:|---:| | Sequential BIGINT, 10,000 rows | 0.13 | 0.02 | 0.13 | 0.03 | 0.13 | 0.16 | | `lineitem`, 10,000 rows | 2.63 | 0.37 | 2.63 | 0.47 | 2.63 | 0.97 | | `lineitem`, 1,000,000 rows | 178.00 | 37.95 | 178.70 | 52.75 | 178.70 | 83.45 | | `lineitem`, 64 MB limit (389,390 rows) | 87.00 | 14.75 | 88.00 | 22.88 | 88.00 | 38.88 | | Text, 64 MB limit (8,348 rows) | 64.06 | 13.97 | 64.28 | 16.03 | 64.28 | 80.25 | | UUID strings, 1,000,000 rows | 68.00 | 20.98 | 38.59 | 38.47 | 38.59 | 38.47 | | Random BIGINT, 1,000,000 rows (raw output) | 8.00 | 7.63 | 8.13 | 8.00 | 8.13 | 8.00 | - `lineitem` is a generated TPC-H-like batch with 16 columns and 37 Arrow buffers. The text batch has about 8 KB of JSON-like text per row. - Held and peak are on top of the vectors. Held is what a batch holds until it is serialized, and peak is the highest point during `getRecordBatch`. - The peak goes up when one big buffer dominates a batch, as in the text batch, because both buffers are alive during the copy. The threshold limits this to half of that buffer, and it lasts only for the copy. - These numbers are Arrow allocator accounting, not RSS. A new buffer above 4 MiB commits only the pages that zstd writes (checked on macOS), so the effect on RSS is smaller. I did not measure RSS or OOMs in a real executor, or anything on Linux. For the CPU cost, I compared the thread CPU time of `getRecordBatch` and `close` with and without the shrink. Each value is the median after/before ratio across rounds: | Batch | 1 thread | 8 threads | 8 threads with jemalloc | |---|---:|---:|---:| | Sequential BIGINT, 100 rows | 1.015 | 1.029 | 1.086 | | `lineitem`, 100 rows | 1.029 | 1.040 | 1.057 | | `lineitem`, 1,000 rows | 1.010 | 1.030 | 1.035 | | `lineitem`, 10,000 rows | 1.016 | 1.012 | 1.015 | | `lineitem`, 1,000,000 rows | 1.006 | 1.003 | 1.016 | | Text, 64 MB limit | 0.997 | 0.998 | 1.021 | - Batches of 1,000 rows or more cost at most 3.5% more. - 100-row batches cost up to 8.6% more, but the absolute cost is small: about 3 us per `lineitem` batch with 1 thread. - Other builds were running on the machine (load average 46 to 80), so wall time was not usable. Thread CPU time does not show lock waits, so I cannot rule out a small wall time increase with 8 threads. ### 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]
