andygrove opened a new issue, #6183: URL: https://github.com/apache/datafusion-comet/issues/6183
### Describe the bug `spark.comet.shuffle.native.writeBufferSize` is declared with `bytesConf(ByteUnit.MiB)` and a default of `1` ([CometConf.scala#L721-L731](https://github.com/apache/datafusion-comet/blob/67803a7a422c44de07af1e5d25c1dbeae8df68d4/spark/src/main/scala/org/apache/comet/CometConf.scala#L721-L731)), so `COMET_SHUFFLE_NATIVE_WRITE_BUFFER_SIZE.get()` returns the size in MiB. `CometNativeShuffleWriter` passes that number straight into the `write_buffer_size` proto field ([CometNativeShuffleWriter.scala#L326-L327](https://github.com/apache/datafusion-comet/blob/67803a7a422c44de07af1e5d25c1dbeae8df68d4/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometNativeShuffleWriter.scala#L326-L327)), and the native planner uses it as a byte count ([planner.rs#L1928](https://github.com/apache/datafusion-comet/blob/67803a7a422c44de07af1e5d25c1dbeae8df68d4/native/core/src/execution/planner.rs#L1928)). With the default configuration the native shuffle writer gets a 1-byte write buffer, not the documented 1 MB. The generated config table shows the default as `1048576b`. Values that users set are off by the same factor: | Setting | `.get()` on the JVM | Bytes the native writer uses | | --- | --- | --- | | default | 1 | 1 | | `8m` | 8 | 8 | | `8388608` | 8388608 | 8388608, correct only because both sides misread the unit | With a 1-byte buffer: - The output data file and the spill file are wrapped in `BufWriter::with_capacity(1, ...)` ([local_partition_writer.rs#L127](https://github.com/apache/datafusion-comet/blob/67803a7a422c44de07af1e5d25c1dbeae8df68d4/native/shuffle/src/writers/local/local_partition_writer.rs#L127), [spill.rs#L222](https://github.com/apache/datafusion-comet/blob/67803a7a422c44de07af1e5d25c1dbeae8df68d4/native/shuffle/src/writers/local/spill.rs#L222)), and `BufBatchWriter` writes out each block as soon as it is encoded ([buf_batch_writer.rs#L171](https://github.com/apache/datafusion-comet/blob/67803a7a422c44de07af1e5d25c1dbeae8df68d4/native/shuffle/src/writers/buf_batch_writer.rs#L171)), so nothing is coalesced. - `BufBatchWriter::flush` shrinks the shared scratch buffer to the configured size ([buf_batch_writer.rs#L212](https://github.com/apache/datafusion-comet/blob/67803a7a422c44de07af1e5d25c1dbeae8df68d4/native/shuffle/src/writers/buf_batch_writer.rs#L212)), which frees it after every partition. That undoes the scratch reuse added in #5568. - The spill copy allocates a 1-byte read buffer ([local_partition_writer.rs#L286](https://github.com/apache/datafusion-comet/blob/67803a7a422c44de07af1e5d25c1dbeae8df68d4/native/shuffle/src/writers/local/local_partition_writer.rs#L286)), so the single `pread` path added in #5916 ([local_partition_writer.rs#L401](https://github.com/apache/datafusion-comet/blob/67803a7a422c44de07af1e5d25c1dbeae8df68d4/native/shuffle/src/writers/local/local_partition_writer.rs#L401)) never applies and every range goes through `io::copy`. The Rust benchmarks and `shuffle_bench` pass `1048576` directly, so they measure the intended configuration rather than what runs under Spark. #2899 added the config with a `max(Int.MaxValue)` clamp, which always sent 2 GiB whatever the setting. #3914 changed the clamp to `min`, and since then the native writer has received the value in MiB. ### Steps to reproduce With default settings, `CometConf.COMET_SHUFFLE_NATIVE_WRITE_BUFFER_SIZE.get()` returns `1`. The conversion can be checked in isolation: `JavaUtils.byteStringAs("1048576b", ByteUnit.MiB)` returns `1`. To see the cost, compare `shuffle_bench --write-buffer-size 1` against `--write-buffer-size 1048576`. On a macOS laptop with TPC-H SF1 `lineitem` (strings read as `Utf8`), 200 hash partitions, lz4, one warmup and five iterations: | Run | 1 byte (current) | 1 MiB (intended) | | --- | --- | --- | | No memory limit, avg time | 2.161s | 2.169s | | 64 MiB memory limit (17 spills), avg time | 2.054s | 1.989s | | 64 MiB memory limit, write time | 0.192s | 0.116s | There is no difference without spilling. With spilling the write is about 3% slower overall, and its write time is about 65% higher. Linux, where `io::copy` can use `copy_file_range`, was not measured. ### Expected behavior The native writer uses 1 MiB by default, and `8m` means 8 MiB. ### Additional context Suggested fix: declare the config with `bytesConf(ByteUnit.BYTE)` and a default of `1024 * 1024`. Bare numbers keep their current meaning, while the default and values with a unit become correct. A test should assert the value that ends up in the shuffle writer proto. After the fix, each shuffle-writing task holds up to about 3 MiB of buffers that the memory pool does not track: the output writer, the spill writer and the spill read buffer. They are allocated once per task, not once per partition. -- 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]
