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

Posted 2023-03-31

【中文标题】Pyspark - 过滤数据框并创建排名列

【英文标题】：Pyspark - filter dataframe and create rank columns

【发布时间】：2019-12-12 21:29:29

【问题描述】：

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

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： 计算 Column1 上 Column2 为 X1 且 Column3 为空并按 Column4 排序的行的排名

Rank2： 计算 Column1 上 Column2 为 X1 且 Column3 为 Test 的行的排名，并按 Column4 排序

所以预期的结果是：

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

我尝试使用 when 过滤数据，但排名不是从 1 开始。

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

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

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

【参考方案1】：

需要在partitionby窗口的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()

【讨论】：

非常感谢，我一直在到处寻找这个。