如何在Polars中将嵌套列表列展平为单层列表列？

问题场景 #

你现有如下Polars代码：

import polars as pl

pl.DataFrame(pl.Series("x", ["1, 0", "2,3", "5 4"])).with_columns(
    pl.col("x").str.split(",").list.eval(pl.element().str.split(" "))
)

执行后得到嵌套列表格式的输出：

shape: (3, 1)
┌────────────────────┐
│ x                  │
│ ---                │
│ list[list[str]]    │
╞════════════════════╡
│ [["1"], ["", "0"]] │
│ [["2"], ["3"]]     │
│ [["5", "4"]]       │
└────────────────────┘

需要将嵌套列表展平为单层列表，目标输出为：

shape: (3, 1)
┌────────────────┐
│ x              │
│ ---            │
│ list[str]      │
╞════════════════╡
│ ["1", "", "0"] │
│ ["2", "3"]     │
│ ["5", "4"]     │
└────────────────┘

解决方案 #

方法一：基于现有代码修改（展平嵌套列表） #

在现有代码的链式调用末尾添加.list.flatten()方法，该方法会将list[list[str]]类型的列直接展平为list[str]：

import polars as pl

pl.DataFrame(pl.Series("x", ["1, 0", "2,3", "5 4"])).with_columns(
    pl.col("x").str.split(",").list.eval(pl.element().str.split(" ")).list.flatten()
)

方法二：更简洁的直接处理方式 #

可以先将所有逗号替换为空格，再按空格分割字符串，一步得到单层列表，避免嵌套操作：

import polars as pl

pl.DataFrame(pl.Series("x", ["1, 0", "2,3", "5 4"])).with_columns(
    pl.col("x").str.replace_all(",", " ").str.split(" ")
)

两种方法都能得到你想要的目标输出，可根据实际场景选择使用。

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