Pyspark 将数组列分解为带有滑动窗口的子列表

Question

我有一个 PySpark 行，我想根据列中的一个值将其分成更小的行。

给定一个df：

input_df = spark.createDataFrame([
    (2,[1,2,3,4,5],),
    ], ("id", "list"))

+---+------------+
| id|        list|
+---+------------+
|  2|[1, 2, 3, 4]|
+---+------------+

我想用固定大小的滑动窗口将每一行分成多个子集。结果 df 将是这样的：

output_df = spark.createDataFrame([
    (2, [0,0], 1), (2, [0,1], 2), (2, [1,2], 3), (2, [2,3], 4), (2, [3,4], 5),
    ], ("id", "past", "future"))

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

关于如何打破列表以使指针查看列表的每个元素的逻辑，使用前面的 N 个元素（在这种情况下 N=2）作为过去（如果没有足够的元素，则用 0 填充）并使用当前指针元素作为未来。对每个元素执行此操作会创建数据框。

我想不出用 pyspark 来做这件事的方法，我会用 pandas 数据框为每一行做一个迭代循环。有没有办法用 pyspark 做到这一点？

Answer 1

处理数组列的最佳方法是使用higher-order functions。

import pyspark.sql.functions as f

output_df = (input_df
             .withColumn('list', f.expr('TRANSFORM(list, (element, i) -> STRUCT(ARRAY(COALESCE(list[i - 2], 0), COALESCE(list[i - 1], 0)) AS past, element AS future))'))
             .selectExpr('id', 'inline(list)'))

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

更新

在创建表达式之前动态传递N：

N = 2

expr = 'TRANSFORM(list, (element, i) -> STRUCT(TRANSFORM(sequence({N}, 1), k -> COALESCE(list[i - k], 0)) AS past, element AS future))'.format(N=N)
output_df = (input_df
             .withColumn('list', f.expr(expr))
             .selectExpr('id', 'inline(list)'))