如何在PyArrow表中按分组计算列众数并优化性能

核心思路 #

放弃循环过滤单组的低效方式，利用PyArrow的自定义聚合函数，将pc.compute.mode整合到分组聚合流程中，和内置的均值、极值聚合一起批量执行，充分利用PyArrow的列存矢量化性能优势。

实现步骤 #

1. 定义自定义众数聚合函数 #

PyArrow的分组聚合支持传入自定义Aggregate对象，基于pc.mode实现矢量化的众数聚合：

import pyarrow as pa
import pyarrow.compute as pc

def mode_aggregate(col_type):
    return pa.Aggregate(
        # 初始化：创建对应类型的空数组存储组内数据
        init=lambda: pa.array([], type=col_type),
        # 更新：追加组内批次数据到存储数组
        update=lambda state, values: pa.concat_arrays([state, values]),
        # 合并：合并多批次的存储数组
        merge=lambda state1, state2: pa.concat_arrays([state1, state2]),
        # 最终计算：取众数的第一个结果（兼容多众数场景），空组返回None
        finalize=lambda state: pc.mode(state)[0] if len(state) > 0 else None
    )

注：传入col_type参数匹配目标列类型（如pa.int64()、pa.string()、pa.float64()），避免隐式类型转换开销。

2. 批量执行分组聚合 #

假设表以ID为分组键，需计算Col1、Col2等列的多指标，一次性定义聚合规则并执行：

# 示例测试表
data = {
    "ID": [1,1,1,2,2,3,3,3,3],
    "Col1": [5,5,3,2,2,7,7,7,9],
    "Col2": [10,10,10,20,30,30,30,40,40]
}
table = pa.table(data)

# 批量生成聚合规则：每列对应均值、最大、最小、众数
agg_rules = {}
target_cols = ["Col1", "Col2"]
for col in target_cols:
    col_type = table.schema.field(col).type
    agg_rules[f"{col}_mean"] = pc.mean(col)
    agg_rules[f"{col}_max"] = pc.max(col)
    agg_rules[f"{col}_min"] = pc.min(col)
    agg_rules[f"{col}_mode"] = mode_aggregate(col_type)

# 一次完成分组聚合
result = table.group_by("ID").aggregate(agg_rules)
print(result.to_pandas())

3. 性能优化细节 #

• 空值处理：若数据含大量空值，可在finalize阶段改为pc.mode(pc.drop_null(state))，排除空值后计算众数。
• 内存控制：对于超大规模表，可启用PyArrow的内存池优化，避免临时数组占用过多内存。
• 类型一致性：确保自定义聚合的初始化类型与目标列严格匹配，杜绝不必要的类型转换。

性能对比 #

原循环过滤方式需多次扫描全表，时间复杂度为O(N*K)（N为行数，K为分组数）；自定义聚合仅需一次全表扫描，时间复杂度为O(N)，真实数据下可将众数计算时间从几十分钟压缩至与均值/极值相当的秒级范围。

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