如何将Polars中的Struct类型展开为行而非列？

我有如下Polars DataFrame：

import polars as pl

df = pl.DataFrame({
    'as_of':    ['2024-08-01', '2024-08-02', '2024-08-03', '2024-08-04'],
    'quantity': [{'A': 10, 'B': 5}, {'A': 11, 'B': 7}, {'A': 9, 'B': 4, 'C': -3},
                 {'A': 15, 'B': 3, 'C': -14, 'D': 50}]
}, schema={'as_of': pl.String, 'quantity': pl.Struct})

对应的DataFrame结构：

shape: (4, 2)
┌────────────┬──────────────────┐
│ as_of      ┆ quantity         │
│ ---        ┆ ---              │
│ str        ┆ struct[4]        │
╞════════════╪══════════════════╡
│ 2024-08-01 ┆ {10,5,null,null} │
│ 2024-08-02 ┆ {11,7,null,null} │
│ 2024-08-03 ┆ {9,4,-3,null}    │
│ 2024-08-04 ┆ {15,3,-14,50}    │
└────────────┴──────────────────┘

执行df.unnest('quantity')会将Struct展开为列，得到如下结果：

shape: (4, 5)
┌────────────┬─────┬─────┬──────┬──────┐
│ as_of      ┆ A   ┆ B   ┆ C    ┆ D    │
│ ---        ┆ --- ┆ --- ┆ ---  ┆ ---  │
│ str        ┆ i64 ┆ i64 ┆ i64  ┆ i64  │
╞════════════╪═════╪═════╪══════╪══════╡
│ 2024-08-01 ┆ 10  ┆ 5   ┆ null ┆ null │
│ 2024-08-02 ┆ 11  ┆ 7   ┆ null ┆ null │
│ 2024-08-03 ┆ 9   ┆ 4   ┆ -3   ┆ null │
│ 2024-08-04 ┆ 15  ┆ 3   ┆ -14  ┆ 50   │
└────────────┴─────┴─────┴──────┴──────┘

但我希望将Struct展开为行，得到如下格式的DataFrame：

shape: (11, 3)
┌────────────┬──────┬──────────┐
│ as_of      ┆ name ┆ quantity │
│ ---        ┆ ---  ┆ ---      │
│ str        ┆ str  ┆ i64      │
╞════════════╪══════╪══════════╡
│ 2024-08-01 ┆ A    ┆ 10       │
│ 2024-08-01 ┆ B    ┆ 5        │
│ 2024-08-02 ┆ A    ┆ 11       │
│ 2024-08-02 ┆ B    ┆ 7        │
│ 2024-08-03 ┆ A    ┆ 9        │
│ …          ┆ …    ┆ …        │
│ 2024-08-03 ┆ C    ┆ -3       │
│ 2024-08-04 ┆ A    ┆ 15       │
│ 2024-08-04 ┆ B    ┆ 3        │
│ 2024-08-04 ┆ C    ┆ -14      │
│ 2024-08-04 ┆ D    ┆ 50       │
└────────────┴──────┴──────────┘

解决方案 #

可以通过struct.to_list()把Struct转换为包含field（键名）和value（对应值）的Struct列表，再用explode()将列表展开为行，最后提取字段并过滤空值：

result = df.select(
    pl.col("as_of"),
    pl.col("quantity").struct.to_list().alias("kv_pairs")
).explode("kv_pairs").select(
    pl.col("as_of"),
    pl.col("kv_pairs").struct.field("field").alias("name"),
    pl.col("kv_pairs").struct.field("value").alias("quantity")
).drop_nulls()

print(result)

也可以用更简洁的写法：

result = df.with_columns(
    pl.col("quantity").struct.to_list().alias("kv_pairs")
).explode("kv_pairs").unnest("kv_pairs").rename({"field": "name", "value": "quantity"}).drop_nulls()

执行后即可得到目标格式的DataFrame。

内容的提问来源于stack exchange，提问作者Phil-ZXX