Repository navigation
Description
Part of #7267.
Problem
The native shuffle writer builds output batches in two places, and both limit only
by row count:
native/shuffle/src/partitioners/partitioned_batch_iterator.rs:227gathers the
rows for one output partition withinterleave_record_batch. For partitioning
that hands out runs of rows instead of single rows (positional round robin,
PartitionIndices::Runs,multi_partition.rs:241-245), theRunIteratoruses
concat_batchesinstead (line 330).native/shuffle/src/writers/buf_batch_writer.rs:133merges small batches with
arrow'sBatchCoalescerbefore serializing them.
Both are capped at the session batch size (shuffle_writer.rs:348, 502):
insert_batch slices input to batch_size rows (multi_partition.rs:677), and the
coalescer only merges batches smaller than batch_size; batches of batch_size
rows or more pass through without a copy (with_biggest_coalesce_batch_size,
buf_batch_writer.rs:133-134).
So an overflow needs one output chunk of batch_size rows to hold more than 2 GiB
in one string/binary column, which is more than ~256 KiB per value at the default
8192. The gather can still hit this when every input batch fits: it collects one
partition's rows from many input batches, so skew can group the large rows. When it
happens, the task fails with OffsetOverflowError. Neither path retries.
There is also a third limit. After compression, a block over i32::MAX bytes is
rejected (shuffle_block_writer.rs:373-378, "Shuffle block size ... exceeds
maximum size"). One block holds one batch, so a batch that is large in total bytes
fails here even if no single column overflows. The single-partition path passes an
input batch through as one block (single_partition.rs:28-31).
DataFusion's sort merge already handles the gather case: retry_interleave
(sorts/builder.rs:251, DataFusion 55.2.0) catches OffsetOverflowError and retries
with half the rows. It is private, so Comet needs its own copy.
Workaround
Lower spark.comet.batchSize. All three limits scale with rows per batch.
Proposed fix
- Around
interleave_record_batch/concat_batchesin
partitioned_batch_iterator.rs: onOffsetOverflowError, split the index chunk
in half and retry, emitting more, smaller batches. Same approach as
retry_interleave. - In
buf_batch_writer.rs: flush the coalescer before pushing a batch that would
take a string/binary column past a byte budget. If sub-issue 4 lands first, use
the coalescer's byte limit instead. - Block size: split a batch whose encoded size would pass
i32::MAXinto several
blocks before writing, rather than failing. - Reader side: no change. A partition can already hold any number of blocks.
Tests
- The i32 limit can't be lowered for a test, so factor the retry into a helper that
takes the gather function. Test it with an injected function that returns
OffsetOverflowErrorabove N rows, and assert all rows come out, in order, across
several batches. - Byte-budget test for
buf_batch_writer.rswith a small configured budget. - End to end: repartition a table with a binary column of values a few hundred KB
in size and compare with Spark.
Notes
- Shuffle changes need the shuffle CI suites.
- Celeborn and the JVM shuffle path are out of scope here.
Activity
[-][Offset overflow] Native shuffle writer: retry with fewer rows on offset overflow[/-][+]Offset overflow- Native shuffle writer: retry with fewer rows on offset overflow[/+]