Polars中如何对列表类型列计算点积？无需展开列表

问题背景 #

有包含等长列表列的Polars DataFrame，需在不使用explode（避免内存占用过高）的前提下计算两列的点积，原执行代码触发报错：

df = pl.DataFrame({
    "b": [[1,2,3], [2,3],[4],[1,2,3],[]],
    "c": [[1,1,1], [2,2],[4],[1,5,6],[]]
})
df.with_columns(pl.col('b').list.eval(pl.element().dot(pl.col('b'))).alias('b.c'))

报错信息：

ComputeError: named columns are not allowed in `list.eval`; consider using `element` or `col("")`

解决方案 #

方法1：Polars原生向量化操作（推荐）

通过list.zip将两列对应元素配对，对每对元素相乘后求和得到点积。完全基于Polars原生函数，性能最优，无需Python lambda：

import polars as pl

df = pl.DataFrame({
    "b": [[1,2,3], [2,3],[4],[1,2,3],[]],
    "c": [[1,1,1], [2,2],[4],[1,5,6],[]]
})

result = df.with_columns(
    dot=(
        pl.col("b").list.zip(pl.col("c"))
        .list.eval(pl.element().first() * pl.element().second())
        .list.sum()
        .fill_null(0)  # 空列表求和返回null，此处转为0
    )
)

print(result)

方法2：使用struct.map_elements

若逻辑更复杂，可将两列打包成struct，通过map_elements遍历每行计算点积，代码更直观：

result = df.with_columns(
    dot=pl.struct(["b", "c"]).map_elements(
        lambda s: sum(x*y for x, y in zip(s["b"], s["c"])),
        return_dtype=pl.Int64
    )
)

原代码报错原因 #

list.eval的作用域局限于当前列表列内部，仅能用pl.element()指代当前列表的元素，无法直接引用DataFrame的其他命名列（如代码中的pl.col('b')），因此触发报错。

最终输出 #

shape: (5, 3)
┌───────────┬───────────┬─────┐
│ b         ┆ c         ┆ dot │
│ ---       ┆ ---       ┆ --- │
│ list[i64] ┆ list[i64] ┆ i64 │
╞═══════════╪═══════════╪═════╡
│ [1, 2, 3] ┆ [1, 1, 1] ┆ 6   │
│ [2, 3]    ┆ [2, 2]    ┆ 10  │
│ [4]       ┆ [4]       ┆ 16  │
│ [1, 2, 3] ┆ [1, 5, 6] ┆ 29  │
│ []        ┆ []        ┆ 0   │
└───────────┴───────────┴─────┘

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