PySpark数组列高阶函数条件转换：将负数转为0

需求说明 #

有PySpark DataFrame包含数组列，示例数据为[0,-1,0,0,1,1,1]，需要将数组中所有负数转换为0，预期结果：[0,0,0,0,1,1,1]。

尝试的错误代码及问题 #

错误代码1（Lambda写法） #

sdf = (sdf
       .withColumn('ArrayCol', f.transform(f.col('ArrayCol'), lambda x: 0 if x < 0 else x))
)

报错信息：

"Cannot convert column into bool: please use '&' for 'and', '|' for 'or', '~' for 'not' when building DataFrame boolean expressions."

问题：直接使用Python原生if x < 0做判断，但x是PySpark的Column对象，无法直接转为Python布尔值，必须使用PySpark专属的条件表达式API。

错误代码2（Expr写法） #

sdf = (sdf
       .withColumn('ArrayCol', f.expr("transform(ArrayCol, x -> CASE WHEN x < 0 THEN 0 ELSE x "))
)

报错信息：

extraneous input 'WHEN' expecting {')', ','}(line 1, pos 37)

问题：Hive SQL语法中CASE语句必须以END收尾，代码遗漏了该关键字。

已验证可行的Expr写法 #

sdf = sdf.withColumn('ArrayCol', f.expr("transform(ArrayCol, x -> CASE WHEN x < 0 THEN 0 ELSE x END)"))

Lambda函数的正确实现方式 #

要在transform中用lambda实现需求，需使用PySpark的when和otherwise函数构建条件逻辑，替代Python原生if-else：

import pyspark.sql.functions as f

sdf = (sdf
       .withColumn('ArrayCol', f.transform(
           f.col('ArrayCol'),
           lambda x: f.when(x < 0, 0).otherwise(x)
       ))
)

解释：f.when(x < 0, 0)会生成合法的Column表达式，当x小于0时返回0，否则返回x本身，完全适配PySpark的Column操作逻辑，不会触发类型转换错误。

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