优化文本国家识别函数 提升大数据集处理速度的技术问询

核心优化方向 #

原函数在处理大数据集时的主要性能瓶颈包括：重复加载文件、逐行文本预处理效率低、逐国匹配的线性搜索、以及Pandas逐行apply的开销。以下是针对性的优化方案：

1. 预加载静态数据，避免重复IO操作 #

每次调用函数都加载country_names.txt是最大的性能浪费，将数据预加载到模块级变量，仅初始化一次：

import json
import string
import re
from typing import List, Union
import pandas as pd

# 模块级预加载，仅执行一次
with open('country_names.txt') as file:
    _COUNTRY_NAMES = json.load(file)

# 提前编译正则表达式、生成标点转换表，避免重复构建
_REPLACE_PUNCT_PAT = re.compile(r"['’-]")
_MY_PUNCT = string.punctuation + '”“'
_PUNCT_TRANS_TABLE = str.maketrans('', '', _MY_PUNCT)

# 预构建国家到所有提及形式的映射（包含国家名本身）
_COUNTRY_TERMS = {
    country.lower(): {country.lower()} | set(alias.lower() for alias in aliases)
    for country, aliases in _COUNTRY_NAMES.items()
}

2. 优化搜索逻辑：用正则批量匹配替代逐国检查 #

原函数循环遍历每个国家并检查子串，效率极低。改为将所有提及形式整合成一个正则表达式，一次匹配所有目标术语：

def find_countries(text: str, exclude: Union[str, List[str]] = [], extra: Union[str, List[str]] = []) -> List[str]:
    # 统一处理参数格式并转小写
    exclude = [exclude.lower()] if isinstance(exclude, str) else [e.lower() for e in exclude]
    extra = [extra.lower()] if isinstance(extra, str) else [e.lower() for e in extra]

    # 构建当前可用的术语集合：排除指定国家，添加额外术语
    current_terms = {}
    for country, terms in _COUNTRY_TERMS.items():
        if country not in exclude:
            current_terms[country] = terms
    for term in extra:
        current_terms[term] = {term}

    if not current_terms:
        return []

    # 文本预处理复用预定义的正则和转换表
    text = _REPLACE_PUNCT_PAT.sub(" ", text)
    text = text.translate(_PUNCT_TRANS_TABLE).lower()

    # 构建正则模式：按术语长度降序排序，避免短术语匹配长术语前缀（如"china"匹配"chinese"）
    all_terms = {term for terms in current_terms.values() for term in terms}
    sorted_terms = sorted(all_terms, key=len, reverse=True)
    pattern = re.compile(r'\b(' + '|'.join(re.escape(term) for term in sorted_terms) + r')\b')

    # 匹配术语并映射回对应的国家/额外项
    matches = pattern.findall(text)
    # 用反向映射加速查找，避免嵌套循环
    term_to_country = {term: country for country, terms in current_terms.items() for term in terms}
    return list({term_to_country[match] for match in matches})

3. Pandas批量处理优化：减少逐行apply开销 #

直接用pd.Series.apply逐行处理数万条数据效率低下，可通过预构建全局正则和反向映射，实现批量处理：

# 假设exclude和extra为全局固定参数，提前构建全局匹配规则
GLOBAL_EXCLUDE = []
GLOBAL_EXTRA = []

# 构建全局术语映射和反向查找表
_global_current_terms = {}
for country, terms in _COUNTRY_TERMS.items():
    if country not in GLOBAL_EXCLUDE:
        _global_current_terms[country] = terms
for term in GLOBAL_EXTRA:
    _global_current_terms[term] = {term}

_term_to_country = {term: country for country, terms in _global_current_terms.items() for term in terms}

# 预编译全局正则
_global_all_terms = sorted(_term_to_country.keys(), key=len, reverse=True)
_global_pattern = re.compile(r'\b(' + '|'.join(re.escape(t) for t in _global_all_terms) + r')\b')

def batch_process_countries(text_series: pd.Series) -> pd.Series:
    # 批量预处理文本
    processed = text_series.str.replace(_REPLACE_PUNCT_PAT, " ")
    processed = processed.str.translate(_PUNCT_TRANS_TABLE).str.lower()

    # 批量提取并映射国家
    def extract(text):
        if pd.isna(text):
            return []
        matches = _global_pattern.findall(text)
        return list({_term_to_country[match] for match in matches})

    return processed.apply(extract)

# 使用方式：
# df['mentioned_countries'] = batch_process_countries(df['text_column'])

4. 额外优化细节 #

• 单词边界匹配：用\b替代原函数的空格包裹方式，避免漏匹配文本开头/结尾的国家术语
• 反向映射表：提前构建术语到国家的映射，避免每次匹配时嵌套循环查找
• 正则转义：对所有术语做re.escape处理，避免特殊字符破坏正则结构
• 空值处理：在批量处理时明确处理NaN值，避免报错

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