viirya opened a new pull request, #58097:
URL: https://github.com/apache/spark/pull/58097
### 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 --
`PipelinedShuffleDependency`
and 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 `core` transport (`org.apache.spark.shuffle.local.pipelined`):
- `ChannelShuffleRendezvous`: a process-wide rendezvous holding one bounded
`LinkedBlockingQueue`
per `(shuffleId, reducePartitionId)`. Every map task writing a reduce
partition shares the
queue with the single reduce task that drains it. Queue elements are
BATCHES of records (an
`Array` of pairs) or an end-of-stream marker, so the queue's per-operation
lock cost is paid
per batch, not per row.
- `ChannelShuffleWriter` / `ChannelShuffleReader`: the
`ShuffleWriter`/`ShuffleReader` for the
transport. 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 `numMaps` end-of-stream markers. Records +
read/write time are
reported to the shuffle metrics (no byte metrics: an in-process transport
serializes nothing,
so there is no wire-byte count).
- `PipelinedChannelShuffleManager`: a `PipelinedShuffleManager` that 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) and
`requiresDetachedRecords = true` (records cross to a concurrent consumer
thread, so the SQL
layer must copy each row off the producer's reused buffer). It `require`s
local mode in its
constructor -- 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) and `AQEEnablePipelinedShuffle` (AQE)
rewrite eligible
`ShuffleExchangeExec` nodes to `pipelined = true`. Both gate on the opt-in
flag, on local
mode, 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).
- `ShuffleExchangeExec` gains a `pipelined` flag and, for the pipelined
path, copies each row
off the producer's reused buffer before it is handed across the channel.
- Two registered configs: `spark.sql.pipelinedShuffle.enabled` (default
false) turns the
rewrite on; `spark.shuffle.pipelined.channel.batchSize` (default 1024)
sets the per-partition
batch size. `spark.shuffle.manager.incremental` selects 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 `DAGScheduler`
also passes the
result 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.enabled` is
set 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 shuffle
exchanges 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.enabled`
and `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 and
routes 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; `ContextCleaner` frees
the channel's queues
for a tracker-less pipelined shuffle; and deterministic unit tests for the
abandon /
end-of-stream-counting logic (abandon marks and drains;
`clearAbandoned`/`removeShuffle` reset;
the reader stops after exactly `numMaps` markers).
- `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 the
correct rows in both AQE modes (an early-stopping reader must not wedge
the writer).
- `DAGSchedulerSuite` (core): the materialized-prefix relaxation accepts a
fully-materialized
mixed 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 standard `SqlBasedBenchmark`
framework, with a checked-in
results 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
--
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]
