計算Spark數據框中分組百分比的更簡單方法？

Question

考慮下面的數據框：

df = sc.parallelize([ 
    ('2017-05-21', 'a'), 
    ('2017-05-21', 'c'), 
    ('2017-05-22', 'b'), 
    ('2017-05-22', 'c'), 
    ('2017-05-23', 'a'), 
    ('2017-05-23', 'b'), 
    ('2017-05-23', 'c'), 
    ('2017-05-23', 'c'), 
]).toDF(['date', 'foo']) 

我想獲得的foo == a每日百分比：

+----------+----------+ 
|  date|percentage| 
+----------+----------+ 
|2017-05-21|  0.5| 
|2017-05-22|  0.0| 
|2017-05-23|  0.25| 
+----------+----------+ 

這是我想出了：

df.withColumn('foo_a', df.foo == 'a') 
    .groupby('date') 
    .agg((func.sum(col('foo_a').cast('integer'))/func.count('*')).alias('percentage')) 
    .sort('date') 

這工作，但我覺得應該有一個更簡單的方法。具體來說，是否有一個用於統計某個值出現的集合函數？

Answer 1

mean/avg結合when：

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

df.groupBy("date").agg(avg(when(col("foo") == "a", 1).otherwise(0))) 

或cast：

df.groupBy("date").agg(avg((col("foo") == "a").cast("integer"))) 

是你所需要的。