python - 使用Polars处理多个csv文件并对它们进行不同的分区

Question

我有一个充满 csv 文件的目录，每个文件有 100 万行或更多。 每个 csv 文件都具有相同的结构，并且有一个“日期”列，这些日期没有明显的顺序。

我想读入 csv 文件，然后按月年组合将它们拆分，然后再次写出。

lazy_df = pl.scan_csv('data/file_*.tsv', separator='\t')
lazy_df = lazy_df.with_columns(pl.col("Date").str.to_datetime("%Y-%m-%d").alias("Date"))
    
lazy_df = lazy_df.with_columns(pl.col("Date").apply(lambda date: date.strftime("%Y-%m")).alias("year_month"))

year_months = lazy_df.select(['year_month']).unique().collect()

for year_month in year_months:
        df_month = lazy_df.filter(col("year_month") == year_month).collect()

        # Write this dataframe to a new TSV file
        output_filename = f"data_{year_month}.tsv"
        df_month.write_csv(output_filename, delimiter='\t')

我的 jupyter 内核不断崩溃，所以我不确定我是否以正确的方式使用 Polars。

我尝试过的代码在上面。

Answer 1

使用 .apply 方法会削弱使用 Polars 所带来的速度性能优势，因此强烈建议不要使用该方法，除非您无法使用“Polars Expression ”实现数据转换

我认为这样的事情应该适合你:

lazy_df = (
    pl.scan_csv("data/file_*.tsv", separator="\t")
    .with_columns(pl.col("Date").str.to_datetime("%Y-%m-%d").alias("Date_dt"))
    .with_columns(
        pl.col("Date_dt")
        .dt.strftime("%Y-%m")
        .alias("year_month")
    )
    .collect()
)
lazy_df

我将 apply 替换为 Polars .dt.strftime 表达式方法。我将整个 Polars 表达式放入 scan_csv 调用中。

在使用 3 个文件的示例中，每个文件包含 100,000 行日期字符串 (YYYY-MM-DD) 和一个值 (random.random())，即一些 float x，其中 0.0 < x < 1.0)300,000 行被懒惰地收集到一个 df 中，在我的本地个人标准台式计算机(4 核)上执行大约花费了 180-190 毫秒。

然后，对于后续的 Polars 操作，您无需应用 .collect 方法，因为它已经如此，因此:

year_months = lazy_df.select(["year_month"]).unique()

year_months = year_months.to_dict()["year_month"].to_list()

for year_month in year_months:
    df_month = lazy_df.filter(pl.col("year_month") == year_month)

    # Write this dataframe to a new TSV file
    output_filename = f"./data/year_months_{year_month}.tsv"
    df_month.write_csv(output_filename, separator="\t")