在PySpark中，如何在groupby操作内获取值的计数？

Question

我有一个包含几列的数据框 df。我想按照一个（或多个）列进行分组，并对每个组，我想知道另外一些列的值的数量。

这是 df:

col1 col2  col3  col4
1     1      a    2
1     1      b    1
1     2      c    1
2     1      a    3
2     1      b    4

我希望按照'col1'和'col2'分组，对于每组，计算一列中独特值的数量，以及另一列的总和/平均值/最小值/最大值。我还希望维护分组后的列。结果应该是：

col1 col2  count_a count_b count_c col4_sum
1     1      1        1      0       3
1     2      0        0      1       1
2     1      1        1      0       7

我该如何实现这个？

- yuser099881232

Answer 1

您有两个解决方案。

第一个是可以在col3上使用数据透视表获取唯一值的计数，然后将您的数据透视表与聚合的数据框连接起来，计算其他列的总和/平均数/最小值/最大值。

代码如下：

from pyspark.sql import functions as F

result = df \
    .groupBy('col1', 'col2') \
    .pivot('col3') \
    .agg(F.count('col3')) \
    .fillna(0) \
    .join(
        df.groupby('col1', 'col2').agg(F.sum('col4').alias('col4_sum')),
        ['col1', 'col2']
    )

使用您提供的数据框，您将获得以下结果：

+----+----+---+---+---+--------+
|col1|col2|a  |b  |c  |col4_sum|
+----+----+---+---+---+--------+
|1   |1   |1  |1  |0  |3       |
|1   |2   |0  |0  |1  |1       |
|2   |1   |1  |1  |0  |7       |
+----+----+---+---+---+--------+

然而，你无法选择从数据透视表中提取的列的名称，它将是值的名称。

如果你真的想要选择列的名称，可以先检索所有不同的值，然后从每个值构建你的汇总列，具体如下：

from pyspark.sql import functions as F

values = map(lambda x: x.col3, df.select("col3").distinct().collect())

count_of_distinct_values = [F.sum((F.col('col3') == i).cast('integer')).alias('count_' + i) for i in values]
other_column_aggregations = [F.sum('col4').alias('col4_sum')]
aggregated = count_of_distinct_values + other_column_aggregations

result = df.groupBy('col1', 'col2').agg(*aggregated)

你将获得以下数据框：

+----+----+-------+-------+-------+--------+
|col1|col2|count_a|count_b|count_c|col4_sum|
+----+----+-------+-------+-------+--------+
|1   |1   |1      |1      |0      |3       |
|1   |2   |0      |0      |1      |1       |
|2   |1   |1      |1      |0      |7       |
+----+----+-------+-------+-------+--------+

- Vincent Doba