具有动态滞后的窗函数

Question

我正在查看 Spark SQL 中 Spark DataFrame 的窗口滑动功能。

我有一个列的数据帧id，month和volume。

id       month   volume new_col

1        201601  100     0

1        201602  120   100

1        201603  450   220

1        201604  200   670

1        201605  121   870

现在我想添加一个名为 name 的新列new_col，其值new_col是当前行volume和new_col之前的总和，如上所示。new_col第一行的值将为零。

我尝试了以下选项来使用lagPySpark使用窗口函数。但是我发现该new_col列会被递归使用。仅使用lag函数的方式无法做到这一点：

window = Window.partitionBy(F.col('id')).orderBy(F.col('month').asc())

df.withColumn('new_col', F.lag(col('volume'), 1).over(window) + F.lag(col('new_col'), 1).over(window))

有没有办法new_col通过使用窗口函数来动态滞后？或者有其他好的解决方案吗？

Answer 1

您可以在窗口上使用lag和sum来实现这一点。如果在窗口上使用，sum将自动计算cumsum。下面的代码将首先滞后volume列，然后取其 cumsum，但也可以以相反的顺序进行操作。window = Window.partitionBy(F.col('id')).orderBy(F.col('month').asc())df.withColumn('new_col', F.sum(F.lag(col('volume'), 1, 0).over(window)).over(window))