如何在PySpark中按多字段分组分区并获取最新记录

原始数据 #

vehicle_coalesce  vehicleNumber  productionNumber pin  checkDate
V123              V123           P123             null 27/08/2023 01:03
P123              null           P123             W123 27/08/2023 01:05
P123              null           P123             W123 27/08/2023 01:05
V234              V234           P234             null 27/08/2023 01:03
V234              V234           null             W234 27/08/2023 01:05
V234              V234           null             W234 27/08/2023 01:05
P456              null           P456             W456 27/08/2023 01:03
v456              V456           null             W456 27/08/2023 01:05
V456              V456           P456             W456 27/08/2023 01:05

需求说明 #

需要将通过vehicleNumber、productionNumber或pin关联的记录归为同一组，筛选每组中checkDate最新的记录：

• V123与P123通过productionNumber关联，取该组最新的01:05记录
• V234的多条记录通过vehicleNumber关联，取该组最新的01:05记录
• P456、v456、V456通过pin关联，取该组最新的01:05记录

预期输出 #

vehicle_coalesce  vehicleNumber  productionNumber pin  checkDate
P123              null           P123             W123 2023-08-27 01:05:00
P123              null           P123             W123 2023-08-27 01:05:00
V234              V234           null             W234 2023-08-27 01:05:00
V234              V234           null             W234 2023-08-27 01:05:00
v456              V456           null             W456 2023-08-27 01:05:00
V456              V456           P456             W456 2023-08-27 01:05:00

实现代码 #

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

# 1. 将字符串类型的checkDate转换为时间戳，方便日期比较
df = df.withColumn("checkDate", F.to_timestamp("checkDate", "dd/MM/yyyy HH:mm"))

# 2. 定义窗口，传播各分组下的关联字段，确保同一实体的记录能获取到所有关联标识
# 按productionNumber分组，传播该组内的非空vehicleNumber和pin
window_prod = Window.partitionBy("productionNumber")
df = df.withColumn("prod_veh", F.first(F.col("vehicleNumber"), ignorenulls=True).over(window_prod))
df = df.withColumn("prod_pin", F.first(F.col("pin"), ignorenulls=True).over(window_prod))

# 按vehicleNumber分组，传播该组内的非空productionNumber和pin
window_veh = Window.partitionBy("vehicleNumber")
df = df.withColumn("veh_prod", F.first(F.col("productionNumber"), ignorenulls=True).over(window_veh))
df = df.withColumn("veh_pin", F.first(F.col("pin"), ignorenulls=True).over(window_veh))

# 按pin分组，传播该组内的非空productionNumber和vehicleNumber
window_pin = Window.partitionBy("pin")
df = df.withColumn("pin_prod", F.first(F.col("productionNumber"), ignorenulls=True).over(window_pin))
df = df.withColumn("pin_veh", F.first(F.col("vehicleNumber"), ignorenulls=True).over(window_pin))

# 3. 生成全局分组键：合并所有关联标识，生成唯一哈希值确保同一实体归为一组
df = df.withColumn(
    "group_key",
    F.sha2(
        F.concat_ws(
            ",",
            F.coalesce("productionNumber", "veh_prod", "pin_prod"),
            F.coalesce("vehicleNumber", "prod_veh", "pin_veh"),
            F.coalesce("pin", "prod_pin", "veh_pin")
        ),
        256
    )
)

# 4. 按分组键计算每组的最新日期，筛选符合条件的记录
window_max_date = Window.partitionBy("group_key")
df = df.withColumn("max_checkDate", F.max("checkDate").over(window_max_date))
result_df = df.filter(F.col("checkDate") == F.col("max_checkDate"))

# 清理临时字段，输出结果
result_df = result_df.drop("group_key", "max_checkDate", "prod_veh", "prod_pin", "veh_prod", "veh_pin", "pin_prod", "pin_veh")
result_df.show(truncate=False)

代码说明 #

• 先转换checkDate为时间戳类型，避免字符串比较的误差
• 通过三次窗口函数，分别按三个关联字段分组，将组内的非空标识传播到每条记录，确保同一实体的记录能获取到所有关联信息
• 合并所有关联标识生成唯一分组键，保证关联记录归为同一组
• 按分组键计算每组的最新日期，筛选出最新记录后清理临时字段

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