在 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

这是一种方法：

• 使用

  pd.DataFrame.explode

  将“b”的列表值放入单独的行中。
• 现在，链接

  pl.DataFrame.with_columns

  将列“a”乘以列“b”，将结果分配给“new”。
• 接下来，我们要取回列表：在“a”列上使用

  pl.DataFrame.group_by

  ，添加

  maintain_order=True

  以使数据保持正确的顺序，然后在“b”和“列”上应用

  groupby.agg

  新的。

df_new = (
    df_test.explode('b')
           .with_columns(new=(pl.col('a') * pl.col('b')))
           .group_by('a', maintain_order=True)
           .agg(pl.col("b"), pl.col("new"))
)

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