Polars struct.field(list[str]) 在处理 list[Struct] 时返回单个列

Question

我的 Polars Dataframe 中的一些列的数据类型为 pl.List(pl.Struct)。我正在尝试替换这些列，以便获得多个标量值列表的列。

这是我尝试更改的列的示例：

'column_0'
[{'field_1': 'a', 'field_2': 1}, {'field_1': 'b', 'field_2':2}]
[{'field_1': 'c', 'field_2':3}]

df.select(
    pl.col(col_name).list.eval(
        pl.element().struct.field("*")
    )
)

我的期望是我会得到这样的东西：

'field_1', 'field_2'
['a', 'b'], [1, 2]
['c'], [3]

相反，我只得到最后一个字段（在本例中为“field_2”）：

'field_2'
[1, 2]
[3]

Answer 1

struct.field("field_name")

提取特定字段，因此您可以使用

field_name

代替

*

来获得结果。

result = df.select([
    pl.col("column_0").list.eval(pl.element().struct.field("field_1")).alias("field_1"),
    pl.col("column_0").list.eval(pl.element().struct.field("field_2")).alias("field_2"),
])

如果你有更多像field_2这样的字段，那么你可以使用通用函数，

def split_struct_column(df: pl.DataFrame, col_name: str) -> pl.DataFrame:
    first_struct = df[col_name].to_list()[0][0]
    field_names = list(first_struct.keys())

    exprs = [
        pl.col(col_name).list.eval(pl.element().struct.field(field)).alias(field)
        for field in field_names
    ]

    return df.select(exprs)


result = split_struct_column(df, "column_0")