Pyspark基于固定大小滑动窗口拆分多列数组并标记末尾窗口的实现问题

实现思路 #

• 首先按行计算数组长度，生成待遍历的指针序列：指针为0-based索引，范围从N到target数组长度减1，对应需求里从第N+1位开始遍历的规则
• 展开指针序列为多行，每行对应一个滑动窗口
• 基于当前指针位置，通过Spark内置数组函数截取对应区间的过往数据、提取未来target值，同时判断当前指针是否为数组最后一位

注意：Spark的slice、element_at数组函数使用1-based索引，代码中会对0-based的指针做转换处理

完整实现代码 #

from pyspark.sql import functions as F

# 定义滑动窗口大小N
N = 2

# 输入DataFrame使用你提供的示例即可
result_df = input_df \
    # 生成所有需要遍历的指针位置（0-based下标）
    .withColumn("pointer_seq", F.sequence(F.lit(N), F.size("target") - 1)) \
    # 展开指针序列得到多行记录
    .select("*", F.explode("pointer_seq").alias("cur_pointer")) \
    # 计算业务字段
    .withColumn("past-target", F.slice("target", F.col("cur_pointer") - N + 1, N)) \
    .withColumn("future-target", F.element_at("target", F.col("cur_pointer") + 1)) \
    .withColumn("past-feature1", F.slice("feature1", F.col("cur_pointer") - N + 1, N)) \
    .withColumn("past-feature2", F.slice("feature2", F.col("cur_pointer") - N + 1, N)) \
    .withColumn("islast", F.col("cur_pointer") == F.size("target") - 1) \
    # 筛选输出需要的字段
    .select("id", "past-target", "future-target", "past-feature1", "past-feature2", "islast")

执行result_df.show()即可得到你示例中的预期输出结果。

业务前置要求：同一行的target、feature1、feature2三个数组列长度保持一致。

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