Shuffle

Shuffle is the slowest, most expensive part of most Spark jobs and the source of most performance work. The default implementation is sort-based; push-based shuffle is the opt-in evolution shipped in Spark 3.x.

The data path

sequenceDiagram
    participant M as Map task
    participant SW as SortShuffleWriter
    participant DM as DiskBlockManager
    participant ESS as External Shuffle Service
    participant R as Reduce task
    participant SBF as ShuffleBlockFetcherIterator

    M->>SW: write(records)
    SW->>SW: sort by partition (and optionally by key)
    SW->>DM: write data file + index file
    Note right of DM: stable on disk;<br/>survives executor death

    R->>SBF: read(shuffle, reduceId)
    alt Same node
        SBF->>DM: read local data file via index
    else Different node
        SBF->>ESS: open block stream
        ESS->>DM: read data file
        ESS-->>SBF: bytes
    end
    SBF-->>R: deserialized records

Map side

SortShuffleManager (core/.../shuffle/sort/SortShuffleManager.scala) is the default. It chooses a writer for each map task:

• BypassMergeSortShuffleWriter - one file per partition, no sort. Fast for small partition counts.
• UnsafeShuffleWriter - serialized records sorted in off-heap memory. Used when the serializer relocates serialized objects (Kryo with the right options).
• SortShuffleWriter - the general path; sorts records then writes.

All three end up writing one data file + one index file per map task, addressed via IndexShuffleBlockResolver (core/.../shuffle/IndexShuffleBlockResolver.scala).

Reduce side

ShuffleBlockFetcherIterator (core/.../storage/ShuffleBlockFetcherIterator.scala, ~77 KB) runs in each reduce task. It:

• Splits requests into local and remote.
• Pipelines remote fetches up to a per-host concurrency cap.
• Spills oversize batches to disk to keep direct-memory bounded.
• Verifies block lengths and falls back when corrupt blocks are detected.

ShuffledRDD (core/.../rdd/ShuffledRDD.scala) is the RDD-level wrapper that triggers a shuffle dependency.

Tracking output

MapOutputTracker (core/.../MapOutputTracker.scala, ~75 KB) is the registry of "which executor has which shuffle blocks". The driver runs MapOutputTrackerMaster, executors run MapOutputTrackerWorker. After a stage completes, MapStatus ( core/.../scheduler/MapStatus.scala) is reported to the master and broadcast to reducers.

Push-based shuffle

Push-based shuffle (Spark 3.2 onward) reduces fetch fan-out by having mappers proactively push their output to mergers. Reducers then read merged blocks - one big block per merger - instead of a fan-out per mapper.

Key files:

• PushBasedFetchHelper (core/.../storage/PushBasedFetchHelper.scala).
• ShufflePushBlockId, ShuffleMergedBlockId (core/.../storage/BlockId.scala).
• MergeStatus (core/.../scheduler/MergeStatus.scala) - the merge-side analogue of MapStatus.
• Server-side merging logic in common/network-shuffle/.../shuffle/protocol/ and common/network-shuffle/.../shuffle/MergeManager.

DAGScheduler coordinates a "finalize merge" RPC after all map tasks succeed; only after finalization can reducers safely fetch merged blocks. The legacy non-merged blocks are still available as a fallback.

External shuffle service

common/network-shuffle/ and the YARN-specific entry point in common/network-yarn/ define a long-running daemon that lives outside the executor. It serves shuffle blocks from disk even after the executor that wrote them dies. This is essential for dynamic allocation: an idle executor can be killed without losing its shuffle output.

The service is also the host for push-based shuffle's mergers.

Pluggable shuffle SPI

core/src/main/java/org/apache/spark/shuffle/api/ defines a Java SPI that allows third parties to drop in alternative shuffle implementations (RDMA, Apache Celeborn, etc.). The SPI covers writers, readers, plugin lifecycle, and the metadata that the driver needs.

Failure handling

Shuffle is the place where partial failure is most visible:

• FetchFailedException - reducer cannot fetch a block. The DAG scheduler marks the upstream ShuffleMapStage as missing, retries the missing map outputs, and then retries the failed ResultStage.
• IndexShuffleBlockResolver repairs missing or corrupt index files when possible.
• Push-based shuffle can fall back to non-merged blocks transparently.

Integration points

• The SQL exchange operators (sql/core/.../execution/exchange/) build ShuffleExchangeExec plans which produce ShuffledRDDs.
• AQE consumes MapStatus info to coalesce/skew-split partitions.
• BlockManager's getRemoteBytes is what ShuffleBlockFetcherIterator ultimately calls.
• The external shuffle service has its own metric set published via the metrics system.

Entry points for modification

• Add a new shuffle writer strategy: implement ShuffleWriter and update SortShuffleManager.getWriter selection logic.
• Add a new shuffle SPI provider: implement org.apache.spark.shuffle.api.ShuffleExecutorComponents and configure spark.shuffle.sort.io.plugin.class.
• Improve push-based shuffle: most action is in PushBasedFetchHelper.scala and the MergeManager in common/network-shuffle/.
• Add a metric: extend core/.../shuffle/ShuffleMetricsSource.scala (and the analogous external shuffle service source).