feat(core): range shuffle reads only the bytes a consumer's range covers

[avantgardnerio]

Draft. The public RangeShuffleWriter / RangeShuffleReader traits and registration point @phillipleblanc asked for on #2204 are still to come, and the numbers below are from a single box where the network is free. Not for review yet.

Two goals

Read less. A stage that routes rows by value range hands each consumer a set of producer files, of which the consumer wants only the rows in its own cut range. Today it reads them whole and the RangeFilterExec above it discards the rest.

Make persistent shuffle possible. Ballista shuffle files live and die with the executor that wrote them. Serving them from object storage instead is what survives an executor rotating out — which matters on spot infrastructure, and is where both Coralogix and Spice already flush to S3. An object store has no process to run a lookup for you: it answers GET with a Range: header and nothing else. So the seek has to be expressible as a byte range decided by the reader, or it does not port at all.

The second goal is what shapes the design. It would be easier to have the serving side resolve a value range to bytes, and it would also be a dead end.

What the waste is

RangeShuffleReaderExec does no filtering, so the ratio between the rows it emits and the rows the halo trim keeps is the over-fetch — no new instrumentation needed:

scale	reader emits	halo trim keeps	read amplification
1e7	20.5 M / 626 MB	12.1 M / 372 MB	1.69×
1e8	275 M / 8.2 GB	121 M / 3.6 GB	2.27×

At 1e8 that is 56% of the shuffle read fetched, decoded, and thrown away.

All numbers here are at K=8 on 2 executors × 4 vcores — one output partition per vcore, which is the shape worth optimising: every core busy through every stage, and no more partitions than that.

What this PR does

Skipping bytes means seeking, and seeking means the data file needs a chunk index. The Arrow IPC stream format the passthrough shuffle writes has none: finding the bytes that hold a given value means walking every message from the head.

Three pieces, each useless without the next:

A writer that emits a seekable format. RangeShuffleWriterExec writes the Arrow IPC file format, whose footer records a Block { offset, metadata_length, body_length } per record batch. A distinct operator rather than a flag on ShuffleWriterExec, because the two IPC framings are not interchangeable on the read side — a mode flag would put every existing reader one config change away from failing to decode what it is handed. The existing shuffle's format never moves.

A value index beside each data file. data-{file_id}.rangeidx.arrow, one row per IPC message in file order:

<sort_expr_0..n>   key types, nullable    first row's key in this batch
is_dict            Boolean                row describes a dictionary
byte_offset        UInt64                 absolute, from the data file's footer
byte_len           UInt64
num_rows           UInt64, nullable

5 KB of index against 19 MB of data. Keys are one typed column per ORDER BY expression rather than an encoded blob, so multi-column keys, DESC, and mixed types work by construction, and the schema names what the file is indexed on.

A reader that fetches only what its range covers. Locally it seeks. Remotely it fetches the index, searches it, and asks for byte ranges — which is the step that ports to object storage unchanged.

Result

K=8, h2o Q8, verified against single-process DataFusion:

scale	fetched before	fetched now	what the window needs	amplification
1e7	626 MB	373 MB	372 MB	1.69× → 1.01×
1e8	8.2 GB	3.6 GB	3.6 GB	2.27× → 1.00×

The reader now fetches what the window consumes and essentially nothing else. At 1e8 that is 4.6 GB of the 8.2 GB no longer crossing the shuffle boundary.

Who does the searching, and why it matters

The serving side could resolve a value range to bytes itself. It deliberately does not.

If it did, the index file would be pointless — a server holding the data can binary-search it directly — and none of this would port to object storage. Byte offsets exist in the index precisely so a reader can turn a value range into a Range: header. Keeping the lookup on the reading side means swapping Arrow Flight for an object store changes the transport and nothing else, because the selection was never on the serving side to begin with.

Concretely, a consumer task does:

1. fetch the index for a source — one small object
2. binary-search it locally for its own [lo, hi)
3. ask for the byte ranges that search produced

Step 3 is a GET with a Range: header wherever the bytes live. The Flight service that answers it opens a file, seeks, and copies: it reads no index, compares no value, and has no notion of what a range means.

The cost is two fetches per source rather than one. Batched across sources that is one extra round trip per read, not one per source, because the dependency is within a file and not across them. On a single box that cost is invisible and the bandwidth saving is fully visible — the reverse of a real network, which is why the caveat at the top matters.

Protocol

BALLISTA_PROTOCOL_VERSION 2 → 3. FetchPartition gains layout, file_kind, and byte_ranges; is_sort_shuffle is gone, because that bool selected a path shape and misreads as data-vs-index once a file_kind sits beside it.

Coordinates address a logical partition's bytes; byte_ranges narrows them:

• SORT + DATA + no ranges → the Flight service resolves the partition through its own index. Unchanged, one fetch.
• PASSTHROUGH + DATA + no ranges → the whole file, as today.
• any layout + ranges → exactly those absolute offsets, resolving nothing.
• any layout + INDEX → the sidecar, whose encoding follows from the layout.

Sort-based shuffle therefore keeps its single fetch and its existing behaviour. Moving its lookup to the reader later needs no further protocol change — fetch the index, read two i64s, ask for the bytes — and would delete stream_sort_shuffle_block rather than add a path beside it. That is also what would make sort shuffle object-store-native, for free.

Scope, and what is untouched

The range writer is planted only where the consumer is a RangeShuffleReaderExec, which the AQE adapter chooses when a stage's child declares an output ordering. In practice that is the parallel-window rewrites and a sort above a shuffle boundary. It requires adaptive planning; the static planner never produces the pair.

Everything else is byte-for-byte as before: the passthrough shuffle still writes IPC stream, sort-based shuffle still writes its own consolidated format and still serves it in one fetch, and no existing reader's path changes.

Within that scope, one property is worth naming: an IPC stream file is readable up to its last complete message, while an IPC file with no footer does not open at all. A range-shuffle task killed mid-write therefore leaves a file that fails on open rather than one that reads partially. Task-level retry covers it, and it applies only to files this PR's writer creates.

Testing

• Round trip driving RangeShuffleWriterExec into fetch_partition_local, rather than fabricating files and asserting the reader agrees with itself — the specific gap called out on feat(core): ValueIndexExec + ValueIndexReader + shuffle FileWriter switch #2204.
• Byte offsets are checked by slicing the file at each recorded range and decoding the message, not by checking the offsets are self-consistent. Arithmetic agreed with itself while being wrong.
• One test covers dictionary-encoded data, the only shape that distinguishes footer-derived offsets from byte-counted ones.
• The writer asserts its footer's record-block count equals the batches it wrote, so a future re-chunking fails loudly instead of producing an index that addresses the wrong batches.
• h2o Q8 verified against single-process DataFusion at both scales.

What this is not

• Wall-clock evidence. Every number here is bytes or rows. The measurements run on one box where the network is free, so a bandwidth saving does not become a time saving. A real-cluster run is owed before any latency or throughput claim, and will replace these figures.
• The public interface. @phillipleblanc's request on feat(core): ValueIndexExec + ValueIndexReader + shuffle FileWriter switch #2204 — public RangeShuffleWriter / RangeShuffleReader traits with a registration point, so a Vortex implementation can live outside this repo — is not in this diff. The format-specific parts are already confined to two modules, so it is a matter of naming the seam rather than moving the work.
• The end of the read amplification. What remains at K=8 is the halo the window genuinely needs plus one batch of granularity at each edge. Larger K makes selection less effective, not more — a cut eventually becomes narrower than a single 8192-row batch, and nothing can be skipped below one batch. The direction that helps is the opposite one: fewer, larger sorted files per task, where every file straddles every consumer's range and the index has room to work.

[phillipleblanc]

This makes sense to me, but I do wonder if we should try to come up with a more generic shuffle writer that could handle a file format that is registered in the ballista system somehow? i.e. there isn't anything in the RangeShuffleWriterExecNode protobut message itself that couldn't be made generic, its just the message type itself.

Given that we already have the RangeShuffleReader though, I don't think we need to change anything on this PR - but something to think about for making this handle different shuffle formats in the future.