如何在PySpark中用withColumn、循环及UDF批量创建统计字段

需要在PySpark DataFrame中新增sum_fields列，统计满足条件的字段组合次数：对每组daysN和codeN，当daysN > 500且codeN = 'B'时加1，否则加0。当前已实现单组字段的逻辑，但需要扩展到30组字段，希望用循环/列表推导式简化代码，避免手动重复编写30次判断逻辑。

给定的DataFrame示例：

df = spark.createDataFrame(
[(1, 100, 100, 'A', 'A'),
  (2, 1000, 200, 'A', 'A'),
  (3, 1000, 300, 'B', 'A'),
  (4, 1000, 1000, 'B', 'B')],  
"id int, days1 int, days2 int, code1 string, code2 string")

期望输出示例：

+---+-----+-----+-----+-----+----------+
| id|days1|days2|code1|code2|sum_fields|
+---+-----+-----+-----+-----+----------+
|  1|  100|  100|    A|    A|         0|
|  2| 1000|  200|    A|    A|         0|
|  3| 1000|  300|    B|    A|         1|
|  4| 1000| 1000|    B|    B|         2|
+---+-----+-----+-----+-----+----------+

方法1：使用Spark原生表达式（推荐） #

利用列表推导式生成每组字段的条件判断表达式，再通过sum()统一求和，无需自定义UDF，性能更优（Spark原生函数经过分布式优化）。

代码实现：

from pyspark.sql import functions as F

# 定义需要处理的字段组数（此处为30组）
num_groups = 30

# 生成每组字段的条件表达式：满足条件返回1，否则返回0
condition_list = [
    F.when(
        (F.col(f"days{i+1}") > 500) & (F.col(f"code{i+1}") == "B"),
        1
    ).otherwise(0)
    for i in range(num_groups)
]

# 将所有条件表达式的结果求和，生成sum_fields列
result_df = df.withColumn("sum_fields", F.sum(*condition_list))

# 查看结果
result_df.show()

方法2：自定义UDF（不推荐，性能低于原生函数） #

如果必须使用UDF，可以将所有days和code字段分别打包为数组传入UDF，在UDF内部遍历数组统计符合条件的次数。

代码实现：

from pyspark.sql import functions as F
from pyspark.sql.types import IntegerType, ArrayType

def count_valid_pairs(days_array, codes_array):
    count = 0
    for day, code in zip(days_array, codes_array):
        if day > 500 and code == "B":
            count += 1
    return count

# 注册UDF
count_udf = F.udf(count_valid_pairs, IntegerType())

num_groups = 30
# 收集所有days和code字段，打包为数组
days_columns = [F.col(f"days{i+1}") for i in range(num_groups)]
code_columns = [F.col(f"code{i+1}") for i in range(num_groups)]

# 调用UDF生成sum_fields列
result_df = df.withColumn(
    "sum_fields",
    count_udf(F.array(*days_columns), F.array(*code_columns))
)

result_df.show()

注意：优先选择方法1，因为UDF需要在Python和JVM之间进行数据序列化/反序列化，在大数据量场景下会显著降低性能，而Spark原生函数是基于JVM执行的，效率更高。

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