apache / apache/datafusion-comet
Investigate: why Comet shuffle files can be larger than Spark's, and whether byte-based block sizing would fix it
- Dominant language
- Scala
- Stars
- 1.3k
- Forks
- 373
- Avg merge
- 2d 4h
- Merged PRs (30d)
- 198
Description
## Summary
apache/datafusion-comet#3882 recorded a real problem — Comet shuffle files can be substantially
larger than Spark's — but its framing has been overtaken by what we have measured since, and two
of the three fixes it proposed have already been investigated and closed. Rather than keep editing
it, I am closing it and restating the current position here, along with what still needs
measuring.
**Nothing in the "what this asks for" section below has been benchmarked.** This is a request for
investigation, not a proposal for an agreed change.
## Current behavior
### Block format
Every shuffle block is a self-contained Arrow IPC stream.
`ShuffleBlockWriter::write_batch` (`native/shuffle/src/writers/shuffle_block_writer.rs:232`)
emits, per block:
- a 20-byte header — 8-byte compressed length, 8-byte field count, 4-byte codec tag
(`LZ4_` / `ZSTD` / `SNAP` / `NONE`), see `try_new_inner:136`
- a compressed payload holding a complete Arrow IPC stream: schema message, dictionary messages,
one record batch message, and the 8-byte end-of-stream marker
So the Arrow schema flatbuffer is on the wire once per *block*, not once per partition or once per
file. That is structural rather than an oversight: Spark shuffle is block-based, and a reducer
fetches blocks from many mappers in arbitrary order, so a block cannot reference a schema written
earlier in a stream. Both apache/datafusion-comet#1186 and apache/datafusion-comet#2928 closed on
that reasoning.
Since apache/datafusion-comet#5006 the schema message is pre-encoded once in
`ShuffleBlockWriter::try_new` and copied verbatim into each block (`SchemaEncoding::Precoded`).
That removed the CPU cost of re-serializing the flatbuffer per block; it did not remove the bytes.
Schemas containing dictionary types still fall back to a real `StreamWriter` per block
(`SchemaEncoding::Fallback`), because dictionary-id bookkeeping ties the schema and record-batch
encoding together, and remote (RSS/Celeborn) writers always take that fallback path.
### Block sizing
Block boundaries are row counts, not byte counts. Native shuffle coalesces to
`session_config().batch_size()` — i.e. `spark.comet.batchSize`, default 8192 rows — through
`BufBatchWriter`'s `BatchCoalescer` (`native/shuffle/src/writers/buf_batch_writer.rs:85`), with a
zero-copy passthrough for batches that already exceed it. The JVM columnar path uses
`spark.comet.shuffle.jvm.batchSize`, also 8192. A narrow schema therefore produces small blocks and
a wide schema produces large ones, while the fixed per-block overhead stays the same.
### Compression
Compression is applied per block and wraps the entire IPC stream, so the schema message is inside
the compressed region. The default codec is lz4 (`spark.comet.shuffle.compression.codec`).
Spark, by contrast, compresses the shuffle stream at the block level with an lz4 block size of
`spark.io.compression.lz4.blockSize` (32 KB by default). At 8192 rows over a few narrow columns a
Comet block can be well below that, so each block becomes its own compression window and the codec
has little history to work with.
## What we already know
- The original report (karuppayya, on apache/datafusion-comet#3882): 204M rows, 8 columns
(7 string, 1 timestamp), scan → repartition → write. Comet wrote 4.77 GB against Spark's
1.58 GB — 25.1 vs 8.3 bytes per record.
- **The repro is strongly data-dependent.** My first attempt produced Comet files *smaller* than
Spark's. It only reproduced once I switched to short, unique strings, i.e. when the payload per
block is small relative to the fixed overhead.
- **Batch size dominates in the one workload we measured.** At the default batch size, Comet files
were 50% larger than Spark's and the query was 10% slower; doubling the batch size took that to
8% larger and 15% faster than Spark.
- **The per-block schema is mostly not the cause.** apache/datafusion-comet#2928 measured exactly
this: reusing a single encoder and stream writer across all batches gave a ~25% size reduction
for a single-column schema, dropping to **~1% at 100 columns of mixed types**. The schema
flatbuffer grows with column count, but the data per batch grows faster, so the relative
overhead shrinks as schemas widen. It was closed as not worth doing for size.
- **Possible exception: deeply nested schemas.** comphead reported ~1.5x on nested data and
observed that the schema can exceed the data there. That case is tracked separately in
apache/datafusion-comet#5355 and was not covered by the #2928 measurement, which used flat wide
schemas.
The residual hypothesis, then, is not "we repeat the schema too often". It is "row-count-based
block sizing produces blocks that are too small for the compressor on narrow schemas."
## What this issue asks for
That hypothesis has never been measured directly, and neither has the cost of the alternative.
Specific things worth doing:
1. **Attribute the bytes.** Take a repro that shows the gap and break a shuffle file down into
header bytes, schema-message bytes, record-batch metadata, and buffer payload — both
compressed and uncompressed. We are currently arguing about which term dominates by reading the
format rather than by measuring it. Note that the schema sits inside the compressed region, so
its uncompressed size overstates its real cost, and repeating near-identical schema bytes in
every block is exactly the pattern a compressor handles well within a block but cannot exploit
across blocks.
2. **Measure block size against compression ratio.** Sweep `spark.comet.batchSize` and the three
codecs across at least one narrow-schema and one wide-schema dataset, and report compressed
bytes per row. The "doubling the batch size" result above is a single data point from a single
workload, and it is currently carrying more weight than it should.
3. **Decide whether byte-based block sizing is worth building.** This was the main proposal in
apache/datafusion-comet#3882 and it remains unimplemented. The hard part is estimating a
batch's serialized size before encoding it; worth evaluating whether
`RecordBatch::get_array_memory_size` is a good enough proxy, or whether an encode-and-split
approach is needed. It also interacts with the reader: `spark.comet.batchSize` is what
downstream native operators expect to receive, so decoupling block size from batch size may
require the reader to re-slice.
4. **Check whether the deeply-nested case is genuinely different**, or whether it is
apache/datafusion-comet#5355 in another guise. If schema bytes really can exceed data bytes
there, that is the one place where revisiting the format might pay off.
5. **Quantify what tuning already buys.** If the practical answer for most users turns out to be
"raise `spark.comet.batchSize`", that belongs in the tuning guide with numbers attached, and
this becomes a documentation change rather than a format change.
## Tooling
- `native/shuffle/src/bin/shuffle_bench.rs` — standalone shuffle write benchmark, added in
apache/datafusion-comet#3752
- `native/shuffle/benches/shuffle_writer.rs` and `native/shuffle/benches/row_columnar.rs`
- `spark/src/test/scala/org/apache/spark/sql/benchmark/CometShuffleBenchmark.scala`
- `benchmarks/pyspark/benchmarks/shuffle.py`
- apache/datafusion-comet#3909 was a shuffle size comparison benchmark I opened for this and closed
unmerged. Worth resurrecting as the starting point for item 1.
## Related
Replaces apache/datafusion-comet#3882.
- apache/datafusion-comet#5002 — native shuffle writer optimizations (umbrella)
- apache/datafusion-comet#5198 — shuffle performance audit
- apache/datafusion-comet#5355 — scan + shuffle-write slower than Spark on deeply nested schemas
- apache/datafusion-comet#5006 — pre-encoded schema per writer (CPU only, landed)
- apache/datafusion-comet#2928, apache/datafusion-comet#1186 — earlier attempts at removing the
per-block schema, both closed
Carried over from apache/datafusion-comet#3882: kazuyukitanimura noted that
`WriteDistributionAndOrderingSuite` is disabled on Spark 3.5+ by apache/datafusion-comet#834 and
should be re-enabled once the underlying issue is addressed.
Contributor guide
Research direction
Start with native/shuffle/src/bin/shuffle_bench.rs and the closed #3909 shuffle size comparison benchmark, then review native/shuffle/benches/shuffle_writer.rs, native/shuffle/benches/row_columnar.rs, and the Spark benchmark. Measure byte attribution and batch-size/codec sweeps on narrow and wide datasets, including the nested case, and report whether tuning, documentation, or byte-based sizing is justified.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, rust, scala, spark
- Domain
- data-engineering, distributed-systems, performance
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 45/100