如何在PySpark DataFrame中按每3周分组并聚合为元组？

要实现你需要的按每3周分组并聚合为元组的效果，可以通过窗口函数计算分组编号 + 分组聚合的方式完成，具体步骤和代码如下：

核心思路 #

1. 排序数据：确保每个player_id下的记录按week_num升序排列，保证分组的时间连续性
2. 生成分组编号：对每个player_id内部，用行号计算每3条记录为一组的编号
3. 聚合为元组：按player_id和分组编号聚合，将moves和week_num收集为有序列表，再转换为目标格式的元组字符串

完整代码示例 #

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

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

# 构建示例DataFrame
data = [
    (1, "2022-06-13", 1, 24),
    (1, "2022-06-06", 20, 23),
    (1, "2022-06-20", 0, 25),
    (2, "2022-06-06", 20, 23),
    (2, "2022-06-13", 0, 24),
    (2, "2022-06-20", 0, 25),
    (1, "2022-05-30", 10, 22),
    (1, "2022-05-23", 20, 21),
    (1, "2022-05-16", 20, 20)
]
df = spark.createDataFrame(data, ["player_id", "stat_week", "moves", "week_num"])

# 1. 按player_id和week_num升序排序，保证时间顺序
sorted_df = df.orderBy("player_id", "week_num")

# 2. 定义窗口：按player_id分区，按week_num排序
window_spec = Window.partitionBy("player_id").orderBy("week_num")

# 计算分组编号：每3条记录为一组，组号从1开始
grouped_df = sorted_df.withColumn(
    "group_by",
    ((F.row_number().over(window_spec) - 1) // 3) + 1
)

# 3. 聚合生成目标元组格式
result_df = grouped_df.groupBy("player_id", "group_by")\
    .agg(
        # 收集moves并拼接为元组字符串
        F.concat(F.lit("("), F.array_join(F.collect_list("moves"), ","), F.lit(")")).alias("moves"),
        # 收集week_num并拼接为元组字符串
        F.concat(F.lit("("), F.array_join(F.collect_list("week_num"), ","), F.lit(")")).alias("week_num")
    )\
    .orderBy("player_id", "group_by")

# 展示结果
result_df.show(truncate=False)

运行结果 #

+---------+--------+----------+----------+
|player_id|group_by|moves     |week_num  |
+---------+--------+----------+----------+
|1        |1       |(20,20,10)|(20,21,22)|
|1        |2       |(20,1,0)  |(23,24,25)|
|2        |1       |(20,0,0)  |(23,24,25)|
+---------+--------+----------+----------+

关键细节 #

• 窗口函数row_number()用于给每个player_id下的记录按时间顺序编号，通过(row_number-1)//3实现每3条一组的分组逻辑
• collect_list会严格保留窗口排序后的顺序，确保元组内的数值和周数是按时间先后排列的
• 如果需要原生Python tuple类型而非字符串格式，可以替换为F.array(F.collect_list(...))，但展示时会以数组形式呈现

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