如何用PySpark按多标识分组并基于checkDate筛选最新记录？

需求说明 #

给定包含vehicle_coalesce、vehicleNumber、productionNumber、pin、checkDate字段的Spark DataFrame，需按以下逻辑处理：

• 若两条记录的vehicleNumber、productionNumber或pin任一字段相同，则归为同一组
• 筛选每组中checkDate最新的所有记录

实现思路 #

由于分组逻辑是多字段关联的连通分组（同一实体可能通过不同字段关联），采用图计算中的连通分量算法识别所有关联的记录组，再对每组筛选最新时间的记录。

代码实现 #

1. 初始化环境与创建示例DataFrame #

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, max as spark_max
from graphframes import GraphFrame

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

# 构建示例数据
data = [
    ("V123", "V123", "P123", None, "27/08/2023 01:03"),
    ("P123", None, "P123", "W123", "27/08/2023 01:05"),
    ("P123", None, "P123", "W123", "27/08/2023 01:05"),
    ("V234", "V234", "P234", None, "27/08/2023 01:03"),
    ("V234", "V234", None, "W234", "27/08/2023 01:05"),
    ("V234", "V234", None, "W234", "27/08/2023 01:05"),
    ("P456", None, "P456", "W456", "27/08/2023 01:03"),
    ("v456", "V456", None, "W456", "27/08/2023 01:05"),
    ("V456", "V456", "P456", "W456", "27/08/2023 01:05")
]

schema = ["vehicle_coalesce", "vehicleNumber", "productionNumber", "pin", "checkDate"]
df = spark.createDataFrame(data, schema)

# 将checkDate转换为时间戳类型，便于时间比较
df = df.withColumn("checkDate", col("checkDate").cast("timestamp"))

2. 构建图并计算连通分量 #

通过GraphFrame将记录与关联字段构建为图，计算连通分量以识别同一组的记录：

# 为每条记录生成唯一ID（若vehicle_coalesce不唯一，建议用monotonically_increasing_id()）
df_with_id = df.withColumn("record_id", col("vehicle_coalesce"))

# 生成边：记录与非空的关联字段建立连接
edges_vehicle = df_with_id.filter(col("vehicleNumber").isNotNull())\
    .select(col("record_id").alias("src"), col("vehicleNumber").alias("dst"))
edges_production = df_with_id.filter(col("productionNumber").isNotNull())\
    .select(col("record_id").alias("src"), col("productionNumber").alias("dst"))
edges_pin = df_with_id.filter(col("pin").isNotNull())\
    .select(col("record_id").alias("src"), col("pin").alias("dst"))

# 合并所有边
edges = edges_vehicle.union(edges_production).union(edges_pin)

# 生成顶点：包含所有记录ID和关联字段值
vertices = df_with_id.select(col("record_id").alias("id"))\
    .union(edges.select(col("dst").alias("id")))\
    .distinct()

# 创建GraphFrame并计算连通分量
g = GraphFrame(vertices, edges)
connected_components = g.connectedComponents()

# 将连通分量ID关联回原DataFrame
df_with_group = df_with_id.join(
    connected_components.select(col("id").alias("record_id"), col("component")),
    on="record_id",
    how="left"
)

3. 筛选每组最新记录 #

# 计算每个连通组的最新时间
group_max_date = df_with_group.groupBy("component")\
    .agg(spark_max("checkDate").alias("max_checkDate"))

# 筛选出组内时间等于最新时间的记录
result = df_with_group.join(
    group_max_date,
    on=["component", "checkDate"],
    how="inner"
).drop("record_id", "component")

# 查看结果
result.show(truncate=False)

注意事项 #

• 若vehicle_coalesce存在重复值，需替换为monotonically_increasing_id()生成唯一记录ID，避免分组错误
• 运行代码需确保Spark环境已安装GraphFrame依赖，可通过pip install graphframes安装
• 若无法使用GraphFrame，可通过迭代合并分组的方式实现，但逻辑复杂度更高

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