8. Shuffle Internals

• 4.5에서 Shuffle을 “Shuffle Write → 네트워크 → Shuffle Read"의 흐름으로 봤고, 비용은 직렬화, 디스크, 네트워크, 정렬, Spill로 나눠 봤음
• 이번 장은 그 비용이 실제로 어디서 생기고, Spark가 그걸 어떻게 줄이는지 구현 수준에서 봄
  • Map 쪽: 어떤 Writer가 선택되고, 메모리가 모자라면 언제 Spill하는가
  • 출력 포맷: Map Task 하나가 무엇을 남기는가
  • Reduce 쪽: 블록을 어떻게 찾고, 얼마나 동시에 가져오는가
• 근거는 Apache Spark 4.2.0 소스 (core/src/main/{scala,java}/org/apache/spark/shuffle/)

flowchart TD
  register["registerShuffle<br/>Driver에서 1회"] --> pick{"Writer 선택"}
  pick -- "combine 없음 · 파티션 ≤ 200" --> bypass["BypassMergeSortShuffleWriter"]
  pick -- "relocatable · combine 없음" --> unsafe["UnsafeShuffleWriter<br/>Tungsten"]
  pick -- "그 외" --> sort["SortShuffleWriter<br/>ExternalSorter"]
  bypass --> output["data 파일 1개 + index 파일 1개"]
  unsafe --> output
  sort --> output
  output --> push["Push-based Shuffle<br/>선택"]
  output --> fetch["Reduce Task<br/>ShuffleBlockFetcherIterator"]
  push --> fetch
  fetch --> agg["Reduce 쪽 정렬 · 집계"]

8.1 Choosing a Shuffle Writer

ShuffleDependency가 생성될 때 Driver에서 SortShuffleManager.registerShuffle()이 딱 한 번 호출되고, 여기서 Writer 종류가 정해짐.

override def registerShuffle[K, V, C](
    shuffleId: Int, dependency: ShuffleDependency[K, V, C]): ShuffleHandle = {
  if (SortShuffleWriter.shouldBypassMergeSort(conf, dependency)) {
    new BypassMergeSortShuffleHandle[K, V](shuffleId, dependency...)
  } else if (SortShuffleManager.canUseSerializedShuffle(dependency)) {
    new SerializedShuffleHandle[K, V](shuffleId, dependency...)
  } else {
    new BaseShuffleHandle(shuffleId, dependency)
  }
}

검사 순서가 곧 우선순위임: Bypass → Serialized → Base

Bypass의 200은 어디서 왔나

• Bypass는 출력 파티션 수만큼 파일 Writer를 동시에 엶
• 파티션이 많으면 열린 스트림과 직렬화 버퍼가 너무 많아짐
  • spark.shuffle.file.buffer 기본 32KB × 200개 ≈ 6.4MB가 Task 하나에 동시에 잡힘
• 그래서 파티션 수가 spark.shuffle.sort.bypassMergeThreshold(기본값 200) 이하일 때만 사용함
• 이어 붙일 때는 가능하면 transferTo(커널 수준 복사)를 씀

“직렬화 객체 재배치 가능"이란

• 직렬화된 레코드 바이트를 다른 위치로 옮겨 이어 붙여도 역직렬화 결과가 같은 성질
• Kryo와 Spark SQL 전용 UnsafeRowSerializer가 이 성질을 가짐
• 이 성질이 있어야 레코드 본문을 해석하지 않고 바이트 단위로 정렬·병합할 수 있음 (8.3)

DataFrame 연산의 Shuffle은 대부분 Map-side 집계를 Shuffle 앞의 별도 연산자(HashAggregate partial)로 처리함. 그래서 Shuffle 자체에는 combine이 없고 SQL 전용 직렬화기를 씀.

• 출력 파티션이 200 이하면(기본 spark.sql.shuffle.partitions = 200을 그대로 쓰는 경우) Bypass
• 200을 넘기면 Serialized 경로를 탐

8.2 Sort-Based Write Path: ExternalSorter

Base 경로(RDD의 reduceByKey처럼 Map-side combine이 있는 경우)의 핵심은 ExternalSorter임.

insertAll: 두 가지 자료구조

def insertAll(records: Iterator[Product2[K, V]]): Unit = {
  val shouldCombine = aggregator.isDefined
  if (shouldCombine) {
    while (records.hasNext) {
      ...
      map.changeValue((partitionId, kv._1), update)   // 제자리 집계
      maybeSpillCollection(usingMap = true)
    }
  } else {
    while (records.hasNext) {
      ...
      buffer.insert(partitionId, kv._1, kv._2)         // 단순 append
      maybeSpillCollection(usingMap = false)
    }
  }
}

• 두 구조 모두 Key를 (partitionId, key)로 저장
• 정렬하거나 Spill할 때 Partitioner를 레코드마다 다시 호출하지 않기 위해서

maybeSpill: 5MB에서 시작해 두 배씩

Spill 여부는 Spillable.maybeSpill이 판단함.

protected def maybeSpill(collection: C, currentMemory: Long): Boolean = {
  val shouldSpill = if (_elementsRead > numElementsForceSpillThreshold
      || currentMemory > maxSizeForceSpillThreshold) {
    true                                               // ① 강제 Spill
  } else if (_elementsRead % 32 == 0 && currentMemory >= myMemoryThreshold) {
    // ② 현재 메모리의 두 배까지 Shuffle 메모리 풀에 요청
    val amountToRequest = 2 * currentMemory - myMemoryThreshold
    val granted = acquireMemory(amountToRequest)
    myMemoryThreshold += granted
    currentMemory >= myMemoryThreshold                 // ③ 못 받았으면 바로 Spill
  } else { false }
  if (shouldSpill) {
    spill(collection); _elementsRead = 0
    _memoryBytesSpilled += currentMemory; releaseMemory()
  }
  shouldSpill
}

읽는 법:

레코드 32개마다 한 번만 크기 추정 (추정 자체가 비싸서)
      ↓
현재 크기 ≥ 임계값?
      ↓ 예
"현재 × 2 − 임계값"만큼 추가 메모리 요청
      ↓
다 받음  → 임계값이 현재의 2배로 커짐, 계속 메모리에서 처리
못 받음 → 지금 컬렉션을 디스크로 Spill
      ↓
Spill 후 releaseMemory() → 임계값을 초기값(5MB)으로 되돌림

• 초기 임계값 spark.shuffle.spill.initialMemoryThreshold = 5MB
• 메모리를 “얼마나 받을 수 있는지"는 9.3의 Task별 배분 규칙이 결정함
• Reduce 쪽 ExternalAppendOnlyMap도 같은 Spillable을 상속하므로 이 알고리즘은 Map과 Reduce 양쪽에 공통 (8.6)

Spill 파일과 병합

• Spill은 레코드 10,000개 단위 배치로 flush함
• Spill 파일 메타데이터에 파티션별 원소 수를 기록하므로, 레코드마다 partitionId를 쓰지 않아도 됨
• 최종 출력 시
  • Spill이 없으면 메모리 컬렉션만 정렬해서 바로 씀
  • Spill이 있으면 모든 Spill 파일과 메모리 데이터를 한 번의 pass로 병합함

merge() — 파티션마다
├─ 집계 필요   → mergeWithAggregation
├─ 정렬만 필요 → mergeSort (PriorityQueue 기반 k-way merge)
└─ 둘 다 아님  → 그냥 이어 붙임

• mergeWithAggregation: Key 순서가 완전한 정렬(total order)이면 같은 Key가 연속으로 오므로 바로 합침. 해시 순서만 있으면 해시가 같은 Key들을 모아 ==로 다시 확인함

8.3 Tungsten Write Path: Sorting Pointers, Not Records

Serialized 경로(UnsafeShuffleWriter)는 레코드를 객체로 다루지 않음.

바로 직렬화해 메모리 페이지에 복사하고 포인터만 저장

// UnsafeShuffleWriter.insertRecordIntoSorter
final int partitionId = partitioner.getPartition(record._1());
serBuffer.reset();
serOutputStream.writeKey(key, OBJECT_CLASS_TAG);
serOutputStream.writeValue(record._2(), OBJECT_CLASS_TAG);
sorter.insertRecord(serBuffer.getBuf(), Platform.BYTE_ARRAY_OFFSET, serializedRecordSize, partitionId);

// ShuffleExternalSorter.insertRecord
acquireNewPageIfNecessary(required);
final long recordAddress = taskMemoryManager.encodePageNumberAndOffset(currentPage, pageCursor);
UnsafeAlignedOffset.putSize(base, pageCursor, length);   // 레코드 길이
Platform.copyMemory(recordBase, recordOffset, base, pageCursor, length);  // 바이트 복사
inMemSorter.insertRecord(recordAddress, partitionId);    // 정렬 배열에는 주소만

• 레코드 본문은 Task가 할당받은 메모리 페이지에 연속으로 쌓임
• 정렬 배열에는 (partitionId, 레코드 주소)를 담은 long 하나만 들어감
• 페이지 주소 체계는 9.4에서 볼 TaskMemoryManager 페이지 테이블을 따름

PackedRecordPointer: 8바이트에 다 넣기

 63            40 39          27 26                    0
┌───────────────┬──────────────┬───────────────────────┐
│ partition id  │ page number  │   offset in page      │
│   24 bit      │   13 bit     │       27 bit          │
└───────────────┴──────────────┴───────────────────────┘

static final int MAXIMUM_PAGE_SIZE_BYTES = 1 << 27;      // 128MB
static final int MAXIMUM_PARTITION_ID = (1 << 24) - 1;   // 16,777,215

이 비트 배치가 그대로 API 제약이 됨.

• partition id가 24비트라서 Serialized 경로의 파티션 상한이 2^24 (8.1의 조건)
• offset이 27비트라서 페이지 하나는 최대 128MB
• page number가 13비트라서 2^13 페이지 × 128MB = Task당 최대 1TB
• partition id가 최상위 비트에 있으므로 long 값을 그대로 비교하면 파티션 순으로 정렬됨

Radix Sort vs Tim Sort

if (useRadixSort) {
  offset = RadixSort.sort(array, pos,
    PackedRecordPointer.PARTITION_ID_START_BYTE_INDEX,   // 5
    PackedRecordPointer.PARTITION_ID_END_BYTE_INDEX,     // 7 → 상위 3바이트만
    false, false);
} else {
  new Sorter<>(new ShuffleSortDataFormat(buffer)).sort(array, 0, pos, SORT_COMPARATOR);
}

• 파티션 순서만 맞으면 되므로 상위 3바이트(24비트)에만 Radix Sort를 돌림
• spark.shuffle.sort.useRadixSort 기본 true
• 대가: Radix Sort는 배열 크기만큼의 보조 버퍼가 필요해서 배열의 절반만 실제 데이터에 씀 (Tim Sort는 2/3)

포인터만 정렬하면 왜 빠른가

1. 캐시 효율: 레코드당 8바이트라 정렬 배열이 CPU 캐시에 많이 들어감. 64바이트 캐시 라인 하나에 포인터 8개
2. 역직렬화 없음: 비교에 필요한 건 long의 상위 24비트뿐
3. Spill도 바이트 복사: 정렬 순서대로 페이지의 바이트를 그대로 파일에 씀 (직렬화기는 아무것도 안 하는 DummySerializerInstance)

Spill 병합: Fast Merge

• 압축 코덱이 스트림 이어 붙이기를 지원하면(lz4, snappy, zstd 등) Spill 파일들의 파티션 구간을 바이트를 해석하지 않고 concat함
• transferTo를 쓸 수 있으면 커널 수준에서 zero-copy로 복사함
• Spill이 하나뿐이면 병합 없이 그 파일을 최종 출력으로 사용

8.4 Shuffle Output Format: Data and Index Files

Writer가 무엇이든 출력 형태는 같음.

Map Task 1개
├─ shuffle_<shuffleId>_<mapId>_0.data   ← 모든 Reduce 파티션 데이터가 순서대로
└─ shuffle_<shuffleId>_<mapId>_0.index  ← 파티션 경계 offset

• 파일 이름의 Reduce ID 자리가 항상 0인데, 여러 Reduce 파티션 출력이 한 파일에 합쳐져 있다는 흔적임
• 예전(Hash Shuffle)처럼 “Map 수 × Reduce 수"개의 파일이 생기지 않음

index 파일 = 파티션 길이의 누적합

val offsets = lengths.scanLeft(0L)(_ + _)

파티션 길이 : [100, 250, 30]
index 파일  : [0, 100, 350, 380]   ← N+1개 Long, 마지막 값 = 파일 크기

Reducer가 자기 몫을 찾는 법

channel.position(startReduceId * 8L)
val startOffset = in.readLong()
channel.position(endReduceId * 8L)
val endOffset = in.readLong()
new FileSegmentManagedBuffer(transportConf, getDataFile(shuffleId, mapId, dirs),
  startOffset, endOffset - startOffset)

• index 파일에서 Long 두 개만 읽으면 구간이 나오므로 O(1)로 랜덤 접근함
• FileSegmentManagedBuffer는 파일을 복사하지 않고 “이 파일의 이 구간"만 가리키므로, 네트워크로 보낼 때도 zero-copy

커밋 원자성

• 임시 파일에 쓰고 동기화 블록 안에서 rename함
• 같은 Map Task의 다른 attempt(7.7의 Speculation 등)가 이미 성공했으면 자기 결과를 버림

8.5 Read Path: Fetching Shuffle Blocks

MapOutputTracker로 위치 조회

Reduce Task
  → mapOutputTracker.getMapSizesByExecutorId(shuffleId, startMapIndex, endMapIndex,
                                               startPartition, endPartition)
  → Executor별로 묶인 (blockId, size, mapIndex) 목록  (크기 0인 블록은 제외)

• 블록 크기를 미리 알기 때문에 한 번에 가져올 양을 조절할 수 있음
• Map 범위(startMapIndex·endMapIndex) 파라미터는 AQE가 Skew Join을 분할할 때 씀 (11.5)

Fetch 대상 4가지

ShuffleBlockFetcherIterator
├─ local              같은 Executor → 파일 직접 읽기
├─ host-local         같은 호스트의 다른 Executor → 디스크 직접 읽기 (네트워크 우회)
├─ push-merged-local  같은 호스트의 Push 병합 결과 (8.7)
└─ remote             FetchRequest로 묶어 네트워크 요청

• 전체 흐름이 lazy Iterator라서 블록이 도착하는 대로 스트림 단위로 역직렬화함

maxBytesInFlight 스로틀링

// Make remote requests at most maxBytesInFlight / 5 in length; the reason to keep them
// smaller than maxBytesInFlight is to allow multiple, parallel fetches from up to 5
// nodes, rather than blocking on reading output from one node.
private val targetRemoteRequestSize = math.max(maxBytesInFlight / 5, 1L)

• 동시에 받는 중인 데이터 총량의 상한은 48MB
• 요청 하나는 그 1/5(약 9.6MB)이라서 최대 5개 노드에서 병렬로 받음
• 한 노드가 느려도 전체가 그 노드 하나에 막히지 않게 하려는 설계

Batch Fetch

• AQE가 연속된 파티션 여러 개를 Reduce Task 하나에 몰아줄 때(11.4)
• 같은 Map 출력의 연속 구간을 요청 한 건으로 가져옴
• 조건: 직렬화기가 재배치를 지원하고, 압축하지 않거나 압축 스트림을 이어 붙일 수 있어야 함

8.6 Reduce-Side Aggregation and Sorting

블록을 받은 뒤의 처리는 세 갈래로 나뉨.

keyOrdering 있음 (sortByKey 등)   → ExternalSorter 재사용 (Write path와 같은 클래스)
aggregator만 있음 (reduceByKey 등) → ExternalAppendOnlyMap
둘 다 없음                         → 그대로 통과

• ExternalAppendOnlyMap도 Spillable을 상속하므로, 8.2의 “5MB 시작, 두 배 요청, 실패 시 Spill” 규칙이 Reduce 쪽에도 그대로 적용됨
• Reduce 쪽 OOM은 대개 캐시가 아니라 이 해시 테이블(Task의 working set) 때문에 남
  • 공식 튜닝 가이드의 처방도 “병렬도를 올려 Task당 입력을 줄여라”

8.7 Push-Based Shuffle

기본 Shuffle의 Reduce 쪽 문제:

Map 1,000개 × Reduce 1,000개
→ Reduce Task 하나가 1,000개 노드에서 작은 블록 1,000개를 가져옴
→ 작은 랜덤 I/O + 요청 수 폭증

Push-based Shuffle(3.2+)은 Map이 끝난 뒤 블록을 External Shuffle Service로 밀어 넣어 같은 Reduce 파티션끼리 미리 합침.

flowchart LR
  m1["Map Task 1"] -- "push" --> ess1["ESS Merger A<br/>파티션 0 병합"]
  m2["Map Task 2"] -- "push" --> ess1
  m1 -- "push" --> ess2["ESS Merger B<br/>파티션 1 병합"]
  m2 -- "push" --> ess2
  ess1 -- "큰 블록 1개" --> r0["Reduce Task 0"]
  ess2 -- "큰 블록 1개" --> r1["Reduce Task 1"]

• Reduce는 흩어진 작은 블록 대신 병합된 큰 블록을 읽으므로 랜덤 I/O와 요청 수가 줄어듦
• 원본 블록도 남겨 두므로, 병합에 실패하면 기존 방식으로 fallback함
• Writer 종류와 무관하게 동작
  • 이미 쓰인 data 파일과 파티션 길이만 재료로 씀
  • 8.4의 공통 출력 포맷 덕분
• Push 순서는 일부러 섞음. 여러 Mapper가 같은 Merger에 동시에 몰리지 않게 하려는 것

• 원래 YARN의 External Shuffle Service 기반 기능이라 환경 지원 여부를 먼저 확인해야 함

8.8 The Whole Shuffle on One Page

[MAP SIDE]
registerShuffle에서 Handle 확정
├─ BypassMergeSortShuffleWriter  (combine X, 파티션 ≤ 200)
│    파티션당 파일 → concat
├─ UnsafeShuffleWriter            (relocatable, combine X, 파티션 ≤ 2^24)
│    직렬화 → 페이지 복사 + 8B packed pointer → radix sort → fast merge
└─ SortShuffleWriter              (그 외)
     ExternalSorter.insertAll (집계? Map : Buffer)
     → maybeSpill (5MB 시작, 두 배 요청, 실패 시 Spill) → k-way merge
          ↓
shuffle_<sid>_<mapId>_0.data + .index (누적 offset N+1개)
          ↓ (선택) Push → ESS 병합

[REDUCE SIDE]
MapOutputTracker 조회 → ShuffleBlockFetcherIterator
  (local / host-local / push-merged-local / remote, 요청당 48m/5)
→ 역직렬화 → 정렬? ExternalSorter : 집계? ExternalAppendOnlyMap : 통과

연결 고리 세 개:

1. PackedRecordPointer의 비트 배치가 Serialized Shuffle의 파티션 상한을 정함
2. Spillable.maybeSpill이 Map·Reduce 양쪽에 공유됨
3. data와 index 포맷이 세 Writer의 공통 출력 계약이라서 Push-based Shuffle이 Writer와 무관하게 동작

실무로 옮기면:

• Shuffle 비용을 줄이는 가장 효과 큰 방법은 여전히 “Shuffle 자체를 없애기”(Broadcast Join, 5.3)와 “고르게 나누기”(Skew 대응, 5.4)
• Spill이 보이면 Task당 데이터가 메모리 몫보다 크다는 뜻이므로, 파티션 수를 늘리는 것이 가장 싼 처방 (12.3)