使用pyspark进行条件聚合

Question

考虑以下作为数据框

a        b  c   d   e  
africa  123 1   10  121.2
africa  123 1   10  321.98
africa  123 2   12  43.92
africa  124 2   12  43.92
usa     121 1   12  825.32
usa     121 1   12  89.78
usa     123 2   10  32.24
usa     123 5   21  43.92
canada  132 2   13  63.21
canada  132 2   13  89.23
canada  132 3   21  85.32
canada  131 3   10  43.92

现在我想使用数据帧将下面的 case 语句转换为 PYSPARK 中的等效语句。

我们可以直接在 case 语句中使用 hivecontex/sqlcontest nut 寻找传统的 pyspark nql 查询

select 
case 
    when c <=10 then sum(e)
    when c between 10 and 20 then avg(e)
else 0.00 end 
from table 
group by a,b,c,d

问候 Anvesh

Answer 1

您可以将 SQL 代码直接转换为DataFrame原语：

from pyspark.sql.functions import when, sum, avg, col

(df
    .groupBy("a", "b", "c", "d")  # group by a,b,c,d
    .agg(  # select 
        when(col("c") < 10, sum("e"))  #  when c <=10 then sum(e)
            .when(col("c").between(10 ,20), avg("c"))  # when c between 10 and 20 then avg(e)
            .otherwise(0))   # else 0.00