PySpark中WHEN函数条件未生效，UDF执行报错问题排查

问题背景 #

通过4条过滤语句验证，确认status为"pending"的行中col1和col2不存在NULL或None值，但使用F.when调用自定义UDF计算new_col时，抛出TypeError: unsupported operand type(s) for *: 'NoneType' and 'NoneType'异常；而先过滤status为"pending"的行再调用UDF则执行正常。使用Spark版本为2.4.8。

验证代码 #

# 验证status="pending"的行中col1/col2无NULL/None
df.filter((df.status == "pending") & ((F.col("col1").isNull()) | (F.col("col2").isNull()))).count() -> 0
df.filter((df.status == "pending") & ((F.col("col1").eqNullSafe(None)) | (F.col("col2").eqNullSafe(None)))).count() -> 0
df.filter((df.status == "pending") & ((isnan("col1")) | (isnan("col2")))).count() -> 0
df.filter((df.status == "pending") & ((F.col("col1") == None) | (F.col("col2") == None))).count() -> 0

报错代码 #

df1 = df.withColumn("new_col", F.when(df.status == "pending", udf_func(df.col1, df.col2)).otherwise(F.lit(0.0)))
df1.show(2, False)  # 触发报错

错误信息 #

Caused by: org.apache.spark.api.python.PythonException: Traceback (most recent call last):
  File "/usr/lib/spark/python/lib/pyspark.zip/pyspark/worker.py", line 377, in main
    process()
  File "/usr/lib/spark/python/lib/pyspark.zip/pyspark/worker.py", line 372, in process
    serializer.dump_stream(func(split_index, iterator), outfile)
  File "/usr/lib/spark/python/lib/pyspark.zip/pyspark/worker.py", line 248, in <lambda>
    func = lambda _, it: map(mapper, it)
  File "<string>", line 1, in <lambda>
  File "/usr/lib/spark/python/lib/pyspark.zip/pyspark/worker.py", line 85, in <lambda>
    return lambda *a: f(*a)
  File "/usr/lib/spark/python/lib/pyspark.zip/pyspark/util.py", line 99, in wrapper
    return f(*args, **kwargs)
  File "<stdin>", line 7, in udf_func_name
TypeError: unsupported operand type(s) for *: 'NoneType' and 'NoneType'

正常执行代码 #

df2 = df.filter(df.status == "pending").withColumn("new_col", udf_func(df.col1, df.col2))
df2.show(2, False)  # 执行正常

原因分析 #

Spark 2.4.x中，F.when的参数执行逻辑存在关键特性：

• F.when的第二个参数是列表达式（此处为UDF调用），Spark会对全量数据计算UDF的输入列（col1、col2），而非仅在满足when条件的行上计算。也就是说，即使status不为"pending"的行，其col1/col2的NULL值也会被传入UDF，触发运算错误。
• 先过滤再调用UDF时，已排除非"pending"行，UDF仅处理符合条件的行，因此不会遇到NULL值。
• 之前的过滤验证仅确认了"pending"行无NULL，但非"pending"行的col1/col2可能存在NULL，这些行在when语句中虽不会使用UDF结果，但UDF的计算过程仍会触发，最终抛出异常。

解决思路 #

1. 在UDF内部添加NULL校验 #

让UDF自身处理NULL输入，避免运算错误：

def udf_func(col1, col2):
    if col1 is None or col2 is None:
        return 0.0  # 根据业务需求设置默认值
    return col1 * col2  # 原业务逻辑

2. 提前处理UDF输入列的NULL值 #

调用UDF前，用F.coalesce将NULL替换为默认值，确保UDF接收非空参数：

df1 = df.withColumn(
    "new_col",
    F.when(
        df.status == "pending",
        udf_func(F.coalesce(df.col1, F.lit(0.0)), F.coalesce(df.col2, F.lit(0.0)))
    ).otherwise(F.lit(0.0))
)

3. 用Spark内置函数替代UDF #

若UDF逻辑简单（如仅相乘），直接使用Spark内置函数，内置函数会自动处理NULL（结果为NULL，但otherwise会覆盖为0.0）：

df1 = df.withColumn(
    "new_col",
    F.when(df.status == "pending", F.col("col1") * F.col("col2")).otherwise(F.lit(0.0))
)

4. 分步骤处理数据 #

先处理pending行，再与非pending行合并，保持数据完整性：

pending_df = df.filter(df.status == "pending").withColumn("new_col", udf_func(df.col1, df.col2))
non_pending_df = df.filter(df.status != "pending").withColumn("new_col", F.lit(0.0))
df1 = pending_df.union(non_pending_df)

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