如何按优先级从多列取值聚合PySpark DataFrame的ID分组？

需求说明 #

现有PySpark DataFrame，每个ID对应多行数据，每行仅一个Type1-Type5列包含有效值；需要将每个ID合并为一行，生成Type列，取值规则为按Type1>Type2>Type3>Type4>Type5的优先级，选取第一个非空且大于0的值。

实现方案 #

以下提供两种可行的实现方式，可根据实际场景选择：

方法一：聚合后按优先级筛选（直观简洁）

先对每个ID聚合各Type列的最大值（因每个ID在单个Type列最多一个有效值），再通过coalesce结合when条件按优先级选取符合要求的值：

from pyspark.sql import SparkSession
from pyspark.sql import functions as F

# 初始化SparkSession
spark = SparkSession.builder.appName("MergeTypePriority").getOrCreate()

# 构造示例DataFrame
data = [
    ("001", None, 100, None, None, None),
    ("001", None, None, None, None, 200),
    ("001", None, None, 300, None, None),
    ("001", None, None, None, 150, None),
    ("002", 0, None, None, None, None),
    ("002", None, None, 300, None, None),
    ("002", None, 200, None, None, None),
    ("002", None, None, None, None, 400)
]
schema = ["ID", "Type1", "Type2", "Type3", "Type4", "Type5"]
df = spark.createDataFrame(data, schema=schema)

# 1. 聚合每个ID的各Type列最大值
agg_df = df.groupBy("ID").agg(
    F.max("Type1").alias("Type1"),
    F.max("Type2").alias("Type2"),
    F.max("Type3").alias("Type3"),
    F.max("Type4").alias("Type4"),
    F.max("Type5").alias("Type5")
)

# 2. 按优先级选取第一个非空且>0的值
result_df = agg_df.select(
    "ID",
    F.coalesce(
        F.when(F.col("Type1") > 0, F.col("Type1")),
        F.when(F.col("Type2") > 0, F.col("Type2")),
        F.when(F.col("Type3") > 0, F.col("Type3")),
        F.when(F.col("Type4") > 0, F.col("Type4")),
        F.when(F.col("Type5") > 0, F.col("Type5"))
    ).alias("Type")
)

# 查看结果
result_df.show()

方法二：长格式转换+排序取首（扩展性强）

将宽表转换为长表，提取有效数据并标记优先级，再按ID分组后按优先级排序取第一个有效值，适合Type列数量较多的场景：

from pyspark.sql import SparkSession
from pyspark.sql import functions as F
from pyspark.sql.window import Window

# 初始化SparkSession
spark = SparkSession.builder.appName("MergeTypePriority").getOrCreate()

# 构造示例DataFrame（同方法一）
data = [
    ("001", None, 100, None, None, None),
    ("001", None, None, None, None, 200),
    ("001", None, None, 300, None, None),
    ("001", None, None, None, 150, None),
    ("002", 0, None, None, None, None),
    ("002", None, None, 300, None, None),
    ("002", None, 200, None, None, None),
    ("002", None, None, None, None, 400)
]
schema = ["ID", "Type1", "Type2", "Type3", "Type4", "Type5"]
df = spark.createDataFrame(data, schema=schema)

# 定义Type列优先级（数字越小优先级越高）
type_priority = {"Type1": 1, "Type2": 2, "Type3": 3, "Type4": 4, "Type5": 5}

# 1. 转换为长表，过滤空值和<=0的无效数据，添加优先级列
long_df = df.select(
    "ID",
    F.explode(
        F.array(
            *[F.struct(F.lit(col).alias("type_name"), F.col(col).alias("value")) for col in schema[1:]]
        )
    ).alias("type_info")
).select(
    "ID", "type_info.type_name", "type_info.value"
).filter(
    (F.col("value").isNotNull()) & (F.col("value") > 0)
).withColumn(
    "priority", F.lit(type_priority).getItem(F.col("type_name"))
)

# 2. 按ID分组，按优先级升序排序后取第一个值
window_spec = Window.partitionBy("ID").orderBy("priority")
result_df = long_df.withColumn(
    "Type", F.first("value", ignorenulls=True).over(window_spec)
).select("ID", "Type").distinct()

# 查看结果
result_df.show()

结果输出 #

两种方法最终都会得到目标DataFrame：

+---+----+
| ID|Type|
+---+----+
|001| 100|
|002| 200|
+---+----+

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