如何从Polars DataFrame中按行选取首个非空struct至新列？

输入数据 #

structs_a = [
    {"a_key_1": 1, "a_key_2": "a" },
    {"a_key_1": None, "a_key_2": None }
]
structs_b = [
    {"b_key_1": None, "b_key_2": None },
    {"b_key_1": "b", "b_key_2": 2 }
]
df = pl.DataFrame([structs_a, structs_b])
df

对应的DataFrame结构：

┌─────────────┬─────────────┐
│ column_0    ┆ column_1    │
│ ---         ┆ ---         │
│ struct[2]   ┆ struct[2]   │
╞═════════════╪═════════════╡
│ {1,"a"}     ┆ {null,null} │
│ {null,null} ┆ {"b",2}     │
└─────────────┴─────────────┘

已尝试方法及结果 #

尝试将值拼接为列表并删除空值，但结果得到struct[4]类型的列表，代码如下：

df.with_columns(
    pl.concat_list(
        "column_0",
        "column_1",
    ).list.drop_nulls().alias('final_column')
)

执行结果：

shape: (2, 3)
┌─────────────┬─────────────┬─────────────────────────────────┐
│ column_0    ┆ column_1    ┆ final_column                    │
│ ---         ┆ ---         ┆ ---                             │
│ struct[2]   ┆ struct[2]   ┆ list[struct[4]]                 │
╞═════════════╪═════════════╪═════════════════════════════════╡
│ {1,"a"}     ┆ {null,null} ┆ [{1,"a",null,null}, {null,null… │
│ {null,null} ┆ {"b",2}     ┆ [{null,null,null,null}, {null,… │
└─────────────┴─────────────┴─────────────────────────────────┘

期望结果 #

┌─────────────┬─────────────┬─────────────────┐
│ column_0    ┆ column_1    ┆ final_column    │
│ ---         ┆ ---         ┆ ---             │
│ struct[2]   ┆ struct[2]   ┆ struct[2]       │
╞═════════════╪═════════════╪═════════════════╡
│ {1,"a"}     ┆ {null,null} ┆ {1,"a"}         │ <- 取自column_0的非空值
│ {null,null} ┆ {"b",2}     ┆ {"b",2}         │ <- 取自column_1的非空值
└─────────────┴─────────────┴─────────────────┘

更新尝试 #

尝试使用coalesce函数，但直接调用未得到预期结果，代码如下：

df.with_columns(
    pl.coalesce(
        "column_0",
        "column_1",
    ).alias('final_column')
)

执行结果：

shape: (2, 3)
┌─────────────┬─────────────┬───────────────────────┐
│ column_0    ┆ column_1    ┆ final_column          │
│ ---         ┆ ---         ┆ ---                   │
│ struct[2]   ┆ struct[2]   ┆ struct[4]             │
╞═════════════╪═════════════╪═══════════════════════╡
│ {1,"a"}     ┆ {null,null} ┆ {1,"a",null,null}     │
│ {null,null} ┆ {"b",2}     ┆ {null,null,null,null} │
└─────────────┴─────────────┴───────────────────────┘

解决方案 #

问题原因 #

之前的方法失效是因为：

• concat_list会自动合并不同struct的字段，导致生成包含所有字段的新struct类型
• coalesce仅判断struct本身是否为null，而你的数据中struct是存在的，只是内部字段全为null，因此它会合并两个struct的所有字段，而非选取第一个非空struct

方法1：通用多列场景（推荐） #

通过列表拼接后，过滤掉内部字段全为null的struct，再取第一个元素：

df.with_columns(
    pl.concat_list("column_0", "column_1")
    # 过滤出至少有一个非空字段的struct
    .list.eval(pl.element().filter(pl.any_horizontal(pl.element().struct.fields().is_not_null())))
    # 取第一个符合条件的struct
    .list.first()
    .alias("final_column")
)

方法2：固定两列场景 #

通过条件判断直接选择第一个非全空的struct：

df.with_columns(
    pl.when(pl.any_horizontal(pl.col("column_0").struct.fields().is_not_null()))
    .then(pl.col("column_0"))
    .otherwise(pl.col("column_1"))
    .alias("final_column")
)

执行结果 #

两种方法都能得到预期的输出：

┌─────────────┬─────────────┬─────────────┐
│ column_0    ┆ column_1    ┆ final_column│
│ ---         ┆ ---         ┆ ---         │
│ struct[2]   ┆ struct[2]   ┆ struct[2]   │
╞═════════════╪═════════════╪═════════════╡
│ {1,"a"}     ┆ {null,null} ┆ {1,"a"}     │
│ {null,null} ┆ {"b",2}     ┆ {"b",2}     │
└─────────────┴─────────────┴─────────────┘

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