Как добавить имя файла в Polars pl.scan_csv()?

Здесь вы можете использовать пару подходов, если работаете с местным файловую систему, то вы можете вручную объединить файлы, добавить столбец самостоятельно и объединить результаты.

Объединение полярностей + сканирование

from pathlib import Path
from tempfile import TemporaryDirectory

from numpy.random import default_rng
import polars as pl

def random_dataframe(rng, size=10):
    return pl.DataFrame({
        'x': rng.random(size),
        'y': rng.integers(10, size=size),
        'z': rng.normal(10, scale=3, size=size),
    })

rng = default_rng(0)
with TemporaryDirectory() as d:
    outdir = Path(d)
    for letter in 'abcd':
        df = random_dataframe(rng)
        df.write_csv(outdir / f'{letter}.csv')

    plan = pl.concat([
        pl.scan_csv(p).with_columns(source=pl.lit(str(p)))
        for p in outdir.glob('*.csv')
    ])
    print(plan.collect())
    # shape: (40, 4)
    # ┌──────────┬─────┬───────────┬────────────────────────┐
    # │ x        ┆ y   ┆ z         ┆ source                 │
    # │ ---      ┆ --- ┆ ---       ┆ ---                    │
    # │ f64      ┆ i64 ┆ f64       ┆ str                    │
    # ╞══════════╪═════╪═══════════╪════════════════════════╡
    # │ 0.636962 ┆ 2   ┆ 7.803198  ┆ /tmp/tmp_c8ysn7b/a.csv │
    # │ 0.269787 ┆ 8   ┆ 8.367223  ┆ /tmp/tmp_c8ysn7b/a.csv │
    # │ 0.040974 ┆ 6   ┆ 9.0511    ┆ /tmp/tmp_c8ysn7b/a.csv │
    # │ 0.016528 ┆ 0   ┆ 11.234892 ┆ /tmp/tmp_c8ysn7b/a.csv │
    # │ 0.81327  ┆ 3   ┆ 13.12754  ┆ /tmp/tmp_c8ysn7b/a.csv │
    # │ …        ┆ …   ┆ …         ┆ …                      │
    # │ 0.529312 ┆ 7   ┆ 13.094359 ┆ /tmp/tmp_c8ysn7b/d.csv │
    # │ 0.785786 ┆ 9   ┆ 10.483029 ┆ /tmp/tmp_c8ysn7b/d.csv │
    # │ 0.414656 ┆ 9   ┆ 8.243414  ┆ /tmp/tmp_c8ysn7b/d.csv │
    # │ 0.734484 ┆ 6   ┆ 5.976341  ┆ /tmp/tmp_c8ysn7b/d.csv │
    # │ 0.711143 ┆ 9   ┆ 5.795439  ┆ /tmp/tmp_c8ysn7b/d.csv │
    # └──────────┴─────┴───────────┴────────────────────────┘

Обратите внимание, что сканирование с помощью glob и объединение списка входных данных создают одни и те же планы:

print(pl.scan_csv(outdir / '*.csv').explain())
# UNION
#   PLAN 0:
#
#       Csv SCAN /tmp/tmp3t7rlzyq/a.csv
#       PROJECT */3 COLUMNS
#   PLAN 1:
#
#       Csv SCAN /tmp/tmp3t7rlzyq/b.csv
#       PROJECT */3 COLUMNS
#   PLAN 2:
#
#       Csv SCAN /tmp/tmp3t7rlzyq/c.csv
#       PROJECT */3 COLUMNS
#   PLAN 3:
#
#       Csv SCAN /tmp/tmp3t7rlzyq/d.csv
#       PROJECT */3 COLUMNS
# END UNION

print(
    pl.concat([pl.scan_csv(p) for p in outdir.glob('*.csv')]).explain()
)
# UNION
#   PLAN 0:
#
#       Csv SCAN /tmp/tmp3t7rlzyq/a.csv
#       PROJECT */3 COLUMNS
#   PLAN 1:
#
#       Csv SCAN /tmp/tmp3t7rlzyq/b.csv
#       PROJECT */3 COLUMNS
#   PLAN 2:
#
#       Csv SCAN /tmp/tmp3t7rlzyq/c.csv
#       PROJECT */3 COLUMNS
#   PLAN 3:
#
#       Csv SCAN /tmp/tmp3t7rlzyq/d.csv
#       PROJECT */3 COLUMNS
# END UNION

Duckdb read_csv(…, имя_файла=True)

Альтернативно можно использовать duckdb.read_csv и передать параметр filename=True. Затем материализуйте результат в polars.DataFrame

from pathlib import Path
from string import ascii_lowercase
from tempfile import TemporaryDirectory

from numpy.random import default_rng
import polars as pl
import duckdb

def random_dataframe(rng, size=10):
    return pl.DataFrame({
        'x': rng.random(size),
        'y': rng.integers(10, size=size),
        'z': rng.normal(10, scale=3, size=size),
    })

rng = default_rng(0)
with TemporaryDirectory() as d:
    outdir = Path(d)
    for letter in 'abcd':
        df = random_dataframe(rng)
        df.write_csv(outdir / f'{letter}.csv')

    result = duckdb.read_csv(outdir / '*.csv', filename=True).pl()
    print(result)
    # shape: (40, 4)
    # ┌──────────┬─────┬───────────┬────────────────────────┐
    # │ x        ┆ y   ┆ z         ┆ filename               │
    # │ ---      ┆ --- ┆ ---       ┆ ---                    │
    # │ f64      ┆ i64 ┆ f64       ┆ str                    │
    # ╞══════════╪═════╪═══════════╪════════════════════════╡
    # │ 0.636962 ┆ 2   ┆ 7.803198  ┆ /tmp/tmp8cyagj76/a.csv │
    # │ 0.269787 ┆ 8   ┆ 8.367223  ┆ /tmp/tmp8cyagj76/a.csv │
    # │ 0.040974 ┆ 6   ┆ 9.0511    ┆ /tmp/tmp8cyagj76/a.csv │
    # │ 0.016528 ┆ 0   ┆ 11.234892 ┆ /tmp/tmp8cyagj76/a.csv │
    # │ 0.81327  ┆ 3   ┆ 13.12754  ┆ /tmp/tmp8cyagj76/a.csv │
    # │ …        ┆ …   ┆ …         ┆ …                      │
    # │ 0.529312 ┆ 7   ┆ 13.094359 ┆ /tmp/tmp8cyagj76/d.csv │
    # │ 0.785786 ┆ 9   ┆ 10.483029 ┆ /tmp/tmp8cyagj76/d.csv │
    # │ 0.414656 ┆ 9   ┆ 8.243414  ┆ /tmp/tmp8cyagj76/d.csv │
    # │ 0.734484 ┆ 6   ┆ 5.976341  ┆ /tmp/tmp8cyagj76/d.csv │
    # │ 0.711143 ┆ 9   ┆ 5.795439  ┆ /tmp/tmp8cyagj76/d.csv │
    # └──────────┴─────┴───────────┴────────────────────────┘

include_file_paths= добавлен в Polars 1.2.0

Это позволяет вам указать имя столбца, который будет заполнен путями.

pl.scan_csv("/tmp/*.csv", include_file_paths = "path").collect()

shape: (2, 3)
┌─────┬─────┬────────────┐
│ a   ┆ b   ┆ path       │
│ --- ┆ --- ┆ ---        │
│ i64 ┆ i64 ┆ str        │
╞═════╪═════╪════════════╡
│ 1   ┆ 2   ┆ /tmp/a.csv │
│ 3   ┆ 4   ┆ /tmp/b.csv │
└─────┴─────┴────────────┘