Polars分组计算嵌入向量均值的实现方法

需求概述 #

对包含分组ID和嵌入向量列表的Polars DataFrame，按ID分组后计算每组内向量的逐元素均值。

原始数据 #

首先构造示例数据：

import polars as pl

df = pl.DataFrame({
    "id": [1, 1, 2, 2],
    "values": [
        [1, 1, 1], [3, 3, 3],
        [1, 1, 1], [2, 2, 2]
    ]
})

原始数据输出：

shape: (4, 2)
id  values
i64 list[i64]
1   [1, 1, 1]
1   [3, 3, 3]
2   [1, 1, 1]
2   [2, 2, 2]

实现方案 #

直接使用Polars内置的分组聚合与数组操作完成计算，无需依赖外部库：

result = df.group_by("id").agg(
    # 将每组的向量转为数组，沿第二个维度求均值，再转回列表
    pl.col("values").list.to_array().mean(axis=1).list.from_array()
)

print(result)

输出结果 #

shape: (2, 2)
id  values
i64 list[f64]
1   [2.0, 2.0, 2.0]
2   [1.5, 1.5, 1.5]

代码说明 #

• group_by("id")：按ID对数据分组
• list.to_array()：将每组的多个向量转为二维数组（形状为「分组内元素数, 向量长度」）
• mean(axis=1)：沿数组的第二个维度（即向量的每个位置）计算均值
• list.from_array()：将计算后的数组转回列表格式，保持数据结构一致性

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