Spark之Shuffle内核完全解析

普通的Shuffle的过程

在Spark处理程序时当运算中当涉及到key聚合或者是key的混乱移动（比如说排序）的时候就会发生Shuffle，RDD之间也就产生了宽依赖。而每有一个宽依赖就会有一个Stage，也就是说每当有一个Shuffle就会有一个Stage产生。而归根结底这些数据的处理是由Task的子类去完成。在最后一个Stage中的Task为ResultTask，而其它的Stage中的Task为ShuffleMapTask。话不多说首先上图：

Map端（Shuffle的写过程）

首先说Shuffle的写过程，我们的目标是要将每一个partition中的key进行聚合分桶。这里使用了Hash的分桶机制将不同的key以及对应的值写入缓冲区（也就是我们所说的bucket）。然后进而溢写到bucket对应的磁盘文件ShuffleBlockFile上。这里需要注意的是如果总共有m个key与n个ShuffleMapTask，那么我们最后产生的文件为O(m*n)个（一个不是很严格的表示方式，因为在一些分区可能不会包含所有的key）。最后ShuffleMapTask会将自己的写出的结果信息封装到MapStatus中进而发送到DAGScheduler的MapOutPutTracker中，可以在reduce端在拉取数据的时候提供数据存储位置信息。

Reduce端（Shuffle的读过程）

在上面的Map端我们把数据按照Hash分桶的逻辑写到了ShuffleBlockFile上面。现在就到了真正key聚合的时候了。假设我们的聚合key所在的Stage为最后一个Stage，这时聚合key的Task就为ResultTask，同时假设我们的Map端与reduce端处于不同的节点上。好了这时我们就要开始拉取数据了，我们要将所有的属于同一个key的数据都聚合到一起，而这些key分布在ShuffleMapTask所产生的ShuffleBlockFile上面。我们要拉取数据就需要知道这些数据具体的物理存储位置，这时也就需要用到我们前边所产生的MapStatus了，它里面封装着输出文件的位置信息。所以我们会去DAGScheduler中的MapTracker中获取这些信息。然后根据这些信息使用BlockStoreShuffleFetcher去实际的物理位置上拉取文件，而它的底层是通过BlockManager来实现的。

写在最后

以上的全部过程可以称之为一个Shuffle过程，它整体发生在两个Stage中。Map端为上一个Stage，Reduce为下一个Stage。而且在这个Shuffle中会产生三个RDD

其中在Map端中数据未处理的时候的RDD为MapPartitionsRDD。
在Shuffle的reduce端将数据拉过来的时候会在内部形成一个RDD，这个RDD为ShuffledRDD注意这里只是将数据拉过来了而没有将拉过来的数据进行聚合。
在将上面的RDD中的数据再次进行聚合处理之后，最后就是我们想要得到的目标RDD它也是一个MapPartitionsRDD。

优化之后的Shuffle过程

优化之后的Shuffle与传统的Shuffle最大的区别就是在Map端的Hash分桶与写的过程。在一般的Shuffle中我们的ShuffleMapTask会针对每一个key创建一个桶也就会对应一个ShuffleBlockFile那么最后的文件数量（在最坏情况下）就会是ShuffleMapTask的数量与key的乘积。而优化过后的Shuffle过程会将文件的数量大大减少。以前是一个Task会分别给每一个key分配一个桶，而现在我们不会分的这么细致而是按照组去给每个key分配桶。一个组中会给一个key分配一个桶。这样的话我们生成的文件数量就会是组的数量与key的数量的乘积。话说话来了我们怎么去分组呢？分组就是依据这个节点的并行度来划分的，比如说有2个CPU那么Spark就会将这个节点上的ShuffleMapTask分为两组，那么最后生成的文件数量就会是2 * key的数量。图解如下：

这也就是Spark中的consolidation机制提出了ShuffleGroup概念。这里需要注意的是在合并的文件中默认没有进行多个Task的预聚合而是每个ShuffleMapTask会在文件中对应自己的一个Segment，可以通过索引去找到对应的数据。

Shuffle核心源码解析

这段文字将Shuffle的一般和优化过的步骤做一个统一的描述。

Map端（Shuffle的写操作）

本地聚合。首先Map写数据数据的时候会考虑Spark是否要求进行预聚合，也就是

def.aggregator.isDefined为true，def.mapSideCombine为true。这时在map端写操作的时候就会进行预聚合。

然后就开始写数据了。怎么确定往哪里写对于每一条数据都会调用partitioner（默认是HashPartitioner）去生成bucketId，这样就决定了每一份数据要写入哪个桶中。接下来调用ShuffleBlockManager.forMapTask()方法来生成bucketId对应的writer然后用writer将数据写入bucket中。

细节：writer的生成策略：

带有优化的生成（Consolidation）

带有优化的生成就是一个组中的一个key只会有一个bucket。也就是为一个bucket获取一个ShuffleGroup的writer。首先用ShuffleId、mapId，bucketId(reduceId)来生成一个唯一的ShuffleBlockId然后使用bucketId来调用ShuffleFileGroup的apply()方法为bucket获取一个ShuffleFileGroup，而底层是通过BlockManager的DiskBlockManager来获取一个针对ShuffleFileGroup的writer。

B. 普通的shuffle

与上面最大的不同就是最后生成的一个Task-key级别的ShuffleBlockFile。

Reduce端（Shuffle的读操作）

Shuffle的读操作发生在RDD的Compute中，这是由于Task在对分区进行计算时会需要这个分区的数据，这些数据必须得拉取过来。这里的获取数据会分为两步：①获取从DAGScheduler中获取目标数据存放的位置，底层为BlockStoreShuffleFetcher来从DAGScheduler的MapOutPutTrackerMaster中来获取自己的数据。②根据拿到的位置信息去拿目标数据。这里会调用ShuffleManager的getReader()方法，获取一个HashShuffleReader来拉取ShuffleMapTask或者是ResultTask需要聚合的数据。

一个重要的方法fetch

fetch()中和核心就是讲究怎样去更加具体的确定如何拿数据。首先拿到一个全局的MapOutPutTrackerMaster的引用，然后调用其方法getServerStatus()方法，这里会传入一些参数ShuffleId与reduceId。

ShuffleId：它可以将数据的获取定位到上一个Stage，具体就是上一个Stage中所有ShuffleMapTask输出的MapStatus这个Status需要走网络与Driver中的DAGScheduler通信。

reduceId：也就是bucketId，它决定了在上一个Stage中去读取哪个块。

拉取数据的过程

核心方法为ShuffleBlockFetcherIterator的initialize()。

将本地的或者是远程的数据切块便于传输。
切分之后的block进行随机排序操作。
循环往复的去拉取数据，如果没有拉完那么就继续发送数据去拉。
MaxBytesInFlight（本地内存大小）这个参数就决定了最多能够拉取多少的数据来本地。
规定的数据拉取完成后就要进行reduce操作了。

SortShuffle与数据的拉取

可能我学习的资料比较老，在这个版本中HashShuffle可以一边进行写操作一边进行数据拉取。而SortShuffle的话需要进行排序所以只有将这些数据排序之后reduce端才能够拉取数据。