Python代码无法从较大CSV文件提取指定行，请求技术支持

问题描述 #

我需要从指定文件夹的不同文件中，根据不同值提取对应行（例如从file1提取value1相关行、file2提取value2相关行等）。以下是我的实现代码，该代码对4KB、9KB的小型CSV文件有效，但无法处理3.8MB的CSV文件，恳请协助解决。

原代码 #

import glob, os
import pandas as pd

def clean_files(input_folder_path,output_folder_path, input_trade):
     
        # 获取指定文件夹下目标CSV文件列表
        files =[file for file in os.listdir(input_folder_path) if file.endswith('INTEREST_RATE_SWAP_OTC_20231019.csv')]
        #print(files)
        
        # 初始化空DataFrame存储合并结果
        rows = pd.DataFrame()
        # 遍历每个文件并读取为DataFrame
        for file in files:
            file_path = os.path.join(input_folder_path,file)
           
            # 跳过前5行无效数据
            df = pd.read_csv(file_path,skiprows = 5,low_memory=False)

            # 数据清洗
            df1 = df.iloc[:-1,:-2].set_axis(df.columns[2:], axis=1)
            clean_data = df1.drop(df1.columns[[0]], axis=1, inplace = False)
          
            # 逐个查找目标trade_id对应的行
            for trades in input_trade:
                trade_rows = clean_data[clean_data['TRADE_ID'] == trades]
                print(trade_rows)
                if not trade_rows.empty:
                   rows = rows._append(trade_rows, ignore_index = True)    
              
        # 将合并结果写入输出文件夹的CSV文件
        output_file_path = os.path.join(output_folder_path, f"FIXCF_NH3.csv")
        rows.to_csv(output_file_path,index = False)
        # 若仅需写入文件可删除return语句
        return clean_data

问题分析 #

原代码无法处理大文件的核心原因：

1. 内存过载：一次性将3.8MB的CSV文件全部加载到内存，若文件包含大量行/列，极易触发内存不足
2. 效率低下：遍历每个trade_id时重复扫描整个DataFrame，时间复杂度高
3. 追加性能差：使用已弃用的_append方法逐次合并数据，频繁生成新DataFrame，损耗性能

优化方案 #

采用分块读取+集合快速查找+批量存储结果的方式，大幅降低内存占用并提升效率：

import os
import pandas as pd

def clean_files(input_folder_path, output_folder_path, input_trade):
    # 将trade_id转为集合，提升查找匹配速度
    trade_set = set(input_trade)
    # 用列表存储筛选结果块，避免频繁DataFrame追加操作
    result_rows = []
    
    # 筛选目标CSV文件
    files = [file for file in os.listdir(input_folder_path) 
             if file.endswith('INTEREST_RATE_SWAP_OTC_20231019.csv')]
    
    for file in files:
        file_path = os.path.join(input_folder_path, file)
        
        # 分块读取大文件，chunksize可根据内存情况调整（示例为1000行/块）
        chunk_iter = pd.read_csv(file_path, skiprows=5, low_memory=False, chunksize=1000)
        
        for chunk in chunk_iter:
            # 先完成数据清洗，再筛选目标行，减少后续处理的数据量
            cleaned_chunk = chunk.iloc[:-1, :-2]
            cleaned_chunk.columns = chunk.columns[2:]
            cleaned_chunk = cleaned_chunk.drop(cleaned_chunk.columns[0], axis=1)
            
            # 批量筛选包含目标trade_id的行，集合in操作比逐个遍历效率更高
            filtered_rows = cleaned_chunk[cleaned_chunk['TRADE_ID'].isin(trade_set)]
            if not filtered_rows.empty:
                result_rows.append(filtered_rows)
    
    # 合并所有结果块并写入输出文件
    if result_rows:
        final_df = pd.concat(result_rows, ignore_index=True)
        output_file_path = os.path.join(output_folder_path, "FIXCF_NH3.csv")
        final_df.to_csv(output_file_path, index=False)
    
    # 返回最后一个文件的清洗后数据（若不需要可删除此语句）
    return cleaned_chunk if 'cleaned_chunk' in locals() else None

优化说明 #

1. 分块读取：通过chunksize参数将大文件拆分成小块处理，每次仅加载部分数据到内存，避免内存过载
2. 集合查找：将input_trade转为集合，isin操作的时间复杂度从O(n)降为O(1)，彻底避免循环遍历每个trade_id的低效操作
3. 批量存储：用列表存储筛选后的DataFrame块，最后用pd.concat一次性合并，比逐次_append效率提升数倍
4. 提前清洗：在每个数据块上先完成清洗再筛选，减少后续需要处理的数据量

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