Spark中如何依据列数据生成指定格式新列并聚合最大值？

Hey there! Let's work through this Spark DataFrame transformation to get your desired output. I'll show you two approaches depending on whether your code values are fixed or dynamic.

First: If your code values are fixed (only BI_ and col_) #

This approach is straightforward—we'll split the data by code, compute the max values per id, then join the results back together.

from pyspark.sql import functions as F

# Sample input data
data = [
    ("BodilyInjured", 1001, 45, 50, "BI_"),
    ("BodilyInjured", 1001, 12, 55, "BI_"),
    ("Collision", 1001, 18, 21, "col_"),
    ("Collision", 1002, 21, 25, "col_")
]
df = spark.createDataFrame(data, ["des", "id", "veh_cd", "veh_p_cd", "code"])

# Handle BodilyInjured (BI_) records: compute max veh_cd/veh_p_cd per id
bi_agg = df.filter(F.col("code") == "BI_")\
           .groupBy("id")\
           .agg(
               F.max("veh_cd").alias("BI_veh_cd"),
               F.max("veh_p_cd").alias("BI_veh_p_cd")
           )

# Handle Collision (col_) records: compute max veh_cd/veh_p_cd per id
col_agg = df.filter(F.col("code") == "col_")\
            .groupBy("id")\
            .agg(
                F.max("veh_cd").alias("col_veh_cd"),
                F.max("veh_p_cd").alias("col_veh_p_cd")
            )

# Full outer join to combine results (preserves all ids, fills NULL where no data exists)
final_df = bi_agg.join(col_agg, on="id", how="full_outer")

# Show the result
final_df.show()

This will output exactly what you're looking for:

+----+----------+------------+-----------+-------------+
|  id|BI_veh_cd |BI_veh_p_cd |col_veh_cd |col_veh_p_cd |
+----+----------+------------+-----------+-------------+
|1001|        45|          55|         18|           21|
|1002|      null|        null|         21|           25|
+----+----------+------------+-----------+-------------+

Second: If your code values are dynamic (might add more codes later) #

This approach uses pivot to automatically handle any number of code values, so you don't have to update code for new codes.

from pyspark.sql import functions as F

# Start with the same sample df
data = [
    ("BodilyInjured", 1001, 45, 50, "BI_"),
    ("BodilyInjured", 1001, 12, 55, "BI_"),
    ("Collision", 1001, 18, 21, "col_"),
    ("Collision", 1002, 21, 25, "col_")
]
df = spark.createDataFrame(data, ["des", "id", "veh_cd", "veh_p_cd", "code"])

# Step 1: Compute max values per id + code group
agg_df = df.groupBy("id", "code")\
           .agg(
               F.max("veh_cd").alias("veh_cd"),
               F.max("veh_p_cd").alias("veh_p_cd")
           )

# Step 2: Pivot the code column to turn codes into column prefixes
pivoted_df = agg_df.groupBy("id").pivot("code").agg(
    F.first("veh_cd").alias("veh_cd"),
    F.first("veh_p_cd").alias("veh_p_cd")
)

# Step 3: Rename columns to match the required format (code + veh_cd/veh_p_cd)
final_df = pivoted_df.select(
    "id",
    *[F.col(col_name).alias(f"{col_name.split('_')[0]}_veh_cd" if "veh_cd" in col_name else f"{col_name.split('_')[0]}_veh_p_cd") 
      for col_name in pivoted_df.columns if col_name != "id"]
)

# Show the result
final_df.show()

This will produce the same output as the first method, but it's scalable—if you add a new code like PD_ later, it will automatically generate PD_veh_cd and PD_veh_p_cd columns without any code changes.

Key Notes #

• Both methods use groupBy + max to get the highest values per id and code group, which handles multiple records for the same id/code pair (like the two BI_ records for id 1001).
• The full_outer join in the first method ensures we don't lose any ids (like id 1002, which only has col_ records).

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