PySpark перечислить в разделе при условии

Question

У меня есть PySpark dataFrame, например:

class   classID Property
1       1       1
1       2       0
1       3       1
1       4       1
2       1       0   
2       2       0
2       3       1

Теперь мне нужно добавить столбец с информацией о количестве строк в текущем разделе, пока эта строка не получит свойство == 1. Как здесь:

class   classID Property    relevantCount
1       1       1           1
1       2       0           1   
1       3       1           2
1       4       1           3
2       1       0           0   
2       2       0           0
2       3       1           1

Например, я попробовал функцию Window:

import pyspark.sql.functions as f
from pyspark.sql.window import Window
windowSpec = Window().partitionBy('class').orderBy(f.col('classID'))
df = df \
.withColumn('relevantCount',(f.when((f.col('rank') == f.lit(1)) & (f.col('Property') == f.lit(0)),0)).otherwise(f.col('Property')+f.col(f.lag('deliveryCountDesc').over(windowSpec))))

Но я не могу ссылаться на предыдущие значения новой строки.

У кого-нибудь есть лучшеидея?

Answer 1

Ваша спецификация окна должна быть изменена, чтобы включить все предыдущие строки в раздел и взять счетчик с соответствующими значениями свойств.Попробуйте это

import pyspark.sql.functions as f
from pyspark.sql.window import Window
data = [(1, 1,1),(1, 2,0),(1, 3,1),(1, 4,1),(2,1 , 0),(2, 2, 0),(2, 3, 1)]
df = spark.createDataFrame(data,['class','classID','Property'])
windowSpec = Window().partitionBy('class').orderBy('classID').rowsBetween(Window.unboundedPreceding, Window.currentRow)
df = df.withColumn('relevantCount',f.count(f.when(f.col('Property') == 1, f.col('Property'))).over(windowSpec))
df.show()

Результат>

+-----+-------+--------+-------------+
|class|classID|Property|relevantCount|
+-----+-------+--------+-------------+
|    1|      1|       1|            1|
|    1|      2|       0|            1|
|    1|      3|       1|            2|
|    1|      4|       1|            3|
|    2|      1|       0|            0|
|    2|      2|       0|            0|
|    2|      3|       1|            1|
+-----+-------+--------+-------------+