Pyspark 跨行保留价值?

Pyspark retain value across rows?

我有一个问题,使用逐行 SAS 方法自然可以解决,但我无法使用 Pyspark。我有一个按时间排序的人的事件数据集,例如:

test_df = pd.DataFrame({'event_list':[["H"], ["H"], ["H","F"], ["F"], ["F"], ["H"], ["W"], ["W"]], 'time_order':[1,2,3,4,5,6,7,8], 'person':[1,1,1,1,1,1,1,1]})
test_df = spark.createDataFrame(test_df)
test_df.show()

+----------+----------+------+
|event_list|time_order|person|
+----------+----------+------+
|       [H]|         1|     1|
|       [H]|         2|     1|
|    [H, F]|         3|     1|
|       [F]|         4|     1|
|       [F]|         5|     1|
|       [H]|         6|     1|
|       [W]|         7|     1|
|       [W]|         8|     1|
+----------+----------+------+

我想将这些事件分组到剧集中,其中初始事件之后的所有事件都是初始事件列表的一部分。因此,在我的 test_df 中,我希望有 3 集:

+----------+----------+------+-------+
|event_list|time_order|person|episode|
+----------+----------+------+-------+
|       [H]|         1|     1|      1|
|       [H]|         2|     1|      1|
|    [H, F]|         3|     1|      2|
|       [F]|         4|     1|      2|
|       [F]|         5|     1|      2|
|       [H]|         6|     1|      2|
|       [W]|         7|     1|      3|
|       [W]|         8|     1|      3|
+----------+----------+------+-------+

在 SAS 中,我会保留前一行的值 event_list,如果当前 event_list 包含在前 event_list 中,我会保留当前 event_list ] 值而不是之前的 event_list。例如。我保留的值是 [null, ["H"], ["H"], ["H","F"], ["H","F"], ["H","F"], [“W”]]。然后我可以通过跟踪保留值的变化来生成剧集。

在 Pyspark 中,我不确定如何跨行操作按顺序保留信息...这甚至可能吗?我尝试使用 window 函数(按 person 分区并按 time_order 排序)失败了。我如何在 Pyspark 中解决这个问题?

如果您使用的 spark 版本 >= 2.4,请在 event_list 列上使用 collect_list window、flatten他们,使用 array_distinct 删除重复项,最后使用 size 来计算一段时间内有多少不同的事件。应该是这样的:

from pyspark.sql.functions import col, collect_list, flatten, array_distinct, size
from pyspark.sql.window import Window

w = Window.partitionBy('person').orderBy('time_order').rowsBetween(Window.unboundedPreceding, 0)

test_df = test_df.withColumn('episode', size(array_distinct(flatten(collect_list(col('event_list')).over(w)))))
test_df.show()

+----------+----------+------+-------+
|event_list|time_order|person|episode|
+----------+----------+------+-------+
|       [H]|         1|     1|      1|
|       [H]|         2|     1|      1|
|    [H, F]|         3|     1|      2|
|       [F]|         4|     1|      2|
|       [F]|         5|     1|      2|
|       [H]|         6|     1|      2|
|       [W]|         7|     1|      3|
|       [W]|         8|     1|      3|
+----------+----------+------+-------+