如何在Polars中实现类似Spark transform的列表嵌套分组？

问题背景 #

有一个名为structures的列，值是包含"a"和"b"字段的结构体组成的列表。要求避免使用explode操作，仅通过列表操作实现按"b"字段分组并聚合为嵌套列表。

示例数据 #

import polars as pl

# 创建包含结构体列表的数据集
list_of_structs = [
    [{"a":1, "b": 2}, {"a":3, "b": 4}, {"a":5, "b": 4}]
]

df = pl.DataFrame({"structures": list_of_structs})

预期结果 #

按b分组后的嵌套列表：

[[{'a': 1, 'b': 2}], [{'a': 3, 'b': 4}, {'a': 5, 'b': 4}]]

遇到的困境 #

Spark中可以通过transform+filter实现类似逻辑，但Polars的list.eval内部无法直接引用外部列表元素，尝试的代码无法正确完成过滤分组。

解决方案 #

利用Polars的列表操作API，通过提取唯一b值后遍历过滤的方式实现需求：

df = df.with_columns(
    grouped_by_b=pl.col("structures")
    # 提取列表中所有结构体的b字段并去重
    .list.eval(pl.element().struct.field("b"))
    .list.unique()
    # 遍历每个唯一b值，过滤原列表中匹配的结构体
    .list.eval(
        pl.col("structures").list.filter(pl.element().struct.field("b") == pl.element())
    )
)

print(df["grouped_by_b"].to_list())

代码说明 #

1. 提取唯一b值：通过list.eval取出每个结构体的b字段，再用list.unique得到当前行列表中所有不重复的b值。
2. 按b分组过滤：遍历每个唯一b值，使用list.filter在原structures列表中筛选出所有b字段匹配的结构体，最终生成分组后的嵌套列表结构。

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