groupBy 聚合函数中的 PySpark 循环答案

【问题标题】：PySpark loop in groupBy aggregate functiongroupBy 聚合函数中的 PySpark 循环
【发布时间】：2021-03-18 13:13:04
【问题描述】：

我有一个大表，我正在尝试计算按位置分组的某些列的总和（带条件）。

我的代码长这样，而且列越来越多

df.groupBy(location_column).agg(
        F.sum(F.when(F.col(col1) == True, F.col(value))).alias("SUM " + col1),
        F.sum(F.when(F.col(col2) == True, F.col(value))).alias("SUM " + col2),
        F.sum(F.when(F.col(col3) == True, F.col(value))).alias("SUM " + col3),
        ....
        # Additional lines for additional columns (around 20)
)

我想重构我的代码，使其看起来不那么愚蠢，基本上做一些类似的事情

cols = [col1, col2, col3, ... , coln]
df.groupBy(location_column).agg([F.sum(F.when(F.col(x) == True, F.col(value))).alias("SUM " + x)] for x in cols)

它不起作用，因为 agg() 函数不接受列表：

assert all(isinstance(c, Column) for c in exprs), "all exprs should be Column"

有重构它的解决方案吗？谢谢

【问题讨论】：

标签： pyspark group-by aggregate

【解决方案1】：

for x in cols 应该在方括号内。您还需要在列表推导之前添加 * 以扩展参数：

df.groupBy(location_column).agg(
    *[F.sum(F.when(F.col(x) == True, F.col(value))).alias("SUM " + x) for x in cols]
)

【讨论】：