使用 Apache Arrow 加速 PySpark

已发布 2017年7月26日
作者 BryanCutler

Bryan Cutler 是 IBM Spark 技术中心 STC 的软件工程师

从 Apache Spark 2.3 版本开始，Apache Arrow 将成为受支持的依赖项，并开始通过列式数据传输提供更高的性能。如果您是喜欢使用 Python 和 Pandas 的 Spark 用户，那么这是一个令人兴奋的消息！最初的工作仅限于使用 toPandas() 收集 Spark DataFrame，我将在下面讨论，但是目前还有许多其他改进正在 进行中。

优化 Spark 到 Pandas 的转换

以前在 PySpark 中使用 DataFrame.toPandas() 将 Spark DataFrame 转换为 Pandas 的方法效率极低。基本上，它的工作原理是首先将所有行收集到 Spark 驱动程序。接下来，每一行都会被序列化为 Python 的 pickle 格式并发送到 Python 工作进程。此子进程将每一行反序列化为一个巨大的元组列表。最后，使用 pandas.DataFrame.from_records() 从列表创建 Pandas DataFrame。

这一切似乎都是标准程序，但存在两个明显的问题：1）即使使用 CPickle，Python 序列化也是一个缓慢的过程；2）使用 from_records 创建 pandas.DataFrame 必须缓慢迭代纯 Python 数据列表并将每个值转换为 Pandas 格式。有关详细分析，请参阅 此处。

这就是 Arrow 真正发挥作用来帮助优化这些步骤的地方：1) 一旦数据采用 Arrow 内存格式，就不再需要序列化/pickle，因为 Arrow 数据可以直接发送到 Python 进程，2) 当 Python 接收到 Arrow 数据时，pyarrow 可以利用零拷贝方法从整个数据块一次性创建 pandas.DataFrame，而不是处理单个标量值。此外，转换为 Arrow 数据可以在 JVM 上完成，并推回 Spark 执行器并行执行，从而大大减少驱动程序的负载。

自 SPARK-13534 合并以来，需要通过将 SQLConf“spark.sql.execution.arrow.enabled”设置为“true”来启用调用 toPandas() 时的 Arrow 使用。让我们看一个简单的用法示例。

Welcome to
      ____              __
     / __/__  ___ _____/ /__
    _\ \/ _ \/ _ `/ __/  '_/
   /__ / .__/\_,_/_/ /_/\_\   version 2.3.0-SNAPSHOT
      /_/

Using Python version 2.7.13 (default, Dec 20 2016 23:09:15)
SparkSession available as 'spark'.

In [1]: from pyspark.sql.functions import rand
   ...: df = spark.range(1 << 22).toDF("id").withColumn("x", rand())
   ...: df.printSchema()
   ...: 
root
 |-- id: long (nullable = false)
 |-- x: double (nullable = false)


In [2]: %time pdf = df.toPandas()
CPU times: user 17.4 s, sys: 792 ms, total: 18.1 s
Wall time: 20.7 s

In [3]: spark.conf.set("spark.sql.execution.arrow.enabled", "true")

In [4]: %time pdf = df.toPandas()
CPU times: user 40 ms, sys: 32 ms, total: 72 ms                                 
Wall time: 737 ms

In [5]: pdf.describe()
Out[5]: 
                 id             x
count  4.194304e+06  4.194304e+06
mean   2.097152e+06  4.998996e-01
std    1.210791e+06  2.887247e-01
min    0.000000e+00  8.291929e-07
25%    1.048576e+06  2.498116e-01
50%    2.097152e+06  4.999210e-01
75%    3.145727e+06  7.498380e-01
max    4.194303e+06  9.999996e-01

此示例在我的笔记本电脑上使用 Spark 默认值本地运行，因此显示的时间不应精确计算。即使如此，很明显性能有了巨大的提升，使用 Arrow 将极其缓慢的操作加速到几乎察觉不到。

使用说明

在使用此新功能之前，请记住以下几点。在撰写本文时，pyarrow 不会与 pyspark 自动安装，需要手动安装，请参阅安装 说明。计划将 pyarrow 添加为 pyspark 依赖项，以便 > pip install pyspark 也会安装 pyarrow。

目前，控制 SQLConf 默认情况下处于禁用状态。这可以通过编程方式启用，如上例所示，或者通过将“spark.sql.execution.arrow.enabled=true”添加到 SPARK_HOME/conf/spark-defaults.conf 中。

此外，并非所有 Spark 数据类型都受支持，目前仅限于基本类型。扩展类型支持正在开发中，预计也将在 Spark 2.3 版本中发布。

未来改进

如前所述，这只是使用 Arrow 为 Spark Python 用户提供便利的第一步。一些正在进行的令人兴奋的计划是允许向量化 UDF 评估（SPARK-21190，SPARK-21404），以及使用 Pandas DataFrame 对分组数据应用函数的能力（SPARK-20396）。正如 Arrow 帮助将 Spark 转换为 Pandas 一样，它也可以在从现有 Pandas DataFrame 创建 Spark DataFrame 时反向工作（SPARK-20791）。敬请期待！

贡献者

达到第一个里程碑是 Apache Arrow 和 Spark 社区的共同努力。感谢 Wes McKinney、Li Jin、Holden Karau、Reynold Xin、Wenchen Fan、Shane Knapp 和许多其他人的辛勤工作，他们帮助推动了这项工作的发展。