You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "andygrove (via GitHub)" <gi...@apache.org> on 2026/07/05 17:22:26 UTC
[PR] feat: bound shuffle fetch with a reduce-side in-flight governor [datafusion-ballista]
andygrove opened a new pull request, #1951:
URL: https://github.com/apache/datafusion-ballista/pull/1951
# Which issue does this PR close?
Closes #1943.
# Rationale for this change
The shuffle reader fetches each remote partition as an open-ended `do_get` stream with **no global in-flight-bytes / in-flight-requests governor**. That forces a lose-lose choice between three failure corners:
| Corner | Symptom |
|---|---|
| Multiplex many streams over one h2 connection | h2 64 KB connection-window deadlock |
| One exclusive connection per stream | connection churn / **ephemeral-port exhaustion** at high `target_partitions` (#1943) |
| Broken pipe mid-body | fatal — a partition is an unbounded, non-retriable stream |
Spark solves the same problem with a **reduce-side in-flight governor** (`maxBytesInFlight` / `maxReqsInFlight` / `maxBlocksInFlightPerAddress`, see `ShuffleBlockFetcherIterator`). That governor is what makes *multiplexing over very few connections* safe: it bounds how much data is in flight at the application layer, so the transport is never flooded and a small, reused connection set is sufficient. This PR ports that model onto Ballista's gRPC/Arrow-Flight transport.
# What changes are included in this PR?
1. **Reduce-side in-flight governor.** Each remote fetch acquires three tokio semaphore permits — an in-flight-**bytes** budget (acquiring `min(block_size, max_bytes)`, charged from `PartitionStats.num_bytes`), an in-flight-**request** count, and a **per-address** slot — and holds them (via a drop-guard on the returned stream) until the body is fully consumed. This replaces the previous count-only semaphore whose permit was released as soon as the stream handle was obtained (before the body was read), so it bounds concurrent *in-flight data*, not just connection establishment. By bounding concurrency it sharply reduces how many connections are opened per peer, which is what resolves the ephemeral-port exhaustion in #1943.
2. **h2 flow-control window sizing.** The shuffle data-plane client's HTTP/2 initial connection and stream windows are made configurable and default `>=` the governor byte budget, so the application-level governor — not the 64 KB transport window — is the binding backpressure (the property Spark gets for free from Netty). Control-plane clients are unchanged.
3. **Retriable, buffered block fetch.** Each partition body is buffered into memory as a discrete, idempotent unit and the whole fetch is retried on a transport error (up to `io_retries_times`), so a mid-body broken pipe refetches the block instead of failing the task. The outer retry owns retry policy (transport-error-gated) and disables the client's inner establish-retry to avoid multiplying attempts.
4. **Configuration** (all with Spark-parity defaults):
- `ballista.shuffle.reader.max_bytes_in_flight` (48 MiB)
- `ballista.shuffle.reader.max_blocks_in_flight_per_address` (128)
- `ballista.shuffle.reader.default_block_size_bytes` (1 MiB, used when partition stats carry no byte count)
- `ballista.client.initial_connection_window_size` (64 MiB) / `ballista.client.initial_stream_window_size` (16 MiB)
- reuses the existing `ballista.shuffle.max_concurrent_read_requests` (64) as the request-count cap and `io_retries_times` (3) for fetch retry.
**Testing.** Unit tests cover the permit helpers, the `GovernedStream` release-on-drop guard, and the transport-gated retry (retriable vs non-retriable, attempt counts). A standalone integration test runs a shuffle query under a deliberately tiny in-flight budget (64 KiB, 2 blocks/address) to prove it completes without deadlock. The full `sort_shuffle` standalone suite passes.
**Preliminary cluster result (SF100, 2×8 = 16 task slots):** with this change, `target_partitions=128` — which previously died very early on shuffle-fetch port exhaustion — runs the full TPC-H set, and `target_partitions=32` completes ~10% faster than before. (Larger-partition runs also exposed an unrelated scheduler memory ceiling that is a separate concern.)
**Deferred (follow-ups, not in this PR):** spilling oversized partition bodies to disk (`spark.maxRemoteBlockSizeFetchToMem` analog — until then a single partition larger than the budget is buffered in full, and the byte budget bounds compressed *wire* bytes, so decoded RAM is larger by the compression ratio), and a hard `connections_per_peer` cap (which would require a shared-connection pool model).
# Are there any user-facing changes?
New configuration keys are added (listed above), all with defaults, so existing deployments behave sensibly with no changes. The shuffle read path now bounds in-flight bytes/requests and buffers each partition body into memory before yielding it downstream. No public API breakage.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: bound shuffle fetch with a reduce-side in-flight governor [datafusion-ballista]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on PR #1951:
URL: https://github.com/apache/datafusion-ballista/pull/1951#issuecomment-4893653213
Merged. Thanks @milenkovicm
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: bound shuffle fetch with a reduce-side in-flight governor [datafusion-ballista]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove merged PR #1951:
URL: https://github.com/apache/datafusion-ballista/pull/1951
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org