Python-polars:如何将列表中的每个元素与不同列中的值相乘？

Question

我有一个包含一定数量组的dataframe，包含一个权重列和一个值列表，这些值可以是任意长度的，例如：

df = pl.DataFrame(
    {
        "Group": ["Group1", "Group2", "Group3"],
        "Weight": [100.0, 200.0, 300.0],
        "Vals": [[0.5, 0.5, 0.8],[0.5, 0.5, 0.8], [0.7, 0.9]]
    }
)

┌────────┬────────┬─────────────────┐
│ Group  ┆ Weight ┆ Vals            │
│ ---    ┆ ---    ┆ ---             │
│ str    ┆ f64    ┆ list[f64]       │
╞════════╪════════╪═════════════════╡
│ Group1 ┆ 100.0  ┆ [0.5, 0.5, 0.8] │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ Group2 ┆ 200.0  ┆ [0.5, 0.5, 0.8] │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ Group3 ┆ 300.0  ┆ [0.7, 0.9]      │
└────────┴────────┴─────────────────┘

我的目标是计算一个“加权”列，它是值列表中每一项的倍数，而权重列中的值是：

┌────────┬────────┬─────────────────┬─────────────────┐
│ Group  ┆ Weight ┆ Vals            ┆ Weighted        │
│ ---    ┆ ---    ┆ ---             ┆ ---             │
│ str    ┆ f64    ┆ list[f64]       ┆ list[i64]       │
╞════════╪════════╪═════════════════╪═════════════════╡
│ Group1 ┆ 100.0  ┆ [0.5, 0.5, 0.8] ┆ [50, 50, 80]    │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ Group2 ┆ 200.0  ┆ [0.5, 0.5, 0.8] ┆ [100, 100, 160] │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ Group3 ┆ 300.0  ┆ [0.7, 0.9]      ┆ [210, 270]      │
└────────┴────────┴─────────────────┴─────────────────┘

我尝试过几种不同的方法：

df.with_columns([
        pl.col("Vals").arr.eval(pl.element() * 3).alias("Weight1"), #Multiplying with literal works
        pl.col("Vals").arr.eval(pl.element() * pl.col("Weight")).alias("Weight2"), #Does not work
        pl.col("Vals").arr.eval(pl.element() * pl.col("Unknown")).alias("Weight3"), #Unknown columns give same value
        pl.col("Vals").arr.eval(pl.col("Vals") * pl.col("Weight")).alias("Weight4"), #Same effect
        # pl.col('Vals') * 3 -> gives an error
]
)

┌────────┬────────┬────────────┬────────────┬──────────────┬──────────────┬────────────────────┐
│ Group  ┆ Weight ┆ Vals       ┆ Weight1    ┆ Weight2      ┆ Weight3      ┆ Weight4            │
│ ---    ┆ ---    ┆ ---        ┆ ---        ┆ ---          ┆ ---          ┆ ---                │
│ str    ┆ f64    ┆ list[f64]  ┆ list[f64]  ┆ list[f64]    ┆ list[f64]    ┆ list[f64]          │
╞════════╪════════╪════════════╪════════════╪══════════════╪══════════════╪════════════════════╡
│ Group1 ┆ 100.0  ┆ [0.5, 0.5, ┆ [1.5, 1.5, ┆ [0.25, 0.25, ┆ [0.25, 0.25, ┆ [0.25, 0.25, 0.64] │
│        ┆        ┆ 0.8]       ┆ 2.4]       ┆ 0.64]        ┆ 0.64]        ┆                    │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ Group2 ┆ 200.0  ┆ [0.5, 0.5, ┆ [1.5, 1.5, ┆ [0.25, 0.25, ┆ [0.25, 0.25, ┆ [0.25, 0.25, 0.64] │
│        ┆        ┆ 0.8]       ┆ 2.4]       ┆ 0.64]        ┆ 0.64]        ┆                    │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ Group3 ┆ 300.0  ┆ [0.7, 0.9] ┆ [2.1, 2.7] ┆ [0.49, 0.81] ┆ [0.49, 0.81] ┆ [0.49, 0.81]       │
└────────┴────────┴────────────┴────────────┴──────────────┴──────────────┴────────────────────┘

除非我没有正确理解它，否则您似乎无法从eval函数中访问列表之外的列。也许有一种方法可以在语句中使用列表理解，但这看起来并不是一个很好的解决方案。

这里推荐的方法是什么？任何帮助都将不胜感激！

Answer 1

我认为这对你有用：

(df
    .explode('Vals')
    .with_columns(Weighted = pl.col('Weight')*pl.col('Vals'))
    .groupby('Group')
    .agg([
        pl.col('Weight').first(),                                                                                                             
        pl.col('Vals').list(),
        pl.col('Weighted').list()
        ])                                                                                                 
)

Answer 2

所以你可以用两种方法解决这个问题。老实说，我不知道哪个更好，我还没有测试性能或RAM的使用情况。

算法1

我们可以将两列都放在一个结构中，然后在它们上应用一个自定义函数。(这里文档中对此进行了解释)

import polars as pl
import numpy as np

def weighted_list(ls, weight):
    return(list(np.array(ls) * weight))

(df.with_columns([pl.struct(["Weight", "Vals"])
                  .apply(lambda x: weighted_list(x["Vals"], x["Weight"]))])
)

算法2

因此，在这种情况下，您必须小心，您的列“组”、“权重”是唯一的。因此，如果您有两个条目，例如Group3和with 300。

(df.explode("Vals")
    .with_column((pl.col("Vals") * pl.col("Weight")).alias("Weighted"))
    .groupby(["Group", "Weight"])
    .agg([pl.list("Vals"), pl.list("Weighted")])
)