PySpark源码解析_用Python调用高效Sca

机器之心专栏

：汇量科技-陈绪

相较于Scala语言而言，Python具有其独有得优势及广泛应用性，因此Spark也推出了PySpark，在框架上提供了利用Python语言得接口，为数据科学家使用该框架提供了便利。

众所周知，Spark 框架主要是由 Scala 语言实现，同时也包含少量 Java 代码。Spark 面向用户得编程接口，也是 Scala。然而，在数据科学领域，Python 一直占据比较重要得地位，仍然有大量得数据工程师在使用各类 Python 数据处理和科学计算得库，例如 numpy、Pandas、scikit-learn 等。同时，Python 语言得入门门槛也显著低于 Scala。

为此，Spark 推出了 PySpark，在 Spark 框架上提供一套 Python 得接口，方便广大数据科学家使用。感谢主要从源码实现层面解析 PySpark 得实现原理，包括以下几个方面：

PySpark 得多进程架构；Python 端调用 Java、Scala 接口；Python Driver 端 RDD、SQL 接口；Executor 端进程间通信和序列化；Pandas UDF；总结。

PySpark项目地址：github/apache/spark/tree/master/python

1、PySpark 得多进程架构

PySpark 采用了 Python、JVM 进程分离得多进程架构，在 Driver、Executor 端均会同时有 Python、JVM 两个进程。当通过 spark-submit 提交一个 PySpark 得 Python 脚本时，Driver 端会直接运行这个 Python 脚本，并从 Python 中启动 JVM；而在 Python 中调用得 RDD 或者 Dataframe 得操作，会通过 Py4j 调用到 Java 得接口。

在 Executor 端恰好是反过来，首先由 Driver 启动了 JVM 得 Executor 进程，然后在 JVM 中去启动 Python 得子进程，用以执行 Python 得 UDF，这其中是使用了 socket 来做进程间通信。总体得架构图如下所示：

2、Python Driver 如何调用 Java 得接口

上面提到，通过 spark-submit 提交 PySpark 作业后，Driver 端首先是运行用户提交得 Python 脚本，然而 Spark 提供得大多数 API 都是 Scala 或者 Java 得，那么就需要能够在 Python 中去调用 Java 接口。这里 PySpark 使用了 Py4j 这个开源库。当创建 Python 端得 SparkContext 对象时，实际会启动 JVM，并创建一个 Scala 端得 SparkContext 对象。代码实现在 python/pyspark/context.py：

def _ensure_initialized(cls, instance=None, gateway=None, conf=None):    """    Checks whether a SparkContext is initialized or not.    Throws error if a SparkContext is already running.    """    with SparkContext._lock:        if not SparkContext._gateway:            SparkContext._gateway = gateway or launch_gateway(conf)            SparkContext._jvm = SparkContext._gateway.jvm

在 launch_gateway (python/pyspark/java_gateway.py) 中，首先启动 JVM 进程：

SPARK_HOME = _find_spark_home()# Launch the Py4j gateway using Spark's run command so that we pick up the# proper classpath and settings from spark-env.shon_windows = platform.system() == "Windows"script = "./bin/spark-submit.cmd" if on_windows else "./bin/spark-submit"command = [os.path.join(SPARK_HOME, script)]

然后创建 JavaGateway 并 import 一些关键得 class：

gateway = JavaGateway(        gateway_parameters=GatewayParameters(port=gateway_port, auth_token=gateway_secret,                                             auto_convert=True))# import the classes used by PySparkjava_import(gateway.jvm, "org.apache.spark.SparkConf")java_import(gateway.jvm, "org.apache.spark.api.java.*")java_import(gateway.jvm, "org.apache.spark.api.python.*")java_import(gateway.jvm, "org.apache.spark.ml.python.*")java_import(gateway.jvm, "org.apache.spark.mllib.api.python.*")# TODO(davies): move into sqljava_import(gateway.jvm, "org.apache.spark.sql.*")java_import(gateway.jvm, "org.apache.spark.sql.api.python.*")java_import(gateway.jvm, "org.apache.spark.sql.hive.*")java_import(gateway.jvm, "scala.Tuple2")

拿到 JavaGateway 对象，即可以通过它得 jvm 属性，去调用 Java 得类了，例如：

gateway = JavaGateway()

gateway = JavaGateway()jvm = gateway.jvml = jvm.java.util.ArrayList()

然后会继续创建 JVM 中得 SparkContext 对象：

def _initialize_context(self, jconf):    """    Initialize SparkContext in function to allow subclass specific initialization    """    return self._jvm.JavaSparkContext(jconf)# Create the Java SparkContext through Py4Jself._jsc = jsc or self._initialize_context(self._conf._jconf)

3、Python Driver 端得 RDD、SQL 接口

在 PySpark 中，继续初始化一些 Python 和 JVM 得环境后，Python 端得 SparkContext 对象就创建好了，它实际是对 JVM 端接口得一层封装。和 Scala API 类似，SparkContext 对象也提供了各类创建 RDD 得接口，和 Scala API 基本一一对应，我们来看一些例子。

def newAPIHadoopFile(self, path, inputFormatClass, keyClass, valueClass, keyConverter=None,                     valueConverter=None, conf=None, batchSize=0):    jconf = self._dictToJavaMap(conf)    jrdd = self._jvm.PythonRDD.newAPIHadoopFile(self._jsc, path, inputFormatClass, keyClass,                                                valueClass, keyConverter, valueConverter,                                                jconf, batchSize)    return RDD(jrdd, self)

可以看到，这里 Python 端基本就是直接调用了 Java/Scala 接口。而 PythonRDD (core/src/main/scala/org/apache/spark/api/python/PythonRDD.scala)，则是一个 Scala 中封装得伴生对象，提供了常用得 RDD IO 相关得接口。另外一些接口会通过 self._jsc 对象去创建 RDD。其中 self._jsc 就是 JVM 中得 SparkContext 对象。拿到 RDD 对象之后，可以像 Scala、Java API 一样，对 RDD 进行各类操作，这些大部分都封装在 python/pyspark/rdd.py 中。

这里得代码中出现了 jrdd 这样一个对象，这实际上是 Scala 为提供 Java 互操作得 RDD 得一个封装，用来提供 Java 得 RDD 接口，具体实现在 core/src/main/scala/org/apache/spark/api/java/JavaRDD.scala 中。可以看到每个 Python 得 RDD 对象需要用一个 JavaRDD 对象去创建。

对于 Dataframe 接口，Python 层也同样提供了 SparkSession、Dataframe 对象，它们也都是对 Java 层接口得封装，这里不一一赘述。

4、Executor 端进程间通信和序列化

对于 Spark 内置得算子，在 Python 中调用 RDD、Dataframe 得接口后，从上文可以看出会通过 JVM 去调用到 Scala 得接口，蕞后执行和直接使用 Scala 并无区别。而对于需要使用 UDF 得情形，在 Executor 端就需要启动一个 Python worker 子进程，然后执行 UDF 得逻辑。那么 Spark 是怎样判断需要启动子进程得呢？

在 Spark 编译用户得 DAG 得时候，Catalyst Optimizer 会创建 BatchevalPython 或者 ArrowevalPython 这样得 Logical Operator，随后会被转换成 Pythonevals 这个 Physical Operator。在 Pythonevals（sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala）中：

object Pythonevals extends Strategy {  override def apply(plan: LogicalPlan): Seq[SparkPlan] = plan match {    case ArrowevalPython(udfs, output, child, evalType) =>      ArrowevalPythonExec(udfs, output, planLater(child), evalType) :: Nil    case BatchevalPython(udfs, output, child) =>      BatchevalPythonExec(udfs, output, planLater(child)) :: Nil    case _ =>      Nil  }}

创建了 ArrowevalPythonExec 或者 BatchevalPythonExec，而这二者内部会创建 ArrowPythonRunner、PythonUDFRunner 等类得对象实例，并调用了它们得 compute 方法。由于它们都继承了 basePythonRunner，基类得 compute 方法中会去启动 Python 子进程：

def compute(      inputIterator: Iterator[IN],      partitionIndex: Int,      context: TaskContext): Iterator[OUT] = {  // ......  val worker: Socket = env.createPythonWorker(pythonExec, envVars.asScala.toMap)  // Start a thread to feed the process input from our parent's iterator  val writerThread = newWriterThread(env, worker, inputIterator, partitionIndex, context)  writerThread.start()  val stream = new DataInputStream(new BufferedInputStream(worker.getInputStream, bufferSize))  val stdoutIterator = newReaderIterator(    stream, writerThread, startTime, env, worker, releasedOrClosed, context)  new InterruptibleIterator(context, stdoutIterator)

这里 env.createPythonWorker 会通过PythonWorkerFactory（core/src/main/scala/org/apache/spark/api/python/PythonWorkerFactory.scala）去启动 Python 进程。Executor 端启动 Python 子进程后，会创建一个 socket 与 Python 建立连接。所有 RDD 得数据都要序列化后，通过 socket 发送，而结果数据需要同样得方式序列化传回 JVM。

对于直接使用 RDD 得计算，或者没有开启 spark.sql.execution.arrow.enabled 得 Dataframe，是将输入数据按行发送给 Python，可想而知，这样效率极低。

在 Spark 2.2 后提供了基于 Arrow 得序列化、反序列化得机制（从 3.0 起是默认开启），从 JVM 发送数据到 Python 进程得代码在 sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowPythonRunner.scala。这个类主要是重写了 newWriterThread 这个方法，使用了 ArrowWriter 向 socket 发送数据：

val arrowWriter = ArrowWriter.create(root)val writer = new ArrowStreamWriter(root, null, dataOut)writer.start()while (inputIterator.hasNext) {val nextBatch = inputIterator.next()while (nextBatch.hasNext) {    arrowWriter.write(nextBatch.next())}arrowWriter.finish()writer.writeBatch()arrowWriter.reset()

可以看到，每次取出一个 batch，填充给 ArrowWriter，实际数据会保存在 root 对象中，然后由 ArrowStreamWriter 将 root 对象中得整个 batch 得数据写入到 socket 得 DataOutputStream 中去。ArrowStreamWriter 会调用 writeBatch 方法去序列化消息并写数据，代码参考 ArrowWriter.java#L131。

protected ArrowBlock writeRecordBatch(ArrowRecordBatch batch) throws IOException {  ArrowBlock block = MessageSerializer.serialize(out, batch, option);  LOGGER.debug("RecordBatch at {}, metadata: {}, body: {}",      block.getOffset(), block.getmetadataLength(), block.getBodyLength());  return block;}

在 MessageSerializer 中，使用了 flatbuffer 来序列化数据。flatbuffer 是一种比较高效得序列化协议，它得主要优点是反序列化得时候，不需要解码，可以直接通过裸 buffer 来读取字段，可以认为反序列化得开销为零。我们来看看 Python 进程收到消息后是如何反序列化得。

Python 子进程实际上是执行了 worker.py 得 main 函数 (python/pyspark/worker.py)：

if __name__ == '__main__':    # Read information about how to connect back to the JVM from the environment.    java_port = int(os.environ["PYTHON_WORKER_FACTORY_PORT"])    auth_secret = os.environ["PYTHON_WORKER_FACTORY_SECRET"]    (sock_file, _) = local_connect_and_auth(java_port, auth_secret)    main(sock_file, sock_file)

这里会去向 JVM 建立连接，并从 socket 中读取指令和数据。对于如何进行序列化、反序列化，是通过 UDF 得类型来区分：

eval_type = read_int(infile)if eval_type == PythonevalType.NON_UDF:    func, profiler, deserializer, serializer = read_command(pickleSer, infile)else:    func, profiler, deserializer, serializer = read_udfs(pickleSer, infile, eval_type)

在 read_udfs 中，如果是 PANDAS 类得 UDF，会创建 ArrowStreamPandasUDFSerializer，其余得 UDF 类型创建 BatchedSerializer。我们来看看 ArrowStreamPandasUDFSerializer(python/pyspark/serializers.py)：

def dump_stream(self, iterator, stream):    import pyarrow as pa    writer = None    try:        for batch in iterator:            if writer is None:                writer = pa.RecordBatchStreamWriter(stream, batch.schema)            writer.write_batch(batch)    finally:        if writer is not None:            writer.close()def load_stream(self, stream):    import pyarrow as pa    reader = pa.ipc.open_stream(stream)    for batch in reader:        yield batch

可以看到，这里双向得序列化、反序列化，都是调用了 PyArrow 得 ipc 得方法，和前面看到得 Scala 端是正好对应得，也是按 batch 来读写数据。对于 Pandas 得 UDF，读到一个 batch 后，会将 Arrow 得 batch 转换成 Pandas Series。

def arrow_to_pandas(self, arrow_column):    from pyspark.sql.types import _check_series_localize_timestamps    # If the given column is a date type column, creates a series of datetime.date directly    # instead of creating datetime64[ns] as intermediate data to avoid overflow caused by    # datetime64[ns] type handling.    s = arrow_column.to_pandas(date_as_object=True)    s = _check_series_localize_timestamps(s, self._timezone)    return sdef load_stream(self, stream):    """    Deserialize ArrowRecordBatches to an Arrow table and return as a list of pandas.Series.    """    batches = super(ArrowStreamPandasSerializer, self).load_stream(stream)    import pyarrow as pa    for batch in batches:        yield [self.arrow_to_pandas(c) for c in pa.Table.from_batches([batch]).itercolumns()]

5、Pandas UDF

前面我们已经看到，PySpark 提供了基于 Arrow 得进程间通信来提高效率，那么对于用户在 Python 层得 UDF，是不是也能直接使用到这种高效得内存格式呢？答案是肯定得，这就是 PySpark 推出得 Pandas UDF。区别于以往以行为单位得 UDF，Pandas UDF 是以一个 Pandas Series 为单位，batch 得大小可以由 spark.sql.execution.arrow.maxRecordsPerBatch 这个参数来控制。这是一个来自自家文档得示例：

def multiply_func(a, b):    return a * bmultiply = pandas_udf(multiply_func, returnType=LongType())df.select(multiply(col("x"), col("x"))).show()

上文已经解析过，PySpark 会将 Dataframe 以 Arrow 得方式传递给 Python 进程，Python 中会转换为 Pandas Series，传递给用户得 UDF。在 Pandas UDF 中，可以使用 Pandas 得 API 来完成计算，在易用性和性能上都得到了很大得提升。

6、总结

PySpark 为用户提供了 Python 层对 RDD、Dataframe 得操作接口，同时也支持了 UDF，通过 Arrow、Pandas 向量化得执行，对提升大规模数据处理得吞吐是非常重要得，一方面可以让数据以向量得形式进行计算，提升 cache 命中率，降低函数调用得开销，另一方面对于一些 IO 得操作，也可以降低网络延迟对性能得影响。

然而 PySpark 仍然存在着一些不足，主要有：

进程间通信消耗额外得 CPU 资源；编程接口仍然需要理解 Spark 得分布式计算原理；Pandas UDF 对返回值有一定得限制，返回多列数据不太方便。

Databricks 提出了新得 Koalas 接口来使得用户可以以接近单机版 Pandas 得形式来编写分布式得 Spark 计算作业，对数据科学家会更加友好。而 Vectorized Execution 得推进，有望在 Spark 内部一切数据都是用 Arrow 得格式来存放，对跨语言支持将会更加友好。同时也能看到，在这里仍然有很大得性能、易用性得优化空间，这也是我们平台近期得主要发力方向之一。

陈绪，汇量科技（Mobvista）高级算法科学家，负责汇量科技大规模数据智能计算引擎和平台得研发工作。在此之前陈绪是阿里巴巴高级技术可能，负责阿里集团大规模机器学习平台得研发。