如何通过Pandas UDF在PySpark中对接SciPy API实现列转换？

原代码存在几个关键问题：

• UDF返回类型定义错误：标注返回"col1 int, col2 int"，但实际仅返回一个计算后的浮点Series，类型不匹配
• 数据类型不兼容：计算s1 + s2*scipy.pi得到的是浮点数，而非整数，原返回类型的int约束会导致报错
• 返回值结构不匹配：函数标注返回pd.Series，但原类型定义是多列结构体，两者无法对应

修正后的实现步骤 #

1. 确保环境依赖

Spark集群的所有节点都需要安装scipy和pandas，可通过pip install scipy pandas完成安装。

2. 创建原始DataFrame

columns = ['N','P']
data = [(1,3), (3,3), (5,3)]
df = spark.createDataFrame(data=data, schema=columns)

3. 定义正确的Pandas UDF

根据需求返回单个浮点列，UDF返回类型定义为"double"，同时确保返回值是匹配的pd.Series：

from pyspark.sql.functions import pandas_udf
import pandas as pd
import scipy

@pandas_udf("double")
def func(s1: pd.Series, s2: pd.Series) -> pd.Series:
    return s1 + s2 * scipy.pi

4. 调用UDF添加新列

df2 = df.withColumn('transformed', func("N", "P"))
df2.show()

执行结果

+---+---+------------------+
|  N|  P|        transformed|
+---+---+------------------+
|  1|  3|10.42477796076938|
|  3|  3|12.42477796076938|
|  5|  3|14.42477796076938|
+---+---+------------------+

扩展：返回多列的场景 #

如果需要返回多个计算列，可将UDF返回类型定义为结构体，函数返回pd.DataFrame：

from pyspark.sql.types import StructType, StructField, DoubleType

# 定义返回的结构体类型
result_schema = StructType([
    StructField("calc1", DoubleType()),
    StructField("calc2", DoubleType())
])

@pandas_udf(result_schema)
def multi_func(s1: pd.Series, s2: pd.Series) -> pd.DataFrame:
    calc1 = s1 + s2 * scipy.pi
    calc2 = s1 * s2 + scipy.pi
    return pd.DataFrame({"calc1": calc1, "calc2": calc2})

# 调用并展开列
df3 = df.withColumn('results', multi_func("N", "P"))
df3 = df3.select("N", "P", "results.calc1", "results.calc2")
df3.show()

内容的提问来源于stack exchange，提问作者FalconX