Pyspark-过滤数据框并创建排名列

Question

我有一种情况，我想根据不同的条件在数据帧中创建排名列，并将第一排名设置为 true，将其他排名设置为 false。下面是一个示例数据帧：

Column1    Column2   Column3   Column4
ABC        X1        null      2016-08-21 11:31:08
ABC        X1        Test      2016-08-22 11:31:08
ABC        X1        null      2016-08-20 11:31:08
PQR        X1        Test      2016-08-23 11:31:08
PQR        X1        Test      2016-08-24 11:31:08
PQR        X1        null      2016-08-24 11:31:08

这里我想根据以下条件创建排名列：

Rank1:计算列2为X1且列3为空的行在列1上的排名，并按列4排序

Rank2：计算列2为X1而列3为测试的行在列1上的排名，并按列4排序

因此，预期结果将是：

Column1    Column2   Column3   Column4                Rank1    Rank2
ABC        X1        null      2016-08-21 11:31:08    2        null
ABC        X1        Test      2016-08-22 11:31:08    null     1
ABC        X1        null      2016-08-20 11:31:08    1        null
PQR        X1        Test      2016-08-23 11:31:08    null     1
PQR        X1        Test      2016-08-24 11:31:08    null     2
PQR        X1        null      2016-08-24 11:31:08    1        null

我尝试使用何时过滤掉数据来执行此操作，但随后排名不是从 1 开始的。

df = df.withColumn("Rank1", F.when((df.Column2 == 'X1') & (df.Column3.isNull()), rank().over(Window.partitionBy('Column1').orderBy('Column4')))

这确实给了我顺序，但顺序是随机的。我需要给第一等级贴上标签，所以对我来说了解它很重要。

我尝试的另一个选项是过滤临时数据帧中的数据并计算排名并将其连接回主数据帧。但是数据帧大小很大，需要计算多个列，因此会给出内存不足错误。任何关于解决这个问题的帮助将不胜感激。

Answer 1

您需要将该条件添加到“分区依据”窗口的order by子句中。

这应该对你有用：

condition_rank1 = (col("column2") == 'X1') & (col("column3").isNull())
condition_rank2 = (col("column2") == 'X1') & (col("column3") == 'Test')

w_rank1 = Window.partitionBy('column1').orderBy(*[when(condition_rank1, lit(1)).desc(), col("column4")])
w_rank2 = Window.partitionBy('column1').orderBy(*[when(condition_rank2, lit(1)).desc(), col("column4")])

df.withColumn("Rank1", when(condition_rank1, rank().over(w_rank1))) \
    .withColumn("Rank2", when(condition_rank2, rank().over(w_rank2))) \
    .show()