PySpark中使用Pandas UDF计算多列行方差报错求解

问题根源 #

你的代码存在两个核心问题：

1. 类型注解错误：*cols:List[pd.Series] 写法不符合语法，*cols 代表可变数量的独立参数，每个参数是 pd.Series，而非一个包含多个Series的列表。
2. 输入处理错误：pd.Series(cols) 会把传入的多个Series作为单个Series的元素，而非按行拼接成表格结构，导致无法正确计算行方差。

修正后的代码 #

from pyspark.sql import SparkSession
from pyspark.sql.functions import pandas_udf
from pyspark.sql.types import DoubleType
import pandas as pd

# 构造测试数据
test = pd.DataFrame({'id': ['a', 'b', 'c', 'd', 'e'],
                     'feat1': [3,4,5,6,7],
                     'feat2': [6,9,2,4,5]
                    })
test['var_pd'] = test[['feat1', 'feat2']].var(axis=1)

# 初始化Spark会话
spark = SparkSession.builder.appName("variance_udf_test").getOrCreate()
test_spark = spark.createDataFrame(test)

# 定义正确的Pandas UDF
@pandas_udf(returnType=DoubleType())
def variance_udf(*cols: pd.Series) -> pd.Series:
    # 将多列Series合并为DataFrame，保证行对齐
    input_df = pd.concat(cols, axis=1)
    # 计算行方差，ddof=1对应样本方差（与Pandas默认行为一致）
    return input_df.var(axis=1, ddof=1)

# 调用UDF并查看结果
test_spark = test_spark.withColumn("variance_udf", variance_udf('feat1', 'feat2'))
test_spark.show()

关键说明 #

• 使用 pd.concat(cols, axis=1) 将传入的多列Series拼接成完整的DataFrame，确保每一行对应原数据的一行记录。
• 显式指定 ddof=1 保证计算的是样本方差，和你之前用Pandas计算的var_pd结果完全一致。
• 类型注解修正为 *cols: pd.Series，符合可变参数的语法规范。

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