PySpark SQL DataFrame按id分组实现条件判断生成bin字段的方法

解决方案 #

核心逻辑：将每条数据中符合code = 03/06的行标记为1，其余为0，分组后取最大值，即可保证只要该id下存在任意一条符合条件的数据，最终结果为1，否则为0。

方式1：DataFrame API 实现

from pyspark.sql import functions as F

result_df = df.groupBy("id") \
    .agg(
        F.max(
            F.when(F.col("code").isin("03", "06"), 1).otherwise(0)
        ).alias("bin")
    )

result_df.show()

方式2：Spark SQL 实现

# 注册临时视图
df.createOrReplaceTempView("source_data")

# 执行SQL查询
result_df = spark.sql("""
SELECT 
    id,
    MAX(CASE WHEN code IN ('03', '06') THEN 1 ELSE 0 END) AS bin
FROM source_data
GROUP BY id
""")

result_df.show()

注意：如果你的code字段为数值类型，将判断条件中的字符串引号去掉即可（即isin(03,06) / IN (03,06)）。

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