如何基于条件创建PySpark Window并按批次分组数据？

针对你描述的场景，可以通过窗口累加生成批次ID的方式实现动态分区，无需依赖固定长度的rowsBetween。以下是具体实现步骤和代码：

核心思路 #

为每个asset下的记录生成唯一批次标识：每次遇到seq_start=True时，批次ID递增，同一批次（从当前True到下一个True前的所有记录）会共享同一个批次ID。之后就可以基于asset和批次ID进行分组聚合。

实现步骤（以PySpark为例） #

1. 定义分区排序窗口：按asset分区，timestamp升序排序，窗口范围覆盖从当前分区起始行到当前行。
2. 生成批次ID：用when标记seq_start=True的行为1，其余为0，再通过窗口内的sum累加该标记值，得到每个行的批次ID。
3. 分组聚合：基于asset和批次ID分组，对reading执行所需的聚合操作。

代码示例 #

from pyspark.sql.window import Window
from pyspark.sql.functions import col, when, sum as spark_sum

# 确保seq_start为布尔类型（null会被自动视为False处理）
df = df.withColumn("seq_start", col("seq_start").cast("boolean"))

# 定义窗口：按asset分区，timestamp升序，窗口范围从分区开头到当前行
window_spec = Window.partitionBy("asset").orderBy("timestamp").rowsBetween(Window.unboundedPreceding, Window.currentRow)

# 生成批次ID：累加seq_start=True的次数
df_with_batch = df.withColumn(
    "batch_id",
    spark_sum(when(col("seq_start") == True, 1).otherwise(0)).over(window_spec)
)

# 按asset和batch_id分组，聚合reading（示例为求和，可替换为avg、max等）
aggregated_result = df_with_batch.groupBy("asset", "batch_id").agg(
    spark_sum("reading").alias("total_reading"),
    # avg("reading").alias("average_reading")  # 可添加其他聚合逻辑
)

# 查看结果
aggregated_result.show()

逻辑说明 #

• 窗口的rowsBetween(Window.unboundedPreceding, Window.currentRow)确保累加计算从当前asset的第一条记录开始，到当前行结束，保证批次ID的连续性。
• 当seq_start=True时，累加值+1，后续行继承该值，直到下一个True出现，自动完成动态批次划分。
• 即使两个True连续（中间无null），每个True所在行也会成为独立批次，符合需求。

注意事项 #

• 必须保证timestamp的排序正确性，这是批次划分的时间基准。
• 若seq_start的null值无需转换，直接使用即可，因为when会将null视为0，不影响累加逻辑。

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