主页 / coding / 问题 / 79107659
Accepted
Oyibo
Asked: 2024-10-21 01:31:07 +0800 CST

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

  • 772

我们如何将聚合函数作为参数传递给 Polars 中的自定义聚合函数？如果按列进行不同的聚合，则应该能够为所有列或字典传递单个函数。

import polars as pl

# Sample DataFrame
df = pl.DataFrame({
    "category": ["A", "A", "B", "B", "B"],
    "value": [1, 2, 3, 4, 5]
})

def agg_with_sum(df: pl.DataFrame | pl.LazyFrame) -> pl.DataFrame | pl.LazyFrame:
    return df.group_by("category").agg(pl.col("*").sum())

# Custom function to perform aggregation
def agg_with_expr(df: pl.DataFrame | pl.LazyFrame,
                  agg_expr: pl.Expr | dict[str, pl.Expr]) -> pl.DataFrame | pl.LazyFrame:
    if isinstance(agg_expr, dict):
        return df.group_by("category").agg([pl.col(col).aggexpr() for col, aggexpr in agg_expr.items()])
    return df.group_by("category").agg(pl.col("*").agg_expr())

# Trying to pass a Polars expression for sum aggregation
print(agg_with_sum(df))
# ┌──────────┬───────┐
# │ category ┆ value │
# │ ---      ┆ ---   │
# │ str      ┆ i64   │
# ╞══════════╪═══════╡
# │ A        ┆ 3     │
# │ B        ┆ 12    │
# └──────────┴───────┘

# Trying to pass a custom Polars expression
print(agg_with_expr(df, pl.sum))
# AttributeError: 'Expr' object has no attribute 'agg_expr'

print(agg_with_expr(df, {'value': pl.sum}))
# AttributeError: 'Expr' object has no attribute 'aggexpr'
python

1 个回答

  1. Best Answer

    您可以将其作为匿名函数传递，并以表达式作为参数（我简化了您的示例只是为了说明这一点）：

    def agg_with_expr(df, agg_expr):
    return df.group_by("category").agg(agg_expr(pl.col("*")))

agg_with_expr(df, lambda x: x.sum())

    shape: (2, 2)
┌──────────┬───────┐
│ category ┆ value │
│ ---      ┆ ---   │
│ str      ┆ i64   │
╞══════════╪═══════╡
│ B        ┆ 12    │
│ A        ┆ 3     │
└──────────┴───────┘

    更新。正如@orlp 在评论中提到的那样，在这种特殊情况下，您可以在不使用匿名函数的情况下完成此操作，而只需简单使用pl.Expr.sum()，这样更加简洁。

    agg_with_expr(df, pl.Expr.sum)

    shape: (2, 2)
┌──────────┬───────┐
│ category ┆ value │
│ ---      ┆ ---   │
│ str      ┆ i64   │
╞══════════╪═══════╡
│ A        ┆ 3     │
│ B        ┆ 12    │
└──────────┴───────┘
    • 2

相关问题