PySpark的实现原理

PySpark实现了Spark对于Python的API，
通过它，用户可以编写运行在Spark之上的Python程序，
从而利用到Spark分布式计算的特点。

基本流程

PySpark的整体架构图如下，
可以看到Python API的实现依赖于Java的API，
Python程序端的SparkContext通过py4j调用JavaSparkContext，
后者是对Scala的SparkContext的一个封装。
对RDD进行转换和操作的函数由用户通过Python程序来定义，
这些函数会被序列化然后发送到各个worker，
然后每一个worker启动一个Python进程来执行反序列化之后的函数，
通过管道拿到执行之后的结果。

下面更加细致地介绍各个细节。

Python程序的启动

和Scala程序一样，Python程序也是通过SparkSubmit提交得以执行，
在SparkSubmit中会判断提交的程序是否为Python，
如果是，则设置mainClass为PythonRunner。

PythonRunner的执行代码如下：

def main(args: Array[String]) {
  val pythonFile = args(0)
  val pyFiles = args(1)
  val otherArgs = args.slice(2, args.length)
  val pythonExec = sys.env.get("PYSPARK_PYTHON").getOrElse("python") // TODO: get this from conf

  // Format python file paths before adding them to the PYTHONPATH
  val formattedPythonFile = formatPath(pythonFile)
  val formattedPyFiles = formatPaths(pyFiles)

  // Launch a Py4J gateway server for the process to connect to; this will let it see our
  // Java system properties and such
  val gatewayServer = new py4j.GatewayServer(null, 0)
  gatewayServer.start()

  // Build up a PYTHONPATH that includes the Spark assembly JAR (where this class is), the
  // python directories in SPARK_HOME (if set), and any files in the pyFiles argument
  val pathElements = new ArrayBuffer[String]
  pathElements ++= formattedPyFiles
  pathElements += PythonUtils.sparkPythonPath
  pathElements += sys.env.getOrElse("PYTHONPATH", "")
  val pythonPath = PythonUtils.mergePythonPaths(pathElements: _*)

  // Launch Python process
  val builder = new ProcessBuilder(Seq(pythonExec, "-u", formattedPythonFile) ++ otherArgs)
  val env = builder.environment()
  env.put("PYTHONPATH", pythonPath)
  env.put("PYSPARK_GATEWAY_PORT", "" + gatewayServer.getListeningPort)
  builder.redirectErrorStream(true) // Ugly but needed for stdout and stderr to synchronize
  val process = builder.start()

  new RedirectThread(process.getInputStream, System.out, "redirect output").start()

  System.exit(process.waitFor())
}

在PythonRunner中，会根据配置选项，以及用户通过命令行提供的--py-files选项，
设置好PYTHONPATH，然后启动一个Java的GatewayServer用来被Python程序调用，
然后以用户配置的PYSPARK_PYTHON选项作为Python解释器，
执行Python文件，至此用户的Python程序得以启动。

SparkContext

和在Scala中一样，SparkContext是调用Spark进行计算的入口。
在Python的context.py中定义了类SparkContext，
它封装了一个JavaSparkContext作为它的_jsc属性。
在初始化SparkContext时，
首先会调用java_gateway.py中定义的launch_gateway方法来初始化JavaGateWay,
在launch_gateway中会引入在Spark中定义的类到SparkContext的属性_jvm，
比如:

java_import(gateway.jvm, "org.apache.spark.SparkConf")

这样在Python中就可以通过SparkContext._jvm.SparkConf引用在Scala中定义的SparkConf这个类,
可以实例化这个类的对象，可以调用对象的方法等。

在初始化完毕之后，用户就可以调用SparkContext中的方法了，比如textFile和parallelize，
下面以这两个方法作为例子看SparkContext的实现。

textFile的实现

textFile的调用和在Scala中一样，提供一个路径，以及一个可选的参数minPartitions，
后者说明最少的Partition的数目，返回一个RDD。textFile的实现如下：

def textFile(self, name, minPartitions=None):
    minPartitions = minPartitions or min(self.defaultParallelism, 2)
    return RDD(self._jsc.textFile(name, minPartitions), self,
               UTF8Deserializer())

Python中的SparkContext调用JavaSparkContext.textFile，
后者返回一个JavaRDD[String]（JavaRDD是对RDD的封装，可以直接把它当做RDD来看待）,
Python把JavaRDD封装成Python的RDD（RDD的细节请看下面的内容）。

parallelize的实现

parallelize把Python里面的list转换为RDD，
调用示例：

>>> sc.parallelize(range(5), 5).collect()
[0, 1, 2, 3, 4]

parallelize的实现代码如下：

def parallelize(self, c, numSlices=None):
  numSlices = numSlices or self.defaultParallelism
  # Calling the Java parallelize() method with an ArrayList is too slow,
  # because it sends O(n) Py4J commands.  As an alternative, serialized
  # objects are written to a file and loaded through textFile().
  tempFile = NamedTemporaryFile(delete=False, dir=self._temp_dir)
  # Make sure we distribute data evenly if it's smaller than self.batchSize
  if "__len__" not in dir(c):
      c = list(c)    # Make it a list so we can compute its length
  batchSize = min(len(c) // numSlices, self._batchSize)
  if batchSize > 1:
      serializer = BatchedSerializer(self._unbatched_serializer,
                                     batchSize)
  else:
      serializer = self._unbatched_serializer
  serializer.dump_stream(c, tempFile)
  tempFile.close()
  readRDDFromFile = self._jvm.PythonRDD.readRDDFromFile
  jrdd = readRDDFromFile(self._jsc, tempFile.name, numSlices)
  return RDD(jrdd, self, serializer)

首先，把数据序列化到临时文件中，
然后调用PythonRDD的readRDDFromFile从文件中读取一个个Byte，
恢复成JavaRDD[Array[Byte]]。最后封装成Python的RDD。

RDD

Python中的RDD对Spark中的RDD进行了一次封装，
每一个RDD都对应了一个反序列化的函数。
这是因为，尽管在Spark中RDD的元素可以具有任意类型，
提供给JavaSparkContext中生成的RDD的只具有Array[Byte]类型，
也就是说JavaSparkContext的函数返回值是JavaRDD[Array[Byte]]，
这样，Python程序需要把对象先序列化成byte数组，
然后把它分布到各个节点进行计算。
计算完之后再反序列化成Python的对象。
（这其中有一个特殊情况，就是JavaSparkContext返回的是JavaRDD[String]，
可以把它当成是不需要序列化和反序列化的对象。）
在Spark中不需要知道Array[Byte]反序列化之后是什么。
如何序列化和反序列化、如何对这些Array[Byte]进行转换和操作都由Python程序来控制，
Spark只是负责资源的调度，负责如何把这些计算分配到各个节点上去执行。

下面以map和collect为例，分析RDD的转换和操作如何实现。

map的实现

map函数把一个函数应用到RDD中的每一个元素，
把它转换为另一个RDD，调用示例：

>>> rdd = sc.parallelize(["b", "a", "c"])
>>> sorted(rdd.map(lambda x: (x, 1)).collect())
[('a', 1), ('b', 1), ('c', 1)]

map的实现如下:

def map(self, f, preservesPartitioning=False):
    def func(split, iterator): return imap(f, iterator)
    return PipelinedRDD(self, func, preservesPartitioning)

调用map时，先对函数f做一个封装，
转换为统一形式的func，
然后返回一个PipelinedRDD，
当前的RDD作为它的prev，
func作为它的func。
在PipelinedRDD中，
会根据它的prev是否为PipelinedRDD组成一条转换的链条，
把转换的函数一层层封装起来，
这样它的属性_prev_jrdd指向第一个非PipelinedRDD的RDD的_jrdd，
它的func属性为一条计算当前PipelinedRDD的函数链。
通过这两个属性就可以计算当前的RDD。
通过Pipeline的方式能够减少序列化和进程间交互的开销。

PipelinedRDD的计算函数如下：

@property
def _jrdd(self):
    if self._jrdd_val:
        return self._jrdd_val
    if self._bypass_serializer:
        self._jrdd_deserializer = NoOpSerializer()
    command = (self.func, self._prev_jrdd_deserializer,
               self._jrdd_deserializer)
    pickled_command = CloudPickleSerializer().dumps(command)
    broadcast_vars = ListConverter().convert(
        [x._jbroadcast for x in self.ctx._pickled_broadcast_vars],
        self.ctx._gateway._gateway_client)
    self.ctx._pickled_broadcast_vars.clear()
    class_tag = self._prev_jrdd.classTag()
    env = MapConverter().convert(self.ctx.environment,
                                 self.ctx._gateway._gateway_client)
    includes = ListConverter().convert(self.ctx._python_includes,
                                 self.ctx._gateway._gateway_client)
    python_rdd = self.ctx._jvm.PythonRDD(self._prev_jrdd.rdd(),
        bytearray(pickled_command), env, includes, self.preservesPartitioning,
        self.ctx.pythonExec, broadcast_vars, self.ctx._javaAccumulator,
        class_tag)
    self._jrdd_val = python_rdd.asJavaRDD()
    return self._jrdd_val

在需要计算PipelinedRDD时，
先把func序列化成commands，
然后调用PythonRDD的asJavaRDD方法计算，
在计算时，在各个worker上会启动一个Python的进程执行反序列化之后的函数，
通过管道和Python进程进行通信，最后得到JavaRDD。

collect的实现

通过collect，把RDD转换为Python的list，
collect的实现代码如下：

def collect(self):
    """
    Return a list that contains all of the elements in this RDD.
    """
    with _JavaStackTrace(self.context) as st:
      bytesInJava = self._jrdd.collect().iterator()
    return list(self._collect_iterator_through_file(bytesInJava))

def _collect_iterator_through_file(self, iterator):
    # Transferring lots of data through Py4J can be slow because
    # socket.readline() is inefficient.  Instead, we'll dump the data to a
    # file and read it back.
    tempFile = NamedTemporaryFile(delete=False, dir=self.ctx._temp_dir)
    tempFile.close()
    self.ctx._writeToFile(iterator, tempFile.name)
    # Read the data into Python and deserialize it:
    with open(tempFile.name, 'rb') as tempFile:
        for item in self._jrdd_deserializer.load_stream(tempFile):
            yield item
    os.unlink(tempFile.name)

首先，调用JavaRDD.collect().iterator()，得到迭代RDD的迭代器，
然后把这个JavaRDD写入到临时文件，
然后从临时文件读取到array of bytes，将其反序列化成一个个Python的对象，
组成一个list就是最终的结果。

MLlib的实现

在Python中提供了MLlib的接口，
用户可以使用Spark中的MLLib模块提供的算法。
下面分析一下实现这个模块的细节。

交互

以Kmeans这个算法为例，用户调用的示例为：

>>> from numpy import array
>>> data = array([0.0,0.0, 1.0,1.0, 9.0,8.0, 8.0,9.0]).reshape(4,2)
>>> model = KMeans.train(
...     sc.parallelize(data), 2, maxIterations=10, runs=30, initializationMode="random")
>>> model.predict(array([0.0, 0.0])) == model.predict(array([1.0, 1.0]))
True

首先，用户调用Kmeans.train得到训练的模型，Kmeans.train实现代码如下：

class KMeans(object):
    @classmethod
    def train(cls, data, k, maxIterations=100, runs=1, initializationMode="k-means||"):
        """Train a k-means clustering model."""
        sc = data.context
        dataBytes = _get_unmangled_double_vector_rdd(data)
        ans = sc._jvm.PythonMLLibAPI().trainKMeansModel(
            dataBytes._jrdd, k, maxIterations, runs, initializationMode)
        if len(ans) != 1:
            raise RuntimeError("JVM call result had unexpected length")
        elif type(ans[0]) != bytearray:
            raise RuntimeError("JVM call result had first element of type "
                               + type(ans[0]) + " which is not bytearray")
        matrix = _deserialize_double_matrix(ans[0])
        return KMeansModel([row for row in matrix])

首先把数据序列化成RDD，
然后调用PythonMLLibAPI提供的trainKMeansModel接口得到训练之后的结果
（类型为java.util.List[java.lang.Object]），
然后把这个结果反序列化为矩阵，封装成KMeansModel模型。

PythonMLlibAPI.trainKMeansModel接口的实现代码如下：

def trainKMeansModel(
      dataBytesJRDD: JavaRDD[Array[Byte]],
      k: Int,
      maxIterations: Int,
      runs: Int,
      initializationMode: String): java.util.List[java.lang.Object] = {
    val data = dataBytesJRDD.rdd.map(bytes => deserializeDoubleVector(bytes))
    val model = KMeans.train(data, k, maxIterations, runs, initializationMode)
    val ret = new java.util.LinkedList[java.lang.Object]()
    ret.add(serializeDoubleMatrix(model.clusterCenters.map(_.toArray)))
    ret
  }

先反序列化输入，然后调用Spark的MLlib中的Kmeans训练模型，
得到结果之后序列化成java.util.List，返回给Python代码。

向量运算的通用接口

PySpark支持下面的向量：

• Numpy中的array
• Python的list
• MLlib中的SparseVector
• Scipy中的单列的csc_matrix

对于这些向量，在_common.py提供两个通用的函数:

• _squared_distance: 用于计算两个向量的距离。
• _dot: 用于计算向量的乘积。

所有类型的向量都可以使用这两个函数来进行计算。

结语

至此，Python API实现的大致流程已经分析完毕，
本文只是通过示例的方式讲了几个基本的过程，
并没有完整的分析全部的实现细节。
不过，以这个为基础，再去理解其他的代码应该会非常容易。
由于水平有限，错误之处在所难免，欢迎提出改进意见。