如何将含嵌套np.where的Pandas代码转换为PySpark嵌套when实现

首先导入PySpark函数依赖：

from pyspark.sql import functions as F

你给出的Pandas代码可以直接一对一转换为PySpark嵌套when的实现，转换后代码如下：

df = df.withColumn(
    "new_col",
    # 对应最外层np.where，判断col1是否为空
    F.when(
        F.col("col1").isNull(),
        # 嵌套第一层：col1为空时，判断col2是否为空
        F.when(
            F.col("col2").isNull(),
            F.col("col3")
        # col2不为空时进入下一层判断
        ).otherwise(
            # 嵌套第二层：判断col3是否为空
            F.when(
                F.col("col3").isNull(),
                F.col("col2")
            # col3也不为空时执行计算逻辑
            ).otherwise(
                F.col("col2") + F.col("col3") / 2
            )
        )
    # col1不为空时直接取col1的值
    ).otherwise(F.col("col1"))
)

嵌套when的核心使用规则 #

• 每一组np.where(条件, 真值, 假值)完全对应PySpark的F.when(条件, 真值).otherwise(假值)，嵌套顺序和原Pandas逻辑完全一致即可，不需要额外调整逻辑顺序
• 空值判断必须使用PySpark内置的isNull()方法，不要使用Python原生的== None或者is None，这两种写法在分布式计算逻辑中无法正确识别Spark DataFrame的空值，是新手最容易踩的错误点
• 嵌套层级没有固定上限，如果嵌套过深可以把内层逻辑拆成单独变量提高可读性，示例如下：

# 拆分子逻辑简化写法
col2_col3_logic = F.when(F.col("col2").isNull(), F.col("col3"))\
                   .otherwise(F.when(F.col("col3").isNull(), F.col("col2"))\
                   .otherwise(F.col("col2") + F.col("col3")/2))

df = df.withColumn("new_col", F.when(F.col("col1").isNull(), col2_col3_logic).otherwise(F.col("col1")))

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