如何用Pandas DataFrame单列值高效搜索无法载入内存的大型文本文件

原代码性能瓶颈分析 #

原代码采用双层嵌套循环，时间复杂度为O(大文件行数 * 匹配列表行数)，且Pandas的iterrows()遍历本身开销极大，1100万行的场景下运算量会爆炸式增长，是速度慢的核心原因。

优化核心思路 #

把需要匹配的单列值转为Python集合（set），集合的存在性判断是O(1)复杂度，直接把时间复杂度降到O(大文件行数)，性能会有数十到上百倍的提升。

优化后代码 #

通用子串匹配版本（和原代码逻辑完全一致，匹配任意位置子串） #

import os
import pandas as pd

mainPath = r'D:\AllFiles\Projects'

# 提前把匹配列转为集合，一次处理完成
act_df = pd.read_csv(os.path.join(mainPath, 'SingleColList.txt'), header=0)
match_set = set(act_df['col1'].values)

# 用上下文管理器统一管理文件句柄，不需要手动关闭
with open(os.path.join(mainPath, 'A.txt'), 'r', encoding='utf-8') as in_f, \
     open(os.path.join(mainPath, 'A_A.csv'), 'w', encoding='utf-8') as out_f:
    for line in in_f:
        # 直接用集合判断，不需要遍历匹配表
        if any(key in line for key in match_set):
            out_f.write(line)

print('Done!')

精确匹配指定列版本（如果你的匹配值是对应大文件某一列的精确值，推荐用这个，更快且不会误匹配） #

因为大文件是分号分隔的3列，如果你要匹配的是第1/2/3列的精确值，可以拆分行之后匹配，性能更高：

import os
import pandas as pd

mainPath = r'D:\AllFiles\Projects'

act_df = pd.read_csv(os.path.join(mainPath, 'SingleColList.txt'), header=0)
match_set = set(act_df['col1'].values)

with open(os.path.join(mainPath, 'A.txt'), 'r', encoding='utf-8') as in_f, \
     open(os.path.join(mainPath, 'A_A.csv'), 'w', encoding='utf-8') as out_f:
    for line in in_f:
        # 拆分大文件行，假设匹配第1列，索引从0开始，匹配第2列就写[1]，第3列写[2]
        cols = line.strip().split(';')
        if len(cols) ==3 and cols[0] in match_set:
            out_f.write(line)

print('Done!')

额外性能提升建议 #

• 如果匹配的关键词数量很大（十万级以上），可以进一步用正则表达式预编译所有关键词，一次性匹配，比逐个判断key in line更快，示例代码如下：

  import re
  # 预编译正则，注意关键词有特殊字符的话用re.escape转义
  pattern = re.compile('|'.join(re.escape(key) for key in match_set))
  # 每行判断改成
  if pattern.search(line):
      out_f.write(line)
  

• 大文件如果是GBK等非utf-8编码，修改open函数的encoding参数对应即可，避免编码报错。

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