[SPARK-57399][CORE][SQL] Local repartition: in-process pipelined channel shuffle for a single executor - #58097
Open
viirya wants to merge 41 commits into
Open
[SPARK-57399][CORE][SQL] Local repartition: in-process pipelined channel shuffle for a single executor#58097viirya wants to merge 41 commits into
viirya wants to merge 41 commits into
Conversation
…ffle prototype Starting point for exploring an alternative local-repartition implementation that reuses the concurrent-stage pipelined-shuffle machinery (merged for Real-Time Mode) but with an in-process bounded-channel transport instead of the RPC streaming shuffle. Proven end-to-end on local[8]: 1000 rows repartitioned through a PipelinedShuffleDependency + a channel-backed ShuffleManager, no Netty / no serialization / no MapOutputTracker. This is a self-contained prototype (manager + rendezvous + writer + reader + test in one file under test/). Next steps: refactor into production shape (main/ code + real tests), and address the two known constraints (slot-check ceiling on single-machine multi-stage queries; the StreamingShuffleManager coupling that a real transport plug-in needs decoupled upstream). See plans/LOCAL_REPARTITION_PIPELINED_SHUFFLE_FEASIBILITY.md. Co-authored-by: Claude Code
…huffle into production shape
Splits the throwaway single-file prototype into proper main/ + test/ code under
org.apache.spark.shuffle.local.pipelined:
- ChannelShuffleRendezvous: JVM-static per-(shuffleId, reducePartitionId) bounded-queue
registry with end-of-stream marker and per-shuffle cleanup.
- ChannelShuffleWriter / ChannelShuffleReader: the in-process transport (map side pushes
copied pairs onto the reduce partition's queue as produced; reduce side drains until
every map task has signalled end-of-stream). No serialization / disk / network.
- PipelinedChannelShuffleManager: subclasses StreamingShuffleManager (so SparkEnv creates
the StreamingShuffleOutputTracker the pipelined DAGScheduler path requires) and
overrides only getWriter/getReader to use the channel transport; unregisterShuffle
drops the shuffle's queues.
Test suite (5 tests, all green): end-to-end repartition correctness + routing, parity with
a regular shuffle, an admissible fan-in, an explicit assertion that a group whose
whole-group demand exceeds the slots is rejected up front (the single-machine ceiling of
the pipelined model, pinned as a documented constraint), and the unregister cleanup
contract.
Still a v2 exploration, not for merge as-is: the StreamingShuffleManager coupling and the
slot-check ceiling remain open (see plans/LOCAL_REPARTITION_PIPELINED_SHUFFLE_FEASIBILITY.md).
Co-authored-by: Claude Code
…om the concrete streaming manager
The StreamingShuffleOutputTracker was created only when the pipelined shuffle manager
isInstanceOf[StreamingShuffleManager] (SparkEnv), and a PipelinedShuffleDependency
unconditionally demanded it (DAGScheduler). But the tracker is purely a transport
directory of writer host/port locations, read only by the RPC StreamingShuffleReader;
pipelined scheduling (co-scheduling, deferred completion, stage availability) runs off
ShuffleMapStage.pipelinedCompletedPartitions and never consults it. So an in-process
pipelined transport that finds its reader/writer pairs within one JVM needs no tracker,
yet the isInstanceOf coupling forced it to subclass the concrete RPC streaming manager
just to get the tracker created.
Decouple it by capability rather than concrete type:
- PipelinedShuffleManager gains `usesStreamingShuffleOutputTracker` (default true, so the
built-in StreamingShuffleManager is unaffected).
- SparkEnv.initializeStreamingShuffleOutputTracker keys creation off that flag instead of
isInstanceOf[StreamingShuffleManager].
- DAGScheduler.outputTrackerMaster returns Option; a pipelined dep whose manager declares
no tracker registers with none (createShuffleMapStage's registration becomes a no-op via
Option.foreach). When a tracker IS expected but absent, it still fails loudly.
This lets PipelinedChannelShuffleManager implement the PipelinedShuffleManager trait
directly (usesStreamingShuffleOutputTracker = false) instead of subclassing
StreamingShuffleManager, removing the workaround the prototype needed.
Regression: RTM pipelined DAGScheduler tests 51/51, streaming shuffle + tracker suites
74/74, PipelinedChannelShuffleSuite 5/5 (now on the clean trait path). No behavior change
for the RPC streaming manager, which keeps the default `usesStreamingShuffleOutputTracker`.
Co-authored-by: Claude Code
… batch query
Adds the minimal SQL wiring to drive a batch query through the pipelined channel shuffle,
and fixes two bugs that only the real SQL (UnsafeRow) path exposes -- the core-RDD tests
missed them because boxed Integers are not buffer-reused.
SQL entry point:
- EnablePipelinedShuffle: opt-in preparation rule (spark.sql.pipelinedShuffle.enabled,
AQE off) that rewrites hash-partitioning ShuffleExchangeExec nodes to pipelined=true.
Conservative all-or-nothing gate: rewrites only when every shuffle in the plan is
HashPartitioning (a SinglePartition/Range shuffle is a regular ShuffleDependency, and
mixing regular + pipelined in one job is rejected) and none is reused (fan-out is
rejected). Runs last in preparations so it sees the final reuse decision.
Bug 1 -- reader ignored the reduce-partition range. ShuffledRowRDD's
CoalescedPartitionSpec can map a range [startPartition, endPartition) to one reduce task;
the reader drained only startPartition, stranding the rest. ChannelShuffleReader now drains
every queue in the range.
Bug 2 -- reused row buffers. SQL producers reuse their output UnsafeRow across iterations,
so enqueuing the row as-is made every pair alias the same buffer (groupBy collapsed to a
couple of keys with halved counts). A regular shuffle avoids this by serializing on write;
the channel skips serialization, so the row must be copied. The copy cannot be done in
core's ChannelShuffleWriter (core cannot reference InternalRow, and UnsafeRowSerializer has
no single-object serialize), so it is done in the SQL layer: createShuffleWriteProcessor
gains copyRows (true only for the pipelined path) and its write() copies each value row
before delegating. This mirrors v1's SQL-layer row.copy() and keeps the transport
zero-serialization.
End-to-end: PipelinedShuffleSqlSuite runs df.repartition($k) and groupBy($k).count() batch
on local[8] through rule -> pipelined dependency -> channel manager -> concurrent-stage
scheduler, asserting the exchange went pipelined and results match. Both green;
PipelinedChannelShuffleSuite still 5/5.
Co-authored-by: Claude Code
…onings, cover join Relaxes EnablePipelinedShuffle from hash-only to rewriting EVERY ShuffleExchangeExec (hash, single-partition, range) to pipelined. Rewriting all of them keeps the job all-pipelined, which the DAGScheduler requires (a mix of pipelined and regular shuffles is rejected); the previous hash-only gate left SinglePartition/Range exchanges regular and so skipped any query with a final ORDER BY / LIMIT / global aggregate. The channel transport does not care about the partitioning kind -- it routes by partitioner.getPartition(key), and SinglePartition is just numPartitions == 1. The no-reuse skip stays (a pipelined producer cannot fan out to more than one consumer). Test coverage now spans all TPC-DS transport shapes, all end-to-end batch on local[16] (high slot cap so the whole-group admission passes; correctness harness, not perf): - repartition (hash) - keyed groupBy (partial -> hash -> final) - groupBy + ORDER BY (hash + SinglePartition, both pipelined, mixed-job rejection avoided) - SortMergeJoin (both inputs hash-exchanged and pipelined concurrently) The join test also confirms the no-reuse gate: an earlier version with structurally identical join inputs triggered exchange reuse, which the rule correctly declined to pipeline; distinct inputs exercise the two-sided pipelined path. core PipelinedChannelShuffleSuite 5/5, sql PipelinedShuffleSqlSuite 4/4. Co-authored-by: Claude Code
…add benchmark Per-row queue hand-off was the transport's bottleneck: LinkedBlockingQueue costs a lock acquisition per operation (~300ns under contention), paid PER ROW, which made an unaggregated 20M-row repartition ~19x SLOWER than a regular shuffle (6.5s vs 0.35s). Isolation confirmed the diagnosis: raising the queue capacity from 64 to 4M changed nothing (not backpressure), and 6.2s / 20M rows = ~310ns/row matches the lock cost; per-row copy is not the driver (v1's object-batch pays the same copy and wins). Fix is the same lesson v1's object-batch transport learned: the writer accumulates rows in a per-output-partition Array[AnyRef] (spark.shuffle.pipelined.channel.batchSize, default 1024) and hands a full batch across the queue in ONE operation; partial batches are trimmed and flushed before the end-of-stream markers. The reader drains batch by batch and iterates rows out of each array. Lock traffic drops by ~batchSize (20M -> 20K operations). Batch arrays are ownership-transferred (writer allocates a fresh one after each put), so no cross-thread buffer reuse. Benchmark (new PipelinedShuffleBenchmark, local[16] on 16 cores, group demand 14 <= 16 so no oversubscription; 20M rows, best of 6): repartition(k) + count regular 326ms pipelined 228ms 1.43x (was 6537ms, 0.05x) groupBy(k).count regular 67ms pipelined 39ms 1.72x core PipelinedChannelShuffleSuite 5/5, sql PipelinedShuffleSqlSuite 4/4. Co-authored-by: Claude Code
…e; cover range/single Testing RangePartitioning end-to-end (repartitionByRange via Dataset.rdd) exposed silent, racy row loss (875/750 of 1000 rows, ~2/3 reproducible) that had nothing to do with range itself: Dataset.rdd builds its RDD inside a SQL execution scope that ends BEFORE any job runs, and spark.sql.classic.shuffleDependency.fileCleanup.enabled (default under testing) removes the shuffle from every manager at scope end. That wiped the channel manager's numMapsByShuffle registry entry between registration and execution; the reader then read the missing entry as numMaps = 0 (null unboxed), stopped at the FIRST end-of-stream marker instead of one per map task, and dropped whatever the other writer had not yet enqueued. Diagnosed by instrumenting the channel (readers logged "seen=1/0") and stack-tracing unregisterShuffle to SQLExecution's execution-end cleanup RPC. Fix: keep no per-shuffle mutable state in the manager at all. The map-task count is stamped into a ChannelShuffleHandle at registration; the handle travels with the dependency into every task, a plain Int survives task serialization (deriving from handle.dependency.rdd NPEs in a deserialized task -- the dependency's rdd reference is @transient), and no later unregister can take it away. unregisterShuffle now only drops the rendezvous queues. Coverage corrections and additions in PipelinedShuffleSqlSuite (7/7, the once-flaky Dataset.rdd regression test now 5/5 stable): - The ORDER BY test's exchange is RangePartitioning, not SinglePartition as previously claimed; it now pins that (and exercises RangePartitioner's sample job running over a pipelined hash shuffle before the main job). - New repartitionByRange test: pure range transport, asserts non-overlapping key ranges. - New global-aggregate test: a true SinglePartition exchange (numPartitions == 1). - New Dataset.rdd regression test for the unregister-before-job scenario. core PipelinedChannelShuffleSuite 5/5. Co-authored-by: Claude Code
…'s sampling penalty Adds two RangePartitioning cases to PipelinedShuffleBenchmark, separating the control from the differential (viirya's hypothesis: the range sample job re-runs the child, which regular shuffle amortizes via materialized map output but a single-shot channel cannot): repartitionByRange(k) + count regular 414ms pipelined 338ms 1.22x groupBy(k).count + orderBy(k) regular 72ms pipelined 73ms 0.99x Range over a plain scan (control): no differential -- BOTH modes re-run the scan for the sample job (nothing materialized below the range exchange for regular to reuse either); pipelined still wins, with the margin narrowed by the shared sample-pass cost. Range above a shuffle (differential): confirmed exactly. Pipelined pays one full extra upstream pass -- its 73ms is 2x the groupBy-alone 37ms -- because the sample job consumes the single-shot hash channels and completed-job stage cleanup makes the main job re-run scan + partial agg + map side. Regular's main run reuses the hash map output (63ms -> 72ms only). The transport's 1.7x per-pass advantage exactly cancels the extra pass here; a heavier upstream would turn this into a net loss. Co-authored-by: Claude Code
…4 queries verified
Adds TPCDSPipelinedRunner: builds SF=1 parquet from dsdgen .dat files (trailing-delimiter
handled), registers views, and runs selected TPC-DS queries twice in one session --
pipelined rule off (regular baseline) then on -- comparing full result sets. Correctness
harness only: local[128] clears the gang slot check by oversubscribing the cores
(spark.sql.files.maxPartitionBytes=512m keeps fact-scan producer partitions low; the first
attempt at local[64] was rejected at demand 68), so timing is meaningless by design.
Results, all rows identical to the regular-shuffle baseline:
q3 q42 q52 q55 q7 q19 OK 1 pipelined hash exchange each (star join + agg;
ORDER BY+LIMIT becomes TakeOrderedAndProject, no range)
q12 q34 q73 q79 q96 q20 OK 1-2 pipelined exchanges
q98 OK 3 pipelined exchanges: RangePartitioning + 2 hash --
a real query exercising the pipelined range path,
sampling included
q68 OK reuse gate fired (ReusedExchange present): rule left the
plan regular, results still correct -- the fallback is safe
14/14 correct; 13 executed the v2 channel transport end-to-end.
Co-authored-by: Claude Code
…s timed fairly
Adds a `bench` mode to TPCDSPipelinedRunner: same off/on structure as `verify` but timed
(1 warm-up + best of 5) on local[N = physical cores], so the pipelined group does not
oversubscribe and the regular-vs-pipelined comparison is fair. Fitting the gang under 16
honest slots required squeezing scan parallelism for BOTH modes
(files.maxPartitionBytes=64m + minPartitionNum=1 -> ~5 store_sales scan partitions;
FilePartition planning otherwise targets defaultParallelism leaves and every query was
rejected at demand 17-26) -- reduced scan parallelism is itself the honest cost of the
slot ceiling on one box.
SF=1, 16 cores, best of 5 (geomean of the 12 that ran: ~1.10x):
q3 1.44x q42 1.21x q52 1.20x q20 1.18x q55/q7/q12 1.13x q19 1.09x
q79 1.04x q96 1.01x (SinglePartition)
q34 0.94x, q73 0.85x -- the two RANGE-bearing plans, confirming the RangePartitioner
sampling double-pass penalty on real queries (sample job re-runs the unmaterialized
pipelined upstream; regular reuses its materialized map output)
q98 REJECTED (demand 18 > 16): the one 3-stage group that still does not fit
Co-authored-by: Claude Code
…ibe instead of squeeze Adds a third runner mode answering viirya's challenge to the bench methodology: instead of squeezing scan parallelism until the gang fits the physical cores, keep near-natural scans (32m -> ~13 store_sales partitions vs the box's natural 16) and raise local[n] to 2*cores so gang admission passes and the pipelined group oversubscribes the cores. The regular baseline is unaffected by the larger n (its per-stage width still fits the physical cores), so the contention cost lands only on the pipelined side -- measuring it is the point. Results (SF=1, local[32] on 16 cores, best of 5): all 13 queries ran, including q98 (0.86x -- previously rejected at demand 18; its loss matches v1's 0.85x on the same query, again pinning the range-sampling penalty on the shared unmaterialized-transport property). Geomean ~1.08x vs the squeezed bench's ~1.10x: mild oversubscription (demand <= 2x cores) costs only a few percent and does not change the qualitative picture -- the earlier claim that oversubscribed timing is meaningless holds for local[128], not for this regime. Small-data wrinkle: at SF=1 BOTH modes ran slightly faster with the squeezed 5-partition scans than with 13 (per-task overhead dominates tiny scans), so the "sacrificed" scan parallelism cost nothing here; at larger scale factors that reverses. Co-authored-by: Claude Code
…2 wins bigger Adds v1's LocalRepartitionBenchmark "prototype workload" (1M rows, UNIQUE key per row, uncached, wide column pruned -> pure transport overhead) to PipelinedShuffleBenchmark, in two variants. v1's historical record on it was 4.1-4.5x over regular shuffle (AQE off, local[32]); reproduced today on the v1 branch at 4.35x under the matching 32-map shape. prototype, 6 maps, local[16]: regular ~52ms v1 33ms (1.55x) v2 26ms (2.08x) prototype, 32 maps, local[48]: regular ~136ms v1 31ms (4.35x) v2 25ms (5.48x) v2 shows the same large-multiple behavior and beats v1 on its own flagship shape. The mechanism the two variants isolate: going 6 -> 32 maps leaves both in-process transports flat (25-33ms) while the REGULAR baseline degrades 52 -> 136ms -- tiny data multiplied by many map tasks inflates the regular shuffle's fixed per-task/per-segment cost, which is what the big multiple actually measures (and why real TPC-DS queries sit at 1.1-1.4x). Caveat: the 32-map gang demand is 41, so v2's 5.48x requires local[48] oversubscription (v1 was run at the same master for symmetry); at honest local[16] only the 6-map variant is admissible. Co-authored-by: Claude Code
…inst subquery reuse
The gate was plan.exists(_.isInstanceOf[ReusedExchangeExec]), which walks the operator
tree only -- reuse landing inside a subquery plan (embedded in an expression) was
formally invisible to it. Probing every SQL route to that shape showed none can currently
produce a reused PIPELINED exchange, each closed by a different layer:
1. Same-tree reuse: the gate catches it (q68, join shapes).
2. Main-vs-subquery reuse: never fires. Each subquery runs its own full preparation
pass (PlanSubqueries -> prepareExecutedPlan, which includes EnablePipelinedShuffle),
so its exchanges are already pipelined=true when the outer ReuseExchangeAndSubquery
compares canonical forms against the outer, not-yet-pipelined exchange -- the
`pipelined` field diverges and reuse does not match. Accidental, but load-bearing.
3. Subquery-vs-subquery duplication: MergeScalarSubqueries / subquery-level reuse
collapses the duplicates into ONE executed subquery before exchange reuse matters.
Mechanisms 2 and 3 are rule-ordering and optimizer accidents, so the gate now checks
collectWithSubqueries instead of relying on them. A regression test pins all observed
facts (no reused exchange materializes, subquery exchanges pipeline via their own
preparation, both probe queries return correct results) so a change in any layer
surfaces. Suite 8/8.
Co-authored-by: Claude Code
…ow a pipelined suffix
Relaxes the all-regular-or-all-pipelined job rule to admit the MATERIALIZED-PREFIX MIXED
shape: pipelined shuffles in the region reachable from the final RDD, where every regular
shuffle boundary at that region's edge is fully materialized (all MAP outputs registered
with the MapOutputTracker) and no pipelined shuffle sits below any regular boundary. The
prefix never re-runs, so the job executes exactly like an all-pipelined job whose leaves
read materialized shuffle data, and gang admission demand (final stage + suffix producers)
is unchanged. This is the shape adaptive execution produces -- prior map-stage jobs
materialize the prefix stages, the final job runs the pipelined tail -- and is the
scheduler-side prerequisite for running the v2 pipelined shuffle under AQE.
Still rejected, fail-fast with no partial scheduler state:
- an UNMATERIALIZED regular boundary in a pipelined job: its stage would have to run
while gang-admitted producers hold slots blocked on transport backpressure, which
admission does not account for and can deadlock (sequencing the prefix before the gang
is future work);
- a pipelined shuffle BELOW a regular boundary.
classifyJobShuffleKinds is replaced by classifyJobShuffleShape: a suffix walk that stops
at regular boundaries, then a materialization check (against the producer RDD's partition
count -- the tracker counts MAP outputs; the first version compared the reducer-side
partitioner and the end-to-end test's asymmetric 2-map/3-reduce prefix caught it) and a
below-boundary pipelined scan per boundary.
Tests: DAGSchedulerSuite gains the accepted-shape test (asymmetric map/reduce counts to
pin the map-side materialization check); the two existing rejection tests keep passing
with comments updated to the refined rule; PipelinedChannelShuffleSuite gains end-to-end
materialized-prefix (runs, no row loss) and unmaterialized-prefix (still rejected) tests.
DAGSchedulerSuite 213/213, streaming 74/74, MultiShuffleManager 4/4, channel 7/7,
PipelinedShuffleSqlSuite 8/8.
Co-authored-by: Claude Code
…der AQE Builds the AQE-side half on top of the materialized-prefix scheduler relaxation (3958627), adapting v1's AQEReplaceWithLocalRepartition design to v2's mechanics: 1. AQEEnablePipelinedShuffle (queryStagePreparationRules, after skew handling): flips eligible exchanges to pipelined using v1's placement policy -- "free" candidates whose path to the root crosses nothing stats-sensitive, ShuffledJoin inputs only as symmetric pairs, duplicated canonical forms skipped (incl. materialized stages and subqueries). Unlike v1 (hash-only operator, SinglePartition treated as a transparent regular wall), a pipelined exchange supports every partitioning, so SinglePartition exchanges in free position are simply candidates. The walk stops below a flipped exchange: everything underneath stays regular and keeps full AQE treatment (coalescing, skew, join switching). 2. AdaptiveSparkPlanExec.createNonResultQueryStages: a pipelined exchange is never promoted to a query stage (a pipelined producer cannot materialize alone -- the scheduler rejects a map-stage job over a pipelined dep); it stays inline and the walk recurses through it, so stages below still materialize as the regular prefix. The final result job is then exactly the scheduler's admitted shape: fully-materialized prefix + pipelined suffix gang. AQEPipelinedShuffleSuite (4/4): single-exchange aggregate (whole final job pipelined, no stage materialized), groupBy + ORDER BY (the canonical shape: hash exchange materialized as a ShuffleQueryStageExec prefix, range exchange pipelined on top, results correct), shuffled-join pair flip, and an off/on baseline comparison. Plan assertions are made on the SAME Dataset that executed (.as[...] creates a fresh QueryExecution; an unexecuted sibling shows the initial isFinalPlan=false plan -- the first version asserted on that and passed two tests vacuously). AdaptiveQueryExecSuite 129/129 (rule is opt-in and inert by default), PipelinedShuffleSqlSuite 8/8, channel suite 7/7. Co-authored-by: Claude Code
…h SinglePartition First AQE-on benchmark exposed that the placement rule accelerated only the topmost exchange: repartition(k)+count flips just the trivial SinglePartition agg exchange while the 20M-row hash exchange below it materializes as a regular stage (1.09x vs v1's 1.83x -- v1's AQE walk treats SinglePartition as transparent and replaces the hash below, its historical 7e10df5 lesson). Adopt it for v2: a flipped SinglePartition candidate keeps the walk going, flipping free candidates below into a pipelined CHAIN (v2's all-pipelined constraint means the single exchange must flip along, where v1 left it regular). Chain flipping surfaced a second bug: transformUp rebuilds children first, so the upper candidate node reaching the pattern is a new instance whose flipped child no longer matches the collected original structurally -- the upper flip silently dropped, leaving a regular exchange above a pipelined one, which the scheduler correctly rejected as pipelined-below-regular. Fixed by transformDown (candidates hit the pattern before their subtrees are rebuilt). v1 never nests replacements, which is why its transformUp is safe. AQE-on mirror benchmark (local[16], 20M rows / 1M prototype, best of 6), after the fix: repartition(k)+count regular ~340ms v1 1.83x v2 1.09x -> 1.73x groupBy(k).count regular ~55ms v1 1.10x v2 1.08x -> 1.51x groupBy+orderBy regular ~68ms v1 1.06x v2 1.08x (range tail is trivial) prototype 1M uniq regular ~56ms v1 1.62x v2 1.17x -> 2.28x AQEPipelinedShuffleSuite 4/4, AdaptiveQueryExecSuite 129/129. Co-authored-by: Claude Code
…on head-to-head
Adds a `benchaqe` runner mode (bench configs + adaptive execution on; plan summaries via
AdaptiveSparkPlanHelper.collect and an executed probe -- TreeNode.collect sees an AQE plan
as exchange-less, and an unexecuted adaptive plan is the initial one).
AQE-on results (SF=1, local[16], best of 5; v1 numbers from its mirrored scratch runner):
v2 fires on 12/13 queries, geomean ~1.12x; v1 fires on 8/13, geomean ~1.09x. Speedups on
shared queries are comparable (v1 slightly ahead on q3/q12, v2 on q7/q19/q55). The
structural differences:
- q98 RUNS at honest local[16] under AQE (1.15x; rejected at demand 18 AQE-off): gang
demand now scopes to the pipelined tail over materialized stages.
- The range-sampling tax vanishes under AQE for aggregate+ORDER BY shapes: q34 went
0.94x (AQE off) -> 1.17x. The hash stage below the range exchange is materialized, so
RangePartitioner's sample job re-runs only the cheap reduce side over materialized
output instead of the whole upstream.
- v1's AQE rule skips q34/q73/q98 (hash below a stats-sensitive range exchange) and q96
(SinglePartition-only): shapes only v2's any-partitioning flip can serve.
- q79 is skipped by both (its remaining exchange feeds a join; no symmetric pair).
Co-authored-by: Claude Code
…g RTM's hot path Auditing our changes for RTM impact found one real (perf-only) regression: the pipelined branch of prepareShuffleDependency passed copyRows = true unconditionally, and RTM's real-time plans take the same branch -- its RPC streaming transport serializes records promptly (serialization IS the detach), so the per-row InternalRow.copy() was pure added overhead on RTM's hot path. Correctness was unaffected, which is why every suite stayed green while the cost crept in. Express the need as a manager capability, the same pattern as the tracker decoupling: PipelinedShuffleManager.requiresDetachedRecords (default false, so the built-in streaming manager keeps its no-copy behavior); the in-process channel manager overrides it to true (records cross the channel as object references read by a concurrent consumer thread); the SQL pipelined branch reads the capability off SparkEnv's pipelined manager (always initialized; defaults to streaming). ExchangeSuite 16/16, PipelinedShuffleSqlSuite 8/8, AQEPipelinedShuffleSuite 4/4, channel 7/7, streaming 74/74. Co-authored-by: Claude Code
…his; v2 did not)
viirya asked whether v2 carries v1's local-mode restriction -- it did not, anywhere. The
channel rendezvous is JVM-local, so on a multi-executor deployment every reader would
block forever on data written in another JVM: a silent hang, the worst failure mode for a
misconfiguration. Three layers, mirroring v1's gates plus a hard floor v1 does not need:
- PipelinedChannelShuffleManager now REQUIRES a local master at construction (fails the
app loudly at startup with a clear message; unit-tested both ways).
- EnablePipelinedShuffle and AQEEnablePipelinedShuffle return the plan untouched off
local mode, matching v1's ReplaceWithLocalRepartition / AQEReplaceWithLocalRepartition
gates. (The pipelined machinery itself is not local-only -- the RPC streaming
transport is cross-executor -- but batch queries over it are unexplored, so the rules
stay conservative.)
channel 8/8, PipelinedShuffleSqlSuite 8/8, AQEPipelinedShuffleSuite 4/4.
Co-authored-by: Claude Code
… concurrency limit Second gap found by the v1-gate audit (after the local-mode gate): v1 refuses to replace an exchange whose partition count exceeds SparkContext.defaultParallelism -- its channel-deadlock precondition, and under AQE additionally an output cap -- degrading gracefully to a regular run. v2 had no such gate and leaned on the scheduler's gang slot check, which FAILS the query instead of skipping the optimization: with the DEFAULT spark.sql.shuffle.partitions=200 on a 16-core box, enabling the flag would fail essentially every query (our suites masked this by setting partitions=4). Both rules now cap candidates at defaultParallelism: the non-AQE rule skips the WHOLE plan if any exchange is wider (a partial flip would be a rejected mixed job); the AQE rule skips per candidate (a skipped candidate materializes as a regular stage, which the prefix shape supports). Regression tests pin the degrade-to-regular behavior at shuffle.partitions=64 on local[16] in both suites. Remaining audit deltas, accepted and documented rather than ported: v1's maxInputPartitionNum soft cap and maxNum replacement budget (v2's whole-group slot check subsumes their deadlock role; the failure-vs-skip asymmetry for INPUT-side width remains a known gap -- the rules cannot see scan partition counts at plan time). PipelinedShuffleSqlSuite 9/9, AQEPipelinedShuffleSuite 5/5, channel 8/8. Co-authored-by: Claude Code
…s at the concurrency limit" This reverts commit 8f94a73.
…ns (design decision) Follows the revert of the concurrency-cap gates (viirya's call): the user opted into pipelined execution explicitly, so a plan whose flipped exchanges cannot fit the local task-concurrency limit should surface the scheduler's explicit CONCURRENT_SCHEDULER_INSUFFICIENT_SLOT error -- actionable (raise local[n] or lower shuffle.partitions) -- rather than silently degrade to a regular run as v1's gates do. Two tests pin the fail-loud behavior at shuffle.partitions=64 on local[16], non-AQE and AQE. The local-mode gates stay (a cluster misconfiguration is a silent HANG, not an explicit error, so it still fails at startup / skips the rules). PipelinedShuffleSqlSuite 9/9, AQEPipelinedShuffleSuite 5/5. Co-authored-by: Claude Code
…a proper SQLConf spark.sql.pipelinedShuffle.enabled existed only as an unregistered string key read via getConfString in both rules: no type validation, no doc, invisible to SET -v. Register it as an internal boolean SQLConf entry (default false, version 4.3.0; doc covers the local-mode requirement, the manager routing, and the explicit insufficient-slot failure semantics) and switch both rules to conf.pipelinedShuffleEnabled. PipelinedShuffleSqlSuite 9/9, AQEPipelinedShuffleSuite 5/5. Co-authored-by: Claude Code
…ke the flag public
Two conf-surface cleanups (viirya):
- spark.shuffle.pipelined.channel.batchSize is now a registered core ConfigEntry
(intConf, > 0, default 1024, version 4.3.0) next to spark.shuffle.manager.incremental,
replacing the raw conf.getInt in PipelinedChannelShuffleManager.
- spark.sql.pipelinedShuffle.enabled drops .internal(): the flag is the feature's
public opt-in surface.
channel 8/8, PipelinedShuffleSqlSuite 9/9, AQEPipelinedShuffleSuite 5/5, SQLConfSuite
41/41.
Co-authored-by: Claude Code
Adds a regular vs streaming-pipelined vs channel-pipelined comparison over the shapes the streaming transport survives (no range sampling: its tracker accumulates writer registrations across the sample job's producer re-run and the reader dies on its writer-count assertion). Same scheduling for the two pipelined modes -- all-pipelined gang -- so the delta is purely the byte path. local[16], AQE off, best of 6: repartition(k)+count 20M regular 346ms streaming 1096ms (0.32x) channel 204ms (1.70x) groupBy(k).count 20M regular 59ms streaming 124ms (0.48x) channel 41ms (1.44x) join 10Mx10M uniq + count regular 622ms streaming 608ms (1.02x) channel 623ms (1.00x) prototype 1M uniq regular 50ms streaming 133ms (0.38x) channel 26ms (1.92x) The headline: RTM's RPC streaming transport LOSES to the regular materialized shuffle by 2-3x on shuffle-dominated batch shapes -- it is latency-optimized (records flow while the producer runs, at per-record serialization + Netty loopback + tracker cost), not throughput-optimized -- while the in-process channel beats regular by 1.4-1.9x and beats streaming head-to-head by 3-5x. This is the quantified justification for v2 carrying its own transport rather than reusing the streaming one in local mode. The join row is a compute-bound control: all three transports tie when sort-merge work dominates. Co-authored-by: Claude Code
…all shuffles not worth it" Adds a `curve` mode sweeping repartition(k)+count over shrinking row counts (regular vs channel, no aggregation so the transport stays on the hot path) to find the crossover where the channel's fixed per-gang cost (extra concurrent-stage scheduling + queue setup) would overtake its transport win -- the premise behind a size-based cost-aware placement gate. There is NO crossover. Small shuffles are where the channel wins MOST (local[16], best of 6): 20M rows regular 326ms channel 204ms 1.60x 1M regular 53ms channel 25ms 2.12x 100K regular 37ms channel 18ms 2.06x 10K regular 34ms channel 15ms 2.27x 1K regular 32ms channel 14ms 2.29x 100 regular 32ms channel 14ms 2.29x The premise was wrong because the two sides are asymmetric: regular's absolute time floors at ~32ms regardless of size (shuffle-file write, block manager, serializer/fetch framework startup), while the channel floors at ~14ms (gang scheduling + queue setup is real but cheaper than the spill framework). As data shrinks the comparison becomes floor-vs-floor and the channel's lower floor makes the speedup RISE toward a ~2.3x asymptote. Small shuffles are exactly where an in-process transport should win -- regular pays the most per-byte spill overhead there. So a size-based skip does not belong in cost-aware placement; if that gate has value at all, its predicate is not row count. Co-authored-by: Claude Code
…ter bug (known bug) Investigating the unregister lifecycle turned up a real, reachable, common-shape bug via a LIMIT counterexample: an early-stopping reader over a pipelined channel shuffle hangs the writer. The channel queue is bounded (64 batches); once a LIMIT's reduce task is satisfied and stops draining while the map task is still producing, the writer blocks on a full queue's put() with no drainer, forever. Verified to hang under a 90s deadline in BOTH AQE modes. This is unlike regular shuffle, which materializes to disk so the writer finishes regardless of whether the reader drains everything -- LIMIT just skips reading the rest. The pipelined transport's writer/reader concurrency + backpressure assumes the reader stays; early stop (LIMIT, first/head/take, an early-terminating join build side) breaks that assumption. Reachable and common, so it outranks the (unreachable) mid-job-unregister hazard that this investigation started from. Adds PipelinedLimitHangSuite as the reproduction / future acceptance test (both cases `ignore` -- they hang to the deadline, so CI-running them would waste 90s each and, sharing a process, a cancelled first corrupts the second; flip to `test` + assert rows==10 once the fix lands). Also reverts the misguided sentinel attempt in ChannelShuffleRendezvous.remove- Shuffle: a non-empty queue at unregister is NORMAL (LIMIT leaves undrained elements), so queue occupancy is not a usable "unregistered mid-job" signal; comment records why. Fix direction is an open design question (writer-side detect reader-gone and stop/discard, vs a channel "reader abandoned" signal that makes put() return/throw) -- deliberately not chosen here. Co-authored-by: Claude Code
…riter hang
The pipelined channel deadlocks whenever a consumer does not drain every reduce partition
to the end: the single-threaded writer interleaves all partitions, fills an undrained
partition's bounded queue, and blocks -- so even the read partitions never get their data
or end-of-stream. Reachable and common (LIMIT / take / head / first / isEmpty / show, and
any multi-exchange plan). Diagnosed by instrumentation after several wrong guesses; the
fix has two halves plus two correctness refinements that later hangs forced out.
Half 1 -- drop writes to dead partitions. The DAGScheduler knows a job's live reduce
partitions (ResultStage.partitions, set before submitStage; executeTake's per-batch runJob
passes exactly the subset it reads). New job property
SPARK_PIPELINED_LIVE_REDUCE_PARTITIONS carries it to the producer's tasks;
ChannelShuffleWriter drops records (and emits no end-of-stream) for a partition not in the
set. Absent property = all live, so full-read jobs (collect/count) are unchanged.
Half 2 -- reader-departure signal for LIVE partitions that stop early. A live partition
whose reader quits (LIMIT pulled enough) still fills and wedges. The reader's
TaskCompletionListener marks its partitions abandoned; the writer uses a polling
offer(100ms) that bails when the partition becomes abandoned, and abandon() drains the
queue to release a parked put -- a cooperative unblock, no reliance on interrupt.
Refinements the SQL suite's later cases forced (each a real design bug, not test-only):
- live set is per-SHUFFLE-EDGE, not per-job. Setting the result stage's partition subset
on EVERY pipelined producer broke multi-exchange chains: a subquery's hash exchange
(4 parts) below a single-partition agg was told live={0} and dropped the parts the agg
still needed -> hang. Now the property is set only on the producer whose shuffle the
result stage reads DIRECTLY (getShuffleDependenciesAndResourceProfiles); middle
exchanges stay fully live.
- abandonment is per-JOB, but a shuffleId is re-run within one query (RangePartitioner
sampling job then main job; executeTake batches). A prior job's abandoned marks must
not make the re-run's fresh writer think partitions are dead, so the writer clears its
partitions' marks at the start of each attempt; only abandonment during that attempt
counts. clearForTesting/removeShuffle also clear the abandoned set.
PipelinedLimitHangSuite's two cases flipped ignore->test and pass (AQE off + on), no longer
hang. Verified in a working shell (this sandbox's sbt could not start): channel 8/8, SQL
9/9, AQE 5/5.
Co-authored-by: Claude Code
…ss channel transport
PipelinedChannelShuffleManager was handed a ShuffleWriteMetricsReporter and a
ShuffleReadMetricsReporter in getWriter/getReader and discarded both, so the Spark
UI showed all-zero shuffle metrics for a pipelined (channel) shuffle even though
rows were flowing.
Wire the reporters through:
- ChannelShuffleWriter takes the write reporter. On each successful hand-off of a
non-empty batch to a partition's queue it records incRecordsWritten(batch rows)
and incWriteTime(nanos since the put began -- so time parked on a full bounded
queue counts as write time). End-of-stream markers and dropped/abandoned batches
(records == 0, or a partition with no reader / a departed reader) record nothing,
since those rows are never shuffled out.
- ChannelShuffleReader takes the read reporter and records incRecordsRead(batch
rows) as each batch is fetched from the queue.
- The manager now passes both reporters into the writer/reader instead of dropping
them.
Deliberately NOT reported:
- Bytes (incBytesWritten / incLocalBytesRead). The channel hands off object batches
with no serialization, so there is no wire-byte count; a fabricated figure would
be misleading. Records + time are the honest metrics for this transport.
- MapStatus stays the all-zero placeholder (a pipelined reducer never reads
partition lengths), matching the RPC streaming writer.
Co-authored-by: Claude Code
…e decoupling relaxed The usesStreamingShuffleOutputTracker decoupling (earlier in this stack) let an in-process pipelined manager declare it needs no StreamingShuffleOutputTracker, so SparkEnv builds none and DAGScheduler.outputTrackerMaster registers the shuffle with no tracker instead of failing. This deliberately RELAXED the older invariant "a pipelined dependency implies a streaming tracker." PipelinedShuffleRoutingSuite (upstream, SPARK-58454) still carried "createShuffleMapStage fails loud for a pipelined dependency when no streaming tracker", which asserted that older invariant. It can no longer describe a reachable state: SparkEnv's build-it decision and outputTrackerMaster's throw-if-absent decision now read the SAME flag, so they are locked consistent -- flag true builds a tracker (no throw), flag false needs none (no throw). No manager reports "needs a tracker" while being trackerless. Remove the test rather than repurpose it (its fail-loud subject is gone) or ignore it (nothing to revive). The behavior the change actually introduced -- a channel-manager pipelined dependency schedules and runs with no tracker and no throw -- is covered end to end by the channel suites (PipelinedChannelShuffleSuite, PipelinedShuffleSqlSuite, ...), so it is not left untested. Also migrate the suite's IncrementalRecordingManager to override usesStreamingShuffleOutputTracker = false. It is a recording mock, not the RPC StreamingShuffleManager; under the old isInstanceOf check it was classified "no tracker", but the flag's trait default is true, so without the override it now builds a tracker and breaks the two "does / does not initialize the tracker" tests. The override restores its intended non-streaming role. Co-authored-by: Claude Code
…achable defense-in-depth The pipelined branch of DAGScheduler.outputTrackerMaster throws when a pipelined dependency's manager declares it needs a StreamingShuffleOutputTracker (usesStreamingShuffleOutputTracker) yet none exists. That state is unreachable by construction: SparkEnv.initializeStreamingShuffleOutputTracker reads the SAME flag to decide whether to create the tracker, so flag true means a tracker was created (the throw's None arm is never taken) and flag false means the throw's own guard is false. With all current managers (RPC streaming true, channel false, test doubles constant) the two decisions are locked consistent. The usesStreamingShuffleOutputTracker decoupling (3ff57cf) intended this fail-loud to survive ("it still fails loudly"), but the single-flag coupling silently made it dead. Rather than change behavior, annotate it: explain it is deliberate defense-in-depth, unreachable under current wiring, costs nothing on the hot path, and would matter only if the two sites ever diverged (a non-constant-flag manager, or a conditional tracker build). The comment also points at the clean-up alternative -- unify the two flag reads so SparkEnv is the sole reader and this arm trusts only tracker presence -- and at the feasibility doc's decoupling section. No behavior change; comment only. Co-authored-by: Claude Code
EnablePipelinedShuffle leaves a plan regular (no pipelined flip) when it contains a
reused exchange, because a pipelined producer cannot fan out to more than one consumer.
This is a normal, expected fallback: exchange reuse is routine optimizer output (self-
joins, repeated subqueries over one table), the query still runs correctly as a regular
shuffle, and the user has nothing to act on. Logging it at WARN fired on every reuse-
bearing query once the pipelined config is on (much of TPC-DS) and read as a fault.
Downgrade to DEBUG -- it is diagnostic ("why this query did not go pipelined"), not a
warning. Gate logic unchanged; log level and an explanatory comment only.
Co-authored-by: Claude Code
…led path and RTM inertness Two scheduler tests pinning properties of the materialized-prefix mixed-shape relaxation (3958627) that the existing tests did not cover: 1. Losing a materialized prefix's shuffle output mid-group aborts the WHOLE group, not a lone prefix resubmit. A materialized prefix needs no slot to run, but if its output is lost while the gang holds all slots (producers blocked on backpressure), a base-scheduler single-stage resubmit would deadlock -- no slot for the prefix to recompute into. This is the one FetchFailed path where the failing consumer reads OUTSIDE its group (the external prefix); the existing group-internal FetchFailed test does not exercise it. It still routes to a whole-group abort because isPipelinedGroupMember keys off the failing STAGE (the pipelined producer reading the prefix is a member) regardless of which shuffle's fetch failed. Asserts job abort + no lone stage resubmitted into the held slots. 2. An all-pipelined group with no regular prefix classifies identically before and after the relaxation (inert for RTM-style jobs). The relaxation keys off dependency type so it is open to any pipelined dependency, but the materialized-prefix shape it admits is produced only by AQE; an RTM-style all-pipelined job never hits the relaxed path. Pins the "strict superset RTM inherits, nothing RTM runs changes" claim so RTM sign-off is "confirm you do not produce this shape," not "did you break me." DAGSchedulerSuite 215/215 (was 213; +2). Co-authored-by: Claude Code
Two defensive fixes (fail-loud and fallback) plus two clarifying comments, from a pre-PR code review. No behavior change for the execution shapes the SQL rules produce today; both fixes close a hazard that is currently unreachable but one refactor away. 1. FAIL-LOUD: ChannelShuffleReader requires exactly one reduce partition per task. The reader drains a partition range strictly sequentially, while the map-side writer interleaves all partitions on one thread and blocks on a full bounded queue. A coalesced multi-partition read could therefore deadlock (the reader drains partition start to completion while the writer is parked filling start+1, which the reader never reaches). Spark never sends a coalesced spec to a pipelined dependency today (createNonResultQueryStages keeps a pipelined exchange out of any ShuffleQueryStage, so CoalesceShufflePartitions cannot coalesce it), but the class doc previously CLAIMED range support it cannot safely honor. Replace that claim with a `require(endPartition - startPartition == 1)` so a future change that lets a coalesced spec through fails loud instead of hanging. (This also makes the LIMIT live-set's result-partition == reduce- partition assumption hold by construction.) 2. FALLBACK: the two SQL rules (EnablePipelinedShuffle, AQEEnablePipelinedShuffle) now require the active pipelined manager to be the in-process channel manager. The SQL flag spark.sql.pipelinedShuffle.enabled does not pick a transport; the manager is set separately by spark.shuffle.manager.incremental, which DEFAULTS to the RPC StreamingShuffleManager. Enabling the flag without also setting the channel manager previously flipped exchanges to pipelined and routed them to the untested RPC-streaming + local-mode + batch combination (also skipping the row copy, since that manager reports requiresDetachedRecords = false). Now such a configuration cleanly leaves the plan regular with a DEBUG log, mirroring the existing reuse fallback. 3. COMMENT: note in classifyJobShuffleShape that the materialized-prefix check is a point-in-time snapshot, safe because the only supported deployment is single-executor local mode and a FetchFailed on a prefix routes to a whole-group abort (already tested), not a lone-stage resubmit into the held gang slots. 4. COMMENT: note in AQEEnablePipelinedShuffle.duplicatedShuffleForms that pipelined participates in canonicalization, so duplicate detection only compares pre-flip forms. Co-authored-by: Claude Code
…und in second-pass review C1 (production memory leak, was masked by a test-only config default): a channel-manager pipelined shuffle registers in NO output tracker (usesStreamingShuffleOutputTracker = false), so ContextCleaner.doCleanupShuffle found it in neither the MapOutputTracker nor the StreamingShuffleOutputTracker and hit the bare "non-existent shuffle" else branch, never calling shuffleDriverComponents.removeShuffle. That RPC is the only thing that reaches SparkEnv.unregisterShuffleFromAllManagers -> PipelinedChannelShuffleManager.unregisterShuffle -> ChannelShuffleRendezvous.removeShuffle, so the process-wide queues (up to 64 batches x 1024 rows of live heap objects per reduce partition) leaked for the JVM's lifetime, one set per pipelined shuffle ever run. The SQL path (SQLExecution) does remove the shuffle, but only under spark.sql.classic.shuffleDependency.fileCleanup.enabled, whose default is Utils.isTesting -- true in tests (leak hidden), false in production (leak live). Fix: make doCleanupShuffle's tracker-less branch call shuffleDriverComponents.removeShuffle unconditionally (plus notify listeners). Safe for a genuinely non-existent id -- the BlockManager finds no blocks and unregisterShuffle for an unknown id is a no-op on every manager. No shuffle-manager type check in the cleaner: it stays transport-agnostic and just balances the registerShuffle the ShuffleDependency constructor issues for every shuffle. Regression test in PipelinedChannelShuffleSuite runs a real pipelined repartition, then drives doCleanupShuffle directly (the production path, minus flaky GC) and asserts the queues are freed. This test fails without the fix. S1 (latent, fail-safe today): in putUnlessAbandoned a successful q.offer can race abandon() (which does add(mark) then q.clear()); a batch offered just after the clear would be stranded in the queue with no reader to drain it. Re-check isAbandoned after a successful offer and clear the queue if so -- the reader has departed, so discarding is correct, and it keeps the queue empty for cleanup. Not a correctness bug today (each reduce partition is read once, no task retry re-reads a drained partition), but it compounded C1 and would bite if a partition were ever read twice. S2: reword the duplicatedShuffleForms comment -- the rule DOES re-run in later AQE rounds on a plan that already contains an earlier flip; it is not a before-any-flip pass. Missing a flipped/unflipped twin pair is still fail-safe (reuse collapses twins; the fan-out check rejects otherwise). Co-authored-by: Claude Code
…dd deterministic abandon/EOS tests
From a fourth review pass focused on test effectiveness and degenerate inputs (prior three
passes cleared correctness/concurrency/scheduler). No product correctness defect found; the
changes are one defensive guard plus test hardening.
- Guard against a zero-map producer (numMaps == 0, e.g. a pipelined shuffle over an empty
RDD): ChannelShuffleReader.advance() blocking-take()s before its end-of-stream count check,
so with numMaps == 0 (no writer ever enqueues, no marker ever arrives) a reduce task that
reached read() would block forever. Add an early return `if (endOfStreamSeen >= numMaps)`
at the top of advance() -> an empty iterator instead of a hang. Unreachable from the SQL
path (an exchange producer always has >= 1 partition) but constructible via the core RDD
entry point; the guard is harmless for numMaps > 0.
- Add fast, deterministic unit tests (no SparkContext, no wall-clock timeout as the assertion)
for the abandon / end-of-stream logic previously covered only by the 90s e2e LIMIT test:
* abandon() sets the mark and drains the queue (releases a parked writer, readies cleanup);
* clearAbandoned() resets a mark (re-run reuse) and removeShuffle() drops both queues and
all marks for a shuffle (executeTake shuffleId reuse);
* ChannelShuffleReader.read() (the REAL reader, via a mock handle + TempShuffleReadMetrics)
stops after exactly numMaps EndOfStream markers and yields every row in order.
- Rewrite the stale PipelinedLimitHangSuite scaladoc: it still described the tests as `ignore`
reproductions that hang; they are enabled acceptance tests asserting completion + rows == 10
in both AQE modes, each with its own isolated SparkSession.
- Wrap one pre-existing >100-char line in the channel suite (linter fix).
Co-authored-by: Claude Code
…ments Comment/scaladoc/test-title cleanup only, no logic change: strip terminology that should not appear in an upstream PR. Remove "WIP" markers; replace the internal codenames "v1"/"v2" (local-repartition v1/v2, two internal implementations) and "RTM"/"RTM-style" (Real-Time Mode) with self-contained wording that keeps the technical explanation. SPARK-57399 JIRA references are kept. Where a comment explained a mechanism via a "v1 did X, we do Y" comparison (the AQE placement rule's SinglePartition-chain reasoning), it is rewritten as a positive, self-contained statement rather than deleted, so the technical content survives. Co-authored-by: Claude Code
viirya
force-pushed
the
local-repartition-v2-pr
branch
from
August 19, 2026 02:57
7b60d48 to
f9e9e93
Compare
CI's SparkConfigBindingPolicySuite requires every config entry to declare a ConfigBindingPolicy. The two configs this PR adds -- spark.sql.pipelinedShuffle.enabled and spark.shuffle.pipelined.channel.batchSize -- lacked it (the check is newer than when they were first written). Both are NOT_APPLICABLE: neither changes the resolved plan of a view/UDF/procedure body (one is a physical-execution flip, the other a runtime batch size), matching the sibling spark.shuffle.manager.incremental. SparkConfigBindingPolicySuite 3/3 passes. Co-authored-by: Claude Code
Remove exploratory benchmarks that answered our own design questions and do not belong
in the PR:
- TPCDSPipelinedRunner.scala: a TPC-DS runner used to verify correctness and time
queries across AQE modes during development. Its findings are captured; the runner
itself is not product.
- The fixed-cost curve in PipelinedShuffleBenchmark (compareChannel / fixedCostCurve
and the `curve` main branch): it swept repartition+count over shrinking row counts to
settle the "are small shuffles worth pipelining" question (answer: yes, no crossover),
which is now decided and needs no standing benchmark.
KEPT in PipelinedShuffleBenchmark: the three-way transport comparison (regular vs RTM
RPC-streaming vs in-process channel), which answers the reviewer question "why not just
use the existing streaming manager?" (the channel beats streaming 3-5x on batch shapes),
plus the AQE / range / prototype compare rows.
sql/Test/compile clean.
Co-authored-by: Claude Code
…SqlBasedBenchmark framework
The benchmark was an ad-hoc object with a hand-rolled main / timing loop. Rework it onto
Spark's standard benchmark framework so it fits repo conventions and can emit a results file
(SPARK_GENERATE_BENCHMARK_FILES=1 -> benchmarks/PipelinedShuffleBenchmark-results.txt):
- `extends SqlBasedBenchmark`, implements `runBenchmarkSuite`, groups cases under
`runBenchmark(...)` blocks (three-way transport; regular-vs-channel AQE off; AQE on;
32-map oversubscribed), using `new Benchmark(...)` + the framework's warm-up / best-of-N.
- All original workloads preserved unchanged.
The one non-standard need: each transport (regular / RPC-streaming / channel) requires its OWN
SparkSession, because the shuffle manager and the pipelined flag are SparkContext-level and
fixed at startup -- they cannot be switched with setConf on a shared session. Handled with
`addTimerCase`: build the session before `startTiming()` and stop it after `stopTiming()`, so
session build/teardown (seconds of fixed cost) is EXCLUDED from the measured time while the
framework still runs its own iterations. Each case also stops any active/default session first,
so `getOrCreate` builds a fresh one rather than silently reusing the base trait's throwaway
session (which would ignore the case's master / manager config).
Verified by running: no session-reuse WARN, per-case stdev ~2-7% of best, and the headline
signal is clean (channel 1.7-7.2x over regular; RPC-streaming loses to regular on batch
shapes, e.g. 0.4x on repartition -- which is why the channel transport exists).
Co-authored-by: Claude Code
…k21) Generated baseline for PipelinedShuffleBenchmark via SPARK_GENERATE_BENCHMARK_FILES=1, per Spark convention (a *-results.txt under sql/core/benchmarks/). Captured on JDK 21 / Apple M4 Max; the standard CI benchmark workflow can regenerate on the reference environment. Highlights: the in-process channel beats the regular shuffle 1.4-7.4x across the batch shapes, while RTM's RPC-streaming transport loses to regular on the same shapes (0.4x on repartition, 2.2x on groupBy) -- the throughput-vs-latency split that justifies the channel carrying its own transport rather than reusing the streaming manager in local mode. Co-authored-by: Claude Code
viirya
force-pushed
the
local-repartition-v2-pr
branch
from
August 19, 2026 06:24
f9e9e93 to
172e003
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What changes were proposed in this pull request?
This PR adds an opt-in, in-process channel transport for pipelined shuffles and wires it
into the SQL layer, so a local-mode batch query can run its shuffle exchanges as pipelined
shuffles served entirely within one JVM -- producer and consumer stages co-scheduled by the
concurrent-stage scheduler, with records flowing through bounded in-memory queues instead of
shuffle files.
It builds on the already-merged pipelined-shuffle infrastructure --
PipelinedShuffleDependencyand dependency-type shuffle routing (SPARK-58185), concurrent-stage scheduling (SPARK-58263),
group-atomic failure and fail-fast rejection (SPARK-58398), and the MapOutputTracker decoupling
(SPARK-58454) -- and adds the pieces specific to a local, in-process transport.
New
coretransport (org.apache.spark.shuffle.local.pipelined):ChannelShuffleRendezvous: a process-wide rendezvous holding one boundedLinkedBlockingQueueper
(shuffleId, reducePartitionId). Every map task writing a reduce partition shares thequeue with the single reduce task that drains it. Queue elements are BATCHES of records (an
Arrayof pairs) or an end-of-stream marker, so the queue's per-operation lock cost is paidper batch, not per row.
ChannelShuffleWriter/ChannelShuffleReader: theShuffleWriter/ShuffleReaderfor thetransport. The writer batches records per reduce partition, pushes full batches onto the
queues, and emits one end-of-stream marker per partition; the reader drains its partition's
queue until it has seen
numMapsend-of-stream markers. Records + read/write time arereported to the shuffle metrics (no byte metrics: an in-process transport serializes nothing,
so there is no wire-byte count).
PipelinedChannelShuffleManager: aPipelinedShuffleManagerthat mints the writer/reader.It declares
usesStreamingShuffleOutputTracker = false(the reader/writer find each other by(shuffleId, partition)in-JVM, so no writer-location directory is needed) andrequiresDetachedRecords = true(records cross to a concurrent consumer thread, so the SQLlayer must copy each row off the producer's reused buffer). It
requires local mode in itsconstructor -- a cross-executor deployment would give each executor its own empty queue map
and hang every reader, so it fails loud at startup instead.
SQL integration:
EnablePipelinedShuffle(non-AQE) andAQEEnablePipelinedShuffle(AQE) rewrite eligibleShuffleExchangeExecnodes topipelined = true. Both gate on the opt-in flag, on localmode, and on the in-process channel manager actually being the configured pipelined manager;
otherwise they leave the plan regular. Both leave a plan regular when it contains a reused
exchange (a pipelined producer cannot fan out to more than one consumer).
ShuffleExchangeExecgains apipelinedflag and, for the pipelined path, copies each rowoff the producer's reused buffer before it is handed across the channel.
spark.sql.pipelinedShuffle.enabled(default false) turns therewrite on;
spark.shuffle.pipelined.channel.batchSize(default 1024) sets the per-partitionbatch size.
spark.shuffle.manager.incrementalselects the channel manager.Scheduler:
DAGScheduler's job-shape classification is relaxed to admit a MATERIALIZED-PREFIX MIXED job:a job may mix regular and pipelined shuffles when every regular boundary reachable from the
final RDD is fully materialized and no pipelined shuffle sits below a regular one. This is the
shape adaptive execution produces (prior map-stage jobs materialize the prefix; the final job
runs the pipelined tail). An unmaterialized regular prefix, and a pipelined shuffle below a
regular boundary, stay rejected fail-fast. The relaxation is a strict superset: previously
rejected shapes now run, and every previously-valid job classifies and schedules identically
(an all-pipelined job is unchanged, pinned by a test). The
DAGScheduleralso passes theresult stage's live reduce partitions to the producer, so a partial-read job (LIMIT /
executeTake reads a subset) does not fill and wedge the queues of partitions no consumer will
drain.
Why are the changes needed?
Local repartition (a shuffle whose producer and consumer are co-located in one JVM) does not
need the durable, file-based, cross-executor machinery of a regular shuffle. Serving it through
an in-process channel -- records handed directly from writer to a concurrently running reader --
avoids the shuffle-file write/read, block-manager, and serialization/fetch startup costs, which
dominate for the small-to-medium shuffles typical of a single-executor deployment.
Crucially, this reuses Spark's own scheduling machinery (the already-merged concurrent-stage /
pipelined-shuffle infrastructure) rather than introducing a parallel mechanism outside the
shuffle framework: the two shuffle sides remain real, separately scheduled stages, so the shuffle
boundary stays visible in the UI and to AQE, and the same scheduler serves both regular and
pipelined shuffles routed by dependency type. Each transport serves the workload it was built
for -- the RPC streaming transport is latency-optimized for cross-executor streaming, while
the in-process channel is throughput-optimized for local batch (a three-way transport benchmark
in this PR shows the channel beating a regular shuffle while the RPC streaming transport loses to
it on batch shapes, which is why local batch needs its own transport rather than reusing the
streaming one).
Does this PR introduce any user-facing change?
No behavior change by default: the feature is off unless
spark.sql.pipelinedShuffle.enabledisset to true (default false) AND the in-process channel manager is configured via
spark.shuffle.manager.incremental, both in local mode. With those set, eligible shuffleexchanges of a batch query run as in-process pipelined shuffles; results are unchanged, and the
shuffle appears in the UI as concurrently scheduled producer/consumer stages rather than a
materialized boundary. Two new configs are added (both documented,
spark.sql.pipelinedShuffle.enabledand
spark.shuffle.pipelined.channel.batchSize).How was this patch tested?
New unit and end-to-end suites, all passing:
PipelinedChannelShuffleSuite(core): the channel transport loses/duplicates no rows androutes correctly; matches a regular shuffle's grouping; slot admission fits/rejects a group;
the manager refuses to construct outside local mode; a materialized regular prefix runs
end-to-end while an unmaterialized one is rejected;
ContextCleanerfrees the channel's queuesfor a tracker-less pipelined shuffle; and deterministic unit tests for the abandon /
end-of-stream-counting logic (abandon marks and drains;
clearAbandoned/removeShufflereset;the reader stops after exactly
numMapsmarkers).PipelinedShuffleSqlSuite,AQEPipelinedShuffleSuite(sql): batch queries -- repartition,keyed groupBy, range partitioning, sort-merge join, single-partition aggregate, chains through
SinglePartition -- run end-to-end through the channel and produce correct results against an
independently computed ground truth, under both the non-AQE and AQE rules; over-wide plans fail
loud at admission; cross-subquery reuse cannot create a shared pipelined exchange.
PipelinedLimitHangSuite(sql): a LIMIT over a pipelined shuffle completes and returns thecorrect rows in both AQE modes (an early-stopping reader must not wedge the writer).
DAGSchedulerSuite(core): the materialized-prefix relaxation accepts a fully-materializedmixed job and rejects unmaterialized-prefix / pipelined-below-regular; a FetchFailed on a
pipelined group member (including one reading an external materialized prefix) aborts the whole
group rather than resubmitting a lone stage; an all-pipelined job classifies identically
under the relaxation (a previously-valid shape is unchanged).
PipelinedShuffleRoutingSuite,ContextCleanerSuite: routing by dependency type and cleanup.A
PipelinedShuffleBenchmark(on the standardSqlBasedBenchmarkframework, with a checked-inresults file) compares the channel against the regular shuffle and against the RPC streaming
transport across batch shapes.
Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code