我遇到一种情况,我想根据不同的条件在数据框中创建等级列,并将第一等级设置为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')))
这确实给了我顺序顺序,但是顺序是随机的。我需要标记第一名,因此对我来说很重要。
我尝试的其他选项是在临时数据帧中过滤数据并计算等级并将其联接回主数据帧。但是数据帧的大小很大,并且要计算多列,因此会出现内存不足错误。对于解决此问题的任何帮助将不胜感激。
答案 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()