如何高效优化Python脚本来处理大型CSV文件？

Question

我一直在开发一个 Python 脚本来处理大型 CSV 文件（500MB 到 1GB）。该脚本执行的任务包括根据列值过滤行、计算新字段以及将处理后的数据导出到另一个 CSV 文件。虽然该功能运行良好，但随着文件大小的增加，脚本会明显变慢，尤其是在处理接近 1GB 的文件时。

这是我当前脚本的结构：

蟒蛇 复制代码

import pandas as pd  

# Step 1: Read the CSV file  
df = pd.read_csv("large_file.csv")  

# Step 2: Filter rows  
filtered_df = df[df["column_name"] > 100]  

# Step 3: Perform calculations  
filtered_df["new_column"] = filtered_df["column_name"] * 2  

# Step 4: Export the results  
filtered_df.to_csv("filtered_output.csv", index=False)  
The main bottlenecks seem to be in the reading and filtering stages when working with large datasets.

我尝试了几种优化技术：

在 pd.read_csv 中使用 chunksize： 这允许一次处理较小的文件块。然而，它使过滤和聚合数据的逻辑变得复杂，因为它需要跨块组合结果。

切换到Dask： 我尝试使用Dask进行并行处理，这在一定程度上提高了性能。然而，对于较小的项目来说，这感觉有点矫枉过正，而且学习曲线比预期的要陡峭。

预先对 CSV 文件进行排序： 我根据过滤标准对文件进行排序，希望它能加快过滤步骤。虽然它在某些情况下有所帮助，但并不总是实用，尤其是当数据源不受我控制时。

我的期望： 实现一个平衡简单性和性能的解决方案。理想情况下，如果可能的话，我想坚持使用 pandas，但不介意探索其他轻量级解决方案。

Answer 1

您可以以可管理的块的形式处理数据

import pandas as pd

CHUNK_SIZE = 100000


first_chunk = True
for chunk in pd.read_csv("large_file.csv", chunksize=CHUNK_SIZE):
    # Filter rows
    filtered = chunk[chunk["column_name"] > 100]
    
    
    filtered["new_column"] = filtered["column_name"] * 2
    
    
    filtered.to_csv("output.csv", 
                   mode='a' if not first_chunk else 'w',
                   header=first_chunk,
                   index=False)
    first_chunk = False