parthchandra opened a new issue, #7269:
URL: https://github.com/apache/datafusion-comet/issues/7269
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:227`
gathers the
rows for one output partition with `interleave_record_batch`. For
partitioning
that hands out runs of rows instead of single rows (positional round robin,
`PartitionIndices::Runs`, `multi_partition.rs:241-245`), the `RunIterator`
uses
`concat_batches` instead (line 330).
- `native/shuffle/src/writers/buf_batch_writer.rs:133` merges small batches
with
arrow's `BatchCoalescer` before 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_batches` in
`partitioned_batch_iterator.rs`: on `OffsetOverflowError`, 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::MAX` into
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
`OffsetOverflowError` above N rows, and assert all rows come out, in
order, across
several batches.
- Byte-budget test for `buf_batch_writer.rs` with 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.
--
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]
