如何在Python Polars中实现类似R tidyverse的嵌套列表列？

在R的tidyverse中，我们可以通过group_by()配合nest()创建包含列表列的数据集，每个列表元素对应一个分组后的小型DataFrame，这对批量拟合模型等场景非常实用。下面介绍如何在Python的Polars中实现相同效果。

示例Polars DataFrame #

import polars as pl

df = pl.DataFrame(
    [
        pl.Series("dataset_id", [1, 1, 2, 2, 1, 1, 2, 2], dtype=pl.Int64),
        pl.Series("day", [1, 2, 1, 2, 1, 2, 1, 2], dtype=pl.Int64),
        pl.Series("recipe", [1, 1, 1, 1, 2, 2, 2, 2], dtype=pl.Int64),
        pl.Series("cum_trials", [1000, 2000, 1000, 2000, 1000, 2000, 1000, 2000], dtype=pl.Int64),
        pl.Series("cum_events", [644, 1287, 643, 1262, 645, 1312, 655, 1301], dtype=pl.Int64),
        pl.Series("cum_rate", [0.644, 0.643, 0.643, 0.619, 0.645, 0.667, 0.655, 0.646], dtype=pl.Float64),
    ]
)

实现分组嵌套（nest） #

在Polars中，通过group_by()配合map_groups()可以精准复刻R中nest()的效果，将每组数据打包为独立的子DataFrame：

# 按dataset_id和day分组，生成包含子DataFrame的列表列
df_list = df.group_by(["dataset_id", "day"]).map_groups(
    lambda group: pl.DataFrame({"data": [group]})
)

查看结果 #

• 查看嵌套后的完整数据集：

print(df_list)

输出包含dataset_id、day和data三列，data列的每个元素都是对应分组的小型DataFrame。

• 访问第一个分组的子DataFrame：

print(df_list["data"][0])

补充：轻量级分组（仅生成列表） #

如果不需要完整的子DataFrame，仅需将分组后的各列转为列表，可使用更简洁的agg方法：

# 将每组的列转为独立列表（非DataFrame列）
df_list_lists = df.group_by(["dataset_id", "day"]).agg(pl.all())

这种方式适合仅需提取单列分组数据的场景。

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