Polars自定义条件求和表达式实现及求和写法差异解析

需求与疑问 #

自定义求和函数需求 #

需要实现一个求和函数（如nullsum），行为与Polars默认sum一致，但当列中全为null时返回null，且支持min_count参数，能在select、with_columns、groupby等场景复用。

四种求和写法的差异疑问 #

刚接触Polars，无法区分以下四种列求和写法的差异：

• i) 调用DataFrame的sum方法

  frame.sum()
  

• ii) select上下文使用单列求和表达式列表

  frame.select([pl.col("A").sum(), pl.col("B").sum()])
  

• iii) 先选列再调用DataFrame的sum方法

  frame.select([pl.col("A"), pl.col("B")]).sum()
  

• iv) select上下文对多列表达式链式调用sum

  frame.select(pl.col(["A", "B"]).sum())
  

适配性疑问 #

自定义求和函数能否适配上述四种写法？

解决方案：基于Expr.pipe实现自定义nullsum函数 #

通过Polars的Expr.pipe方法，可将自定义求和逻辑封装为可复用的表达式函数，完美适配各类上下文场景：

def nullsum(expr: pl.Expr, min_count: int = 0) -> pl.Expr:
    if min_count > 0:
        return pl.when(expr.is_not_null().sum() >= pl.lit(min_count)).then(expr.sum())
    else:
        return expr.sum()

使用示例 #

result = frame.select(pl.col("A", "B").pipe(nullsum, min_count=1))
print(result)

输出示例：

shape: (1, 2)
┌─────┬─────┐
│ A   ┆ B   │
│ --- ┆ --- │
│ f32 ┆ f64 │
╞═════╪═════╡
│ null┆ 0.0 │
└─────┴─────┘

四种求和写法的差异解析 #

i) frame.sum() #

• 属于DataFrame级别的聚合方法，默认对所有数值型列执行求和，返回单行DataFrame（单列为Series）。
• 可通过numeric_only参数控制是否仅计算数值列，自动跳过非数值列。

ii) frame.select([pl.col("A").sum(), pl.col("B").sum()]) #

• 属于表达式级别的聚合，在select上下文明确指定每列的求和逻辑。
• 灵活性极强，可对不同列应用差异化聚合规则（如添加过滤条件），结果列名默认沿用原列名（可通过.alias()修改）。

iii) frame.select([pl.col("A"), pl.col("B")]).sum() #

• 是列选择+DataFrame聚合的组合操作：先提取指定列生成中间DataFrame，再调用其sum()方法。
• 效果与frame.sum(["A", "B"])基本一致，但多了中间DataFrame创建步骤，性能略逊于直接指定列的写法。

iv) frame.select(pl.col(["A", "B"]).sum()) #

• 属于多列批量聚合表达式，选中多列后链式调用sum()实现逐列求和。
• 写法简洁高效，适合对多列执行相同聚合逻辑的场景，无需逐个列编写表达式。

自定义nullsum的适配性分析 #

• 完全适配写法ii和iv：通过pl.col(...).pipe(nullsum)的方式，可直接在select、with_columns的表达式上下文使用。例如：
  • 写法ii适配：frame.select([pl.col("A").pipe(nullsum), pl.col("B").pipe(nullsum, min_count=2)])
  • 写法iv适配：frame.select(pl.col(["A", "B"]).pipe(nullsum, min_count=1))
• 写法i和iii的适配方式：这两种是DataFrame级别的聚合方法，无法直接替换内置sum。可改为表达式方式实现，比如将frame.sum()替换为frame.select(pl.col(pl.numeric()).pipe(nullsum))。
• groupby场景完美兼容：可直接在聚合阶段使用，例如：

  frame.groupby("category").agg(pl.col("value").pipe(nullsum, min_count=1))
  

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