30万行DataFrame编辑距离计算性能优化方法咨询

问题描述 #

我正在计算两个DataFrame之间的编辑距离，两个DataFrame均包含约30万行数据，由于数据量庞大，计算耗时极长，请问有什么方法可以提升性能？

原始代码 #

主循环代码 #

import pandas as pd
import datetime as dt

for i in range(0,len(targets1)):
    if i % 100 == 0:
        pct = (i/len(targets1)) * 100
        print("(" + str(dt.datetime.now()) + ") completed: " + str(round(pct, 2)) + "%")
    sr1_new=sr1[(sr2==targets2[i]) & (len_sr2>=(len_targets2[i]-10)) & (len_sr2 <=(len_targets2[i]+10))]
    if len(sr1_new) > 0:
        ee=sr1_new.str.edit_distance(targets1[i])
        ee=ee.sort_values()
        output_final = output_final.append({'Name': targets1[i],'Matched_Name': sr1[ee.index[0]],'score':ee[ee.index[0]],'score_final':(len(sr1[ee.index[0]])+len(targets1[i])-ee[ee.index[0]])/(len(sr1[ee.index[0]])+len(targets1[i]))*100}, ignore_index=True)
    else:
        output_final = output_final.append({'Name': targets1[i],'Matched_Name': '','Matched_REF': "0",'score':0,'score_final':0}, ignore_index=True)

测试数据 #

targets1 = pd.Series(['ABBSHHCH','ABBSAJSJAHDKAJKJ', 'BASJBASJASH', 'KJSAKASJAS', 'KJSAIUBDAKS',
                'KAJSNDSAX', 'JASANXAJSKJ', 'NASNXHY', 'AIUSSHXBAHSJASHJ'])

targets2 = pd.Series(['AB','AB', 'BA', 'KJ', 'KJ','KA', 'JA', 'NA', 'AI'])

sr1 = pd.Series(['ABBSHHSJAKX','ABBMNASASJKKLASAHDKAJKJ', 'BASSAMSAJASH', 'KJSMSANMAASJAS', 'KJSSMNASBDAKS',
                'KASKJADSAX', 'JASAKJKJSKJ', 'NASAKXHY', 'AIUSSANMASSJASHJ','NSAASJNCXA','ABBSASMNKAJKJ', 'ASNASNXJASH', 
                'KJSKJSAKSJAS', 'KJASKJSDAKS', 'KAJSAKJSAX', 'JAKJASXAJSKJ', 'NADADHY', 'AIUSNASSASJASHJ'])

sr2 = pd.Series(['AB','AB','BA','KJ','KJ','KA','JA','NA','AI','NS','AB','AS','KJ','KJ','KA','JA','NA','AI'])

len_sr2 = pd.Series([11,23,12,14,13,10,11,8,16,10,13,11,12,11,10,12,7,15])

len_targets2 = pd.Series([8,16,11,10,11,9,11,7,16])

性能优化方案 #

1. 替换低效的DataFrame追加操作 #

pd.DataFrame.append()每次都会创建新的DataFrame对象，30万次循环会产生大量内存开销和时间损耗。改用列表存储结果，最后一次性转为DataFrame：

results = []
for i in range(len(targets1)):
    if i % 100 == 0:
        pct = (i/len(targets1)) * 100
        print(f"({dt.datetime.now()}) completed: {round(pct, 2)}%")
    
    sr1_new = sr1[(sr2 == targets2[i]) & (len_sr2 >= (len_targets2[i]-10)) & (len_sr2 <= (len_targets2[i]+10))]
    if len(sr1_new) > 0:
        ee = sr1_new.str.edit_distance(targets1[i])
        min_idx = ee.idxmin()  # 直接取最小值索引，比排序更高效
        min_dist = ee[min_idx]
        matched_name = sr1[min_idx]
        total_len = len(matched_name) + len(targets1[i])
        score_final = (total_len - min_dist) / total_len * 100
        
        results.append({
            'Name': targets1[i],
            'Matched_Name': matched_name,
            'score': min_dist,
            'score_final': score_final
        })
    else:
        results.append({
            'Name': targets1[i],
            'Matched_Name': '',
            'Matched_REF': "0",
            'score': 0,
            'score_final': 0
        })

output_final = pd.DataFrame(results)

2. 预分组减少重复筛选 #

循环中每次都执行sr2==targets2[i]的筛选是重复操作，预先按sr2的值分组，直接复用分组结果：

from collections import defaultdict

# 预先按sr2分组，存储对应的sr1和长度信息
grouped_data = defaultdict(list)
for s2_val, s1_val, length in zip(sr2, sr1, len_sr2):
    grouped_data[s2_val].append((s1_val, length))

results = []
for i in range(len(targets1)):
    t1_val = targets1[i]
    t2_val = targets2[i]
    t_len = len_targets2[i]
    
    # 直接获取同前缀的候选
    candidates = grouped_data.get(t2_val, [])
    # 过滤长度范围
    filtered_candidates = [(s1, l) for s1, l in candidates if (t_len -10) <= l <= (t_len +10)]
    
    if filtered_candidates:
        min_dist = float('inf')
        best_match = None
        for s1_val, _ in filtered_candidates:
            dist = s1_val.edit_distance(t1_val)  # 后续替换为更快的库
            if dist < min_dist:
                min_dist = dist
                best_match = s1_val
        
        total_len = len(best_match) + len(t1_val)
        score_final = (total_len - min_dist) / total_len * 100
        results.append({
            'Name': t1_val,
            'Matched_Name': best_match,
            'score': min_dist,
            'score_final': score_final
        })
    else:
        results.append({
            'Name': t1_val,
            'Matched_Name': '',
            'Matched_REF': "0",
            'score': 0,
            'score_final': 0
        })

output_final = pd.DataFrame(results)

3. 用C实现的编辑距离库替代pandas原生方法 #

pandas的str.edit_distance是纯Python实现，速度较慢。推荐使用python-Levenshtein库（C语言底层），速度提升数倍：

pip install python-Levenshtein

替换后的计算代码：

import Levenshtein

# 在循环中替换编辑距离计算
dist = Levenshtein.distance(s1_val, t1_val)

4. 并行计算分散任务 #

利用多进程并行处理每个目标的匹配任务，充分利用CPU多核：

from concurrent.futures import ProcessPoolExecutor
import Levenshtein

def process_single_target(args):
    t1_val, t2_val, t_len, grouped_data = args
    candidates = grouped_data.get(t2_val, [])
    filtered_candidates = [(s1, l) for s1, l in candidates if (t_len -10) <= l <= (t_len +10)]
    
    if not filtered_candidates:
        return {
            'Name': t1_val,
            'Matched_Name': '',
            'Matched_REF': "0",
            'score': 0,
            'score_final': 0
        }
    
    min_dist = float('inf')
    best_match = None
    for s1_val, _ in filtered_candidates:
        dist = Levenshtein.distance(s1_val, t1_val)
        if dist < min_dist:
            min_dist = dist
            best_match = s1_val
    
    total_len = len(best_match) + len(t1_val)
    score_final = (total_len - min_dist) / total_len * 100
    return {
        'Name': t1_val,
        'Matched_Name': best_match,
        'score': min_dist,
        'score_final': score_final
    }

# 准备任务参数
task_args = [
    (t1, t2, t_len, grouped_data) 
    for t1, t2, t_len in zip(targets1, targets2, len_targets2)
]

# 并行执行
with ProcessPoolExecutor() as executor:
    results = list(executor.map(process_single_target, task_args))

output_final = pd.DataFrame(results)

5. 进一步缩小候选范围 #

除了前缀和长度过滤，还可以添加字符频率过滤：

• 预先计算每个候选字符串的字符频率字典
• 与目标字符串的字符频率做相似度比较，排除差异过大的候选，减少后续编辑距离计算量

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