如何将Polars数据框指定列横向合并为单个列表列？

我正尝试使用Python Polars将多列合并为一列，但尚未找到一种优雅的方法将列合并为列表。我仅需将b至e列合并为一列，a列需保持原有状态。我曾尝试使用map_elements来实现此功能，但该方法不仅无效，运行速度还很慢，显然并非最佳方案。

现有数据框：

df = pl.from_repr("""
┌─────┬─────┬─────┬─────┬─────┐
│ a   ┆ b   ┆ c   ┆ d   ┆ e   │
│ --- ┆ --- ┆ --- ┆ --- ┆ --- │
│ f64 ┆ f64 ┆ f64 ┆ f64 ┆ f64 │
╞═════╪═════╪═════╪═════╪═════╡
│ 0.1 ┆ 1.1 ┆ 2.1 ┆ 3.1 ┆ 4.1 │
│ 0.2 ┆ 1.2 ┆ 2.2 ┆ 3.2 ┆ 4.2 │
│ 0.3 ┆ 1.3 ┆ 2.3 ┆ 3.3 ┆ 4.3 │
└─────┴─────┴─────┴─────┴─────┘
""")

期望结果：

shape: (3, 2)
┌─────┬──────────────────────┐
│ a   ┆ value                │
│ --- ┆ ---                  │
│ f64 ┆ list[f64]            │
╞═════╪══════════════════════╡
│ 0.1 ┆ [1.1, 2.1, 3.1, 4.1] │
│ 0.2 ┆ [1.2, 2.2, 3.2, 4.2] │
│ 0.3 ┆ [1.3, 2.3, 3.3, 4.3] │
└─────┴──────────────────────┘

高效实现方法 #

不需要使用map_elements，Polars提供了矢量化的pl.concat_list函数，能直接将指定列合并为列表列，性能远优于逐行处理的方法。

代码示例：

result_df = df.select(
    'a',
    pl.concat_list(pl.col('b')..pl.col('e')).alias('value')
)

print(result_df)

说明：

• pl.col('b')..pl.col('e') 选中从b到e的连续列，也可以直接指定列名列表：pl.col(['b', 'c', 'd', 'e'])
• pl.concat_list() 将选中的列按行合并为列表类型的列
• alias('value') 为新生成的列表列设置名称

运行上述代码即可得到期望结果，该方法基于Polars的矢量化引擎实现，处理大规模数据集时的效率会比map_elements高几个数量级。

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