优化大规模数据集上的Pandas向量化性能

原代码的核心问题 #

• 逐行操作效率极低：df.apply(axis=1)本质是逐行遍历，面对20万+行数据时性能会急剧下降，远不如Pandas的矢量化操作高效。
• 内存碎片化：循环中反复给DataFrame新增列，且每次都通过apply生成新列，容易导致内存碎片化，触发insert错误。
• 存在性判断低效：用列表存储匹配ID，str(id) in id_list是O(n)复杂度，多次查询会累积大量耗时。

优化方案 #

1. 用矢量化操作替代逐行apply

直接利用Pandas的布尔掩码特性，通过eval或query生成匹配结果，一次性完成列赋值，完全避免逐行处理。

2. 用集合优化存在性判断

将匹配的ID转换为集合，in操作的复杂度降为O(1)，大幅提升查询速度。

3. 处理多规则叠加的情况

如果需要保留之前分组的"Yes"结果（即同一行满足多个规则时，始终显示"Yes"），可以用逻辑或|=来累积布尔值。

4. 从CSV加载动态规则集

针对动态规则集来自另一个CSV的需求，直接读取规则文件批量处理，无需硬编码规则。

完整优化代码 #

import pandas as pd
import numpy as np

# ----------------------
# 1. 加载大规模数据（实际用pd.read_csv）
# ----------------------
data = {
    'ID': range(1, 200001),  # 20万行模拟数据
    'Gender': ['Male', 'Female'] * 100000,
    'Age': np.random.randint(18, 70, size=200000),
    'Country': ['USA', 'Canada', 'UK', 'Australia'] * 50000
}
df = pd.DataFrame(data)
# 设置dtype减少内存占用，降低碎片化风险
df['Gender'] = df['Gender'].astype('category')
df['Country'] = df['Country'].astype('category')

# ----------------------
# 2. 从CSV加载规则集（模拟规则CSV结构）
# 实际使用：rules_df = pd.read_csv('rules.csv')
# ----------------------
rules_data = {
    'query': [
        '(Gender in ["Male","Female"]) & (Country=="USA")',
        '(Country in ["USA", "UK"]) & (Gender=="Male")',
        '(Age > 40) & (Gender=="Male")'
    ],
    'group_name': ['USA', 'MALE_USA_UK', 'MALE_OVER_40']
}
rules_df = pd.DataFrame(rules_data)

# ----------------------
# 3. 批量处理规则，生成分组列
# ----------------------
for _, row in rules_df.iterrows():
    query = row['query']
    group_name = row['group_name']
    
    # 用eval生成布尔掩码，直接完成赋值
    df[group_name] = np.where(df.eval(query), 'Yes', 'No')
    
    # 若需保留历史匹配结果（同一行满足多规则时保持Yes），替换为以下逻辑：
    # if group_name in df.columns:
    #     df[group_name] = np.where(df[group_name] == 'Yes', 'Yes', np.where(df.eval(query), 'Yes', 'No'))
    # else:
    #     df[group_name] = np.where(df.eval(query), 'Yes', 'No')

关键优化点解释 #

• 矢量化赋值：df.eval(query)直接生成整列布尔数组，配合np.where一次性完成赋值，速度比apply快几十倍以上。
• 内存优化：将字符串列设为category类型，大幅减少内存占用，避免内存碎片化问题。
• 动态规则适配：直接遍历规则DataFrame，无需手动维护规则列表，完美适配动态规则需求。
• 减少冗余操作：不再复制临时DataFrame，直接在原数据上操作，降低内存开销。

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