如何在Polars中基于键列拆分LazyFrame数据集？

方案1：用group_by+map_groups实现无预收集的分组处理 #

不需要提前拉取所有唯一键，直接在LazyFrame上通过分组映射完成流式处理，全程保持惰性执行，不会一次性加载全量数据。

示例代码：

import polars as pl

def process_single_group(df: pl.DataFrame) -> pl.DataFrame:
    # 替换成你对每个子表的实际处理逻辑，比如导出、计算等
    return df.write_parquet(f"group_{df['键列'][0]}.parquet")

# 链式调用完成流式分组处理
(
    pl.scan_parquet("your_large_data.parquet")
    .group_by("键列")
    .map_groups(process_single_group)
    .collect(streaming=True)
)

核心逻辑是map_groups会在流式读取数据时，逐个处理每个键对应的分组，完全跳过预收集唯一键的步骤，内存占用可控。

方案2：自定义流式迭代器手动控内存 #

如果map_groups的封装满足不了你的灵活需求，可以自己写一个流式迭代器，逐批读取数据并按键缓存，按需处理分组：

示例代码：

import polars as pl

def stream_split_by_key(lf: pl.LazyFrame, key_col: str, batch_size: int = 10000):
    # 按批次流式读取数据
    for batch in lf.collect_stream(batch_size=batch_size):
        # 拆分当前批次的分组
        batch_groups = batch.partition_by(key_col)
        # 这里可以加入缓存逻辑，比如把同键的批次合并后处理
        # 示例：直接处理当前批次的分组（也可缓存多批后再处理）
        for key, group_df in batch_groups.items():
            # 执行你的处理逻辑
            group_df.write_parquet(f"batch_group_{key}.parquet")

# 调用示例
large_lf = pl.scan_parquet("your_large_data.parquet")
stream_split_by_key(large_lf, "键列")

这种方式可以自主控制每个分组的处理时机（比如积累N批再处理），适合对内存使用有严格限制的场景。

关键注意点 #

• 确保Polars版本在0.19.0以上，这个版本对分组流式处理的支持更稳定。
• 处理分组时尽量直接输出结果（比如写入文件），不要把所有分组数据留在内存中，避免内存溢出。
• 避免在分组处理中执行全局聚合操作，这类操作会强制加载全量数据，失去流式处理的意义。

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