Python中groupby分组聚合的高效替代方案（字符串列表场景）

你当前的核心问题是Pandas groupby 在处理大量分类组合+字符串列表聚合时的性能瓶颈，以下是几个适配字符串列表场景的高效替代方案：

1. 使用Polars替代Pandas #

Polars是基于Rust的DataFrame库，分组聚合性能远优于Pandas，尤其适合列表这类复杂类型的处理：

import polars as pl
from itertools import combinations

# 转换为Polars DataFrame
pl_df = pl.DataFrame(df)

categories = ['cat1', 'cat2', 'cat3','cat4','cat5']
groups = [list(combo) for combo in combinations(categories, 2)]

for group in groups:
    # 直接扁平化分组内的所有列表，比Pandas的sum()效率更高
    aggregated = pl_df.group_by(group).agg(pl.col("Reports").flatten().alias("Reports"))
    # 后续TF-IDF处理逻辑

2. 预计算分组键+矢量化聚合 #

跳过Pandas groupby 的内部开销，直接用Numpy索引实现分组聚合：

import numpy as np
from itertools import combinations

categories = ['cat1', 'cat2', 'cat3','cat4','cat5']
groups = [list(combo) for combo in combinations(categories, 2)]

for group in groups:
    # 生成分组键的元组数组
    keys = df[group].apply(tuple, axis=1).to_numpy()
    # 获取唯一分组键及对应索引
    unique_keys, indices = np.unique(keys, return_inverse=True)
    
    # 按索引批量聚合列表
    aggregated_reports = []
    for idx in range(len(unique_keys)):
        mask = indices == idx
        combined = [item for sublist in df.loc[mask, 'Reports'] for item in sublist]
        aggregated_reports.append(combined)
    
    # 生成结果DataFrame
    df1 = pd.DataFrame(unique_keys, columns=group)
    df1['Reports'] = aggregated_reports
    # 后续TF-IDF处理逻辑

3. 用Dask做并行分组处理 #

Dask支持数据分片并行计算，适合超大规模数据集：

import dask.dataframe as dd
from itertools import combinations

# 转换为Dask DataFrame，分区数建议匹配CPU核心数
ddf = dd.from_pandas(df, npartitions=4)

categories = ['cat1', 'cat2', 'cat3','cat4','cat5']
groups = [list(combo) for combo in combinations(categories, 2)]

for group in groups:
    # 自动并行执行分组聚合
    aggregated = ddf.groupby(group).agg({'Reports': lambda x: [item for sublist in x for item in sublist]}).compute()
    aggregated = aggregated.reset_index()
    # 后续TF-IDF处理逻辑

4. 预扁平化数据+直接统计词频 #

如果后续要计算TF-IDF，可以跳过列表合并步骤，直接统计每个分组的词频：

from itertools import combinations
from collections import defaultdict

categories = ['cat1', 'cat2', 'cat3','cat4','cat5']
groups = [list(combo) for combo in combinations(categories, 2)]

for group in groups:
    group_word_counts = defaultdict(lambda: defaultdict(int))
    
    # 遍历行直接统计每个分组的词频
    for _, row in df.iterrows():
        key = tuple(row[g] for g in group)
        for word in row['Reports']:
            group_word_counts[key][word] += 1
    
    # 转换为结构化数据
    df1 = pd.DataFrame([
        (*key, dict(counts)) for key, counts in group_word_counts.items()
    ], columns=group + ['WordCounts'])
    # 后续直接基于WordCounts计算TF-IDF

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