Apache spark PySpark-在值交叉上增加秩

我有一个这样的数据帧

data = [(("ID1", 'False', '2021-05-01', 1)), 
    (("ID1", 'False', '2021-05-02', 1)), 
    (("ID1", 'True', '2021-05-03', 2)), 
    (("ID1", 'False', '2021-05-04', 3)), 
    (("ID1", 'False', '2021-05-05', 3)),
   (("ID2", 'True', '2021-05-01', 1)), 
    (("ID2", 'True', '2021-05-02', 1)), 
    (("ID2", 'False', '2021-05-03', 2)), 
    (("ID2", 'True', '2021-05-04', 3))]
df = spark.createDataFrame(data, ["ID", "colA", 'Date', 'Rank'])
df.show()

+---+-----+----------+
| ID| colA|      Date|
+---+-----+----------+
|ID1|False|2021-05-01|
|ID1|False|2021-05-02|
|ID1| True|2021-05-03|
|ID1|False|2021-05-04|
|ID1|False|2021-05-05|
|ID2| True|2021-05-01|
|ID2| True|2021-05-02|
|ID2|False|2021-05-03|
|ID2| True|2021-05-04|
+---+-----+----------+

+---+-----+----------+----+
| ID| colA|      Date|Rank|
+---+-----+----------+----+
|ID1|False|2021-05-01|   1|
|ID1|False|2021-05-02|   1|
|ID1| True|2021-05-03|   2|
|ID1|False|2021-05-04|   3|
|ID1|False|2021-05-05|   3|
|ID2| True|2021-05-01|   1|
|ID2| True|2021-05-02|   1|
|ID2|False|2021-05-03|   2|
|ID2| True|2021-05-04|   3|
+---+-----+----------+----+

我想为每个用户添加一个新的列

rank

，它将从1开始，并且只在colA中的交叉处增加值

因此，新的数据帧是这样的

data = [(("ID1", 'False', '2021-05-01', 1)), 
    (("ID1", 'False', '2021-05-02', 1)), 
    (("ID1", 'True', '2021-05-03', 2)), 
    (("ID1", 'False', '2021-05-04', 3)), 
    (("ID1", 'False', '2021-05-05', 3)),
   (("ID2", 'True', '2021-05-01', 1)), 
    (("ID2", 'True', '2021-05-02', 1)), 
    (("ID2", 'False', '2021-05-03', 2)), 
    (("ID2", 'True', '2021-05-04', 3))]
df = spark.createDataFrame(data, ["ID", "colA", 'Date', 'Rank'])
df.show()

+---+-----+----------+
| ID| colA|      Date|
+---+-----+----------+
|ID1|False|2021-05-01|
|ID1|False|2021-05-02|
|ID1| True|2021-05-03|
|ID1|False|2021-05-04|
|ID1|False|2021-05-05|
|ID2| True|2021-05-01|
|ID2| True|2021-05-02|
|ID2|False|2021-05-03|
|ID2| True|2021-05-04|
+---+-----+----------+

+---+-----+----------+----+
| ID| colA|      Date|Rank|
+---+-----+----------+----+
|ID1|False|2021-05-01|   1|
|ID1|False|2021-05-02|   1|
|ID1| True|2021-05-03|   2|
|ID1|False|2021-05-04|   3|
|ID1|False|2021-05-05|   3|
|ID2| True|2021-05-01|   1|
|ID2| True|2021-05-02|   1|
|ID2|False|2021-05-03|   2|
|ID2| True|2021-05-04|   3|
+---+-----+----------+----+

---

这可以在Pyspark中实现吗？

您可以添加一列

change

，以指定何时存在“交叉”，并对该列进行滚动求和：

from pyspark.sql import functions as F, Window

w = Window.partitionBy('ID').orderBy('Date')

df2 = df.withColumn(
    'change', 
    F.coalesce(
        F.lag('colA').over(w) != F.col('colA'), 
        F.lit(True)    # Take care of first row in each group
    ).cast('int')
).withColumn(
    'rank', 
    F.sum('change').over(w)
)

df2.show()
+---+-----+----------+------+----+
| ID| colA|      Date|change|rank|
+---+-----+----------+------+----+
|ID1|False|2021-05-01|     1|   1|
|ID1|False|2021-05-02|     0|   1|
|ID1| True|2021-05-03|     1|   2|
|ID1|False|2021-05-04|     1|   3|
|ID1|False|2021-05-05|     0|   3|
|ID2| True|2021-05-01|     1|   1|
|ID2| True|2021-05-02|     0|   1|
|ID2|False|2021-05-03|     1|   2|
|ID2| True|2021-05-04|     1|   3|
+---+-----+----------+------+----+