如何在PySpark中根据多列条件生成漏斗步骤列？

需求说明 #

根据date_lead→date_login→date_access→date_buy→date_del的顺序，取每行中最后一个非空日期列对应的步骤名称，生成step列。

方法一：嵌套WHEN逻辑（对应Pandas覆盖式赋值） #

和Pandas中从前往后覆盖的逻辑对应，这里反向从最后一个步骤开始判断，优先匹配后续非空列：

from pyspark.sql import functions as F

# 新增step列
df = df.withColumn(
    "step",
    F.when(F.col("date_del").isNotNull(), "del")
    .when(F.col("date_buy").isNotNull(), "buy")
    .when(F.col("date_access").isNotNull(), "access")
    .when(F.col("date_login").isNotNull(), "login")
    .when(F.col("date_lead").isNotNull(), "lead")
    .otherwise(None)  # 处理所有日期列都为空的情况
)

方法二：数组+内置函数（灵活适配多列场景） #

通过构造日期列与步骤名的映射数组，过滤空值后取最后一个有效步骤，适合列数量较多的场景：

from pyspark.sql import functions as F

# 定义日期列和对应步骤的映射关系
date_step_mapping = [
    ("date_lead", "lead"),
    ("date_login", "login"),
    ("date_access", "access"),
    ("date_buy", "buy"),
    ("date_del", "del")
]

# 构造数组并处理空值，提取最后一个有效步骤
df = df.withColumn(
    "step",
    F.element_at(
        F.filter(
            F.array(*[F.struct(F.col(date_col).alias("dt"), F.lit(step).alias("step_name")) 
                      for date_col, step in date_step_mapping]),
            lambda x: x.dt.isNotNull()
        ),
        -1  # 取数组最后一个元素
    ).step_name
)

两种方法执行后，都能得到预期的输出结果。

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