Polars惰性链中提取中间信息的惯用方法及分支操作实现

可以使用Polars的pl.collect_all()方法实现流式模式下的分支操作，只需读取一次数据源即可完成多个查询的计算，既保留流式处理的内存优势，又避免重复读取的低效问题。

针对计数需求的示例 #

import polars as pl

# 初始化流式LazyFrame
df = pl.scan_csv("my.csv")

# 定义两个分支查询：分别统计总行数和过滤后的行数
branch_total = df.select(pl.len().alias("total_count"))
branch_filtered = df.filter(pl.col("A") != 0).select(pl.len().alias("filtered_count"))

# 一次性执行所有分支查询，仅读取一次CSV
results = pl.collect_all([branch_total, branch_filtered])

# 提取结果
l1 = results[0]["total_count"][0]
l2 = results[1]["filtered_count"][0]

print(f"{l1}, {l2}")

通用分支操作方案 #

如果需要更复杂的中间操作（如多维度聚合、不同过滤条件的统计等），同样可以用collect_all()实现，每个分支对应独立的LazyFrame操作链：

import polars as pl

df = pl.scan_csv("my.csv")

# 定义多个不同的分支查询
branch_total_stats = df.select(
    pl.len().alias("total_rows"),
    pl.col("A").mean().alias("mean_A_total")
)
branch_filtered_stats = df.filter(pl.col("A") != 0).select(
    pl.len().alias("filtered_rows"),
    pl.col("A").median().alias("median_A_filtered")
)
branch_high_A_stats = df.filter(pl.col("A") > 100).select(
    pl.col("B").sum().alias("sum_B_high_A"),
    pl.col("C").n_unique().alias("unique_C_high_A")
)

# 一次性执行所有分支
all_results = pl.collect_all([branch_total_stats, branch_filtered_stats, branch_high_A_stats])

# 按需提取各分支结果
total_stats = all_results[0].to_dict(as_series=False)
filtered_stats = all_results[1].to_dict(as_series=False)
high_A_stats = all_results[2].to_dict(as_series=False)

print("总行数与A列均值:", total_stats)
print("过滤后行数与A列中位数:", filtered_stats)
print("A>100时的B列总和与C列唯一值数:", high_A_stats)

collect_all()会自动合并所有分支的查询计划，共享初始的数据源读取步骤，Polars的查询优化器会确保操作的高效性，全程无需将完整数据集加载到内存，完美适配大文件的流式处理场景。

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