PySpark如何将数组列通过固定大小滑动窗口拆分为多行子列表

实现方案 #

可以完全用PySpark内置函数实现需求，无需转pandas处理也不需要自定义UDF，性能更高，适配分布式场景。

核心思路

1. 提前给原数组头部拼接N个0，解决前置元素不足补0的问题
2. 用posexplode将数组炸开为多行，同时拿到每个元素的索引和取值（即future字段）
3. 按照索引从补0后的数组中切片取N个元素，即可得到past字段

完整实现代码

from pyspark.sql import SparkSession
from pyspark.sql.functions import posexplode, concat, array_repeat, slice, col

# 构造测试数据
spark = SparkSession.builder.appName("array_split").getOrCreate()
input_df = spark.createDataFrame([
    (2,[1,2,3,4,5],),
    ], ("id", "list"))

# 定义past字段的长度
N = 2

result_df = input_df \
    .withColumn("padded_list", concat(array_repeat(0, N), col("list"))) \
    .select("id", "padded_list", posexplode("list").alias("pos", "future")) \
    .withColumn("past", slice(col("padded_list"), col("pos") + 1, N)) \
    .select("id", "past", "future")

# 查看结果
result_df.show()

输出结果和你要求的完全一致：

+---+------+------+
| id|  past|future|
+---+------+------+
|  2|[0, 0]|     1|
|  2|[0, 1]|     2|
|  2|[1, 2]|     3|
|  2|[2, 3]|     4|
|  2|[3, 4]|     5|
+---+------+------+

如果需要调整past的长度，直接修改N的取值即可。

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