如何使用Polars高效加载多文件并添加自定义源信息

import polars as pl
from pathlib import Path

# 假设当前目录下有 data_product_1.csv, data_product_2.csv 等文件
# 为了演示，我们先创建一些模拟文件
file_contents = """data,value
2000-01-01,1
2000-01-02,2
"""
Path("data_product_1.csv").write_text(file_contents)

file_contents_2 = """data,value
2000-01-01,3
2000-01-02,4
"""
Path("data_product_2.csv").write_text(file_contents_2)

file_contents_3 = """data,value
2000-01-01,4
2000-01-02,5
"""
Path("data_product_3.csv").write_text(file_contents_3)

# 遍历所有匹配的文件，并为每个文件创建一个LazyFrame
csv_lazyframes = []
for f_path in Path().glob("data_*.csv"):
    # 使用 scan_csv 惰性读取文件
    lazy_df = pl.scan_csv(f_path)
    csv_lazyframes.append(lazy_df)

# 此时，数据尚未被实际读取
print(f"创建了 {len(csv_lazyframes)} 个 LazyFrame 对象。")

import polars as pl
from pathlib import Path

# (省略模拟文件创建部分，假设文件已存在)

# 遍历所有匹配的文件，并为每个文件创建一个LazyFrame，同时添加product_code列
csv_lazyframes_with_product_code = [
    pl.scan_csv(f_path).with_columns(product_code=pl.lit(f_path.name))
    for f_path in Path().glob("data_*.csv")
]

# 此时，每个LazyFrame都包含一个添加product_code列的指令，但数据仍未加载
print(f"创建了 {len(csv_lazyframes_with_product_code)} 个包含 'product_code' 列指令的 LazyFrame 对象。")

import polars as pl
from pathlib import Path

# 1. 创建模拟数据文件 (如果尚未创建)
file_contents_1 = """data,value
2000-01-01,1
2000-01-02,2
"""
Path("data_product_1.csv").write_text(file_contents_1)

file_contents_2 = """data,value
2000-01-01,3
2000-01-02,4
"""
Path("data_product_2.csv").write_text(file_contents_2)

file_contents_3 = """data,value
2000-01-01,4
2000-01-02,5
"""
Path("data_product_3.csv").write_text(file_contents_3)


# 2. 核心解决方案：使用LazyFrames处理和合并文件
# 遍历文件，创建LazyFrame，并添加文件名作为product_code列
lazy_frames = [
    pl.scan_csv(f_path).with_columns(product_code=pl.lit(f_path.name))
    for f_path in Path().glob("data_*.csv")
]

# 使用pl.concat合并所有LazyFrames，并调用.collect()执行计算
# pl.concat在处理LazyFrames时会默认尝试并行化读取和转换操作
final_df = pl.concat(lazy_frames).collect()

# 3. 打印结果
print("最终合并的DataFrame:")
print(final_df)

# 清理模拟文件
Path("data_product_1.csv").unlink()
Path("data_product_2.csv").unlink()
Path("data_product_3.csv").unlink()

最终合并的DataFrame:
shape: (6, 3)
┌────────────┬───────┬────────────────────┐
│ data       ┆ value ┆ product_code       │
│ ---        ┆ ---   ┆ ---                │
│ str        ┆ i64   ┆ str                │
╞════════════╪═══════╪════════════════════╡
│ 2000-01-01 ┆ 1     ┆ data_product_1.csv │
│ 2000-01-02 ┆ 2     ┆ data_product_1.csv │
│ 2000-01-01 ┆ 3     ┆ data_product_2.csv │
│ 2000-01-02 ┆ 4     ┆ data_product_2.csv │
│ 2000-01-01 ┆ 4     ┆ data_product_3.csv │
│ 2000-01-02 ┆ 5     ┆ data_product_3.csv │
└────────────┴───────┴────────────────────┘