如何在Polars中不使用map_rows实现浮点列与列表列相乘？

给定如下Polars DataFrame，列b包含长度一致的列表：

import polars as pl

df_test = pl.DataFrame({'a': [1., 2., 3.], 'b': [[2,2,2], [3,3,3], [4,4,4]]})

其输出为：

shape: (3, 2)
┌─────┬───────────┐
│ a   ┆ b         │
│ --- ┆ ---       │
│ f64 ┆ list[i64] │
╞═════╪═══════════╡
│ 1.0 ┆ [2, 2, 2] │
│ 2.0 ┆ [3, 3, 3] │
│ 3.0 ┆ [4, 4, 4] │
└─────┴───────────┘

需要在不使用map_rows的前提下，得到如下结果（新增列new，值为列a的数值与列b中每个元素相乘后的列表）：

shape: (3, 3)
┌─────┬───────────┬────────────────────┐
│ a   ┆ b         ┆ new                │
│ --- ┆ ---       ┆ ---                │
│ f64 ┆ list[i64] ┆ list[f64]          │
╞═════╪═══════════╪════════════════════╡
│ 1.0 ┆ [2, 2, 2] ┆ [2.0, 2.0, 2.0]    │
│ 2.0 ┆ [3, 3, 3] ┆ [6.0, 6.0, 6.0]    │
│ 3.0 ┆ [4, 4, 4] ┆ [12.0, 12.0, 12.0] │
└─────┴───────────┴────────────────────┘

目前已知用map_rows可以实现，但效率较低，代码如下：

df_temp = df_test.map_rows(lambda x: ([x[0] * i for i in x[1]],))
df_temp.columns = ['new']
df_test = df_test.hstack(df_temp)

寻求更高效的替代方法。

高效解决方案 #

方法1：使用list.eval（推荐） #

Polars的列表表达式支持list.eval方法，可直接对列表内元素进行矢量化运算，完全避免逐行循环：

df_result = df_test.with_columns(
    pl.col("b").list.eval(pl.element() * pl.col("a")).alias("new")
)

pl.element()指代列表中的单个元素，pl.col("a")会自动广播到列表的每个位置完成乘法，最终返回列表类型结果。

方法2：展开列表计算后重新聚合 #

先将列表展开为多行，计算乘积后再按原行索引重新聚合为列表，同样是矢量化操作：

df_result = df_test.with_row_index() \
    .explode("b") \
    .with_columns(new=pl.col("a") * pl.col("b")) \
    .group_by("index") \
    .agg(
        pl.col("a").first(),
        pl.col("b").alias("b"),
        pl.col("new").alias("new")
    ) \
    .drop("index")

这种方法适合更复杂的列表运算场景，性能优于逐行映射。

方法3：转换为数组后运算 #

如果列b的列表长度固定，可转为数组类型利用Polars的数组运算能力：

df_result = df_test.with_columns(
    (pl.col("b").cast(pl.Array(pl.Int64, 3)) * pl.col("a")).cast(pl.List(pl.Float64)).alias("new")
)

先将列表转为固定长度数组，乘法运算自动广播，最后转回列表类型即可。

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