如何获取 pyspark 中 groupBy 之后每个计数的总数百分比？

Question

给定以下数据框：

import findspark
findspark.init()
from pyspark.sql import SparkSession

spark = SparkSession.builder.master("local").appName("test").getOrCreate()
df = spark.createDataFrame([['a',1],['b', 2],['a', 3]], ['category', 'value'])
df.show()


+--------+-----+
|category|value|
+--------+-----+
|       a|    1|
|       b|    2|
|       a|    3|
+--------+-----+

我想计算每个类别中的项目数量，并提供每个计数的总数百分比，如下所示

+--------+-----+----------+
|category|count|percentage|
+--------+-----+----------+
|       b|    1|     0.333|
|       a|    2|     0.667|
+--------+-----+----------+

Answer 1

您可以通过以下方式获取总数的计数和百分比/比率

import pyspark.sql.functions as f
from pyspark.sql.window import Window
df.groupBy('category').count()\
  .withColumn('percentage', f.round(f.col('count') / f.sum('count')\
  .over(Window.partitionBy()),3)).show()

+--------+-----+----------+
|category|count|percentage|
+--------+-----+----------+
|       b|    1|     0.333|
|       a|    2|     0.667|
+--------+-----+----------+

前面的陈述可以分为步骤。df.groupBy('category').count() 产生count：

+--------+-----+
|category|count|
+--------+-----+
|       b|    1|
|       a|    2|
+--------+-----+

然后通过应用窗口函数，我们可以获得每行的总计数：

df.groupBy('category').count().withColumn('total', f.sum('count').over(Window.partitionBy())).show()

+--------+-----+-----+
|category|count|total|
+--------+-----+-----+
|       b|    1|    3|
|       a|    2|    3|
+--------+-----+-----+

其中total列是通过将分区（包含所有行的单个分区）中的所有计数加在一起来计算的。

一旦我们有了count 和 ，total 我们就可以计算每一行的比率：

df.groupBy('category')\
  .count()\
  .withColumn('total', f.sum('count').over(Window.partitionBy()))\
  .withColumn('percentage',f.col('count')/f.col('total'))\
  .show()

+--------+-----+-----+------------------+
|category|count|total|        percentage|
+--------+-----+-----+------------------+
|       b|    1|    3|0.3333333333333333|
|       a|    2|    3|0.6666666666666666|
+--------+-----+-----+------------------+

Answer 2

您可以groupby并聚合agg：

import pyspark.sql.functions as F

df.groupby('category').agg(F.count('value') / df.count()).show()

输出：

+--------+------------------+
|category|(count(value) / 3)|
+--------+------------------+
|       b|0.3333333333333333|
|       a|0.6666666666666666|
+--------+------------------+

为了让它更好，你可以使用：

df.groupby('category').agg(
    (
        F.round(F.count('value') / df.count(), 2)
    ).alias('ratio')
).show()

输出：

+--------+-----+
|category|ratio|
+--------+-----+
|       b| 0.33|
|       a| 0.67|
+--------+-----+

您还可以使用 SQL：

df.createOrReplaceTempView('df')

spark.sql(
    """
    SELECT category, COUNT(*) / (SELECT COUNT(*) FROM df) AS ratio
    FROM df
    GROUP BY category
    """
).show()