如何在Polars中将聚合函数作为函数参数传递？

要实现支持单个聚合函数应用到所有列或按列指定不同聚合函数的自定义聚合函数，核心是正确处理聚合函数与列表达式的结合，而非使用不存在的.agg_expr()方法。以下是修正后的实现方案：

错误原因分析 #

你代码中的pl.col(col).aggexpr()或.agg_expr()是无效语法——Polars中聚合函数（如pl.sum）是接收列表达式并返回聚合表达式的函数，需通过聚合函数(列表达式)或列表达式.聚合方法()的形式调用。

正确实现代码 #

import polars as pl
from typing import Callable, Dict, Union

# 示例DataFrame
df = pl.DataFrame({
    "category": ["A", "A", "B", "B", "B"],
    "value": [1, 2, 3, 4, 5],
    "score": [10, 20, 30, 40, 50]
})

# 基础聚合示例（作为对比）
def agg_with_sum(df: pl.DataFrame | pl.LazyFrame) -> pl.DataFrame | pl.LazyFrame:
    return df.group_by("category").agg(pl.col("*").sum())

# 支持灵活传递聚合函数的自定义函数
def agg_with_custom(df: pl.DataFrame | pl.LazyFrame,
                  agg_spec: Union[Callable[[pl.Expr], pl.Expr], Dict[str, Callable[[pl.Expr], pl.Expr]]]) -> pl.DataFrame | pl.LazyFrame:
    if isinstance(agg_spec, dict):
        # 按列映射生成聚合表达式
        agg_exprs = [func(pl.col(col)) for col, func in agg_spec.items()]
    else:
        # 单个函数应用到所有非分组列
        agg_exprs = agg_spec(pl.col("*"))
    
    return df.group_by("category").agg(agg_exprs)

# 测试1：传递单个聚合函数（所有列求和）
print("=== 单个聚合函数测试 ===")
print(agg_with_custom(df, pl.sum))
# 输出：
# ┌──────────┬───────┬───────┐
# │ category ┆ value ┆ score │
# │ ---      ┆ ---   ┆ ---   │
# │ str      ┆ i64   ┆ i64   │
# ╞══════════╪═══════╪═══════╡
# │ A        ┆ 3     ┆ 30    │
# │ B        ┆ 12    ┆ 120   │
# └──────────┴───────┴───────┘

# 测试2：按列指定不同聚合函数（value求和，score求均值）
print("\n=== 按列聚合测试 ===")
print(agg_with_custom(df, {"value": pl.sum, "score": pl.mean}))
# 输出：
# ┌──────────┬───────┬───────┐
# │ category ┆ value ┆ score │
# │ ---      ┆ ---   ┆ ---   │
# │ str      ┆ i64   ┆ f64   │
# ╞══════════╪═══════╪═══════╡
# │ A        ┆ 3     ┆ 15.0  │
# │ B        ┆ 12    ┆ 40.0  │
# └──────────┴───────┴───────┘

# 测试3：使用LazyFrame也能正常工作
lazy_df = df.lazy()
print("\n=== LazyFrame测试 ===")
print(agg_with_custom(lazy_df, {"value": pl.max, "score": pl.min}).collect())
# 输出：
# ┌──────────┬───────┬───────┐
# │ category ┆ value ┆ score │
# │ ---      ┆ ---   ┆ ---   │
# │ str      ┆ i64   ┆ i64   │
# ╞══════════╪═══════╪═══════╡
# │ A        ┆ 2     ┆ 10    │
# │ B        ┆ 5     ┆ 30    │
# └──────────┴───────┴───────┘

关键说明 #

1. 类型标注：使用Union[Callable[[pl.Expr], pl.Expr], Dict[str, Callable[[pl.Expr], pl.Expr]]]明确参数类型，聚合函数是接收列表达式并返回聚合表达式的可调用对象。
2. 单个函数处理：agg_spec(pl.col("*"))会将聚合函数应用到所有非分组列（category是分组列，不会被包含）。
3. 字典映射处理：遍历字典的键值对，为每个列生成对应的聚合表达式。
4. 兼容LazyFrame：代码同时支持DataFrame和LazyFrame，无需额外修改。

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