PySpark新增列实现两列相加 仅双NULL时保留NULL不转0

需求说明 #

对DataFrame的cola、colb两列做加法生成新列colc，NULL处理规则如下：

• 两列均为非NULL值：返回两列之和
• 仅单列值为NULL：返回另一列的非NULL值
• 两列均为NULL：返回NULL
  *注意：直接将所有NULL替换为0再相加的方案不满足需求，该方案会在两列全为NULL时错误返回0。

输入输出示例 #

输入数据 #

期望输出 #

实现代码 #

使用PySpark内置的coalesce函数即可实现，无需写复杂的分支判断，逻辑简洁且性能更好：

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, coalesce

# 初始化Spark会话、构造测试数据
spark = SparkSession.builder.appName("add_cols_demo").getOrCreate()
source_data = [
    (100, 100),
    (None, 300),
    (200, None),
    (400, 200),
    (None, None)
]
df = spark.createDataFrame(source_data, schema=["cola", "colb"])

# 核心计算逻辑
df_result = df.withColumn(
    "colc",
    coalesce(col("cola") + col("colb"), col("cola"), col("colb"))
)

# 打印结果验证
df_result.show()

逻辑说明 #

coalesce函数会按传入参数的顺序，返回第一个非NULL的值，对应三种场景的计算逻辑完全匹配需求：

1. 两列都有值：col("cola") + col("colb")返回正常求和结果，直接作为colc的值
2. 单列值为NULL：加法运算结果为NULL，coalesce会向后匹配，取到非NULL的那列的值作为结果
3. 两列全为NULL：加法结果、cola、colb全为NULL，最终返回NULL
   运行上述代码后输出的结果和预期完全一致：

+----+----+----+
|cola|colb|colc|
+----+----+----+
| 100| 100| 200|
|null| 300| 300|
| 200|null| 200|
| 400| 200| 600|
|null|null|null|
+----+----+----+

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