PySpark中使用UDF处理DataFrame：为df1添加gamma列的需求

我明白你的需求：要给df1新增一列gamma，每一行的gamma值等于df2中所有满足x1等于当前df1.n或者x2等于当前df1.n的w值之和。用UDF确实容易遇到性能或者序列化的问题，这里推荐一种更高效的Spark原生方法，完全不需要写UDF：

步骤1：将df2转换为长格式（Unpivot） #

首先把df2里的x1和x2列合并成一个单列n，这样每个x值都对应一个w，方便后续分组求和。

Scala版本：

import org.apache.spark.sql.functions._

// 将x1和x2转成一列n，同时保留w值
val df2Unpivoted = df2.selectExpr("w", "stack(2, x1, x2) as n")

Python版本：

from pyspark.sql import functions as F

# 用explode+array把x1和x2拆成多行
df2_unpivoted = df2.select(F.col("w"), F.explode(F.array(F.col("x1"), F.col("x2"))).alias("n"))

步骤2：按n分组计算w的总和 #

现在我们可以按n分组，对w求和，得到每个n对应的总w值：

Scala版本：

val sumWByN = df2Unpivoted.groupBy("n").agg(sum("w").alias("gamma"))

Python版本：

sum_w_by_n = df2_unpivoted.groupBy("n").agg(F.sum("w").alias("gamma"))

步骤3：关联df1和求和结果 #

最后把求和后的结果和df1做左关联，这样每个df1.n就能匹配到对应的gamma值：

Scala版本：

// 左关联保证df1的所有行都保留，没有匹配的n会得到null
val resultDf = df1.join(sumWByN, Seq("n"), "left")
// 如果需要把null替换成0，可以加coalesce
// val resultDf = df1.join(sumWByN, Seq("n"), "left").withColumn("gamma", coalesce(col("gamma"), lit(0)))

Python版本：

result_df = df1.join(sum_w_by_n, on="n", how="left")
# 替换null为0的话：
# result_df = df1.join(sum_w_by_n, on="n", how="left").withColumn("gamma", F.coalesce(F.col("gamma"), F.lit(0)))

结果示例 #

按照你的数据，最终的结果会是这样：
| n|val| distances| gamma|
+---+---+----------+----------+
| 1| 1|0.27308652|0.75747334|
| 2| 1|0.24969208|0.03103427|
| 3| 1|0.21314497| null|

为什么不推荐用UDF？ #

UDF需要对df1的每一行单独去查询df2，不仅性能极低（尤其数据量大时），还可能因为没有正确广播df2导致大量的shuffle操作，甚至出现序列化错误。而上面的方法用Spark原生的分布式计算逻辑，效率和稳定性都要好很多。

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