Skip to content

fix: bound a shuffle read's coalesced batches by bytes as well as rows - #24

Merged
msafonov merged 1 commit into
branch-1.1from
joom/1.1-shuffle-read-coalesce-bytes
Oct 18, 2026
Merged

msafonov merged 1 commit into
branch-1.1from
joom/1.1-shuffle-read-coalesce-bytes

Conversation

@msafonov

Copy link
Copy Markdown
Collaborator

Which issue does this PR close?

Closes #. (Joom fork, OOM diagnosis of night 10-17 -> 10-18, fix A)

Rationale for this change

ShuffleReadCoalescer joins decoded shuffle blocks until spark.comet.batchSize rows, with no bound on bytes, outside any memory pool. In prepare_device_events (stage 52: shuffle read of HLL sketches, about 190 KB per row, about 6 rows per block) that is about 1.5 GB per batch per task plus the concat_batches copy: allocated reached 16 GiB with 0 reserved within 4 s of the stage start, and 12 executors were OOMKilled (em32, limit 26g).

What changes are included in this PR?

  • ShuffleReadCoalescer::new(target_rows, target_bytes): a batch is passed on at the row target or the byte target, whichever comes first; a block that alone exceeds the byte target passes through unchanged (flushing what is pending first).
  • Bytes are the sliced size of each column (ArrayData::get_slice_memory_size), because a decoded block slices every column out of one IPC body and buffer capacity would count the body once per column.
  • New config spark.comet.shuffle.read.coalesce.maxBytes, default 64 MiB, 0 = no byte bound. It reaches the JVM decoder through createShuffleReadCoalescer(batchSize, maxBytes) and the native ShuffleScanExec through the new ShuffleScan.coalesce_max_bytes proto field.
  • Narrow rows are unaffected: 8192 rows stay under 64 MiB up to about 8 KB per row.

How are these changes tested?

  • New Rust tests in read_coalescer.rs: stops_joining_wide_rows_at_the_byte_target (60 blocks of 6 x 190 KB rows: one batch of >60 MiB without the byte bound, batches <= 4 MiB with it, same rows in order), passes_a_block_larger_than_the_byte_target_through, counts_a_body_shared_by_all_columns_once. With the byte checks disabled all three fail.
  • Local (aarch64 container): cargo fmt --check, cargo clippy --all-targets --workspace -D warnings, workspace tests green (core 661, shuffle 193, spark-expr 972, jni-bridge 28, common 70).
  • Cluster validation on the datasci-ab playground follows (prepare_device_events vs vanilla).

🤖 Generated with Claude Code

ShuffleReadCoalescer joined decoded blocks until spark.comet.batchSize rows. With wide rows
(an HLL sketch of about 190 KB per row, a map task writing about 6 rows per reduce block)
8192 rows is about 1.5 GB per batch, held outside any memory pool, plus the copy
concat_batches makes. prepare_device_events reached 16 GiB allocated with 0 reserved
within seconds of its shuffle read and its executors were OOMKilled.

The coalescer now passes a batch on at spark.comet.batchSize rows or
spark.comet.shuffle.read.coalesce.maxBytes (default 64 MiB) of column data, whichever comes
first; a block larger than the bound passes through as it is. The bytes are the sliced size
of each column, because a decoded block slices all of its columns out of one IPC body and
the buffer capacity would count that body once per column. Both readers take the bound: the
JVM decoder through createShuffleReadCoalescer and the native ShuffleScanExec through the
new ShuffleScan.coalesce_max_bytes field.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
@github-actions github-actions Bot added bug Something isn't working area:shuffle area:writer labels Oct 18, 2026
@msafonov
msafonov merged commit cbee655 into branch-1.1 Oct 18, 2026
36 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:shuffle area:writer bug Something isn't working

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant