在Polars中，如何将浮点列与列表列相乘？

Question

假设我们有一个数据表，其中有一列'b'，这列里包含了多个列表，而且每个列表的长度都是一样的（所以这些列表也可以转换成数组）。

df_test = pl.DataFrame({'a': [1., 2., 3.], 'b': [[2,2,2], [3,3,3], [4,4,4]]})
df_test
shape: (3, 2)
┌─────┬───────────┐
│ a   ┆ b         │
│ --- ┆ ---       │
│ f64 ┆ list[i64] │
╞═════╪═══════════╡
│ 1.0 ┆ [2, 2, 2] │
│ 2.0 ┆ [3, 3, 3] │
│ 3.0 ┆ [4, 4, 4] │
└─────┴───────────┘

我想得到的结果是

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 这个方法。

我能想到的最好办法是用 map_rows，这个方法在pandas中类似于 apply。根据文档来看，这个方法并不是最有效率的，但它确实能实现我的需求：

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)

Answer 1

编辑: 调整了答案，以确保它可以处理列'a'中的重复值。

---

这里有一种方法：

数据

注意：下面将'a'从 [1., 2., 3.] 改为 [1., 1., 3.]，以示例说明需要额外的临时列'idx'来进行分组。

import polars as pl

# changing 'a' from `[1., 2., 3.]` to `[1., 1., 3.]` to exemplify need temp `idx`
df_test = pl.DataFrame({'a': [1., 1., 3.], 
                        'b': [[2,2,2], [3,3,3], [4,4,4]]})
df_test

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

代码

df_new = (
    df_test.with_columns(idx=pl.arange(0, pl.len()))
           .explode('b')
           .with_columns(new=(pl.col('a') * pl.col('b')))
           .group_by(['idx', 'a'], maintain_order=True)
           .agg(pl.col("b"), pl.col("new"))
           .drop('idx')
)

df_new

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

解释

• 首先，创建一个名为'idx'的列（使用 pl.DataFrame.with_columns, pl.arange, 和 pl.len），用来跟踪每一行。也就是说，我们用这个列来区分在'a'中有相同值的行。
• 接下来，使用 pd.DataFrame.explode 将'b'中的列表值分到不同的行中。
• 然后，继续使用 pl.DataFrame.with_columns 将'a'列和'b'列相乘，并把结果赋值给'new'。
• 最后，我们想要恢复列表：使用 pl.DataFrame.group_by 对'idx'和'a'列进行分组，添加 maintain_order=True 来保持数据的正确顺序，并对'b'和'new'列应用 groupby.agg。
• 最后，清理一下，删除'idx'列（使用 pl.DataFrame.drop）。

Answer 2

很遗憾，polars 不支持 在 pl.Expr.list.eval 中引用命名列。要是支持的话，那就能很方便地解决问题了。

我觉得，@ouroboros1 提出的解决方案已经在正确的方向上了，方法是先把列“炸开”，进行操作后再“合并”回去。不过，这个过程可以简化得更好，像下面这样。

(
    df_test
    .with_columns(
        (
            pl.col("b").explode() * pl.col("a")
        )
        .implode().over(pl.int_range(pl.len()))
        .alias("new")
    )
)

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] │
└─────┴───────────┴────────────────────┘

特别是，我们可以避免显式地创建和删除索引列，以及使用 pl.DataFrame.group_by().agg() 这个结构。如果数据表的列更多，这样做会变得更加麻烦。