groupBy聚合函数中的PySpark循环

我有一个很大的表，我试图计算一些按位置分组的列的和（带条件）

我的代码是这样的，我有越来越多的列

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"

有重构它的解决方案吗？

---

谢谢

，cols中的x应该在方括号内。您还需要在列表前放置一个
*
，以展开参数：

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