并行正则采样排序算法及在 Mars 中的应用

相信大家对排序算法都非常熟悉了，快速排序、堆排序、归并排序等等。如果我们想在一个很大的数据集上进行排序，能利用上多核，甚至是分布式集群，有什么办法么？

本文就介绍一种并行排序算法：并行正则采样排序算法（Parallel Sorting by Regular Sampling），简称 PSRS 算法。

PSRS 算法过程

PSRS 算法的整个过程分为四步，如图所示，我们拆解开来讲。

现在假设我们有一个数组，有 48 个元素，现在数据被分成4份，即有4个分区。

阶段1，每个分区分别排序，并正则采样

我们对每个分区的数据调用快速排序，这样每个分区都是排好的数据。接着，我们从排好序的数据里正则采样4个数据。所谓正则，即有规律的，这里我们就每隔4个元素采样一个元素。

阶段2，归并采样数据，选出关键点

前面四个分区产生了4份采样数据，收集之，然后调用归并排序让他们有序。接着，我们从中选出 p - 1 （p 是分区个数）个关键点，这里还是正则采样的方式。

阶段3，数据分区

此时将关键点数据广播给每个分区，每个分区就可以根据关键点，将数据分成4份，使得每个数据落在各自的范围内。

阶段4，合并数据，归并排序

最后一个阶段是一个 shuffle 阶段，即每个下游都依赖前面的所有上游。此时每个分区将上游分好的数据收集起来，最终再进行一个归并排序。这样，我们最终的结果就是整体排序的了。

整个过程中，阶段1、阶段3 和 阶段4 可以并行。

MPI 的实现可以参考这里。

PSRS 算法在 Mars 中的应用

Mars 以并行和分布式化 Python 数据科学栈为目标，PSRS 算法能很好解决并行排序问题，因此，Mars 中和排序有关的操作都是基于 PSRS 算法实现的。

以张量排序为例。

首先我们通过 Numpy 创建 1 亿个元素的数组。

In [1]: import numpy as np In [2]: a = np.random.rand(1_0000_0000) In [3]: a.nbytes Out[3]: 800000000

我们来看看使用 Numpy 的排序需要多久。

In [4]: %time np.sort(a) CPU times: user 10.8 s, sys: 394 ms, total: 11.2 s Wall time: 9.4 s Out[4]: array([1.05764619e-10, 5.86309734e-09, 1.76225879e-08, ..., 9.99999976e-01, 9.99999983e-01, 9.99999998e-01])

接着，我们来看看基于 PSRS 算法的 Mars tensor 排序需要多长时间。

In [10]: t = mt.tensor(a, chunk_size=1500_0000) In [12]: %time mt.sort(t).execute() CPU times: user 18.7 s, sys: 7.03 s, total: 25.7 s Wall time: 2.66 s

在我的笔记本上，可以看到 Numpy 的排序时长是 Mars 的 3.53 倍。

总结

本文介绍了并行正则排序算法，这个算法也在 Mars 项目里得到了广泛的使用。

如果对 Mars 感兴趣，可以关注 Mars 团队专栏，或者钉钉扫二维码加入 Mars 讨论群。