Polars中数组列分组逐元素求和的优化实现方法

问题背景 #

刚接触Python版Polars，现有一个DataFrame，其中series列是长度固定为18的numpy数组。需要按group列分组后，对series列执行逐元素求和的聚合操作。目前已通过map_elements结合np.sum实现需求，但该方式效率较低，想寻求更优的Polars表达式实现方案（暂不考虑Rust）。

当前实现代码：

import polars as pl
import numpy as np
data = [
{'group': 1,
  'series': np.array([ 2398,  2590,  3000,  3731,  3986,  4603,  4146,  4325,  6068,
          6028,  7486,  7759,  8323,  8961,  9598, 10236, 10873, 11511])},
{'group': 1,
  'series': np.array([ 2398,  2590,  3000,  3731,  3986,  4603,  4146,  4325,  6068,
          6028,  7486,  7759,  8323,  8961,  9598, 10236, 10873, 11511])},
 {'group': 2,
  'series': np.array([1132, 1269, 1452, 1687, 1389, 1655, 1532, 1661, 1711, 1528, 1582,
         1638, 1603, 1600, 1597, 1594, 1591, 1588])},
 {'group': 3,
  'series': np.array([ 2802,  3065,  3811,  4823,  4571,  4817,  4668,  5110,  6920,
          7131, 10154, 11138, 11699, 12840, 13981, 15123, 16264, 17405])},
]
df = pl.DataFrame(data)
# 实现所需聚合（series数组逐元素求和）
# 将前两行（group 1）相加，其余行保持不变
df.group_by('group').agg(
  pl.col('series').map_elements(lambda x: np.sum(x.to_list(), axis=0))
).to_dicts()

期望输出：

group    series
i64    object
2    [1132 1269 1452 1687 1389 1655 1532 1661 1711 1528 1582 1638 1603 1600
 1597 1594 1591 1588]
1    [ 4796  5180  6000  7462  7972  9206  8292  8650 12136 12056 14972 15518
 16646 17922 19196 20472 21746 23022]
3    [ 2802  3065  3811  4823  4571  4817  4668  5110  6920  7131 10154 11138
 11699 12840 13981 15123 16264 17405]

优化方案 #

方案一：数组展开+聚合+重组（兼容全版本Polars） #

核心思路是避开Python层面的逐行遍历，利用Polars向量化操作完成：

1. 将numpy数组展开为18个独立列
2. 按group分组对这些列分别求和
3. 将求和后的列重新组合为numpy数组

实现代码：

import polars as pl
import numpy as np

data = [
    {'group': 1, 'series': np.array([2398, 2590, 3000, 3731, 3986, 4603, 4146, 4325, 6068, 6028, 7486, 7759, 8323, 8961, 9598, 10236, 10873, 11511])},
    {'group': 1, 'series': np.array([2398, 2590, 3000, 3731, 3986, 4603, 4146, 4325, 6068, 6028, 7486, 7759, 8323, 8961, 9598, 10236, 10873, 11511])},
    {'group': 2, 'series': np.array([1132, 1269, 1452, 1687, 1389, 1655, 1532, 1661, 1711, 1528, 1582, 1638, 1603, 1600, 1597, 1594, 1591, 1588])},
    {'group': 3, 'series': np.array([2802, 3065, 3811, 4823, 4571, 4817, 4668, 5110, 6920, 7131, 10154, 11138, 11699, 12840, 13981, 15123, 16264, 17405])},
]
df = pl.DataFrame(data)

result = (
    df
    # 展开数组为18个独立列
    .with_columns(
        pl.col('series').list.to_struct(fields=[f"s{i}" for i in range(18)]).unnest()
    )
    # 分组求和
    .group_by('group')
    .agg(pl.col('s*').sum())
    # 重组为numpy数组
    .with_columns(
        pl.struct(pl.col('s*')).list.to_array().map_elements(np.array)
    )
    # 整理列名
    .select('group', pl.col('array').alias('series'))
)

print(result)

方案二：利用arr.eval简化写法（Polars ≥0.19.0） #

新版本Polars提供了数组命名空间的聚合函数，写法更简洁，且支持并行计算：

result = (
    df.group_by('group')
    .agg(
        pl.col('series').arr.eval(pl.element().sum(), parallel=True).first()
    )
)

• arr.eval(pl.element().sum())会对数组每个位置的元素执行分组求和
• parallel=True可开启并行计算进一步提升效率
• .first()用于取出arr.eval返回的单元素列表，得到最终求和数组

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