使用Polars处理多份大TSV文件并按年月分区时内核崩溃求助

你的代码出现内核崩溃主要是两个核心问题：

• 用map_elements+lambda做逐行日期格式化，这不是Polars的矢量化操作逻辑，大文件下会疯狂消耗内存和CPU
• 循环里每次过滤年月都重新collect()整个数据集，等于把所有文件重复读取N次（N是不同年月的数量），直接撑爆IO和内存

另外代码里col("year_month")要改成pl.col("year_month")，否则会报未定义错误。

下面是优化后的代码，两种方案都能高效处理百万行级别的TSV文件：

方案一：用partition_by一键拆分（Polars 0.19+版本推荐） #

import polars as pl

# 扫描文件时直接解析日期，比后续转换更高效
lazy_df = pl.scan_csv(
    "data/file_*.tsv",
    separator="\t",
    parse_dates={"Date": "%Y-%m-%d"}
)

# 用矢量化方法生成年月列，替代低效的map_elements
lazy_df = lazy_df.with_columns(
    pl.col("Date").dt.strftime("%Y-%m").alias("year_month")
)

# 按年月分组拆分并写入文件，仅需一次扫描
lazy_df.partition_by("year_month").map_groups(
    lambda df: df.drop("year_month").write_csv(f"data_{df['year_month'][0]}.tsv", separator="\t")
).collect()

方案二：流式过滤写入（兼容旧版本Polars） #

import polars as pl

lazy_df = pl.scan_csv(
    "data/file_*.tsv",
    separator="\t",
    parse_dates={"Date": "%Y-%m-%d"}
).with_columns(
    pl.col("Date").dt.strftime("%Y-%m").alias("year_month")
)

# 先获取所有唯一年月值
year_months = lazy_df.select("year_month").unique().collect().to_series().to_list()

# 用sink_csv流式写入，无需加载全量数据到内存
for ym in year_months:
    (
        lazy_df.filter(pl.col("year_month") == ym)
        .drop("year_month")
        .sink_csv(f"data_{ym}.tsv", separator="\t")
    )

优化点说明 #

1. 读取时解析日期：通过parse_dates参数直接在扫描阶段完成日期转换，减少后续处理步骤
2. 矢量化日期格式化：dt.strftime是Polars内置的矢量化函数，比逐行处理的map_elements效率提升几个数量级
3. 避免重复扫描：两种方案都只扫描一次原始文件，要么分组批量写入，要么流式过滤写入，彻底解决重复IO的问题

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