識別與 PySpark dataframe 中的當前值不同的最新記錄

Question

我有一個 PySpark dataframe，每個用戶在某個時間點都有一定的狀態，如下面的虛擬數據

    --------------------------
    |user_id| status| month  |
    --------------------------
    | 1     | A     | 12/2020|
    | 1     | A     | 11/2020|
    | 1     | B     | 10/2020|
    | 1     | B     | 09/2020|
    | 1     | A     | 08/2020|
    | 1     | C     | 07/2020|
    | 2     | A     | 12/2020|
    | 2     | A     | 11/2020|
    | 2     | A     | 10/2020|
    | 2     | B     | 09/2020|

我想在我的 PySpark dataframe 中創建另外兩列（previous_status_value 和 previous_status_month），對於每條記錄，它們指示用戶狀態與記錄中的狀態不同的最近日期，以及該值是多少。 使用上述虛擬數據，結果將是

    ------------------------------------------------------------------------
    |user_id| status| month  | previous_status_value| previous_status_month|
    ------------------------------------------------------------------------
    | 1     | A     | 12/2020| B                    | 10/2020              |
    | 1     | A     | 11/2020| B                    | 10/2020              |
    | 1     | B     | 10/2020| A                    | 08/2020              |
    | 1     | B     | 09/2020| A                    | 08/2020              |
    | 1     | A     | 08/2020| C                    | 07/2020              |
    | 1     | C     | 07/2020| Null                 | Null                 |
    | 2     | A     | 12/2020| B                    | 09/2020              |
    | 2     | A     | 11/2020| B                    | 09/2020              |
    | 2     | A     | 10/2020| B                    | 09/2020              |
    | 2     | B     | 09/2020| Null                 | Null                 |

dataframe 有數百萬條記錄，所以我試圖使用 Window 函數（類似於這個答案）來解決這個問題，但沒有到達那里。

Answer 1

使用lead查找狀態變化的位置，只保留與狀態變化相對應的status和month ，並使用 null 進行掩碼，否則使用when(F.col('begin'), F.col('status')) ，並獲取上一個使用F.last(..., ignorenulls=True)的值。

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

w = Window.partitionBy('user_id').orderBy('month')
begin = F.lead('status').over(w) != F.col('status')
df = df.select('*', begin.alias('begin'))

w = w.rowsBetween(Window.unboundedPreceding, -1)
previous_status_value = F.last(F.when(F.col('begin'), F.col('status')), ignorenulls=True).over(w).alias('previous_status_value')
previous_status_month = F.last(F.when(F.col('begin'), F.col('month')), ignorenulls=True).over(w).alias('previous_status_month ')

df = df.select('*', previous_status_value, previous_status_month).drop('begin').orderBy('user_id', F.col('month').desc())

df.show()
+-------+------+-------+---------------------+----------------------+
|user_id|status|  month|previous_status_value|previous_status_month |
+-------+------+-------+---------------------+----------------------+
|      1|     A|12/2020|                    B|               10/2020|
|      1|     A|11/2020|                    B|               10/2020|
|      1|     B|10/2020|                    A|               08/2020|
|      1|     B|09/2020|                    A|               08/2020|
|      1|     A|08/2020|                    C|               07/2020|
|      1|     C|07/2020|                 null|                  null|
|      2|     A|12/2020|                    B|               09/2020|
|      2|     A|11/2020|                    B|               09/2020|
|      2|     A|10/2020|                    B|               09/2020|
|      2|     B|09/2020|                 null|                  null|
+-------+------+-------+---------------------+----------------------+