Python 性能优化：从 30 秒到 3 秒的实战记录

背景

最近接手了一个数据处理脚本，处理 100 万行 CSV 数据需要 30 秒。业务方要求优化到 5 秒以内。经过一系列优化，最终做到了 2.8 秒。记录整个过程。

原始代码问题

import csv

def process_data(filepath):
    results = []
    with open(filepath, "r") as f:
        reader = csv.reader(f)
        for row in reader:
            # 逐行处理，每次都创建新对象
            data = {"id": row[0], "value": float(row[1])}
            if data["value"] > 100:
                data["category"] = "high"
            elif data["value"] > 50:
                data["category"] = "medium"
            else:
                data["category"] = "low"
            results.append(data)
    return results

问题分析：

• 逐行读取，I/O 开销大
• 每行都创建字典对象，内存分配频繁
• 纯 Python 循环，没有利用向量化

优化方案一：使用 Pandas（15 秒）

import pandas as pd

def process_data_v2(filepath):
    df = pd.read_csv(filepath)
    df["category"] = pd.cut(
        df["value"],
        bins=[0, 50, 100, float("inf")],
        labels=["low", "medium", "high"]
    )
    return df.to_dict("records")

效果：30 秒 → 15 秒

Pandas 的向量化操作比纯 Python 循环快，但仍有优化空间。

优化方案二：Pandas + 数据类型优化（8 秒）

def process_data_v3(filepath):
    # 指定数据类型，减少内存占用
    dtypes = {"id": "string", "value": "float32"}
    df = pd.read_csv(filepath, dtype=dtypes)
    
    # 使用 numpy 的 where 进行条件判断
    import numpy as np
    df["category"] = np.where(
        df["value"] > 100, "high",
        np.where(df["value"] > 50, "medium", "low")
    )
    
    return df.to_dict("records")

效果：15 秒 → 8 秒

关键点：

• 明确指定数据类型（float32 而非 float64）
• 使用 numpy 的条件判断

优化方案三：分块处理 + 并行（4 秒）

from concurrent.futures import ProcessPoolExecutor

def process_chunk(df_chunk):
    import numpy as np
    df_chunk["category"] = np.where(
        df_chunk["value"] > 100, "high",
        np.where(df_chunk["value"] > 50, "medium", "low")
    )
    return df_chunk

def process_data_v4(filepath):
    # 分块读取
    chunks = pd.read_csv(filepath, chunksize=250000)
    
    # 并行处理
    with ProcessPoolExecutor() as executor:
        results = list(executor.map(process_chunk, chunks))
    
    return pd.concat(results).to_dict("records")

效果：8 秒 → 4 秒

利用多核 CPU 并行处理，适合大数据量场景。

终极优化：Polars（2.8 秒）

import polars as pl

def process_data_v5(filepath):
    df = pl.read_csv(filepath)
    
    df = df.with_columns(
        pl.when(pl.col("value") > 100)
        .then(pl.lit("high"))
        .when(pl.col("value") > 50)
        .then(pl.lit("medium"))
        .otherwise(pl.lit("low"))
        .alias("category")
    )
    
    return df.to_dicts()

效果：4 秒 → 2.8 秒

Polars 是 Rust 编写的高性能 DataFrame 库，比 Pandas 快 3-10 倍。

性能对比总结

方案	耗时	提升倍数
原始代码	30 秒	1x
Pandas	15 秒	2x
Pandas + 类型优化	8 秒	3.75x
分块并行	4 秒	7.5x
Polars	2.8 秒	10.7x

关键经验

1. 优先选择向量化操作：避免 Python 循环
2. 数据类型很重要：float32 比 float64 快且省内存
3. 考虑替代方案：Polars 在大数据场景下优势明显
4. 并行要适度：小数据量并行反而更慢
5. ** profiling 先行**：用 cProfile 找出真正的瓶颈

工具推荐

• 性能分析：cProfile, line_profiler
• 内存分析：memory_profiler, tracemalloc
• 可视化：snakeviz, py-spy

结语

性能优化不是一蹴而就，而是循序渐进的过程。从最简单的改进开始，逐步深入，往往能获得意想不到的效果。

---

你的 Python 代码有过类似的优化经历吗？欢迎分享！