如何在预测阶段使用pandas.get_dummies并高效处理未知类别

原有实现的性能问题核心在于两层嵌套循环，尤其是iterrows逐行遍历的操作，Python层的循环在处理百万级数据时会产生极高的开销，完全可以用Pandas原生向量化操作替换，性能提升可达数百倍。

优化后实现逻辑 #

先把每列的未知类别统一替换为NaN，再调用原生get_dummies编码，最后对齐训练阶段的列即可，代码如下：

import pandas as pd
import numpy as np

def populate_data_frame_in_prediction_time(data, train_cols_dict):
    # 1. 替换每列的未知类别为NaN
    processed = data.copy()
    for col in processed.columns:
        # 提取该列训练时的已知类别（去掉列名前缀）
        known_cats = [c.split(f"{col}_", 1)[1] for c in train_cols_dict[col] if c != f"{col}_nan"]
        # 不在已知类别里的值全部替换为NaN
        processed[col] = np.where(processed[col].isin(known_cats), processed[col], np.nan)
    
    # 2. 调用原生get_dummies编码，开启空值编码
    encoded = pd.get_dummies(processed, dummy_na=True, dtype='uint8')
    
    # 3. 对齐训练阶段的列，缺失列补0
    all_train_cols = [c for cols in train_cols_dict.values() for c in cols]
    encoded = encoded.reindex(columns=all_train_cols, fill_value=0)
    
    return encoded

测试示例 #

针对你给出的场景调用方式如下：

training_data = pd.DataFrame({'City': ['Delhi', None, 'Hyderabad', 'Delhi', 'Hyderabad'],
                     'Country': ['D', 'H', 'D', 'H', None]})
real_time_data = pd.DataFrame({'City': ['Delhi', None, 'Dhaka'],  
                               'Country': [None, "H", "B"]})
# 训练阶段保存的列名映射
train_cols_dict = {'City': ['City_Delhi', 'City_Hyderabad', 'City_nan'], 'Country': ['Country_D', 'Country_H', 'Country_nan']}

result = populate_data_frame_in_prediction_time(real_time_data, train_cols_dict)
print(result)

输出和原有逻辑完全一致：

City_Delhi  City_Hyderabad  City_nan  Country_D  Country_H  Country_nan
0           1               0         0          0          0            1
1           0               0         1          0          1            0
2           0               0         1          0          0            1

性能说明 #

优化后的实现所有核心操作都在Pandas底层的C语言层面执行，没有Python层的循环开销，处理100万行数据通常仅需百毫秒级耗时，远快于原有实现。

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