使用Polars处理大数据集聚合时内核崩溃的原因与解决方法

核心问题分析 #

代码中虽启用了streaming=True，但两次嵌套的group_by聚合可能导致中间结果内存占用过高，加上循环执行查询的累积效应，最终触发内核崩溃。以下是具体问题和对应解决方案：

1. 冗余聚合导致中间结果过大 #

两次聚合逻辑（先按ix+iy分组求和，再按ix分组求和）存在冗余——由于求和操作的可加性，最终结果等价于直接按ix分组对目标列求和。这种冗余会生成不必要的大中间结果，直接占用大量内存。

解决办法：
简化查询逻辑，合并两次聚合为单次操作，彻底避免中间大结果：

# 直接按ix分组求和，等价于原逻辑的结果
res = df.group_by(ix).agg(pl.col(col).sum()).collect(streaming=True)

2. 流式处理未覆盖全流程 #

仅在最后一步collect时启用流式处理，第一次group_by生成的中间数据集仍会在内存中完整加载。如果ix+iy的组合基数极大（如千万级），内存会被瞬间占满。

解决办法：
对中间聚合结果也启用流式处理，或通过窗口函数优化查询：

# 用窗口函数实现嵌套聚合，全程流式处理
res = df.group_by(ix).agg(
    pl.col(col).sum().over([ix, iy]).sum()
).collect(streaming=True)

3. 循环查询的内存累积 #

循环中多次执行聚合查询，每次的结果对象res若未及时释放，会导致内存持续累积，最终触发内核崩溃。

解决办法：

• 每次循环结束后手动清理内存：

import gc
# 循环内处理完结果后
del res
gc.collect()

• 合并多列聚合逻辑，减少重复扫描CSV的次数：

for ix, iy in combinations:
    query_index += 1
    t1 = time.time()
    # 一次性处理两个目标列的聚合
    res = df.group_by([ix, iy]).agg(
        pl.col('request_io_size_bytes').sum(),
        pl.col('disk_time').sum()
    ).group_by(ix).agg(
        pl.col('request_io_size_bytes').sum(),
        pl.col('disk_time').sum()
    ).collect(streaming=True)
    # 分别记录两列的耗时与内存
    time_elapsed = time.time() - t1
    log_results_to_file(result_file, time_elapsed, res.select('request_io_size_bytes').estimated_size())
    log_results_to_file(result_file, time_elapsed, res.select('disk_time').estimated_size())

4. 流式配置与读取优化 #

Polars的流式处理默认配置可能不匹配你的内存环境，或CSV读取时的类型推断、加载策略导致额外内存开销。

解决办法：

• 调整流式块大小，适配可用内存：

import polars as pl
# 设置流式处理的块大小为1GB（根据实际内存调整）
pl.Config.set_streaming_chunk_size(1024 * 1024 * 1024)

• 显式指定CSV列类型，避免自动推断的内存浪费：

df = pl.scan_csv(
    log_dir,
    dtypes={
        'request_io_size_bytes': pl.UInt64,
        'disk_time': pl.Float64
        # 其他列也按需指定类型
    },
    low_memory=True  # 启用低内存读取模式
)

内容的提问来源于stack exchange，提问作者Kamen Petkov