Skip to content

Offset overflow- Native shuffle writer: retry with fewer rows on offset overflow #7269

Open

Description

@parthchandra

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.

Activity

  1. changed the title [-][Offset overflow] Native shuffle writer: retry with fewer rows on offset overflow[/-] [+]Offset overflow- Native shuffle writer: retry with fewer rows on offset overflow[/+] on Oct 27, 2026
  2. self-assigned this
    on Oct 27, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

Labels

area:shuffleShuffle (JVM and native)enhancementNew feature or request

Type

No type

Projects

No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions