PySpark - 對於每個唯一 ID 和列條件設置值為 1

Question

您好 - 當 PurchasePrice > 0 時，我正在嘗試在 dataframe 中為 ID 的第一個實例分配 1。對於以下實例或在此之前的實例，列的值將為 0。例如，在下面的屏幕截圖中，ID 123，'Wanted Column' 將在 MonYer = 201909 中設置為 1，因為這是 PurchasePrice > 0 時的第一個實例。 對於201911的下一次觀察，該值將是0。我想過使用.groupBy或rank（），dense_rank（），但實在想不出如何做到這一點。

任何形式的指導或幫助表示贊賞！

Answer 1

您可以將sum與window結合使用。 在 window 中，您只需匯總前幾行的價格。 使用結果列，您可以檢查記錄是否是第一個非零條目。 前面行的總和應該為零，記錄本身的價格非零。

我創建了一個與您的略有不同的示例數據集，但它應該向您展示該方法。

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

spark = sql.SparkSession.builder.master('local').getOrCreate()

df = spark.createDataFrame([[123,201902,0],[123,201903,0],[123,201904,100],[123,201905,100],[123,201906,0]], ['ID', 'MonYer', 'Price'])

w = Window.partitionBy('ID').orderBy('MonYer').rangeBetween(Window.unboundedPreceding, -1)

df = (df
    .withColumn('sum', f.sum('Price').over(w) )
    .withColumn('wanted', f.when((f.col('Price') > 0) & (f.col('Sum') == 0), 1).otherwise(0))
    .drop('sum')
)

df.show()

+---+------+-----+------+                                                       
| ID|MonYer|Price|wanted|
+---+------+-----+------+
|123|201902|    0|     0|
|123|201903|    0|     0|
|123|201904|  100|     1|
|123|201905|  100|     0|
|123|201906|    0|     0|
+---+------+-----+------+

注意：此解決方案假定所有價格值都 >= 0