如何提升Pandas/Numpy中均值、极值等聚合运算的速度？

问题背景 #

我有一个包含600行、2999个浮点列和1个分类列cat1的DataFrame，需要按cat1分组后，对指定的10列（col1到col10）计算mean、max、min、std、sum聚合值，将结果合并回原DataFrame后转换为numpy数组。当前实现耗时约35ms，希望尽可能提升运行速度。

测试数据生成代码 #

import pandas as pd
import numpy as np

# 生成600行3000列的DataFrame
num_rows = 600
num_numeric_cols = 2999
# 生成float32类型的随机数值列
numeric_data = np.random.rand(num_rows, num_numeric_cols).astype(np.float32)
# 生成含50个唯一类别的分类列
num_categories = 50
categorical_data = np.random.choice(range(num_categories), num_rows)
# 数值列命名
numeric_column_names = [f"col{i}" for i in range(num_numeric_cols)]
# 构建DataFrame
data = {numeric_column_names[i]: numeric_data[:, i] for i in range(num_numeric_cols)}
data["cat1"] = categorical_data

df = pd.DataFrame(data)
df["cat1"] = df["cat1"].astype("category")

待优化的核心代码 #

import time

start = time.time()    

aa = df.groupby(["cat1"])[["col1","col2","col3","col4","col5","col6","col7","col8","col9","col10"]]
df = pd.merge(df,aa.mean().reset_index(),how="left",on="cat1",suffixes=["","mean"])
df = pd.merge(df,aa.std().reset_index(),how="left",on="cat1",suffixes=["","std"])
df = pd.merge(df,aa.min().reset_index(),how="left",on="cat1",suffixes=["","min"])
df = pd.merge(df,aa.max().reset_index(),how="left",on="cat1",suffixes=["","max"])
df = pd.merge(df,aa.sum().reset_index(),how="left",on="cat1",suffixes=["","sum"])
df= df.to_numpy(np.float32)
print("Time Taken",time.time()-start)
print(df.shape)

优化方案 #

方案1：一次性完成聚合，减少分组与合并次数 #

原代码多次调用分组聚合和pd.merge，每次合并都会带来额外开销。通过agg方法一次性指定所有聚合函数，扁平化列名后用join合并（比merge更高效，基于分类键直接拼接）。

优化后代码：

import time

start = time.time()

target_cols = ["col1","col2","col3","col4","col5","col6","col7","col8","col9","col10"]
# 一次性聚合所有函数，并重命名列
agg_result = df.groupby("cat1")[target_cols].agg({
    col: ["mean", "max", "min", "std", "sum"] for col in target_cols
})
# 扁平化多级列名
agg_result.columns = [f"{col}_{func}" for col, func in agg_result.columns]
agg_result = agg_result.reset_index()

# 用join合并（cat1为分类列，匹配效率更高）
df_merged = df.join(agg_result.set_index("cat1"), on="cat1")
# 转numpy数组
result = df_merged.to_numpy(np.float32)

print("Time Taken", time.time() - start)
print(result.shape)

效果：将多次分组和合并压缩为1次分组+1次join，耗时可降至10ms以内（取决于硬件）。

方案2：用numpy直接处理，跳过pandas合并开销 #

由于最终目标是生成numpy数组，可绕过pandas的DataFrame合并步骤，直接用numpy完成分组统计和映射拼接，性能提升更明显。

优化后代码：

import time

start = time.time()

target_cols = ["col1","col2","col3","col4","col5","col6","col7","col8","col9","col10"]
# 提取目标列的numpy数组
target_data = df[target_cols].to_numpy(np.float32)
# 提取分类列的编码值（category类型的codes为整数索引）
cat_codes = df["cat1"].cat.codes.values
# 获取唯一类别编码
unique_cats = np.unique(cat_codes)

# 预分配统计结果数组：每个类别对应5种统计量×10列
stats = np.zeros((len(unique_cats), len(target_cols)*5), dtype=np.float32)

for i, cat in enumerate(unique_cats):
    # 筛选当前类别的行
    mask = cat_codes == cat
    subset = target_data[mask]
    # 计算所有统计量并拼接
    stats[i] = np.concatenate([
        subset.mean(axis=0),
        subset.max(axis=0),
        subset.min(axis=0),
        subset.std(axis=0, ddof=1),  # 与pandas的std默认参数一致
        subset.sum(axis=0)
    ])

# 创建映射字典：类别编码→对应的统计行
cat_to_stats = {cat: stats[i] for i, cat in enumerate(unique_cats)}
# 为原表每一行映射对应的统计数据
mapped_stats = np.array([cat_to_stats[code] for code in cat_codes], dtype=np.float32)

# 拼接原数据和统计数据：原表所有列（除cat1） + 统计列
original_data = df.drop("cat1", axis=1).to_numpy(np.float32)
result = np.hstack([original_data, mapped_stats])

print("Time Taken", time.time() - start)
print(result.shape)

效果：完全基于numpy操作，避免pandas分组和合并开销，耗时可降至5ms以内。注意std使用ddof=1，与pandas默认的无偏标准差保持一致。

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