python - Pyspark 列在查找前几行和计算时生成

Question

需要通过查找以前的行值动态生成列。

目前已尝试使用下面共享的代码

Spark Dataframe 如下

cat a    b    c
1   null 0    0
1   0    9    0 
2   0    126  0
2   7    0    0
2   36   0    0
3   35   0    0

输出应该如下所示

cat a    b    c    d
1   null 0    0    null
1   0    9    0    0
2   0    126  0    9
2   7    0    0    119
2   36   0    0    83 
2   35   0    0    48

此处 d 列是动态生成的，下面是使用的代码，但无法正常工作

w=Window().partitionBy().orderBy('cat')
df=df.withColumn('d',lag("b").over(w)-df.a+df.c)
df.show()

问题是公式没有得到正确的复制 prperly 公式是 prev(b)-a+c

Answer 1

import pyspark.sql.functions as f
from pyspark.sql.window import Window

df = sc.parallelize([
    [1,None, 0, 0],
    [1,0,9,0],
    [2,0,126,0],
    [2,7,0,0],
    [2,36,0,0],
    [2,35,0,0]
]).toDF(('cat','a', 'b', 'c'))

df1 = df.withColumn("row_id", f.monotonically_increasing_id())
w = Window.partitionBy().orderBy(f.col("row_id"))
df1 = df1.withColumn("d", f.lag("b").over(w)- f.col("a") + f.col("c")).drop("row_id")
df1.show()

输出是:

+---+----+---+---+----+
|cat|   a|  b|  c|   d|
+---+----+---+---+----+
|  1|null|  0|  0|null|
|  1|   0|  9|  0|   0|
|  2|   0|126|  0|   9|
|  2|   7|  0|  0| 119|
|  2|  36|  0|  0| -36|
|  2|  35|  0|  0| -35|
+---+----+---+---+----+