在 PySpark 中计算 groupby 后的 sum 和 countDistinct

Tito 发表于 Dev

铁托

我有一个 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 和 count distinct ？我知道，我可以用sum列做一次，用列做一次，countDistinct然后加入两个数据帧，但应该有一个解决方案可以一步完成......

沃纳

不要使用agg的字典版本，而是使用带有列列表的版本：

from pyspark.sql import functions as F
df = ...
exprs1 = [F.sum(c) for c in sum_cols]
exprs2 = [F.countDistinct(c) for c in count_cols]

df_aggregated = df.groupby('month_product').agg(*(exprs1+exprs2))

如果你想保持当前的逻辑，你可以切换到approx_count_distinct。与countDistinct此功能不同，该功能可作为 SQL 功能使用。

本文收集自互联网，转载请注明来源。

如有侵权，请联系 [email protected] 删除。