使用 PySpark 并行化自定义函数

Question

我熟悉使用 UDF 将自定义函数逐行应用于 DataFrame。但是，我想知道如何将自定义函数并行应用于我的 DataFrame 的不同子集。

这是一个简化的例子：

import numpy as np
import pandas as pd

dummy_data = pd.DataFrame({'id':np.random.choice(['a','b','c'],size=100),
                           'val':np.random.normal(size=100)})

我的自定义函数将一组数字作为输入。对于每个唯一的“id”，我想将我的函数应用于与该 id 关联的“val”值数组。

我现在做的最简单的方法是循环我的 PySpark DataFrame，并将每个“id”的数据转换为 pandas DataFrame，然后应用该函数。它可以工作，但显然它很慢并且没有使用火花。

我怎样才能并行化这个？

Answer 1

这个答案太短了，应该是评论，但没有足够的声誉来评论。

Spark 2.3 引入了 pandas 矢量化 UDF，这正是您正在寻找的：在分组的 Spark DataFrame 上以分布式方式执行自定义 pandas 转换，并且由于 PyArrow 序列化而具有出色的性能。

见

• https://databricks.com/blog/2017/10/30/introducing-vectorized-udfs-for-pyspark.html
• http://spark.apache.org/docs/latest/api/python/pyspark.sql.html?pyspark.sql.functions.pandas_udf#pyspark.sql.functions.pandas_udf
• Using Collect_set after exploding in a groupedBy object in Pyspark

了解更多信息和示例。