PySpark中如何处理输入长度不同的DataFrame的Pandas UDF？

问题描述 #

我是PySpark新手，希望编写一个函数来处理两个列结构相同但行数不同的DataFrame，请问最佳实现方式是什么？我假设Pandas UDF只能作用于单个PySpark DataFrame（若我的理解有误，请予以指正），目前想到的一种方法是将其中一个DataFrame追加到另一个中，再在UDF里使用iloc进行拆分，但这是否是最优方案？

相关参考代码：

import pandas as pd
from pyspark.sql.types import StructType, StructField, IntegerType, DoubleType
from pyspark.sql import SparkSession
from pyspark.sql.functions import pandas_udf, struct, lit

spark = SparkSession.builder.appName("Example").getOrCreate()

schema = StructType([StructField("num", IntegerType(), True)])

data1 = [(1,), (2,), (3,)]
df1 = spark.createDataFrame(data1, schema)
data2 = [(4,), (5,), (6,), (7,), (8,)]
df2 = spark.createDataFrame(data2, schema)

@pandas_udf(DoubleType())
def udf_test(df1, df2):
    val1, _ = some_func(df1['num'], df2['num'])
    return pd.Series(val1)

解答 #

纠正你的误解：Pandas UDF可以处理多个DataFrame #

你的假设不成立，Pandas UDF并非只能作用于单个DataFrame。可以通过将多个DataFrame的字段打包成struct，或结合广播变量的方式，在Pandas UDF中同时处理两个DataFrame的数据。

你的追加拆分方案并非最优 #

把两个DF追加后用iloc拆分的方法存在明显问题：

• 破坏Spark分布式计算特性，追加后的DF在UDF处理时会被当成单批数据，无法利用集群并行能力
• 行数不同时，拆分逻辑易出错，且无法保证数据对应关系（若原本需要按规则关联）
• 数据量较大时性能低下，会导致大量数据在单节点处理

推荐的最佳实现方式 #

根据实际业务场景，选择以下方案：

1. 按关联键处理（最常见场景）

如果需要将两个DF中对应键的行配对处理，先通过join操作合并两个DF，再用Pandas UDF处理合并后的行：

# 示例：给每个DF加标识列，实际替换为业务关联键
df1_with_id = df1.withColumn("id", lit(1))
df2_with_id = df2.withColumn("id", lit(1))

# 按关联键合并DF
joined_df = df1_with_id.join(df2_with_id, on="id", how="outer")

@pandas_udf(DoubleType())
def process_pair(num1, num2):
    # 替换为你的some_func逻辑
    result = num1.add(num2).fillna(0)
    return result

result_df = joined_df.withColumn("result", process_pair("num", "num"))
result_df.show()

2. 全局处理两个DF的整体数据（无关联键场景）

如果需要对两个DF的全部数据进行整体计算，可将小DF广播，在Pandas UDF中使用广播的数据：

from pyspark.sql.functions import broadcast

# 广播较小的DF，减少数据传输
broadcast_df1 = broadcast(df1)
pd_df1 = broadcast_df1.toPandas()

@pandas_udf(DoubleType())
def process_with_broadcast(num_series):
    val1, _ = some_func(pd_df1['num'], num_series)
    return pd.Series(val1)

result_df = df2.withColumn("result", process_with_broadcast("num"))
result_df.show()

3. 使用Grouped Map进行分组处理

如果需要按分组分别处理两个DF的数据，先合并两个DF并标记来源，再分组后用Grouped Map UDF处理：

# 合并两个DF并标记来源
df1_tagged = df1.withColumn("source", lit("df1"))
df2_tagged = df2.withColumn("source", lit("df2"))
combined_df = df1_tagged.unionByName(df2_tagged)

@pandas_udf(schema)
def process_group(pdf):
    df1_data = pdf[pdf['source'] == 'df1']['num']
    df2_data = pdf[pdf['source'] == 'df2']['num']
    val1, _ = some_func(df1_data, df2_data)
    return pd.DataFrame({'num': val1})

result_df = combined_df.groupBy().applyInPandas(process_group, schema=schema)
result_df.show()

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