高效聚合大量Polars DataFrame实现逐行累加/均值并控内存

核心实现思路 #

要严格控制内存占用，核心是避免同时加载所有DataFrame到内存，采用逐列、逐帧的增量累加方式。以下是一个鲁棒的实现：

import polars as pl

def compute_rowwise_accumulation(df_generator, compute_mean=False):
    # 从生成器中获取第一个DataFrame，确定列名与初始类型
    first_df = next(df_generator)
    col_names = first_df.columns
    # 初始化累加器：与第一帧行数相同的零值Series，匹配原列类型
    accumulator = {col: pl.Series([0] * len(first_df), dtype=first_df[col].dtype) for col in col_names}
    frame_count = 1

    # 遍历剩余DataFrame，逐列增量累加
    for df in df_generator:
        # 严格按预定义列名访问，确保列对齐
        for col in col_names:
            accumulator[col] = accumulator[col] + df[col]
        frame_count += 1

    # 组合结果DataFrame
    result = pl.DataFrame({col: accumulator[col] for col in col_names})
    
    # 若需计算均值，除以总帧数
    if compute_mean:
        result = result / frame_count
    
    return result

使用时，只需将DataFrame的生成器（比如从文件逐次加载的迭代器）传入即可，全程仅需保留当前处理的DataFrame和累加器在内存中。

疑问解答 #

1. Schema完全一致时，iter_columns()的顺序是否可靠？ #

当两个DataFrame的schema完全一致（包括列顺序、列名、数据类型三者均匹配），Polars保证iter_columns()的输出顺序与DataFrame的列顺序完全一致，此时依赖它是安全的。

但如果要做到绝对鲁棒（比如避免某些边缘场景下的列顺序意外变动，比如从外部数据源加载时的隐式列重排），更稳妥的方式是显式使用固定列名列表对齐：

• 从第一个DataFrame中提取列名列表col_names = first_df.columns
• 后续所有DataFrame均通过df[col]的方式按该列表顺序访问列

这种方式彻底规避了列顺序不一致的风险，代码可读性也更强。

2. LazyFrames能否进一步优化性能？ #

是的，LazyFrames在性能和内存控制上有显著优势，适合处理大规模数据场景：

• 延迟执行：LazyFrames不会立即加载数据到内存，直到调用collect()，全程仅维护执行计划，内存占用极低。
• 查询优化：Polars的查询优化器会自动合并、简化操作逻辑，减少不必要的计算步骤。
• 分块并行处理：配合pl.scan_csv()/pl.scan_parquet()等扫描函数，Polars会分块读取数据，并行处理后释放内存，进一步降低峰值内存占用。

以下是LazyFrames的实现示例（以CSV文件为例）：

import polars as pl
from functools import reduce

def lazy_rowwise_accumulation(file_paths, compute_mean=False):
    # 扫描第一个文件获取列名
    first_scan = pl.scan_csv(file_paths[0])
    col_names = first_scan.columns

    # 定义累加逻辑：对齐列名后逐列相加
    def accumulate(lf_a, lf_b):
        return pl.DataFrame({col: lf_a[col] + lf_b[col] for col in col_names}).lazy()

    # 将所有文件转为LazyFrame
    lazy_frames = [pl.scan_csv(path) for path in file_paths]

    # 累加所有LazyFrame
    accumulated_lf = reduce(accumulate, lazy_frames[1:], lazy_frames[0])

    # 计算均值（若需要）
    if compute_mean:
        accumulated_lf = accumulated_lf / len(file_paths)

    # 执行并返回结果
    return accumulated_lf.collect()

额外优化建议 #

• 数据类型优化：在精度允许的情况下，使用float32替代float64、int32替代int64，大幅降低内存占用。
• 分块加载：对于单个超大文件，使用pl.read_csv(batch_size=...)分块加载，逐块累加，进一步控制峰值内存。
• 并行配置：通过pl.Config.set_threads(n)设置线程数，充分利用多核CPU提升处理速度。

内容的提问来源于stack exchange，提问作者Benjamin Trendelkamp-Schroer