Pyspark 使用sparksql进行条件滞后求和

下面是我的数据框架

df = sqlContext.createDataFrame(
    [("0", "0"), ("1", "2"), ("2", "3"), ("3", "4"), ("4", "0"), ("5", "5"), ("6", "5")],
    ["id", "value"])

+---+-----+
| id|value|
+---+-----+
|  0|    0|
|  1|    2|
|  2|    3|
|  3|    4|
|  4|    0|
|  5|    5|
|  6|    5|
+---+-----+

我想得到的是：

+---+-----+---+-----+
| id|value|masterid|partsum|
+---+-----|---+-----+
|  0|    0|  0|    0| 
|  1|    2|  0|    2|
|  2|    3|  0|    5|
|  3|    4|  0|    9|
|  4|    0|  4|    0|
|  5|    5|  4|    5|
|  6|    5|  4|   10|
+---+-----+---+-----+

因此，我尝试使用SparkSQL来实现这一点：

df=df.withColumn("masterid", F.when( df.value !=0 , F.lag(df.id)).otherwise(df.id))

我原以为lag函数可以帮助我在下一次迭代之前进行处理，以获得masterid col。不幸的是，在我查看手册后，它帮不上忙

---

所以，我想问一下，是否有任何特殊的功能，我可以用来做我想做的？或者我可以使用“条件滞后”函数吗？所以，当我看到非零项时，我可以使用滞后，直到找到一个零数

IIUC，您可以尝试定义一个子组标签（

g

，在下面的代码中）和两个窗口规格：

from pyspark.sql import Window, functions as F

w1 = Window.orderBy('id')
w2 = Window.partitionBy('g').orderBy('id')

df.withColumn('g', F.sum(F.expr('if(value=0,1,0)')).over(w1)).select(
    'id'
  , 'value'
  , F.first('id').over(w2).alias('masterid')
  , F.sum('value').over(w2).alias('partsum')
).show()
#+---+-----+--------+-------+
#| id|value|masterid|partsum|
#+---+-----+--------+-------+
#|  0|    0|       0|    0.0|
#|  1|    2|       0|    2.0|
#|  2|    3|       0|    5.0|
#|  3|    4|       0|    9.0|
#|  4|    0|       4|    0.0|
#|  5|    5|       4|    5.0|
#|  6|    5|       4|   10.0|
#+---+-----+--------+-------+