如何使用Polars LazyFrame实现按分组值分别保存为Parquet文件

嘿，我来帮你把这段Pandas的分组保存逻辑转换成Polars LazyFrame的实现方式～

首先先回顾一下你原来用Pandas实现的代码，就是按列a分组后，把每个子数据集存成对应分组值的Parquet文件：

import numpy as np
import pandas as pd

df = pd.DataFrame({'a': np.random.randint(0, 5, 1000), 'b': np.random.random(1000)})

for aval, subdf in df.groupby('a'):
    subdf.to_parquet(f'/tmp/{aval}.parquet')

换成Polars LazyFrame的话，因为LazyFrame是延迟执行的，我们不能直接像Pandas那样循环groupby后的结果，得换个思路来实现，具体代码如下：

import polars as pl
import numpy as np

# 生成LazyFrame格式的数据，也可以用pl.from_pandas(你的PandasDataFrame).lazy()转换现有数据
lazy_df = pl.DataFrame({
    'a': np.random.randint(0, 5, 1000),
    'b': np.random.random(1000)
}).lazy()

# 第一步：获取所有唯一的分组值，因为LazyFrame延迟执行，需要先collect拿到实际结果
group_values = lazy_df.select(pl.col('a').unique()).collect().to_series().to_list()

# 循环每个分组值，过滤并保存为Parquet
for aval in group_values:
    (
        lazy_df
        .filter(pl.col('a') == aval)  # 过滤出当前分组的数据
        .collect()  # 触发LazyFrame的计算，拿到实际数据
        .write_parquet(f'/tmp/{aval}.parquet')  # 写入Parquet文件
    )

这里有两个关键点要注意：

• 因为LazyFrame不会立刻执行计算，所以我们得先通过collect()获取到所有唯一的分组键值，这样才能循环处理每个分组
• 每次循环里，我们先过滤出对应分组的数据，然后collect()触发执行，最后调用write_parquet保存文件

如果你是从已有的Pandas DataFrame转换过来，只需要把生成LazyFrame的那行换成lazy_df = pl.from_pandas(你的PandasDF).lazy()就可以顺利运行啦～

备注：内容来源于stack exchange，提问作者Stephen