andygrove opened a new pull request, #3824: URL: https://github.com/apache/datafusion-comet/pull/3824
## Which issue does this PR close? Closes #3821. ## Rationale for this change The shuffle writer's memory reservation only tracks buffered `RecordBatch`es and `partition_indices`, but significant untracked allocations occur in the write path: - **BufBatchWriter buffers** — up to `write_buffer_size` (1 MB) per partition - **BatchCoalescer** — holds up to `batch_size` rows of data internally - **Arrow IPC StreamWriter** — internal serialization buffers - **Compression encoder state** — lz4/zstd/snappy frame buffers - **`interleave_record_batch()` output** — new RecordBatch per partition iteration These untracked allocations cause actual RSS to significantly exceed the configured memory pool limit, because the pool doesn't know about them when making spill decisions. ## What changes are included in this PR? Apply a 2x multiplier to `mem_growth` when reserving memory in `buffer_partitioned_batch_may_spill`. This triggers spills earlier, keeping total process memory closer to the configured limit. ## How are these changes tested? Benchmarked with `shuffle_bench` on TPCH SF100 lineitem (100M rows, 16 columns, hash partitioning with 200 partitions, lz4 compression) at memory limits of 2 GB, 4 GB, 8 GB, and 16 GB. ### Before | Memory Limit | Peak RSS | RSS / Limit | Write Time | Throughput | |---|---|---|---|---| | 2 GB | 3.8 GB | 1.78x | 43.4s | 2.30M rows/s | | 4 GB | 6.7 GB | 1.56x | 45.1s | 2.22M rows/s | | 8 GB | 11.0 GB | 1.28x | 46.4s | 2.15M rows/s | | 16 GB | 13.1 GB | 0.76x | 51.9s | 1.93M rows/s | ### After (2x multiplier) | Memory Limit | Peak RSS | RSS / Limit | Write Time | Throughput | |---|---|---|---|---| | 2 GB | 2.8 GB | 1.37x | 43.8s | 2.34M rows/s | | 4 GB | 4.0 GB | 0.94x | 43.2s | 2.33M rows/s | | 8 GB | 6.7 GB | 0.78x | 45.7s | 2.20M rows/s | | 16 GB | 10.8 GB | 0.63x | 47.9s | 2.10M rows/s | Peak RSS is reduced by 26-40% across all configurations while throughput is equal or better. The 4 GB case now stays within the configured limit (0.94x). All existing tests pass. -- 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]
