PySpark实现带填充的行级Sliding Window构建问题

问题说明 #

需要对分组后的行生成带填充的滑动窗口，将窗口内的行收集为向量形式。当前代码仅能生成无填充的滚动窗口，尝试过ROWS N PRECEDING等SQL方法未成功，已明确填充规则，需实现滑动窗口的行组织逻辑。

输入示例 #

+---+-----+-----+
| id|Label|group|
+---+-----+-----+
|  A|    T|    1|
|  B|    T|    1|
|  C|    F|    2|
|  D|    F|    2|
|  E|    F|    3|
|  F|    T|    3|
|  G|    F|    3|
|  H|    T|    3|
+---+-----+-----+

预期输出（参数配置） #

windows_size = 3
stride = 1
id_padding = ''
label_padding = 'f'

+-----+-------------+-------------+
|group|      Windows|       Labels|
+-----+-------------+-------------+
|    1|   [A, B, '']|    [T, T, f]|
|    2|   [C, D, '']|    [F, F, f]|
|    3|    [E, F, G]|    [F, T, F]|
|    3|    [F, G, H]|    [T, F, T]|
+-----+-------------+-------------+

当前代码（仅支持滚动窗口） #

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

data = [
    ("A", "T", 1),
    ("B", "T", 1),
    ("C", "F", 2),
    ("D", "F", 2),
    ("E", "F", 3),
    ("F", "T", 3),
    ("G", "F", 3),
    ("H", "T", 3),
]
df = spark.createDataFrame(data, ['id', 'label', 'group'])

grouping = 3

w2 = Window.partitionBy('group').orderBy('id')

df = df.withColumn("rows",((F.row_number().over(w2)-1) / grouping).astype('int') )
df.groupBy('group', 'rows')\
  .agg(F.collect_list('id').alias("Windows"), F.collect_list('Label').alias("Labels"))\
  .drop('rows') \
  .orderBy('group').show()

解决方案代码 #

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

# 定义滑动窗口参数
window_size = 3
stride = 1
id_pad = ''
label_pad = 'f'

# 初始化数据
data = [
    ("A", "T", 1),
    ("B", "T", 1),
    ("C", "F", 2),
    ("D", "F", 2),
    ("E", "F", 3),
    ("F", "T", 3),
    ("G", "F", 3),
    ("H", "T", 3),
]
df = spark.createDataFrame(data, ['id', 'label', 'group'])

# 步骤1：给每个组内的行分配行号
w = Window.partitionBy('group').orderBy('id')
df = df.withColumn('rn', F.row_number().over(w))

# 步骤2：计算每个组的总行数，生成所有滑动窗口的起始行索引
group_counts = df.groupBy('group').agg(F.max('rn').alias('total_rows')).cache()
df = df.join(group_counts, on='group')

# 生成每个窗口的起始行（从1开始，按步长递增）
df = df.withColumn('start_rn', F.explode(F.sequence(F.lit(1), F.col('total_rows'), F.lit(stride))))

# 步骤3：标记当前行是否属于对应窗口，然后收集并填充窗口数据
windowed_df = df.groupBy('group', 'start_rn')\
    .agg(
        # 收集窗口内的id，用填充值补全到窗口大小
        F.array_pad(
            F.collect_list(F.when(F.col('rn').between(F.col('start_rn'), F.col('start_rn') + window_size -1), F.col('id'))),
            window_size,
            id_pad
        ).alias('Windows'),
        # 收集窗口内的label，用填充值补全到窗口大小
        F.array_pad(
            F.collect_list(F.when(F.col('rn').between(F.col('start_rn'), F.col('start_rn') + window_size -1), F.col('label'))),
            window_size,
            label_pad
        ).alias('Labels')
    )\
    .orderBy('group', 'start_rn')\
    .drop('start_rn')

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

代码说明 #

1. 行号分配：通过窗口函数给每个组内的行按id排序后分配行号，用于确定窗口范围。
2. 窗口起始行生成：利用sequence函数生成每个组内所有滑动窗口的起始行索引，步长由参数控制。
3. 窗口数据收集与填充：通过between判断行是否属于当前窗口，用collect_list收集有效数据，最后用array_pad将数组补全到指定窗口大小，填充缺失值。

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