使用Polars并行处理超内存Parquet数据集的高效安全实现方案咨询

我来帮你梳理这个场景下的Polars使用问题，结合你的代码和疑问逐一解答，再给出更优的实现思路：

你当前的核心场景是：用Polars polars~=1.24.0 的scan_parquet加载35个超内存级Parquet文件得到LazyFrame，想要通过ThreadPoolExecutor多线程批量处理数据，同时避免本地OOM问题，还编写了一个分片处理的parse函数。下面针对你的三个问题逐一解答，并给出优化建议：

问题1：是否安全将pl.LazyFrame传递到其他线程？ #

完全安全。Polars的LazyFrame本质是不可变的查询计划，它本身不存储任何实际数据，仅记录了要执行的操作序列。多个线程同时基于同一个LazyFrame生成各自的查询分支（比如slice操作）时，每个操作都会生成新的查询计划节点，不会修改原LazyFrame的内容。只要你只做只读操作（如slice、collect），就不会有线程安全问题。

问题2：用slice只读访问并覆盖所有行是否安全？ #

安全，前提是你的分片逻辑准确。

• LazyFrame的slice仅在查询计划中添加分片操作，不会提前加载数据，属于纯只读的逻辑操作，不会对原数据或查询计划造成修改。
• 你的分片逻辑（基于divmod分配范围）只要确保size是准确的总行数，就能覆盖所有行。需要注意：如果查询计划中包含filter、join等会改变行数的操作，建议用df.estimated_row_count()或df.collect_estimated_size()获取更精准的行数参考，避免漏行或越界。

问题3：collect导致OOM的问题与优化方向 #

你在htop中观察到的内存占用，是因为当前代码会一次性将整个分片的数据加载到内存中。要避免OOM，核心是流式处理小批量数据，而非一次性加载大分片：

1. 不要给线程分配大分片后全量collect，而是用collect_streaming（或collect(new_streaming=True)）结合batch_size参数，让Polars流式返回小批量数据，每处理完一批就释放内存。
2. 优先按单个Parquet文件分片：你有35个独立的Parquet文件，这是天然的内存友好分片，每个线程处理一个文件的流式数据，逻辑更简单，内存可控性更强。

优化后的实现示例 #

方案1：基于LazyFrame分片的流式处理

import polars as pl
from concurrent.futures import ThreadPoolExecutor

def parse(df: pl.LazyFrame, rank: int, total_rows: int, batch_size: int = 1000):
    # 计算当前线程的分片范围
    split, remain = divmod(total_rows, rank)
    start = rank * split + min(rank, remain)
    end = start + split + (1 if rank < remain else 0)
    chunk_size = end - start

    # 生成当前线程的分片LazyFrame
    chunk_df = df.slice(start, chunk_size)

    # 流式处理小批次，避免一次性加载全部分片
    for batch in chunk_df.collect_streaming(batch_size=batch_size):
        # 推荐用named=True获取带字段名的行，更易处理
        for row in batch.iter_rows(named=True):
            # 处理单条数据
            pass

# 加载LazyFrame
df = pl.scan_parquet("base-name-*.parquet")
# 获取总行数（若有数据变换操作，estimated_row_count可能需替换为collect统计）
total_rows = df.estimated_row_count()

# 启动线程池处理
num_threads = 4
with ThreadPoolExecutor(max_workers=num_threads) as executor:
    futures = [
        executor.submit(parse, df, rank, total_rows, batch_size=2000)
        for rank in range(num_threads)
    ]
    # 等待所有任务完成
    for future in futures:
        future.result()

方案2：按单个Parquet文件分片（更简单安全）

import polars as pl
from concurrent.futures import ThreadPoolExecutor
import glob

def process_file(file_path: str, batch_size: int = 1000):
    # 流式读取单个Parquet文件的小批次数据
    for batch in pl.scan_parquet(file_path).collect_streaming(batch_size=batch_size):
        # 批量处理示例：直接用Polars表达式处理列，比逐行更高效
        processed_batch = batch.with_columns(
            pl.col("some_column").str.to_upper()
        )
        # 再导出处理后的数据
        for row in processed_batch.iter_rows(named=True):
            # 处理单条数据
            pass

# 获取所有Parquet文件路径
file_paths = glob.glob("base-name-*.parquet")

# 启动线程池处理文件
with ThreadPoolExecutor(max_workers=4) as executor:
    executor.map(process_file, file_paths)

额外OOM规避建议 #

1. 优先批量处理：尽量用Polars的表达式在DataFrame层面批量处理数据，而非逐行处理，既能提升效率，也能降低内存波动。
2. 监控内存：通过pl.Config.set_tracker_memory(True)跟踪Polars的内存使用，或用psutil库动态调整batch_size。
3. 避免不必要的数据复制：用iter_rows(named=True)或直接操作列，减少数据转换的内存开销。

备注：内容来源于stack exchange，提问作者Arda Aytekin