PySpark中如何简洁实现多列求和并生成条件标记列？

问题背景 #

你有一个包含20列（列名col1到col20）的Spark DataFrame，需要将这些列求和后生成一个新列：求和结果大于0时为1，否则为0。目前你通过两步实现：

df = df.withColumn("temp_col", col("col1")+col("col2")+...+col("col20"))
df = df.withColumn("new_col_2", when(col("temp_col") > 0, 1).otherwise(0))

尝试一步实现时遇到报错：

df.na.fill(0).withColumn("new_col" ,reduce(add, [col(col(f'{x}') for x in range(0,20))]))

报错信息：

name 'add' is not defined
Traceback (most recent call last):
NameError: name 'add' is not defined

正确的一步实现代码 #

首先需要导入必要的模块，然后直接在withColumn中完成求和与二值转换：

导入依赖 #

from pyspark.sql.functions import col, when, add
from functools import reduce

针对col1到col20的实现 #

df = df.na.fill(0).withColumn(
    "new_col",
    when(reduce(add, [col(f'col{x}') for x in range(1, 21)]) > 0, 1).otherwise(0)
)

如果是对所有列求和的简化写法 #

如果DataFrame中只有col1到col20这20列，可以直接用df.columns代替手动生成列名：

df = df.na.fill(0).withColumn(
    "new_col",
    when(reduce(add, [col(c) for c in df.columns]) > 0, 1).otherwise(0)
)

你之前代码的错误原因 #

1. 缺少必要导入：add函数需要从pyspark.sql.functions导入，reduce需要从functools导入，未导入就会触发NameError。
2. 列名生成错误：
   • 你用了range(0,20)，生成的是0到19的数字，但列名是col1到col20，应该用range(1,21)。
   • 代码中col(col(f'{x}'))是嵌套错误，正确写法是col(f'col{x}')，不需要嵌套调用col函数。

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