Python大规模文本匹配性能优化：4M段落与3W句子匹配提速方案

数据集情况 #

• 400万个段落，每个段落长度为10-60词：

paragraphs = ['yarrow heart lamium bleeding daisy pea',  'sweet shasta sedum daisy yarrow rhododendron',  'rhododendron pea shasta daisy gladiolus heart',  'gladiolus lamium rhododendron heart pea shasta',  'heart daisy yarrow gladiolus rhododendron sedum',  'pea sedum sweet shasta yarrow bleeding', 
 'yarrow sedum lamium sweet daisy gladiolus',  'heart daisy sweet bleeding pea shasta',  'daisy sweet lamium rhododendron pea bleeding',  'daisy lamium rhododendron gladiolus yarrow',  'rhododendron daisy lamium yarrow sedum',  'bleeding sedum pea heart daisy yarrow', 
 'pea sweet yarrow gladiolus lamium shasta',  'pea rhododendron sweet daisy bleeding yarrow',  'gladiolus daisy bleeding lamium sedum shasta',  'bleeding yarrow pea sedum daisy sweet',  'lamium sweet gladiolus heart rhododendron daisy',  'lamium yarrow sedum pea heart shasta',  'shasta lamium pea heart sedum yarrow',  'lamium bleeding daisy rhododendron gladiolus pea', 
 'lamium yarrow shasta heart sweet gladiolus',  'pea shasta heart sweet yarrow gladiolus',  'sedum shasta rhododendron daisy pea bleeding',  'sedum rhododendron shasta daisy lamium sweet',  'sweet rhododendron yarrow heart sedum daisy',  'bleeding sedum heart gladiolus daisy', 
 'lamium yarrow gladiolus pea sweet rhododendron',  'pea sedum bleeding daisy rhododendron',  'shasta pea rhododendron daisy sedum sweet',  'lamium yarrow bleeding pea shasta sedum']

• 包含30000个唯一句子的集合：

set_sentences = {'bleeding daisy yarrow',  'bleeding lamium shasta',  'bleeding sweet daisy',  'daisy lamium',  'daisy shasta yarrow',  'gladiolus lamium daisy',  'gladiolus shasta',  'heart daisy lamium',  'heart shasta lamium', 
 'heart sweet daisy',  'heart sweet lamium',  'lamium daisy shasta',  'lamium sweet pea',  'lamium yarrow',  
 'pea daisy rhododendron',  'pea shasta sweet',  'pea sweet gladiolus',
  'rhododendron bleeding sedum',  'rhododendron daisy',  'rhododendron gladiolus shasta',  'sedum bleeding yarrow',  'sedum lamium bleeding',  'sweet bleeding pea',  'sweet lamium daisy',  'sweet shasta',  'yarrow gladiolus',  'yarrow sedum heart',  'yarrow sedum rhododendron',  'yarrow sedum shasta',  'yarrow sedum sweet'}

需求 #

筛选出包含上述集合中任意句子的段落，保留符合条件的段落，丢弃其余段落。

当前实现及问题 #

现有基于in关键字的实现可正常运行，但处理400万条数据时速度极慢，用swifter优化后预估仍需约5小时：

def membership_testing(para, set_item):
    for item in set_item:
        if item in para:
            return 'VALID'

df = pd.DataFrame(data={'PARAGRAPH': paragraphs})
df['VALIDITY'] = df['PARAGRAPH'].apply(lambda x: membership_testing(x, set_sentences))
df['VALIDITY'] = df['VALIDITY'].fillna('INVALID')
df = df[df['VALIDITY'] == 'VALID'].reset_index(drop=True)

优化方案建议 #

1. 正则表达式预编译批量匹配 #

把所有待匹配句子合并成单个正则表达式，利用正则引擎的高效匹配能力，避免逐句循环：

import re
import pandas as pd

# 转义特殊字符，避免正则语法冲突
escaped_sentences = [re.escape(s) for s in set_sentences]
# 用|连接所有句子，构建正则模式
pattern = re.compile('|'.join(escaped_sentences))

# 利用pandas向量化字符串操作，比apply快数倍
df = pd.DataFrame(data={'PARAGRAPH': paragraphs})
df = df[df['PARAGRAPH'].str.contains(pattern)].reset_index(drop=True)

2. Dask并行处理 #

如果单进程还是慢，用Dask实现多进程/分布式并行处理：

import dask.dataframe as dd
import re

escaped_sentences = [re.escape(s) for s in set_sentences]
pattern = re.compile('|'.join(escaped_sentences))

# 用Dask读取数据（假设数据存储在CSV中）
ddf = dd.read_csv('your_paragraphs.csv')
# 并行筛选符合条件的段落
valid_ddf = ddf[ddf['PARAGRAPH'].str.contains(pattern)]
# 计算结果并转为pandas DataFrame
result_df = valid_ddf.compute()
result_df = result_df.reset_index(drop=True)

3. 文件格式优化 #

把CSV换成Parquet/Feather列式存储格式，大幅提升读取速度：

• 格式转换（仅需执行一次）：

df.to_parquet('paragraphs.parquet')

• 后续读取：

ddf = dd.read_parquet('paragraphs.parquet')

这类格式不仅读取更快，还支持压缩，节省磁盘空间。

4. AC自动机多模式匹配 #

针对3万条句子的规模，用AC自动机实现线性时间复杂度的多模式匹配，比正则更高效：

import ahocorasick
import pandas as pd

# 构建AC自动机
automaton = ahocorasick.Automaton()
for idx, sentence in enumerate(set_sentences):
    automaton.add_word(sentence, (idx, sentence))
automaton.make_automaton()

# 定义匹配函数
def has_match(para):
    for _ in automaton.iter(para):
        return True
    return False

# 批量筛选
df = pd.DataFrame(data={'PARAGRAPH': paragraphs})
df = df[df['PARAGRAPH'].apply(has_match)].reset_index(drop=True)

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