在Polars DataFrame中高效实现列表列与Python普通列表元素逐位相乘的方法

嘿，刚好碰到过类似需求，我来给你说个高效的实现方式～

你要的是把roc列里每个列表的元素，和weights列表对应位置的元素逐一相乘对吧？Polars本身有很棒的向量化列表操作支持，完全不用写循环，效率拉满。

这里有两种简洁又高效的写法：

方法一：利用list.eval结合位置索引 #

这种方法直接针对列表里的每个元素，匹配权重列表对应位置的值进行相乘：

import polars as pl

df = pl.DataFrame(
    {
        "symbol": ["A", "A", "B", "B"],
        "roc": [[0.1, 0.2], [0.3, 0.4], [0.5, 0.6], [0.7, 0.8]],
    }
)
weights = [0.3, 0.7]

# 计算加权后的roc列
df = df.with_columns(
    roc_wgt=pl.col("roc").list.eval(pl.element() * pl.lit(weights).list.get(pl.int_range(0, pl.count())))
)

print(df)

方法二：用list.zip配对元素后相乘 #

这种思路是先把roc的每个列表和weights列表配对成元组，再对每个元组里的两个元素做乘法：

df = df.with_columns(
    roc_wgt=pl.col("roc").list.zip(pl.lit(weights)).list.eval(pl.element()[0] * pl.element()[1])
)

两种方法都能得到你要的结果，运行后输出如下：

shape: (4, 3)
┌────────┬────────────┬──────────────┐
│ symbol ┆ roc        ┆ roc_wgt      │
│ ---    ┆ ---        ┆ ---          │
│ str    ┆ list[f64]  ┆ list[f64]    │
╞════════╪════════════╪══════════════╡
│ A      ┆ [0.1, 0.2] ┆ [0.03, 0.14] │
│ A      ┆ [0.3, 0.4] ┆ [0.09, 0.28] │
│ B      ┆ [0.5, 0.6] ┆ [0.15, 0.42] │
│ B      ┆ [0.7, 0.8] ┆ [0.21, 0.56] │
└────────┴────────────┴──────────────┘

这两种方式都是完全利用Polars的内部向量化引擎，比手动遍历行高效得多，尤其是处理大数据量的时候优势更明显～

备注：内容来源于stack exchange，提问作者Andi