Spark 1.6.2：如何无循环对所有列执行聚合操作

我完全理解你想避免多次groupby带来性能开销的诉求——循环里每次groupby都会触发一次shuffle，数据量大的时候确实会拖慢整个流程。你之前尝试的写法之所以不符合预期，问题出在countDistinct(*df.schema.names)的用法上：这个函数接收多个列时，计算的是这些列组合在一起的唯一值数量，而不是每个列单独的distinct count。

要实现一次性按colX分组，计算所有列各自的distinct count，你需要为每一列单独创建一个countDistinct聚合表达式，然后把这些表达式批量传给agg方法。具体代码如下：

from pyspark.sql import functions as F

# 为每个列生成单独的countDistinct聚合，并用别名标记对应的列
agg_exprs = [F.countDistinct(col).alias(f"distinct_{col}") for col in df.schema.names]

# 只执行一次groupby和agg，然后收集结果
uniqs_df = df.groupby("colX").agg(*agg_exprs)
uniqs = uniqs_df.collect()

这样处理的话，Spark只会执行一次groupby和shuffle操作，性能和循环版本相比会有明显提升。

如果你需要把结果整理成和原来循环版本类似的字典结构（键是列名，值是对应分组的distinct count列表），可以在收集结果后做简单的转换：

# 将收集到的Row转换成字典格式
uniqs_dict = {}
for col in df.schema.names:
    col_alias = f"distinct_{col}"
    uniqs_dict[col] = [(row["colX"], row[col_alias]) for row in uniqs]

这样就能得到和你原来循环逻辑一致的结果，但只触发了一次Spark作业，避免了重复shuffle的开销。

内容的提问来源于stack exchange，提问作者Thomas