如何基于自定义函数为Pandas on Spark DataFrame添加列？

问题场景 #

给定Pandas on Spark DataFrame（简称psdf）：

name     p1     p2
0     AAA    1.0    1.0
1     BBB    1.0    1.0

定义了支持动态参数的函数（示例为双参数，实际还有三参数版本）：

import math

def f(a: float, b: float) -> float:
    return math.sqrt(a**2 + b**2)

期望得到新增计算列R的结果：

name     p1     p2  R
0     AAA    1.0    1.0  1.4
1     BBB    1.0    1.0  1.4

尝试使用apply时触发compute.ops_on_diff_frames参数错误，且该操作开销高，希望避免；同时不确定转普通Spark DataFrame用withColumn是否有性能损失。

可行解决方案 #

方法1：使用map_in_pandas（官方推荐高效方案） #

这是Pandas on Spark原生支持的批量行处理API，适配动态参数场景，无需开启高开销的跨帧操作。

双参数场景示例：

import pandas as pd

def process_batch(batch: pd.DataFrame) -> pd.DataFrame:
    batch['R'] = batch.apply(lambda x: f(x['p1'], x['p2']), axis=1)
    return batch

# 应用到psdf，需指定新增列的schema
result_psdf = psdf.map_in_pandas(process_batch, schema=psdf.schema + "R double")

三参数场景扩展：

先定义三参数函数：

def f3(a: float, b: float, c: float) -> float:
    return math.sqrt(a**2 + b**2 + c**2)

再修改批量处理逻辑：

def process_batch_3col(batch: pd.DataFrame) -> pd.DataFrame:
    batch['R'] = batch.apply(lambda x: f3(x['p1'], x['p2'], x['p3']), axis=1)
    return batch

result_psdf_3col = psdf.map_in_pandas(process_batch_3col, schema=psdf.schema + "R double")

优势：基于批量处理Pandas DataFrame，比逐行apply效率高，完全适配Pandas on Spark执行模型，无需切换数据结构。

方法2：转换为Spark DataFrame使用Pandas UDF（性能可控） #

转换为普通Spark DataFrame后，使用矢量化Pandas UDF而非普通UDF，性能与Pandas on Spark原生操作差异极小。

双参数场景示例：

from pyspark.sql.functions import pandas_udf

# 定义Pandas UDF
@pandas_udf('double')
def f_udf(a: pd.Series, b: pd.Series) -> pd.Series:
    return (a**2 + b**2).apply(math.sqrt)

# 转换为Spark DataFrame并计算
spark_df = psdf.to_spark()
result_spark_df = spark_df.withColumn('R', f_udf(spark_df['p1'], spark_df['p2']))

# 如需转回Pandas on Spark DataFrame
result_psdf = result_spark_df.to_pandas_on_spark()

三参数场景扩展：

@pandas_udf('double')
def f3_udf(a: pd.Series, b: pd.Series, c: pd.Series) -> pd.Series:
    return (a**2 + b**2 + c**2).apply(math.sqrt)

result_spark_df_3col = spark_df.withColumn('R', f3_udf(spark_df['p1'], spark_df['p2'], spark_df['p3']))

说明：Spark的Pandas UDF基于Arrow传输数据，批量处理，性能接近Pandas on Spark原生操作，适合需要复用Spark生态工具的场景。

方法3：开启compute.ops_on_diff_frames（仅适合小数据集） #

如果一定要用apply，可开启该配置，但会触发全量数据收集到Driver，大数据集下易导致内存溢出，不推荐生产环境使用：

from pyspark.pandas.config import set_option

set_option("compute.ops_on_diff_frames", True)
psdf['R'] = psdf[['p1','p2']].apply(lambda x: f(*x), axis=1)

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