如何从Polars DataFrame中提取带关联ID的嵌套表？

从GraphQL接口获取的实体列表响应中，每个实体包含嵌入式关联实体，对应的Polars DataFrame如下：

import polars as pl

df = pl.from_dicts([
    {"id":0, "z":0.0, "right": [{"x":0.0, "y":0.0}]},
    {"id":1, "z":1.0, "right": [{"x":0.0, "y":0.0}, {"x":1.0, "y":1.0}]},
    {"id":2, "z":2.0, "right": [{"x":0.0, "y":0.0}, {"x":1.0, "y":1.0}, {"x":2.0, "y":2.0}]},
])
print(df.schema)

其schema为：

OrderedDict([('id', Int64), ('z', Float64), ('right', List(Struct({'x': Float64, 'y': Float64})))])

需要将该DataFrame拆分为left和right两个DataFrame：

• left包含id和z字段，已完成提取：

left = df.select("id", "z")
left

输出结果：

shape: (3, 2)
┌─────┬─────┐
│ id  ┆ z   │
│ --- ┆ --- │
│ i64 ┆ f64 │
╞═════╪═════╡
│ 0   ┆ 0.0 │
│ 1   ┆ 1.0 │
│ 2   ┆ 2.0 │
└─────┴─────┘

• right需要包含id、x和y字段，但当前提取的right缺少关联id：

right = df.select(pl.col("right").list.explode()).unnest('right')
right

输出结果：

shape: (6, 2)
┌─────┬─────┐
│ x   ┆ y   │
│ --- ┆ --- │
│ f64 ┆ f64 │
╞═════╪═════╡
│ 0.0 ┆ 0.0 │
│ 0.0 ┆ 0.0 │
│ 1.0 ┆ 1.0 │
│ 0.0 ┆ 0.0 │
│ 1.0 ┆ 1.0 │
│ 2.0 ┆ 2.0 │
└─────┴─────┘

目标right DataFrame格式如下：

right = right.with_columns(id = pl.Series([0,1,1,2,2,2]))
right

输出结果：

shape: (6, 3)
┌─────┬─────┬─────┐
│ x   ┆ y   ┆ id  │
│ --- ┆ --- ┆ --- │
│ f64 ┆ f64 ┆ i64 │
╞═════╪═════╪═════╡
│ 0.0 ┆ 0.0 ┆ 0   │
│ 0.0 ┆ 0.0 ┆ 1   │
│ 1.0 ┆ 1.0 ┆ 1   │
│ 0.0 ┆ 0.0 ┆ 2   │
│ 1.0 ┆ 1.0 ┆ 2   │
│ 2.0 ┆ 2.0 ┆ 2   │
└─────┴─────┴─────┘

解决方案 #

只需在提取时同时保留id字段，再对right列执行list.explode()，最后unnest展开结构字段即可：

right = df.select("id", pl.col("right").list.explode()).unnest("right")

执行后得到的right即为目标格式：

shape: (6, 3)
┌─────┬─────┬─────┐
│ id  ┆ x   ┆ y   │
│ --- ┆ --- ┆ --- │
│ i64 ┆ f64 ┆ f64 │
╞═════╪═════╪═════╡
│ 0   ┆ 0.0 ┆ 0.0 │
│ 1   ┆ 0.0 ┆ 0.0 │
│ 1   ┆ 1.0 ┆ 1.0 │
│ 2   ┆ 0.0 ┆ 0.0 │
│ 2   ┆ 1.0 ┆ 1.0 │
│ 2   ┆ 2.0 ┆ 2.0 │
└─────┴─────┴─────┘

原理说明 #

在select中同时指定id和处理后的right列，list.explode()会将每个列表元素展开为单独行，且每行都会保留对应的id值，最后通过unnest把right结构体中的x、y字段拆分为独立列，最终得到关联了id的right DataFrame。

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