使用 spark.sparkContext.addPyFile 导入 Pandas UDF

Question

我有以下玩具示例。我正在使用 bootstrap 脚本将 pandas 和 pyarrow 安装到我的工作节点上。当我运行在 jupyter notebook 中使用以下代码时，它运行没有错误。

# Declare the function and create the UDF
def multiply_func(a: pd.Series, b: pd.Series) -> pd.Series:
    return a * b

@f.pandas_udf("float")
def udf_multiply(a: pd.Series, b: pd.Series) -> pd.Series:
    df = pd.DataFrame({'a': a, 'b': b})
    df['product'] = df.apply(lambda x : multiply_func(x['a'], x['b']), axis = 1)
    return df['product']

x = pd.Series([1, 2, 3])
#print(multiply_func(x, x))
# 0    1
# 1    4
# 2    9
# dtype: int64

# Create a Spark DataFrame, 'spark' is an existing SparkSession
df = spark.createDataFrame(pd.DataFrame(x, columns=["x"]))

# Execute function as a Spark vectorized UDF
df.select(udf_multiply(f.col("x"), f.col("x"))).show()

但是，我有很多 pandas_udf 想要导入到我的工作区，我不想将它们中的每一个都复制粘贴到我的 Jupyter Notebook 的顶部。我想要的目录结构如下所示：

eda.ipynb
helpful_pandas_udfs/toy_example.py

我查看了其他 SO 帖子并确定我应该能够像这样添加 Python 文件：

spark.sparkContext.addPyFile("helpful_pandas_udfs/toy_example.py")
from toy_example import udf_multiply

然而，当我尝试运行这段代码时，出现以下错误：

AttributeError: 'NoneType' object has no attribute '_jvm'

请帮忙！我完全被这个难住了。

Answer 1

我可以通过在创建 spark 会话后将我的 UDF 复制为文本来解决这个问题。这不是我满意的解决方案，但它确实有效。

使用 spark.sparkContext.addPyFile 导入 Pandas UDF

Import Pandas UDF With spark.sparkContext.addPyFile

python

pandas

pyspark

databricks