Polars中struct与col实现加权求和的性能差异及自定义聚合问询

测试背景与现象 #

我测试了Polars中两种加权求和的实现方式，发现直接基于列的实现比用pl.struct的实现性能高出3-6倍（1000万行快3倍，10亿行快6倍）。测试代码如下：

import polars as pl
import numpy as np
import time

def weighted_sum(expr: pl.Expr) -> pl.Expr:
    return (expr.struct.field("value") * expr.struct.field("weight")).sum() / expr.struct.field("weight").sum()

if __name__ == "__main__":
    size = 10_000_000
    frame = pl.from_dict({"weight": np.random.rand(size), "value": np.random.randn(size)}).lazy()

    t0 = time.time()
    frame.select(pl.struct("weight", "value").pipe(weighted_sum)).collect()
    t = time.time() - t0
    print(f"Time for struct based weighted sum {t:.3f}s.")

    t0 = time.time()
    frame.select((pl.col("value") * pl.col("weight")).sum() / pl.col("weight").sum()).collect()
    t = time.time() - t0
    print(f"Time for col based weighted sum {t:.3f}s.")

性能差异的原因 #

• 额外的内存拷贝与打包开销：使用pl.struct需要先将两列的数据打包成struct类型，这一步会产生额外的内存分配和数据拷贝，数据量越大，拷贝的时间和内存成本越高。
• 连续内存访问被打断：Polars对单列操作的优化依赖于连续的内存布局，直接操作列时可以充分利用向量化计算的优势；而访问struct的字段需要额外的解引用步骤，破坏了内存连续性，降低了计算效率。
• 执行计划复杂度更高：struct方式需要先提取字段再进行计算，增加了执行计划的中间步骤，Polars的查询优化器对这种嵌套操作的优化空间有限，不如直接列操作的执行计划高效。

自定义加权求和聚合函数的标准方式 #

自定义聚合函数的参数类型应该是pl.Expr，因为Polars的表达式API是懒执行的，使用Expr可以让函数融入Polars的查询优化流程，支持懒执行和高效的向量化计算。如果用pl.Series则是即时执行，无法参与查询优化，也不适用于groupby.agg等懒操作场景。

标准实现与调用 #

def weighted_sum(values: pl.Expr, weights: pl.Expr) -> pl.Expr:
    return (values * weights).sum() / weights.sum()

# 在groupby聚合中调用
frame.groupby("group_col").agg(weighted_sum(pl.col("value"), pl.col("weight")))

支持pipe调用的实现 #

如果希望用pl.col("value", "weight").pipe(...)的方式提升可读性，可以调整函数接收列的列表：

def weighted_sum_from_cols(cols: list[pl.Expr]) -> pl.Expr:
    values, weights = cols
    return (values * weights).sum() / weights.sum()

# 调用方式
frame.groupby("group_col").agg(pl.col("value", "weight").pipe(weighted_sum_from_cols))

兼顾可读性与性能的方案 #

上面的pipe方式既保持了代码的简洁可读性（类似pl.col("value", "weight").pipe(weighted_sum)的写法），又避免了pl.struct带来的性能损耗，因为它直接操作原始列，没有额外的struct打包和解包步骤，完全利用Polars对列操作的优化能力。

内容的提问来源于stack exchange，提问作者Benjamin Trendelkamp-Schroer