PySpark聚合代码优化咨询：多列组最新值高效提取方案

原始DataFrame #

+---+----------------------------------+----------+----------+
|id |timestamp                         |Fname     |Lname     |
+---+----------------------------------+----------+----------+
|1  |2024-01-19T11:52:44.775205Z       |Robert    |Albert    |
|1  |2024-01-20T11:52:44.775205Z       |Remo      |Lergos    |
|2  |2024-01-21T11:52:44.775205Z       |Charlie   |Jameson   |
|2  |2024-01-22T11:52:44.775205Z       |Anastacio |Sporer    |
|2  |2024-01-23T11:52:44.775205Z       |Luz       |Toy       |
|3  |2024-01-24T11:52:44.775205Z       |Crystal   |Hills     |
|3  |2024-01-25T11:52:44.775205Z       |Nicholas  |Johnson   |
+---+----------------------------------+----------+----------+

处理需求 #

• 按id字段分组
• 将所有姓名信息（Fname、Lname）收集为字典列表
• 根据每个id的最新timestamp提取对应姓名，存储为单独的字典列
• 将步骤3中使用的最新timestamp存储为单独列

期望结果DataFrame #

+----+--------------------------------------+----------------------------+-----------------------------------------------------------------------------------------------------+
|id  |latest_names                          |latest_timestamp            |all_names                                                                                            |
+----+--------------------------------------+----------------------------+-----------------------------------------------------------------------------------------------------+
|1   |{"Fname":"Remo","Lname":"Lergos"}     |2024-01-20T11:52:44.775205Z |[{"Fname":"Remo","Lname":"Lergos"},{"Fname":"Remo","Lname":"Lergos"}]                                |
|2   |{"Fname":"Luz","Lname":"Toy"}         |2024-01-23T11:52:44.775205Z |[{"Fname":"Luz","Lname":"Toy"},{"Fname":"Remo","Lname":"Lergos"},{"Fname":"Remo","Lname":"Lergos"}]  |     
|3   |{"Fname":"Nicholas","Lname":"Johnson"}|2024-01-25T11:52:44.775205Z |[{"Fname":"Nicholas","Lname":"Johnson"},{"Fname":"Remo","Lname":"Lergos"}]                           |
+----+--------------------------------------+----------------------------+-----------------------------------------------------------------------------------------------------+

现有实现代码 #

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

windowspec = Window.partitionBy("id").orderBy(df["timestamp"].desc())
columns_names = ["Fname","Lname"]
df.withColumn(
    "all_names",
    F.to_json(
        F.struct(
            "Fname","Lname"
        )
    ),
)
.withColumn(
    "latest_names",
    F.to_json(
        F.struct(*[F.first(field).over(windowspec).alias(field) for field in columns_names])
    ),
)
.withColumn("latest_timestamp", F.first("timestamp").over(windowspec).alias("timestamp"))
.groupBy("id")
.agg(
    F.collect_set("all_names").alias("all_names"),
    F.first("latest_names").alias("latest_names"),
    F.first("l_timestamp").alias("latest_timestamp"),  # 注：原代码存在笔误，应为latest_timestamp
)

当前疑问 #

现在需要处理多组类似列（如address1、address2、address3），每组都要执行「提取最新timestamp对应值为字典列」的操作，目前用单个WindowSpec处理，想知道有没有更优的实现方式？

针对多组列的场景，可以通过封装逻辑、批量处理列组减少重复代码，同时保持Spark执行效率，以下是几种可行的优化思路：

1. 封装列组处理逻辑，批量生成所需列 #

把每组列的处理逻辑写成可复用逻辑，传入列组名称和输出列名，批量生成latest_xxx列，避免重复编写struct和first函数代码。

示例代码：

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

# 定义窗口（仅创建一次，全局复用）
win_spec = Window.partitionBy("id").orderBy(F.col("timestamp").desc())

# 配置列组：键为输出列前缀，值为对应列集合
column_groups = {
    "names": ["Fname", "Lname"],
    "address": ["address1", "address2", "address3"]
}

# 第一步：生成all_xxx列（收集所有行的列组字典）
processed_df = df
for output_col, cols in column_groups.items():
    processed_df = processed_df.withColumn(
        f"all_{output_col}",
        F.to_json(F.struct(*cols))
    )

# 第二步：生成latest_timestamp和各列组的latest_xxx列
processed_df = processed_df.withColumn(
    "latest_timestamp",
    F.first("timestamp").over(win_spec)
)

for output_col, cols in column_groups.items():
    # 生成该列组的最新字段struct，再转成json
    latest_struct = F.struct(*[F.first(c).over(win_spec).alias(c) for c in cols])
    processed_df = processed_df.withColumn(
        f"latest_{output_col}",
        F.to_json(latest_struct)
    )

# 第三步：分组聚合
final_df = processed_df.groupBy("id").agg(
    # 聚合all_xxx列
    *[F.collect_set(f"all_{col}").alias(f"all_{col}") for col in column_groups.keys()],
    # 聚合latest_xxx和latest_timestamp（同一id下值重复，用first取任意一个即可）
    *[F.first(f"latest_{col}").alias(f"latest_{col}") for col in column_groups.keys()],
    F.first("latest_timestamp").alias("latest_timestamp")
)

2. 先获取最新记录再关联聚合结果（大数据量场景优先） #

这种方式避免在窗口中重复计算first，通过groupBy+max(timestamp)+join获取最新记录，减少窗口函数计算量，适合数据量较大的场景：

# 第一步：获取每个id的最新记录
latest_records_df = df.groupBy("id").agg(
    F.max("timestamp").alias("latest_timestamp")
).join(df, on=["id", "timestamp"], how="inner")

# 第二步：处理最新记录的列组字典
for output_col, cols in column_groups.items():
    latest_records_df = latest_records_df.withColumn(
        f"latest_{output_col}",
        F.to_json(F.struct(*cols))
    )

# 第三步：处理原始数据的all_xxx列聚合
all_data_agg = df
for output_col, cols in column_groups.items():
    all_data_agg = all_data_agg.withColumn(
        f"all_{output_col}",
        F.to_json(F.struct(*cols))
    )
all_data_agg = all_data_agg.groupBy("id").agg(
    *[F.collect_set(f"all_{col}").alias(f"all_{col}") for col in column_groups.keys()]
)

# 第四步：关联最新记录和聚合结果
final_df = all_data_agg.join(
    latest_records_df.select("id", "latest_timestamp", *[f"latest_{col}" for col in column_groups.keys()]),
    on="id",
    how="inner"
)

3. 用struct统一打包列组（简化代码） #

如果列组命名有规律，可以先把同一组的列打包成struct，再统一处理，减少代码量：

# 先把列组转成struct
processed_df = df.withColumn("names_struct", F.struct("Fname", "Lname")) \
                 .withColumn("address_struct", F.struct("address1", "address2", "address3"))

# 生成all_xxx列
processed_df = processed_df.withColumn("all_names", F.to_json("names_struct")) \
                           .withColumn("all_address", F.to_json("address_struct"))

# 窗口处理最新值
win_spec = Window.partitionBy("id").orderBy(F.col("timestamp").desc())
processed_df = processed_df.withColumn("latest_timestamp", F.first("timestamp").over(win_spec)) \
                           .withColumn("latest_names", F.to_json(F.first("names_struct").over(win_spec))) \
                           .withColumn("latest_address", F.to_json(F.first("address_struct").over(win_spec)))

# 聚合
final_df = processed_df.groupBy("id").agg(
    F.collect_set("all_names").alias("all_names"),
    F.collect_set("all_address").alias("all_address"),
    F.first("latest_names").alias("latest_names"),
    F.first("latest_address").alias("latest_address"),
    F.first("latest_timestamp").alias("latest_timestamp")
)

关键优化点总结 #

• 复用WindowSpec：仅创建一次窗口，避免重复定义带来的冗余
• 批量处理列组：用循环封装逻辑，减少重复代码，便于后续新增列组维护
• 减少窗口计算量：大数据量场景优先用groupBy+max(timestamp)+join获取最新记录，替代多次调用first().over(win)
• 规避代码笔误：批量处理模式可以避免原代码中l_timestamp这类拼写错误

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