Início / coding / Perguntas / 79107659
Accepted
Oyibo
Asked: 2024-10-21 01:31:07 +0800 CST

Como passar funções de agregação como argumento de função em Polars?

  • 772

Como podemos passar funções de agregação como argumento para uma função de agregação personalizada no Polars? Você deve conseguir passar uma única função para todas as colunas ou um dicionário se tiver diferentes agregações por coluna.

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 respostas

  1. Best Answer

    Você pode passar isso como uma função anônima com expressão como parâmetro (simplifiquei seu exemplo apenas para ilustrar o ponto):

    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     │
└──────────┴───────┘

    atualizar . como @orlp mencionou nos comentários, neste caso específico você poderia fazer isso sem a função anônima, com o uso simples de pl.Expr.sum(), o que é muito mais legal.

    agg_with_expr(df, pl.Expr.sum)

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

relate perguntas