Python十亿行数据处理性能优化实战

1. 十亿行挑战：Python性能优化实战

最近在开发者社区掀起了一股"十亿行挑战"(1BRC)的热潮，这个挑战要求参与者处理一个包含十亿行气象站温度数据的文本文件(约13.8GB)，为每个站点计算最小值、平均值和最大值。作为一名长期使用Python的数据工程师，我决定接受这个挑战，看看Python在这种极限场景下的表现如何。

1.1 挑战背景与规则解析

十亿行挑战的核心规则相当明确：

• 输入是一个13.8GB的文本文件，每行格式为"站点名称;温度值"
• 需要为每个站点计算温度的最小值、平均值和最大值
• 官方规则限制只能使用语言的标准库

我的测试环境是一台16英寸M3 Pro Macbook Pro，配备12核CPU和36GB内存。为了确保结果可靠，每个实现我都运行了5次并取平均值。

2. 纯Python实现方案

2.1 单核基础实现

最直接的实现方式是单线程逐行读取文件，使用字典记录每个站点的统计信息：

python复制def process_file(file_name: str):
    result = dict()
    with open(file_name, "rb") as f:
        for line in f:
            location, measurement = line.split(b";")
            measurement = float(measurement)
            if location not in result:
                result[location] = [
                    measurement,  # min
                    measurement,  # max
                    measurement,  # sum
                    1,           # count
                ]
            else:
                _result = result[location]
                _result[0] = min(_result[0], measurement)
                _result[1] = max(_result[1], measurement)
                _result[2] += measurement
                _result[3] += 1
    # 输出结果...

这个实现简单直接，但性能显然不会太好。在我的测试中，处理完整数据集需要约10分钟。

注意：使用二进制模式('rb')读取文件比文本模式稍快，因为避免了编码解码的开销。对于纯ASCII数据，这是安全的优化。

2.2 多核并行处理

为了利用多核CPU，我们可以将文件分割成多个块，并行处理：

python复制import os
import multiprocessing as mp

def get_file_chunks(file_name: str, max_cpu: int = 8):
    """将文件分割为适合多进程处理的块"""
    cpu_count = min(max_cpu, mp.cpu_count())
    file_size = os.path.getsize(file_name)
    chunk_size = file_size // cpu_count
    
    start_end = []
    with open(file_name, "r+b") as f:
        chunk_start = 0
        while chunk_start < file_size:
            chunk_end = min(file_size, chunk_start + chunk_size)
            # 确保块结束在行边界
            while chunk_end > chunk_start:
                f.seek(chunk_end)
                if f.read(1) == b"\n":
                    break
                chunk_end -= 1
            start_end.append((file_name, chunk_start, chunk_end))
            chunk_start = chunk_end + 1
    return cpu_count, start_end

def _process_file_chunk(file_name, chunk_start, chunk_end):
    """处理单个文件块的函数"""
    result = {}
    with open(file_name, "rb") as f:
        f.seek(chunk_start)
        while f.tell() < chunk_end:
            line = f.readline()
            # 处理逻辑与单核版本相同...
    return result

def process_file(cpu_count, start_end):
    """主处理函数"""
    with mp.Pool(cpu_count) as pool:
        chunk_results = pool.starmap(_process_file_chunk, start_end)
    # 合并各块结果...

多核实现将处理时间从10分钟缩短到约2分钟，提升显著。关键在于：

1. 合理分割文件，确保每个块在行边界结束
2. 使用multiprocessing.Pool管理进程池
3. 最后合并各进程的统计结果

2.3 PyPy加速方案

PyPy是Python的即时编译实现，对纯Python代码通常有显著加速效果。我们只需用PyPy解释器运行相同的多核代码：

bash复制pypy multiprocessing_impl.py

PyPy版本将处理时间进一步缩短到约90秒，比CPython快约33%。PyPy的优势在于：

• 即时编译热点代码
• 更高效的内存管理
• 优化的内置类型实现

实测心得：PyPy对纯Python代码优化效果显著，但如果使用了C扩展模块，可能反而会变慢。在这个挑战中，由于只使用标准库，PyPy是最简单的优化手段。

3. 第三方库性能对比

虽然使用第三方库违反官方规则，但实际工作中我们完全可以利用这些工具。我测试了四个主流数据处理库的表现。

3.1 Pandas实现

Pandas是最流行的数据分析库，虽然不以速度见长：

python复制import pandas as pd

df = pd.read_csv("data/measurements.txt", sep=";", 
                header=None, names=["station", "temp"],
                engine="pyarrow")  # 使用PyArrow引擎加速

result = df.groupby("station").agg(["min", "mean", "max"])
# 格式化输出...

Pandas实现耗时约45秒，比纯Python快很多，但内存占用较高。关键点：

• 使用PyArrow引擎替代默认的C引擎，读取速度提升2-3倍
• 避免链式操作，减少中间DataFrame创建
• 明确指定dtype可以减少内存使用

3.2 Polars实现

Polars是基于Rust的多线程DataFrame库：

python复制import polars as pl

df = pl.scan_csv("data/measurements.txt", separator=";",
                has_header=False, new_columns=["station", "temp"])
        .group_by("station")
        .agg([
            pl.min("temp").alias("min"),
            pl.mean("temp").alias("mean"),
            pl.max("temp").alias("max")
        ])
        .sort("station")
        .collect(streaming=True)  # 流式处理降低内存占用

Polars仅用约15秒就完成了处理，优势在于：

• 多线程并行处理
• 查询优化器重新排列操作顺序
• 零拷贝数据处理
• 流式处理支持大数据集

3.3 DuckDB实现

DuckDB是嵌入式OLAP数据库，擅长分析查询：

python复制import duckdb

conn = duckdb.connect()
result = conn.execute("""
    SELECT station,
        min(temp) as min,
        avg(temp) as mean,
        max(temp) as max
    FROM read_csv('data/measurements.txt',
                 delim=';',
                 header=false,
                 columns={'station':'VARCHAR','temp':'FLOAT'},
                 parallel=true)
    GROUP BY station
    ORDER BY station
""").fetchall()

DuckDB表现最佳，仅需约12秒。其优势包括：

• 向量化查询执行
• 多核并行处理
• 高效的聚合算法
• 可以直接查询CSV文件而无需导入

3.4 性能对比总结

各方案在13.8GB文本数据上的表现：

从结果看，DuckDB和Polars是性能最好的解决方案，而Pandas在易用性和生态系统方面仍有优势。

4. 文件格式优化：Parquet的威力

原始文本格式效率低下，我们可以将数据转换为列式存储的Parquet格式：

python复制# 转换脚本
import pandas as pd

df = pd.read_csv("measurements.txt", sep=";", header=None,
                names=["station", "temp"], engine="pyarrow")
df.to_parquet("measurements.parquet", compression="zstd")

转换后文件从13.8GB缩小到2.5GB。各库处理Parquet的表现：

DuckDB处理Parquet仅需3.8秒，比处理文本快3倍多！Parquet的优势：

• 列式存储，只读取需要的列
• 内置压缩，减少I/O
• 元数据丰富，优化查询
• 谓词下推，减少数据扫描量

5. 实战经验与避坑指南

5.1 内存管理技巧

处理大数据时，内存是关键限制因素：

• 使用迭代或流式处理替代全量加载
• 明确指定数据类型减少内存占用
• 及时释放不再需要的数据
• 考虑使用内存映射文件

Polars示例：

python复制# 流式处理避免内存溢出
df = pl.scan_csv("data.txt").group_by("station").agg(
    pl.col("temp").min().alias("min"),
    # ...
).collect(streaming=True)

5.2 性能调优要点

• I/O优化：使用更快的存储设备，或者将数据放在内存文件系统中
• 并行度设置：根据CPU核心数和内存带宽调整并行度
• 批处理大小：适当调整批处理大小平衡内存和CPU利用率
• 预热：对于JIT编译的系统，先运行热身查询

5.3 常见问题排查

问题1：内存不足错误

• 解决方案：使用流式处理，或者增加机器内存
• DuckDB示例：SET memory_limit='16GB'

问题2：性能不如预期

• 检查CPU利用率是否达到100%
• 确认是否使用了所有核心
• 检查磁盘I/O是否成为瓶颈

问题3：结果不正确

• 验证文件分割是否破坏了数据完整性
• 检查聚合逻辑是否正确
• 确保数据类型转换无误

6. 总结与选型建议

经过全面测试，我的结论是：

1. 遵守规则场景：PyPy多核实现是最佳选择，约90秒完成
2. 追求极致性能：DuckDB+Parquet组合仅需3.8秒
3. 开发效率优先：Polars提供了良好的平衡，约8秒完成且API友好

Python可能永远无法在纯性能上击败Java等编译语言，但通过合理选择工具链，我们完全可以在保证开发效率的同时，获得足够好的性能表现。对于大多数实际应用场景，3.8秒处理十亿行数据已经绰绰有余。

最后的小技巧：DuckDB可以直接查询远程文件，这对于分析存储在云存储中的数据特别方便：

python复制# 直接查询S3上的Parquet文件
duckdb.sql("""
    SELECT * FROM read_parquet('s3://bucket/data.parquet')
    WHERE temperature > 30
""")