Spark的Shuffle机制

Shuffle机制

1 Shuffle的意义及设计挑战

运行在不同stage、不同节点上的task间如何进行数据传递，这种传递过程称为Shuffle机制。Shuffle解决的问题是如何将数据重新组织，使其能够在上游和下游task之间进行传递和计算。

如果只是单纯的数据传输，只需要对数据进行分区，通过网络传输即可，但是Shuffle还需要应对各种类型的计算（比如聚合、排序），而且数据量的问题还需要解决。

2 Shuffle的设计思想

2.1 解决数据分区和数据聚合问题

2.1.1 数据分区问题

数据分区问题针对Shuffle Write阶段，这个阶段要解决两个问题，一个是如何确定分区个数？在Spark中，分区个数和task个数一致。分区个数可以由用户自定义，如groupByKey(numPartitions)中的numPartitions一般定义为集群中可用CPU个数的1-2倍，即对每个map task的输出数据划分为numParition份。如果用户没有自定义，默认分区个数是parent RDD分区个数的最大值，第二个问题是如何对map task输出数据进行分区？解决方法是对map task输出的每一个<K,V>record，根据Key计算出partitionId，具有不同partitionId的record被输出到不同的分区。Hask(Key)%numParirion，这种方法操作简单，但是不支持Shuffle Write端的combine()操作。

2.1.2 数据聚合问题

该问题针对Shuffle Read阶段，即如何获取上游不同task的输出数据并按照Key进行聚合？

groupByKey()	<K,V>record> -> <K,CompactBuffer(V)>
reduceByKey()	<K,V>record> -> <K,func(list(V))>

数据聚合的本质是将相同Key的record放在一起，并进行必要的计算，这个过程可以利用HashMap实现，方法是使用two-phase aggregation(两步聚合)，先将不同tasks获取到的<K,V>record存放到HashMap中，HashMap<K,list(V)>。然后对HashMap中每一个<K,list(V)>record，使用func计算得到<K,func(list(V))>record。

两步聚合方案的优点是可以解决数据聚合问题，逻辑清晰，容易实现，缺点是所有Shuffle的record都会被放到HashMap中，占用内存空间较大，另外，对于包含聚合函数的操作，如reduceByKey(func)，需要先将数据聚合到HashMap中以后在执行func()聚合函数，效率较低。

优化方案：对于包含聚合函数的操作来说，我们可以采用Online aggregation(在线聚合)的方法来减少空间的占用。

2.2 解决map()端的combine问题

进行combine()操作的目的是减少Shuffle的数据量，只有包含聚合函数的数据操作需要进行map()端的combine，如reduceByKey()、foldByKey()、aggregateByKey()、combineByKey()、distinct()等。对于不包含聚合函数的操作，combine并不能减少数据量。

解决方案：从本质上将，combine和Shuffle Read端的聚合过程没有区别，区别是数据来源不一样，Shuffle Read端聚合的是来自所有map task输出的结果，而combine只需要对单一的task进行处理。因此可以采用SHuffle Read端基于HashMap的解决方案。

2.3 解决sort问题

理论上，在Shuffle Write端不需要排序，但是如果进行了排序，Shuffle Read获取到的数据已经是部分有序的数据，可以减少Shuffle Read端的复杂度。

对于排序的时机，有以下3种方案：

1. 先排序再聚合。这种方案需要先使用线性数据结构如Array，存储Shuffle Read的<K,V>record，然后对Key进行排序，排序后的数据可以直接从前到后进行扫描聚合，不需要再使用HashMap进行hash-based聚合。这也是Hadoop MR引擎采用的方案，有点是既可以排序又可以聚合，缺点是需要较大内存空间来存储线性数据结构，而且排序和聚合过程不能同时进行，即不能使用在线聚合，效率较低。
2. 排序和聚合同时进行。可以使用带有排序功能的Map，比如TreeMap来对中间数据进行聚合，每次Shuffle Read获取到一个record，就将其放入TreeMap中与现有的record进行聚合，过程与HashMap类似，只是TreeMap自带排序功能。优点是排序和聚合可以同时进行，缺点是相比HashMap，TreeMap的排序复杂度较高，TreeMap的插入时间复杂度是O(nlogn)，而且需要不断调整树的结构，不适合数据量很大的情况。
3. 先聚合再排序。即维持现有的基于HashMap的聚合方案不变，将HashMap中的record或record引用放入线性数据结构中进行排序。优点是聚合和排序过程独立，灵活度高，而且之前的在线聚合的方案不需要改动，缺点是需要赋值数据或引用，空间占用较大。

Spark中使用的第3种方案。

2.4 解决内存不足的问题

解决方案：使用内存+磁盘混合存储方案

3 Spark中Shuffle框架的设计

3.1 Shuffle Write框架设计和实现

在Shuffle Write阶段，数据需要分区、聚合和排序3个功能。但在每个数据操作只需要满足其中的一个或两个功能，Spark为了支持所有情况，设计了一个通用的Shuffle Write框架，map()输出 -> 聚合 -> 排序 -> 分区。

map task每计算出一个record及其partitionId，将record放入类似HashMap的数据结构中进行聚合，聚合完成后，再将HashMap中的数据放入类似Array的数据结构进行排序，即可以按照partitionId，也可以按照partitionId+Key进行排序，最后根据partiionId将数据写入不同的数据分区，存放到本地磁盘上，其中，aggregate和sort过程是可选的。

在实现过程中，Spark对不同的情况进行分类，以及针对性的优化调整，形成不同的Shuffle Write方式。

3.1.1 不需要map()端combine和sort

这种情况只需要实现分区功能，根据record计算出PID，Spark根据PID，将record一次输出到不同的buffer中，每当buffer填满就将record溢写到磁盘上的分区文件中。分配buffer的原因是map()输出record的速度很快，需要进行缓冲来减少IO。

BypassMergeSortShuffleWriter(不需要进行排序的Shuffle Write方式)

优点是速度快，直接将record输出到不同的分区文件中。缺点是资源消耗过高，每个分区需要一个buffer(spark.Shuffle.file.buffer，默认为32KB)，且同时需要建立多个分区文件进行溢写，分区个数太大时(如10000，map task需要消耗320MB内存)，此时内存消耗过大，而且每个task需要建立和打开10000个文件，造成资源不足。因此，该Shuffle方案适合分区个数较少的情况(<200)。

适合groupByKey(100)，partitionBy(100)，sortByKey(100)

3.1.2 不需要map()端combine，但需要sort

这种情况需要按照PID+Key进行排序，Spark的实现方法是建立一个Array来存放map()输出的record，并对Array中元素的Key进行精心设计，将<K,V>record -> <(PID,K),V>record存储，然后按照PID+Key进行排序，最后将所有record写入一个文件中，通过建立索引来标识每个分区。

如果Array放不下，则会先扩容，如果还是存放不下，将Array中的record排序后spill到磁盘上，等待map()输出完以后，再将Array中的record与磁盘上已排序的record进行全局排序，得到最终有序的record，并写入文件中。

Shuffle模式命名为SortShuffleWriter(KeyOrder=true)

使用的Array被命名为PartionedPairBuffer

有点是只需要一个Array结构就可以支持按照PID+Key进行排序，Array的大小可控，而且还有扩容和spill到磁盘上的功能，支持从小规模到大规模的数据排序，同时，输出的数据已经按照PId进行排序，因此只需要一个分区文件进行存储，即可标识不同的分区数据，客服了BypassMergeSortShuffleWriter中建立文件过多的问题，适用于分区个数很大的情况。缺点是排序增加计算时延。

目前Spark官方没有支持这种模式的算子，用户可以自定义使用。

这种方法给我们提供一个解决BypassMergeSortShuffleWriter的buffer分配过多的问题，只需要将PID+Key排序该位只按PID排序，就可以实现不需要map()端combine，不需要按照Key进行排序，分区个数过大的问题。

3.1.3 需要map()端combine，需要或者不需要按Key排序

Spark采用的实现方法是建立一个类似HashMap的数据结构对map()输出的record进行聚合，HashMap中的Key是PID+Key，Value是经过相同combine 的聚合结果。聚合完成后，Spark对HashMap中的record进行排序，如果需要按Key排序，那么按照PID+Key进行排序，如果不需要按照Key排序，那么只按PID排序，最后将排序结果写入一个分区文件中。

如果HashMap存放不下，则会先扩容为两倍大小，如果还是存放不下，将HashMap中的record排序后spill到磁盘上。此时HashMap被清空，可以继续对map()输出的record进行聚合，如果内存再次不够，那么继续spill到磁盘上，此过程可以重复多次。当map()输出完成后，将此时HashMap中的record和磁盘上已排序的record进行merge，得到的record输出到分区文件中。

优点是只需要一个HashMap结构就可以支持map()端的combine功能，HashMap具有扩容和spill功能，支持小规模到大规模数据的聚合，也适合分区个数很大的情况，在聚合后使用Array排序，可以灵活支持不同的排序需求。缺点是在内存中进行聚合，内存消耗较大，需要额外的数组进行排序，而且如果有spill，还需要进行再次数据。

Spark在Shuffle Write端使用一个经过特殊设计和优化的HashMap，命名为PartitionedAppendOnlyMap，可以同时支持聚合和排序操作，相当于HashMap和Array的合体。

SortShuffleWriter

适用的操作比如reduceByKey()、aggregeteByKey().

总结 combine -> sort -> partition

1. 既不需要combine，也不需要sort，分区个数也少，直接输出 ByPassMergeSortShuffleWriter
2. 为了克服ByPassMergeSortShuffleWriter打开文件过多，Buffer分配过多的问题，也为了支持按照Key进行排序的操作，使用SortShuffleWriter，使用基于Array的方案按照PID或PID+Key进行排序，只输出单一的分区文件即可
3. 为了支持map端combine操作，Spark提供了基于HashMap的SortShuffleWriter，将Array替换为类HashMap的操作来支持combine操作，聚合后根据PID或PID+Key进行排序，并输出分区文件。因为SortShuffleWriter按PID进行了排序，所以被称为sort-based Shuffle Writer

3.2 Shuffle Read框架设计和实现

Shuffle Read阶段，数据操作要实现3个功能：跨节点数据获取、聚合和排序。Spark的通用Shuffle Read框架：数据获取 -> 聚合 -> 排序

根据不同的需求，也分为以下3种情况：

3.2.1 不需要聚合，也不需要按Key排序

这种情况只需要实现诗句获取功能。等待上游map task结束后，reduce task开始从各个map task获取<K,V>record，并将record输出到一个buffer中(大小为spark.reducer.maxSizeInFlight=48M)，下一步操作直接从buffer中获取数据即可。

优点是逻辑和实现简单，内存消耗极小。缺点是不支持聚合排序等功能。

适合partitionBy()

3.2.2 不需要聚合，但需要按Key排序

这种情况需要实现数据获取和按Key排序的功能。获取数据后，将buffer中record一次输出到一个Array结构(PartiitonedPairBuffer)中。由于这里采用了本来用于Shuffle Write端的PartitionedPairBuffer结构，所以还保留了每个record的PID。然后对Array的record按Key进行排序并将排序结果输出或者传递给下一步操作。

当内存无法存下所有record时，PartitionedPairBuffer将record排序后spill到磁盘上，最后将内存中和磁盘上的record进行全部排序，得到最终结果。

优点是只需要一个Array结构就可以支持按Key排序，Array大小可控，而且具有扩容和spill到磁盘的功能，不受数据规模限制，缺点是排序增加计算时延。

适合sortByKey()、sortBy()

3.2.3 需要聚合，需要或不需要按Key排序

获取record后，Spark会建立一个类似HashMap的数据结构(ExternalAppendOnlyMap)对buffer中的record进行聚合，HashMap中的Key是record中的Key，HashMap中的Value是经过相同聚合函数(func())计算后的结果。不如不需要排序，直接将aggregate结果输出或传递给下一步操作，如果需要排序，则建立一个Array结构，读取HashMap中的record，并按Key排序，排序完成后，得到最终结果。

如果HashMap放不下，则会先扩容为两倍，如果还存不下，就将HashMap中的record排序后spill到磁盘上。此时HashMap被清空，可以继续对buffer中的record进行聚合，如果内存再次不够用，那么继续spill到磁盘上，此过程可以重复很多次，当聚合完成后，将此时HashMap中record与磁盘上已排序的record进行再次聚合，得到最终的record，输出到分区中。

优点是只需要一个HashMap和一个Array结构就可以实现reduce端聚合和排序功能，HashMap具有扩容和spill到磁盘上的功能支持小规模到大规模数据的聚合，边获取数据边聚合，效率较高。缺点是需要在内存中进行聚合，内存消耗较大，如果有数据spill到磁盘后，还需要进行再次聚合。另外，经过HashMap聚合后的数据仍然需要拷贝到Array中进行排序，内存消耗较大。在实现中，Spark使用的HashMap是经过特殊优化的HashMap，命名为ExternalAppendOnlyMap，可以同时支持聚合和排序操作，相当于HashMap和Array的合体。

适合reduceByKey()、aggregateByKey()

对Spark的Shuffle Read框架，如果应用中的数据不需要聚合，也不需要排序，那么获取数据后直接输出。对于需要按Key进行排序的操作，Spark使用基于Array的方法来对Key进行排序，对于需要聚合的操作，Spark提供了基于HashMap的聚合方法，同时可以再次使用Array来支持按照Key进行排序。

4 支持高效聚合和排序的数据结构

数据结构类型	名称	功能
类HashMap + Array	PartitionedAppendOnlyMap	用户map端聚合和排序，包含PID
类HashMap + Array	ExternalAppendOnlyMap	用于reduce()端聚合和排序
类Array	PartitionedPairBuffer	仅用于map和reduce端数据排序，包含PID

对于Shuffle Write/Read过程，Shuffle机制中使用的数据结构的两个特征：一是只需要支持record的插入和更新操作，不需要支持删除操作，这样我们可以对数据结构进行优化，减少内存消耗；二是只有内存放不下时才需要spill到磁盘上，因此数据结构设计以内存为主，磁盘为辅。Spark中的PartitionedAppendOnlyMap和ExternalAppendOnlyMap都基于AppendOnlyMap实现。

4.1 AppendOnlyMap原理

AppendOnlyMap实际上是一个只支持record添加和对Value进行更新的HashMap。与Java的HashMap采用“数组+链表”实现不同，AppendOnlyMap只使用数组来存储元素，根据元素的Hash值来确定存储位置，如果存储元素时发生Hash值冲突，则使用二次地址探测法(Quadratic pribing)来解决Hash值冲突。

对于每个新来的<K,V>record，先使用Hash(K)计算其存放位置，如果存放位置为空，就把record存放到该位置。如果该位置已经被占用，如K6，则向后选择Hash(K6)+1*2，如果还是被占用，那就继续向后寻址Hash(K6)+2*2。

假设又有<K6,V6>record，需要与刚存放进行AppendOnlyMap中的<K6,V6>record进行聚合，聚合函数是func()，那么首先查找K6所在的位置，查找过程与刚才插入的位置类似，经过3此查找取出<K6,V6>record中V6，进行V’=func(V6,V7)运算，最后将V‘写入V6的位置。

扩容：AppendOnlymap使用数组实现的问题是，如果插入的record太多，则很快就会被填满，此时则需要扩容。在Spark中，如果AppendOnlyMap的利用率达到70%，那么就扩张一倍，扩张意味着原来的Hash()失效，因此所有Key进行reHash，重新排列每个Key的位置。

排序：由于AppendOnlyMap采用了数组作为底层存储结构，可以支持快排等排序算法。先将数组中所有的<K,V>record转移到数组的前端，用begin和end来标记起始位置，然后调用算法对[begin,end]中的record进行排序。对于需要按Key进行排序的操作，如sortBykey()，可以按照Key值进行排序，对于其他操作，按照Key的Hash值进行排序即可。

输出：迭代AppendOnlyMap数组中的record，从前到后扫描输出即可。

4.2 ExternalAppendOnlyMap

AppendOnlyMap的优点是能够将聚合和排序功能很好地结合在一起，缺点是只能使用内存，难以使用与内存空间不足的情况，为了解决这个问题，Spark基于AppendOnlyMap设计实现了基于内存+磁盘的ExternalAppendOnlyMap，用于Shuffle Read端大规模数据聚合。同时，由于Shuffle Write端聚合需要考虑PID，Spark也设计了带有PID的ExternalAppendOnlyMap，命名为PartitionedAppendOnlyHashMap。这两个数据结构功能类似，我们先介绍ExternalAppendOnlyMap。

ExternalAppendOnlyMap的工作原理是，先持有一个AppendOnlyMap来不断接受和聚合新来的record，AppendOnlyMap快被装满是检查一下内存剩余空间是否可以扩展，可以直接在内存中扩展，不可以对AppendOnlyMap中的record进行排序，然后将record对spill到磁盘上。因为record不断到来，可能会多次填满AppendOnlyMap，所以这个spill过程可以出现多次，最终形成多个spill文件。等record都处理完，此时AppendOnlyMap中可能还留存一些聚合后的record磁盘上也有多个spill文件。因为这些数据对经过了部分聚合，还需要进行全局聚合(merge)。因此ExternalAppendOnlyMap的最后一步是将内存AppendOnlyMap的数据与磁盘上spill文件中的数据进行全局聚合，得到最终结果。

这其中有3个核心的问题：

1. AppendOnlyMap的大小估计。虽然我们知道AppendOnlyMap中持有的数组长度和大小，但数组里面存放的是Key和Value的引用，并不是他们的object大小，而且Value会不断被更新，实际大小不断变化，因此，想准确得到AppendOnlyMap的大小比较困难。一种简单的解决方法是在每次插入record或对现有record的Value进行更新后，都扫描一下AppendOnlyMap中存放的record，计算每个record的实际对象大小并相加，但这样非常耗时。而且一般AppendOnlyMap会插入几万甚至几百万的record，如果每次都计算一遍，开销非常大。Spark设计了一个增量式的高效估算算法，在每个record插入或更新时根据历史统计值和当前变化量直接估算当前AppendOnlyMap的大小，算法的复杂度时O(1)，开销很小。在record插入和聚合过程中会定期对当前AppendOnlyMap中的record进行抽样，然后精确计算这些record的总大小、总个数，更新个数及平均值等，并作为历史统计值。之后，每当有record插入或更新时，会根据历史统计值和历史平均值，增量估计AppendOnlyMap的总大小，Spark源码的**SizeTracker.estimateSize()**方法。持之以恒也会定期进行，更新统计值以获得更高的精度。
2. Spill过程与排序。当AppendOnlyMap达到内存限制时，会将record排序后写入磁盘中。排序是为了方便进行下一步聚合(聚合内存和磁盘上的record)可以采用更高效的merge-sort(外部排序+聚合)。Spark采用按照Key的Hash值进行排序的方法，这样既可以进行merge-sort，又不要求操作定义Key排序的方法。然而，这种方法会出现Hash碰撞的问题，因此，Spark在merge-sort时，会同时比较Key的Hash值是否相等，以及Key的实际值是否相等。
3. 全局聚合(merge-sort)。由于内存中的AppendOnlyMap和spill到磁盘上的AppendOnlyMap都是经过部分聚合后的结果，其中可能存在相同Key的record，因此还需要一个sort-merge阶段将这两部的数据进行聚合，得到最终的聚合结果。merge-sort的方法是建立一个最小堆和最大堆，每次从各个spill文件中读取前几个具有相同Key(或者相同Key的Hash值)的record，然后与AppendOnlyMap中的record进行聚合，直到所有record处理完成，输出最终结果。由于每个spill璺中的record是经过排序的，按顺序读取和聚合可以保证能够对每个record得到全局聚合结果。

ExternalAppendOnlyMap是一个高性能的HashMap，只支持数据插入和更新，但可以同时利用内存和磁盘对大规模数据进行聚合和排序，满足了Shuffle Read阶段数据聚合、排序的需求。

4.3 PartitionedAppendOnlyMap

PartitionedAppendOnlyMap用于在Shuffle Write阶段对record进行聚合(combine)。PartitionedAppendOnlyMap的功能和实现与ExternalAppendOnlyMap的功能和实现基本一样，唯一的区别是PartitionedAppendOnlyMap中的Key是PID+Key，这样既可以按照PID进行排序(对于不需要按Key排序的场景)，又可以根据PID+Key进行排序(对于需要按Key排序的场景)，从而在Shuffle Write阶段可以实现聚合、排序和分区。

4.4 PartitionedPairBuffer

PartitionPairBuffer本质上是一个基于内存+磁盘的Array，随着数据的添加，不断扩容，当达到内存限制时，将Array中的数据按PID或PID+Key进行排序，然后spill到磁盘上，该过程可以进行多次，最后对内存和磁盘上的数据进行全局排序，输出或提供给下一步操作。

4.5 与Hadoop MapReduce的Shuffle机制对比

Hadoop MapReduce的Shuffle机制有明显的两个阶段，map stage和reduce stage。

在map stage中，每个map tash首先执行map(K,V)函数，再读取每个record，并输出新的<K,V>record。这些record首先被输出到一个固定大小的spill buffer(默认100MB)，spill buffer如果被填满就将spill buffer中的record按照Key排序后输出到磁盘上，这个过程类似Spark将map task输出的record放到一个排序数组中(PartitionedPairBuffer)，不同的是Hadoop MapReduce是严格按照Key进行排序的，而PartitionedPairBuffer的排序更加灵活(按PID或PID+Key)。另外，由于spill buffer中的record之进行排序，不能完成聚合(combine)功能，所以HadoopMapReduce在晚场map()，等待所有record都spill到磁盘上后，启动一个专门的聚合阶段，使用combine()将所有spill文件中的record进行全局聚合，得到最终聚合结果。这里需要多次全局聚合，因为每次只针对某个分区的spill文件进行聚合。

在Shuffle Read阶段，Hadoop MapReduce先将每个map task输出的相应分区文件通过网络获取，然后再放入内存，如果内存放不下，就先对当前内存中的record进行聚合和排序，再spill到磁盘上。由于每个分区文件中包含的record已经按照Key进行排序，聚合时只需要一个最小堆或最大堆保存当前文件中的前几个record即可，聚合效率较高，但是需要占用大量内存空间来存储这些分区文件。当获取到所有的分区文件时，此时可能存在多个spill文件和内存中剩余的分区文件，这是再启动一个专门的reduce阶段来将这些内存和磁盘上的数据进行全局聚合，这个过长与Spark的全局聚合过程没有区别。

Hadoop MapReduce的Shuffle机制优点：

1. Hadoop MapReduce的Shuffle流程固定，阶段分明，每个阶段读取什么数据，进行什么操作，输出什么数据都是确定的。因此实现起来比较容易。
2. 框架消耗的内存也是确定的，map阶段框架只需要一个大的spill buffer，reduce阶段框架只需要一个大的数组(Merge Queue)来存放获取的分区稳健中的record，这样，什么时候将数据spill到磁盘上是确定的，也易于实现和内存管理。当然，用户自定义的聚合函数，如combine()和reduce()的内存消耗是不确定的。
3. 对Key进行严格排序，是的可以使用最小堆或最大堆进行聚合，非常高效。而且可以原生支持sortByKey()。
4. 按Key进行排序和spill到磁盘上的功能，而快车在Shuffle大规模数据时仍可以顺利运行

缺点：

1. 强制按Key进行排序，大多数引用其实不需要严格按Key进行排序，如groupByKey()，排序增加计算量。
2. 不能在线聚合。不论是map阶段还是reduce阶段，对是先将数据存放到内存或磁盘上后，再执行聚合操作，存储这些数据需要消耗大量的内存和磁盘空间。采用在线聚合，可以有效减少存储空间，并减少时延。
3. 产生的临时文件过多。如果map task个数为M，reduce task个数为N，那么Map阶段会产生M*N个分区文件，当M*N较大时，临时文件过多。