如何在PySpark DataFrame上快速计算多个不同条件的计数？

核心思路是使用条件聚合将所有计数逻辑合并到单次Spark查询中，仅触发1次作业调度，从根源上消除多次count带来的调度开销。

实现代码 #

你可以用sum + when的组合完成任意条件的计数，针对你的示例代码改造如下：

import pyspark.sql.functions as F

# 单次查询完成所有计数
count_result = data.select(
    F.sum(F.when(F.length(F.col('Country')) == 2, 1).otherwise(0)).alias('count_short'),
    F.sum(F.when(F.length(F.col('Country')) > 2, 1).otherwise(0)).alias('count_long')
).collect()[0]

# 直接取出对应统计值
count_short = count_result['count_short']
count_long = count_result['count_long']

100+项统计的批量写法 #

当需要统计大量条件时，可以提前把所有条件整理成字典，批量生成聚合表达式，避免代码冗余：

# 定义所有统计条件，key为统计项名称，value为对应的判断表达式
count_conditions = {
    'count_short': F.length(F.col('Country')) == 2,
    'count_long': F.length(F.col('Country')) > 2,
    # 继续添加其他100+个条件即可，支持任意复杂的判断逻辑
    # 'count_xxx': (F.col('xxx') > 10) & F.col('yyy').isNotNull()
}

# 批量生成聚合列
agg_exprs = [F.sum(F.when(cond, 1).otherwise(0)).alias(name) for name, cond in count_conditions.items()]
# 单次查询得到所有结果
result_row = data.select(*agg_exprs).collect()[0]
# 转成字典方便后续读取
all_counts = result_row.asDict()

极小数据集额外优化 #

如果你的数据集规模始终极小（比如小于10万行），可以直接把数据拉取到Driver端用Pandas处理，完全规避Spark调度开销：

import pandas as pd

pdf = data.toPandas()
all_counts = {
    'count_short': len(pdf[pdf['Country'].str.len() == 2]),
    'count_long': len(pdf[pdf['Country'].str.len() > 2]),
    # 其他统计项
}

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