pandas多列groupby聚合处理百万行数据集的提速优化方法

核心性能瓶颈 #

你当前代码的性能瓶颈完全来自names/details/text三列的自定义lambda聚合：lambda x: " ".join(set(x))。pandas的groupby调用自定义Python函数时，会逐组触发Python层循环，没有向量化优化，百万行数据下会产生极高的函数调用开销；而debit/credit的sum、probables的first都是Cython实现的内置函数，本身开销极低。

可选优化方案 #

方案1：纯Pandas原生逻辑优化 #

通过前置去重+内置聚合函数替代自定义lambda，性能可提升5~10倍：

# 1. 分组键转category类型，大幅降低分组时的哈希计算开销
result_frame["source"] = result_frame["source"].astype("category")
result_frame["target"] = result_frame["target"].astype("category")

# 2. 聚合数值类字段（原生Cython优化，速度极快）
num_agg = result_frame.groupby(["source", "target"], observed=True).agg(
    debit=("debit", "sum"),
    credit=("credit", "sum"),
    probables=("probables", "first")
).reset_index()

# 3. 单独处理字符串字段：前置全局去重替代组内转set，再用内置join聚合
def agg_str_col(col_name):
    # 先按分组键+当前字符串列去重，避免每组内重复做set计算
    unique_df = result_frame[["source", "target", col_name]].drop_duplicates()
    # 用内置字符串join聚合，没有Python层循环开销
    return unique_df.groupby(["source", "target"], observed=True)[col_name].apply(" ".join).reset_index()

# 聚合三个字符串字段
names_agg = agg_str_col("names")
details_agg = agg_str_col("details")
text_agg = agg_str_col("text")

# 4. 合并所有聚合结果
grouped = num_agg.merge(names_agg, on=["source", "target"])\
                 .merge(details_agg, on=["source", "target"])\
                 .merge(text_agg, on=["source", "target"])

方案2：向量化引擎DuckDB优化（推荐） #

使用列式向量化执行引擎DuckDB执行聚合，逻辑完全对齐需求，性能比优化后的Pandas还要高2~5倍，百万行数据可秒级返回：

1. 先安装依赖：pip install duckdb
2. 执行代码：

import duckdb

# 直接用SQL执行聚合，自动优化distinct字符串拼接逻辑
grouped = duckdb.sql("""
    SELECT 
        source,
        target,
        SUM(debit) as debit,
        SUM(credit) as credit,
        STRING_AGG(DISTINCT names, ' ') as names,
        STRING_AGG(DISTINCT details, ' ') as details,
        STRING_AGG(DISTINCT text, ' ') as text,
        FIRST(probables) as probables
    FROM result_frame
    GROUP BY source, target
""").df()

额外优化提示 #

• 如果字符串列重复值占比超过30%，可提前转成category类型进一步提升性能
• 避免在groupby.agg中使用任何自定义Python lambda/函数，优先使用pandas内置算子或向量化引擎

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