如何向量化优化文本匹配循环以提升Python数据处理速度？

我有个Python函数check_segment，用来给包含name列的DataFrame打分类标签。流程是先清洗name字段文本，再依次循环匹配各类预设列表。但第一个循环用textdistance.jaro_winkler计算文本距离，导致整体运行耗时数小时，想通过向量化改造提速。

原核心函数代码 #

def check_segment(name):
    name = re.sub('\W+', ' ', name)
    name = re.sub('\*', '', name)
    
    for i in top_1k:
        x = textdistance.jaro_winkler(i, name)
        if x > 0.8:
            return 'Top 1K'
            
    for i in school:
        if i in name:
            return 'School'
    
    for i in hospital:
        if i in name:
            return 'Hospital'
        
    for i in govt:
        if i in name:
            return 'Government'
    
    for i in coop:
        if i in name:
            return 'Cooperative'
    
    for i in banks:
        if i in name:
            return 'Bank'
    
    for i in sarisari:
        if i in name:
            return 'Sari-Sari Store'
        
    for i in malls:
        if i in name:
            return 'Malls'
    
    for i in remittance_center:
        if i in name:
            return 'Remittance Center'
    
    for i in hotel:
        if i in name:
            return 'Hotels'
    
    for i in foundation:
        if i in name:
            return 'Foundation'
    
    for i in embassy:
        if i in name:
            return 'Embassy'
    
    return 'SME'

原分类匹配列表 #

school = ["UNIVERSITY","ACADEMY","COLLEGE","ACADEME","SCHOOL","MONTESSORI","ELEMENTARY","HIGH SCHOOL","COLLEGIO","INSTITUTE"]
hospital = ["HOSPITAL","LABORATORY","CLINIC","MEDICAL","DIAGNOSTIC","HEALTH","DOCTOR", "HEALTHCARE"]
govt = ["DEPARTMENT OF","CITY GOVERNMENT","OFFICE OF THE","PROVINCE OF","PROVINCIAL","CITY TREASURER","REGISTRY OF","REGISTER OF",
          "BUREAU OF","MUNICIPAL","COMMISSION","PEZA","HDMF","WATER DISTRICT","HOME DEVELOPMENT MUTUAL FUND","CLERK OF COURT", 
          "CITY OF","BARANGAY", "GOVERNMENT"]
coop = ["COOP", "COOPERATIVE"]
hotel = ["HOTEL","RESORT", "CONDOTEL", "HOTELIERS", "INN"]
foundation = ["FOUNDATION"]
embassy = ["EMBASSY"]

原调用方式 #

df['segment'] = df['name'].apply(check_segment)

输入输出示例 #

输入DataFrame #

输出DataFrame #

向量化优化方案 #

原函数的性能瓶颈主要在两个地方：逐个循环的文本清洗，以及Top1K列表的逐元素Jaro-Winkler距离计算。以下是针对性的优化：

1. 向量化文本清洗 #

用Pandas内置的字符串方法批量处理，比apply里逐个调用re.sub快数倍：

import pandas as pd

# 统一清洗所有name字段，转大写避免大小写匹配问题
df['clean_name'] = df['name'].str.replace('\W+', ' ', regex=True) \
                             .str.replace('*', '', regex=False) \
                             .str.upper()
# 同步把top_1k转成大写，保持匹配一致性
top_1k = [item.upper() for item in top_1k]

2. 批量处理子串匹配类分类 #

原函数中除了Top1K，其他分类都是子串匹配，可通过正则表达式批量生成匹配掩码，按优先级依次赋值，彻底摆脱循环：

# 按原函数的优先级定义分类规则（类别名 → 匹配正则）
category_rules = [
    ("School", "|".join(school)),
    ("Hospital", "|".join(hospital)),
    ("Government", "|".join(govt)),
    ("Cooperative", "|".join(coop)),
    ("Bank", "|".join(banks)),
    ("Sari-Sari Store", "|".join(sarisari)),
    ("Malls", "|".join(malls)),
    ("Remittance Center", "|".join(remittance_center)),
    ("Hotels", "|".join(hotel)),
    ("Foundation", "|".join(foundation)),
    ("Embassy", "|".join(embassy))
]

# 初始化所有分类为SME
df['segment'] = "SME"

# 按优先级依次匹配，仅对未分类的行更新结果
for category, pattern in category_rules:
    mask = (df['segment'] == "SME") & df['clean_name'].str.contains(pattern, regex=True)
    df.loc[mask, 'segment'] = category

3. 优化Top1K的Jaro-Winkler匹配 #

这部分是原函数最大的性能瓶颈（时间复杂度O(N*M)），通过过滤长度差异过大的元素减少计算量，再结合swifter自动加速apply：

import textdistance
import swifter

def match_top1k(name):
    name_len = len(name)
    # 只保留长度在当前name的80%-120%之间的Top1K元素，减少无效计算
    filtered_top1k = [item for item in top_1k if 0.8 * name_len <= len(item) <= 1.2 * name_len]
    # 检查是否有匹配度超过0.8的元素
    return any(textdistance.jaro_winkler(item, name) > 0.8 for item in filtered_top1k)

# 对未分类的行检查Top1K匹配
mask = df['segment'] == "SME"
df.loc[mask, 'segment'] = df.loc[mask, 'clean_name'].swifter.apply(
    lambda x: "Top 1K" if match_top1k(x) else "SME"
)

优化后调用说明 #

直接运行上述代码即可完成分类，无需再调用原check_segment函数。

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