如何在Polars中操作向量/矩阵？嵌套结构处理建议

测试用DataFrame #

import polars as pl
df = pl.DataFrame({
   "a": [[1,2], [3,4]],
   "b": [[10, 20], [30, 40]]
})

各方法分析与推荐方案 #

1. 直接列表运算（不支持） #

Polars 0.19.9不支持直接对List类型执行算术操作，会抛出异常：

df.with_columns(pl.col("a") + pl.col("b"))
# PanicException: `add` operation not supported for dtype `list[i64]`

此方法不可行。

2. 列表转结构体运算（推荐） #

这是最符合Polars向量化风格的方案，利用Polars原生支持的结构体运算能力，无需拆分聚合或Python层面的操作：

# 运算后转回List类型，保持字段一致性
df.with_columns(
    (pl.col("a").list.to_struct() + pl.col("b").list.to_struct())
    .struct.to_list()
    .alias("sum")
)

优势：

• 完全基于Polars向量化引擎，性能最优
• 代码简洁，无冗余的索引、拆分、聚合逻辑
• 适用于固定长度的向量场景（结构体要求字段数量固定）

3. explode+group_by+join/hstack（不推荐） #

这类方法需要拆分列表再聚合，代码冗余，维护成本高，且大规模数据下会带来额外性能开销：

# join版本
df = df.with_row_index("i")
c = (
  df
    .select("a", "b", "i")
    .explode("a", "b")
    .group_by("i")
    .agg(c=pl.col("a")+pl.col("b"))
    .select("i", "c")
)
df = df.join(c, on="i")

# hstack版本
df = df.with_row_index("i")
c = (
  df
    .select("a", "b", "i")
    .explode("a", "b")
    .group_by("i", maintain_order=True)
    .agg(c=pl.col("a")+pl.col("b"))
    .select("c")
)
df = df.hstack(c)

仅建议在向量长度不固定且无其他方案时作为备选。

4. map_elements结合numpy（不推荐） #

依赖Python UDF，脱离Polars向量化引擎，单核心运行且存在数据拷贝，性能最差：

import numpy as np
df.with_columns(
   pl.struct("a", "b")
   .map_elements(lambda x: (np.array(x["a"]) + np.array(x["b"])).tolist())
   .alias("c")
)

仅适用于极小数据集或测试场景。

嵌套列表（矩阵）的处理方案 #

对于矩阵（List(List)）这类嵌套结构，优先使用list.eval()实现嵌套向量化运算，避免拆分聚合：

示例：矩阵逐元素相加 #

df_matrix = pl.DataFrame({
    "mat1": [[[1,2], [3,4]], [[5,6], [7,8]]],
    "mat2": [[[10,20], [30,40]], [[50,60], [70,80]]]
})

# 嵌套list.eval实现逐元素相加
df_matrix.with_columns(
    pl.col("mat1").list.eval(
        pl.element().list.eval(
            pl.element() + pl.col("mat2").list.get(pl.int_range(0, pl.len())).list.get(pl.int_range(0, pl.len()))
        )
    ).alias("mat_sum")
)

进阶：使用固定长度Array类型 #

如果向量/矩阵的元素长度固定，优先使用Polars的Array类型，原生支持算术操作，代码最简洁：

# 定义Array类型的DataFrame
df_array = pl.DataFrame({
    "a": pl.Series([[1,2], [3,4]], dtype=pl.Array(pl.Int64, 2)),
    "b": pl.Series([[10,20], [30,40]], dtype=pl.Array(pl.Int64, 2))
})

# 直接执行算术操作
df_array.with_columns((pl.col("a") + pl.col("b")).alias("sum"))

总结 #

1. 固定长度场景：优先使用pl.Array类型，原生支持算术操作，代码简洁高效
2. List类型向量：用list.to_struct()转结构体运算后转回List，完全向量化
3. 嵌套列表/矩阵：使用list.eval()实现嵌套向量化运算
4. 避免：explode+group_by的冗余逻辑、Python UDF的低性能方案

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