Polars处理30GB大文件内存分配失败：脚本或版本问题？

你的问题大概率是脚本内存效率不足+旧版本Polars的内存管理缺陷共同导致的，同时VS Code与Jupyter的环境差异也可能是触发因素。以下是具体的排查和优化方向：

一、脚本层面的内存优化 #

你的agg操作中，list的concat、flatten、unique组合会生成大量中间数据，在处理全量2600万条记录时会快速耗尽内存。可以从以下几点优化：

1. 提前过滤无效数据，减少中间list的大小 #

原逻辑是先concat所有内容再过滤，改成先对每个字段过滤再concat，大幅减少后续list处理的数据量：

(
    df
    .group_by(person_id_key, person_source_id_key, "person_source_type")
    .agg(        
        # 优化后的Name字段处理
        pl.concat_list(
            # 先过滤full_name
            pl.col(full_name_key).fill_null("").filter(~pl.element().is_in(special_removal)),
            # 拆分nickname后直接过滤
            pl.col(nickname_key).fill_null("").str.split("|").list.eval(pl.element().filter(~pl.element().is_in(special_removal))),
            # 先过滤拼接后的全名
            (pl.col(first_name_key).fill_null("") + " " + pl.col(last_name_key).fill_null("")).filter(~pl.element().is_in(special_removal))
        )
        .flatten()
        .unique()
        .alias("name"),

        # 优化Company字段：提前去重再concat
        pl.concat_list(
            pl.col(org_name_key).unique(),
            pl.col(master_org_name_key).unique()
        )
        .flatten()
        .unique()
        .alias("company")
    )
    .with_columns(pl.col('name').list.len().alias("size_of_name"))
    .filter(pl.col("size_of_name") >= 1)
    .drop("size_of_name")
)

2. 切换到Polars Lazy API #

Lazy模式会自动优化执行计划，减少不必要的内存开销（比如延迟计算、合并操作步骤），对于大文件处理效率提升明显：

(
    df.lazy()  # 启用lazy模式
    .group_by(person_id_key, person_source_id_key, "person_source_type")
    .agg(        
        pl.concat_list(
            pl.col(full_name_key).fill_null("").filter(~pl.element().is_in(special_removal)),
            pl.col(nickname_key).fill_null("").str.split("|").list.eval(pl.element().filter(~pl.element().is_in(special_removal))),
            (pl.col(first_name_key).fill_null("") + " " + pl.col(last_name_key).fill_null("")).filter(~pl.element().is_in(special_removal))
        ).flatten().unique().alias("name"),
        pl.concat_list(pl.col(org_name_key).unique(), pl.col(master_org_name_key).unique()).flatten().unique().alias("company")
    )
    .with_columns(pl.col('name').list.len().alias("size_of_name"))
    .filter(pl.col("size_of_name") >= 1)
    .drop("size_of_name")
    .collect()  # 最终触发计算
)

3. 排查分组基数 #

如果person_id_key等分组键的唯一值接近2600万（几乎每条记录一个分组），agg后的list数据会直接占用大量内存。可以先执行以下代码查看分组数量：

print(df.select(person_id_key, person_source_id_key, "person_source_type").n_unique())

如果分组基数过大，需要考虑调整分组策略，或者对分组结果做分批处理。

二、环境差异排查 #

Jupyter能运行但VS Code失败，可能是环境配置不同：

• 确认两者使用的Python环境一致：分别在Jupyter和VS Code的终端执行import polars; print(polars.__version__)和import sys; print(sys.executable)，检查版本和Python路径是否相同。
• 关闭VS Code的调试模式：调试模式会额外占用内存，尝试直接在终端运行脚本而非通过调试按钮启动。
• 检查Windows虚拟内存设置：打开「系统属性-高级-性能设置-高级-虚拟内存」，确保虚拟内存已开启且分配足够空间（建议设置为物理内存的1.5-2倍，或选择「系统管理的大小」）。

三、升级Polars版本 #

Polars 0.17.2是较旧的版本（发布于2023年初），后续版本（0.18+）对group_by、list操作的内存管理做了大量优化，修复了多个内存泄漏和低效问题。直接升级到最新稳定版本：

pip install --upgrade polars

补充建议 #

如果上述优化后仍有问题，可以：

• 用Windows任务管理器监控Python进程的内存占用，定位到内存暴涨的具体操作步骤（比如concat_list还是unique）。
• 生成与原数据结构一致的合成数据集，复现问题后进一步排查。

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