将可变长度的变量传递给PySpark groupby().agg函数

Question

我正在将不同长度的列名列表传递给PySpark的groupby().agg函数？我编写的代码检查列表的长度，例如，如果长度为1，它将对一个元素执行.agg(计数(。如果列表的长度为2，它将执行两个单独的.agg(计数(操作，生成两个新的.agg列。

有没有比通过if语句更简洁的方法来写这篇文章，因为随着列名列表的变长，我将不得不添加更多的elif语句。

例如：

agg_fields：列名列表

if len(agg_fields) == 1:
df = df.groupBy(col1, col2).agg(count(agg_fields[0]))
elif len(agg_fields) == 2:
df = df.groupBy(col1, col2).agg(count(agg_fields[0]), 
count(agg_fields[1]))

Answer 1

是的，您可以简单地循环创建聚合语句：

agg_df = df.groupBy("col1","col2").agg(*[count(i).alias(i) for i in agg_fields])