Python 大数据量分片处理：生成器 + 流式写入的生产实践

Python 大数据量分片处理生成器 流式写入的生产实践一、一次 OOM 让 200 万条数据要重新跑 6 小时周五下班前启动的数据处理任务周一早上发现已经在凌晨崩溃。日志里一行MemoryError进程被 OOM Killer 杀掉了。任务是从数据库读取 200 万条用户行为日志。做格式转换后写入 Parquet 文件。代码看起来很常规。rows db.fetch_all(SELECT * FROM user_behavior WHERE date 2026-07-07) processed [transform(row) for row in rows] write_parquet(processed, output.parquet)问题在于fetch_all一次性加载了所有数据到内存。200 万条数据每条约 2KB总计约 4GB。加上 transform 过程中的临时对象内存直接飙到 8GB。而部署容器的内存限制是 4GB。解决方案不是加内存而是分片处理。二、流式处理管道的核心模式大数据量处理的核心思想是流。数据像水流一样经过处理管道在内存中只保留当前片段。flowchart LR A[数据库游标] --|流式读取 10K 条/批| B[生成器管道] B -- C[数据清洗] C -- D[格式转换] D -- E[数据校验] E -- F[流式写入 Parquet] G[错误批次] -- H[死信队列] H -- I[人工或自动重试] style A fill:#e1f5fe style F fill:#c8e6c9 style H fill:#ffcdd2三个核心模式生成器惰性产出数据不预加载。分片迭代每次处理固定数量的记录。流式写入边处理边写不积攒。三、完整的生产级实现 batch_processor.py - 大数据量分片处理 import json import logging import time from dataclasses import dataclass from typing import ( Any, Callable, Generator, Iterator, List, Optional ) from pathlib import Path import pyarrow as pa import pyarrow.parquet as pq logging.basicConfig(levellogging.INFO) logger logging.getLogger(__name__) dataclass class BatchConfig: 分片配置 batch_size: int 10000 # 每批记录数 max_memory_mb: int 500 # 内存限制 output_dir: str ./output dead_letter_dir: str ./dead_letter progress_interval: int 100000 # 每处理 N 条输出进度 class DataPipeline: 数据管线 - 生成器 流式写入 def __init__( self, config: BatchConfig, transformer: Callable[[dict], Optional[dict]], validator: Optional[Callable[[dict], bool]] None, ): self.config config self.transformer transformer self.validator validator or (lambda _: True) self.metrics PipelineMetrics() def source( self, query: str, db_conn ) - Generator[dict, None, None]: 流式数据源使用服务端游标 使用 SSCursor 或 stream_results 避免客户端缓存全部结果 # MySQL 示例 # conn pymysql.connect(..., cursorclasspymysql.cursors.SSCursor) cursor db_conn.cursor() cursor.execute(query) while True: rows cursor.fetchmany(self.config.batch_size) if not rows: break for row in rows: yield row def batched( self, source: Iterator[dict] ) - Generator[List[dict], None, None]: 将流式数据分片 batch [] for item in source: batch.append(item) if len(batch) self.config.batch_size: yield batch batch [] if batch: yield batch def transform_batch( self, batch: List[dict] ) - tuple[List[dict], List[dict]]: 转换一个批次分离正常数据和异常数据 success [] failed [] for record in batch: try: # 校验 if not self.validator(record): failed.append({ record: record, error: validation_failed, }) continue # 转换 result self.transformer(record) if result is not None: success.append(result) except Exception as e: failed.append({ record: record, error: str(e), }) logger.warning(f记录处理失败: {e}) return success, failed def write_partition( self, batch: List[dict], partition_key: str, schema: pa.Schema, ): 流式写入追加到分区 Parquet 文件 output_dir Path(self.config.output_dir) / partition_key output_dir.mkdir(parentsTrue, exist_okTrue) table pa.Table.from_pylist(batch, schemaschema) # 追加写入同一分区 pq.write_to_dataset( table, root_pathstr(output_dir), partition_cols[], existing_data_behavioroverwrite_or_ignore, ) def handle_failed(self, failed_batch: List[dict]): 将失败记录写入死信队列 dl_dir Path(self.config.dead_letter_dir) dl_dir.mkdir(parentsTrue, exist_okTrue) timestamp int(time.time()) dl_file dl_dir / ffailed_{timestamp}.json with open(dl_file, w) as f: json.dump(failed_batch, f, ensure_asciiFalse, defaultstr) logger.warning( f{len(failed_batch)} 条记录写入死信: {dl_file} ) def run( self, source: Iterator[dict], schema: pa.Schema, partition_fn: Callable[[dict], str], ) - PipelineMetrics: 执行完整管线 partition_fn: 根据记录返回分区键如日期 start_time time.time() logger.info(开始数据处理...) for batch in self.batched(source): success, failed self.transform_batch(batch) if success: # 按分区键分组写入 partitions {} for record in success: key partition_fn(record) partitions.setdefault(key, []).append(record) for key, records in partitions.items(): self.write_partition(records, key, schema) self.metrics.success_count len(success) self.metrics.output_records len(success) if failed: self.handle_failed(failed) self.metrics.failed_count len(failed) self.metrics.total_processed len(batch) # 定时输出进度 if self.metrics.total_processed % self.config.progress_interval 0: elapsed time.time() - start_time rate self.metrics.total_processed / max(elapsed, 0.001) logger.info( f进度: {self.metrics.total_processed} 条 f(成功 {self.metrics.success_count}, f失败 {self.metrics.failed_count}) f速率 {rate:.0f} 条/s ) elapsed time.time() - start_time self.metrics.duration elapsed logger.info( f完成! 共处理 {self.metrics.total_processed} 条, f耗时 {elapsed:.1f}s, f成功率 {self.metrics.success_rate():.1%} ) return self.metrics dataclass class PipelineMetrics: 管线指标 total_processed: int 0 success_count: int 0 failed_count: int 0 output_records: int 0 duration: float 0.0 def success_rate(self) - float: if self.total_processed 0: return 0.0 return self.success_count / self.total_processed # ---- 使用示例 ---- def demo(): # 定义数据 Schema output_schema pa.schema([ (user_id, pa.string()), (event_type, pa.string()), (event_time, pa.timestamp(ms)), (processed_at, pa.string()), ]) # 转换函数 def transform_row(row: dict) - Optional[dict]: return { user_id: str(row[user_id]), event_type: row[event_type].lower(), event_time: row[timestamp], processed_at: time.strftime(%Y-%m-%d %H:%M:%S), } # 校验函数 def validate_row(row: dict) - bool: return bool(row.get(user_id)) and bool(row.get(event_type)) # 分区函数 def partition_by_date(row: dict) - str: return row[event_time].strftime(%Y%m%d) config BatchConfig( batch_size5000, output_dir./output/events, ) pipeline DataPipeline(config, transform_row, validate_row) # 启动管线需要实际的数据库连接 # source pipeline.source(SELECT * FROM events, db_conn) # metrics pipeline.run(source, output_schema, partition_by_date)四、分片批次的调优边界批次大小影响 IO 效率和内存占用。批太小100 条数据库往返次数多吞吐低。批太大10 万条单批内存占用高GC 频繁。经验值每条记录 2KB 时5000-10000 条/批较合适。死信队列需要监控。失败率超过 1% 应触发告警。死信数据定期检查修复上游数据质量。不适合流式处理的场景需要对数据进行全局排序的任务需要 JOIN 多表且无法按分片键关联的场景数据量 10 万条的小数据直接用 Pandas 更简单。五、总结大数据量处理的核心是流式 分片。生成器惰性读取避免一次性加载。批次迭代控制内存占用。流式写入实现边处理边存储。死信队列保障异常数据不丢失。通过批次大小调优可以平衡吞吐和内存。