使用Polars的scan_csv和fetch处理大CSV时如何避免内存错误？

核心问题分析 #

你的代码存在两个关键问题导致内存溢出：

1. 未正确处理fetch()结果：fetch()返回的是一个内存中的DataFrame，但你没有将其赋值给变量，后续调用write_csv()时仍在执行完整的LazyFrame查询（对所有数据排序、去重），而非仅处理子集。
2. 全量排序开销过大：对所有行进行排序需要加载全部数据到内存，这是内存占用过高的主要原因，而你的需求仅需保留每个account_id对应最高account_balance的行。

优化方案代码 #

以下是两种高效处理的实现方式，均避免全量排序并仅加载必要数据：

方式1：使用分组聚合获取每个账户的最高余额行

import os
import polars as pl

# 仅加载需要的列，减少内存占用
needed_columns = ["account_id", "account_balance"]  # 添加你实际需要的其他列

# 扫描CSV文件，建议指定schema以提升性能（可选但推荐）
# schema = {"account_id": pl.UInt64, "account_balance": pl.Float64}
df = pl.scan_csv(os.path.join("data", "*.csv")) \
       .select(needed_columns)

# 按account_id分组，保留每组中account_balance最高的行
df = df.group_by("account_id").agg(
    pl.all().sort_by("account_balance", descending=True).first()
)

# 对结果按余额降序排序，取前100万行并流式收集
result_df = df.sort("account_balance", descending=True) \
              .head(1000000) \
              .collect(streaming=True)

# 写入CSV
result_df.write_csv("df.csv")

方式2：使用窗口函数过滤每个账户的最高余额行

import os
import polars as pl

needed_columns = ["account_id", "account_balance"]  # 按需添加列

df = pl.scan_csv(os.path.join("data", "*.csv")) \
       .select(needed_columns)

# 为每个account_id内的行按余额降序排名
df = df.with_columns(
    pl.col("account_balance")
    .rank(descending=True, method="first")
    .over("account_id")
    .alias("rank")
)

# 保留每个账户的排名第一的行
df = df.filter(pl.col("rank") == 1).drop("rank")

# 排序并取前100万行
result_df = df.sort("account_balance", descending=True) \
              .head(1000000) \
              .collect(streaming=True)

result_df.write_csv("df.csv")

关键优化点说明 #

• 仅加载必要列：通过select()指定需要的列，避免加载无关数据，大幅降低内存占用。
• 避免全量排序：通过分组聚合或窗口函数，仅在每个account_id内部排序，而非对所有数据排序，内存开销骤减。
• 流式收集：collect(streaming=True)将数据分块处理，不一次性加载全部数据到内存。
• 提前限制行数：head(1000000)让Polars在查询计划早期就停止处理多余数据，提升效率。

额外建议 #

• 指定Schema：如果已知CSV文件的列类型，通过schema参数传递给scan_csv()，避免Polars自动推断Schema时的内存开销。
• 检查重复数据：若数据中account_id重复率高，分组后的数据量会大幅减少，后续操作更高效。

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