在 PySpark 中计算 groupby 之后的 sum 和 countDistinct

Question

我有一个 PySpark 数据框，想要对几列进行分组，然后计算一些列的总和并计算另一列的不同值。由于countDistinct不是内置聚合函数，我无法使用像我在这里尝试的那样的简单表达式：

sum_cols = ['a', 'b']
count_cols = ['id']
exprs1 = {x: "sum" for x in sum_cols}
exprs2 = {x: "countDistinct" for x in count_cols}
exprs = {**exprs1, **exprs2}

df_aggregated = df.groupby('month','product').agg(exprs)

我也尝试了这个答案中的方法，但是当我仅针对聚合列尝试时，exprs2 = [countDistinct(x) for x in count_cols]我收到了一条错误消息 。AssertionError: all exprs should be Column

如何将总和和不同计数合并到一个聚合中？我知道，我可以对sum列执行一次，然后对列执行一次countDistinct，然后加入两个数据框，但应该有一种解决方案可以一步完成此操作......

Answer 1

不知道为什么你必须使用expr，但正常的聚合应该可以工作。countDistinct是一个聚合函数。

(df
    .groupBy('month','product')
    .agg(
        F.sum('a', 'b'),
        F.countDistinct('id')
    )
).show()

# +----+-----------+-------------+
# |name|sum(field1)|count(field1)|
# +----+-----------+-------------+
# |   d|          0|            1|
# |   c|         10|            1|
# |   b|          5|            1|
# |   a|          4|            1|
# +----+-----------+-------------+