Pyspark windows 函数不适用于整个数据帧

Question

在有条件的情况下将滞后/Windows 函数应用于整个数据帧时遇到问题。

我想从第 2 周开始将前一行的值 (value1) 与当前行的值 (value2) 相乘。

这是我的数据：

'''

from pyspark.sql import functions as f

data = [
    (1, 1, 1),
    (2, 0, 5),
    (3, 0, 10),
    (4, 0, 20),
    (5, 0, 30),
    (6, 0, 40)
]
columns = ["week", "value1", "value2"]
df = spark.createDataFrame(data, columns)

'''

这是我进行计算的逻辑：

'''

w=Window.orderBy("week")
                          
df2 = df.withColumn('value1',
                    f.when((f.col('week') > 1),
                           f.lag(df['value1']).over(w) * df['value2'] 
                          ).otherwise(
                        f.col('value1')
                    )
                   )

'''

我的输出如下所示：

您可以看到只有第 2 周遵循了逻辑。

第 3 周应该是 50 (5 * 10)，而不是零...

有人可以帮我吗？

Answer 1

查看此解决方案：

w=Window.orderBy("week")
df = (
    df
    .withColumn('value1', f.when(f.col('week') > 1, f.product(f.col('value2').cast(DoubleType())).over(w)).otherwise(f.col('value1')))
)
df.show()

输出：

+----+---------+------+                                                         
|week|   value1|value2|
+----+---------+------+
|   1|      1.0|     1|
|   2|      5.0|     5|
|   3|     50.0|    10|
|   4|   1000.0|    20|
|   5|  30000.0|    30|
|   6|1200000.0|    40|
+----+---------+------+