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

时间:2019-12-12 21:29:29

标签: apache-spark pyspark apache-spark-sql pyspark-sql

我遇到一种情况,我想根据不同的条件在数据框中创建等级列,并将第一等级设置为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

在这里,我想根据以下条件创建排名列:

排名1: 为Column2为X1且Column3为null并按Column4排序的行计算Column1上的排名

排名2: 为列2计算X1且列3测试和按列4排序的行计算列1的排名

所以预期结果将是:

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 个答案:

答案 0 :(得分:3)

您需要将条件添加到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()