Python pandas循环性能与内存优化：批量列分位数均值计算

1. 合并统计计算，避免重复分组开销 #

原代码对单列重复执行4次groupby，完全可以通过agg方法一次性批量计算所有需要的统计量，大幅减少重复计算的性能损耗：

def newsummary_optimized(final_per, grp_lvl, col):
    # 定义要计算的统计规则
    stat_map = {
        '25%': lambda x: x.quantile(0.25),
        '50%': lambda x: x.quantile(0.5),
        '75%': lambda x: x.quantile(0.75),
        'mean': 'mean'
    }
    # 一次分组完成所有统计计算
    result = final_per.groupby(grp_lvl)[col].agg(stat_map).reset_index()
    # 统一重命名列
    new_col_names = [f"{col}_{'_'.join(grp_lvl)}_{stat}" for stat in stat_map.keys()]
    result.columns = grp_lvl + new_col_names
    return result

2. 批量合并结果，减少内存碎片化 #

原循环中逐列处理并拼接的方式会产生大量临时DataFrame，建议先收集所有列的统计结果，再一次性合并：

grp_lvl = ['ZIP_CODE', 'year']
# 收集所有列的统计结果
all_stats = []
for j in list_col:
    col_stats = newsummary_optimized(final_per, grp_lvl, j)
    all_stats.append(col_stats)

# 以分组列为键合并所有结果
final_per2 = all_stats[0]
for stats_df in all_stats[1:]:
    final_per2 = final_per2.merge(stats_df, on=grp_lvl, how='left')

3. 切换到Spark DataFrame适配Databricks集群 #

既然在Azure Databricks环境运行，直接使用Spark DataFrame能充分利用集群分布式计算能力，避免单机Pandas的内存瓶颈：

from pyspark.sql import functions as F

# 假设final_per是Spark DataFrame
grp_lvl = ['ZIP_CODE', 'year']
stats_exprs = []

# 批量生成所有统计量的计算表达式
for col in list_col:
    prefix = f"{col}_{'_'.join(grp_lvl)}"
    stats_exprs.extend([
        F.percentile_approx(col, 0.25).alias(f"{prefix}_25%"),
        F.percentile_approx(col, 0.5).alias(f"{prefix}_50%"),
        F.percentile_approx(col, 0.75).alias(f"{prefix}_75%"),
        F.mean(col).alias(f"{prefix}_mean")
    ])

# 一次分组计算所有统计量
final_per2 = final_per.groupBy(grp_lvl).agg(*stats_exprs)

4. 基础内存优化细节 #

• 清理冗余列：只保留grp_lvl和待计算的list_col列，避免无关数据占用内存
• 优化数据类型：将数值列转换为更小的存储类型，比如Pandas中final_per = final_per.astype({col: 'float32' for col in list_col})
• 移除冗余copy()：原代码中copy()操作完全多余，直接赋值即可，reset_index()返回的是新DataFrame，无需额外复制

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