Pandas大数据集下高效创建表间外键的最优方案

我正在处理一个包含400万行的大型数据集，需要将CSV文件转换为SQL格式，在多个DataFrame间建立关联——把一个表的索引作为另一表的外键。目前用df.apply()的方案速度极慢（仅1800行/秒），寻求更高效的实现方式。

需求说明 #

如何以最快的方式关联以下两个DataFrame？要求在df1.street == df2.street且df1.number == df2.number时，用df2.index替换df1.street字段。

接受任何能提升速度的方案，包括多进程（曾尝试但未成功），同时需尽可能节省内存。也曾尝试df.merge()等函数，但未取得理想效果。

数据集 #

df1 #

import pandas as pd

df1 = pd.DataFrame({
    'street': {'qr16ef3677886a44f8b9c5bc37dd660688a': 'quai de la Tournelle', 'qr112e28085e3c84c41b6b6a5e13ecf15ac': 'r. AlexandreDumas', 'qr1a213d2a5cbf64718892b3dbb3a9024f1': 'pass. Brunoy', 'qr1fb0760cd0fab4c71a4297d006ec3d119': 'Montmartre', 'qr167fce4c23d5c4b879ca6423cea15e742': 'Martel'}
    , 'number': {'qr16ef3677886a44f8b9c5bc37dd660688a': '33', 'qr112e28085e3c84c41b6b6a5e13ecf15ac': '99', 'qr1a213d2a5cbf64718892b3dbb3a9024f1': '18', 'qr1fb0760cd0fab4c71a4297d006ec3d119': '123', 'qr167fce4c23d5c4b879ca6423cea15e742': '4'}
    , 'date': {'qr16ef3677886a44f8b9c5bc37dd660688a': 1914, 'qr112e28085e3c84c41b6b6a5e13ecf15ac': 1900, 'qr1a213d2a5cbf64718892b3dbb3a9024f1': 1921, 'qr1fb0760cd0fab4c71a4297d006ec3d119': 1858, 'qr167fce4c23d5c4b879ca6423cea15e742': 1896}
})

数据展示：

street number  date
qr16ef3677886a44f8b9c5bc37dd660688a  quai de la Tournelle     33  1914
qr112e28085e3c84c41b6b6a5e13ecf15ac     r. AlexandreDumas     99  1900
qr1a213d2a5cbf64718892b3dbb3a9024f1          pass. Brunoy     18  1921
qr1fb0760cd0fab4c71a4297d006ec3d119            Montmartre    123  1858
qr167fce4c23d5c4b879ca6423cea15e742                Martel      4  1896 

df2 #

df2 = pd.DataFrame({
    'number': {'qr152f8de48daa64cf098f44fb3d9e7e145': '123', 'qr18ae0099b6afb48a78d466e5ed6871bec': '18', 'qr183daee61fb98489ebd05556968027a0d': '18', 'qr1e0ee6ec37dbd4e799905db721592ba48': '33', 'qr148505eca183c4fb38f844c35130b92f0': '4'}
    , 'street': {'qr152f8de48daa64cf098f44fb3d9e7e145': 'Montmartre', 'qr18ae0099b6afb48a78d466e5ed6871bec': 'Montmartre', 'qr183daee61fb98489ebd05556968027a0d': 'pass. Brunoy', 'qr1e0ee6ec37dbd4e799905db721592ba48': 'quai de la Tournelle', 'qr148505eca183c4fb38f844c35130b92f0': 'Martel'}
    , 'date': {'qr152f8de48daa64cf098f44fb3d9e7e145': ['1858', '1858'], 'qr18ae0099b6afb48a78d466e5ed6871bec': ['1876', '1881'], 'qr183daee61fb98489ebd05556968027a0d': ['1921', '1921'], 'qr1e0ee6ec37dbd4e799905db721592ba48': ['1914', '1914'], 'qr148505eca183c4fb38f844c35130b92f0': ['1896', '1896']}
}) 

数据展示：

number                street          date
qr152f8de48daa64cf098f44fb3d9e7e145    123            Montmartre  [1858, 1858]
qr18ae0099b6afb48a78d466e5ed6871bec     18            Montmartre  [1876, 1881]
qr183daee61fb98489ebd05556968027a0d     18          pass. Brunoy  [1921, 1921]
qr1e0ee6ec37dbd4e799905db721592ba48     33  quai de la Tournelle  [1914, 1914]
qr148505eca183c4fb38f844c35130b92f0      4                Martel  [1896, 1896]

现有方案 #

当前方案依赖df.apply()调用自定义函数，本质是双重循环，速度极慢：

def foreignkey(ro: pd.core.series.Series) -> pd.core.series.Series:
    """
    replace the address in `ro` of `df1` by a foreign key 
    pointing to `df2`. the key is inserted in `ro.name`
    """
    ro.street = df2.loc[
        ( df2.street == ro.street )    # address has the same street full name
        & ( df2.number == ro.number )  # address has the same street number
    ].index[0]
    return ro

df1 = df1.progress_apply( lambda x: foreignkey(x), axis=1 )

方案1：矢量化merge操作（最快实现） #

apply逐行遍历效率极低，改用Pandas原生矢量化merge，先将df2的索引转为映射列，再关联替换：

# 给df2添加索引列用于映射
df2['df2_index'] = df2.index

# 仅保留关联所需列，减少内存开销
merged = df1.merge(df2[['street', 'number', 'df2_index']], on=['street', 'number'], how='left')

# 替换df1的street字段为df2的索引
df1['street'] = merged['df2_index']

如果df2中存在同一street+number对应多个索引的情况，先去重保留首次匹配项：

# 按street+number去重，保留第一个出现的索引
df2_unique = df2.drop_duplicates(subset=['street', 'number'], keep='first').reset_index(drop=False).rename(columns={'index':'df2_index'})

merged = df1.merge(df2_unique[['street', 'number', 'df2_index']], on=['street', 'number'], how='left')
df1['street'] = merged['df2_index']

方案2：字典映射（内存友好） #

若df2的street+number组合唯一，可构建字典实现O(1)快速查找：

# 构建street+number到df2索引的映射字典
mapping = {(row['street'], row['number']): idx for idx, row in df2.iterrows()}

# 批量映射替换
df1['street'] = df1.apply(lambda x: mapping.get((x['street'], x['number']), x['street']), axis=1)

此方案的apply虽仍逐行，但字典查找比原方案的全表遍历快数倍，适合df2规模较小的场景。

内存与速度优化建议 #

1. 优化数据类型：将street和number转为分类类型，大幅减少内存占用并加速匹配：

df1['street'] = df1['street'].astype('category')
df1['number'] = df1['number'].astype('category')
df2['street'] = df2['street'].astype('category')
df2['number'] = df2['number'].astype('category')

2. 分块处理：内存不足时，将df1分块读取处理，避免一次性加载全量数据：

chunk_size = 100000
# 分块读取原CSV
for i, chunk in enumerate(pd.read_csv('df1.csv', chunksize=chunk_size)):
    merged_chunk = chunk.merge(df2_unique[['street', 'number', 'df2_index']], on=['street', 'number'], how='left')
    merged_chunk['street'] = merged_chunk['df2_index']
    # 保存处理后的块
    merged_chunk.to_csv(f'df1_processed_chunk_{i}.csv', index=False)
# 合并所有处理后的块
final_df = pd.concat([pd.read_csv(f'df1_processed_chunk_{i}.csv') for i in range(num_chunks)])

3. 多进程加速：用swifter自动适配矢量化/多进程，或手动用multiprocessing分块处理：

# 方式1：用swifter自动优化
import swifter
df1['street'] = df1.swifter.apply(lambda x: mapping.get((x['street'], x['number']), x['street']), axis=1)

# 方式2：手动多进程分块
from multiprocessing import Pool

def process_chunk(chunk):
    chunk['street'] = chunk.apply(lambda x: mapping.get((x['street'], x['number']), x['street']), axis=1)
    return chunk

with Pool(processes=4) as pool:
    processed_chunks = pool.map(process_chunk, list(pd.read_csv('df1.csv', chunksize=chunk_size)))
final_df = pd.concat(processed_chunks)

原merge效果不佳的原因 #

大概率是未提前处理df2的重复项、未指定how='left'保留df1全量行，或未筛选关联所需列导致内存过载。按上述方案优化后，merge速度可达到数万行/秒，完全适配400万行规模。

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