Настраиваемая агрегация в group_by с полярами ⇐ Python

[Anonymous]

Мне нужно вычислить средневзвешенное значение процента после группировки по полярностям. Поскольку мне априори нужна общая сумма в качестве знаменателя среднего значения, я думаю, что необходимо создать специальную функцию, которая будет применяться в методе agg().
Это пример того, что я пробовал (не работает):
import polars as pl

df_group = df.group_by(pl.col("entity"))
.agg(
total_amount=pl.col("amount").sum(),
pct_avg=(
pl.col("pct")
* pl.col("amount")
/ pl.col("total_amount")
).sum(),
)

Вот минимальный пример:
a = pl.DataFrame(
{
"entity": ["x", "x", "y", "y"],
"pct": [0.5, 0.4, 0.4, 0.2],
"amount": [1, 10, 100, 1000],
}
)

┌────────┬─────┬────────┐
│ entity ┆ pct ┆ amount │
│ --- ┆ --- ┆ --- │
│ str ┆ f64 ┆ i64 │
╞════════╪═════╪════════╡
│ x ┆ 0.5 ┆ 1 │
│ x ┆ 0.4 ┆ 10 │
│ y ┆ 0.4 ┆ 100 │
│ y ┆ 0.2 ┆ 1000 │
└────────┴─────┴────────┘

И ожидаемый результат:
a_group = pl.DataFrame(
{
"entity": ["x", "y"],
"total_amount": [11, 1100],
"pct_avg": ["(0.5 * 1 + 0.4 * 10)/11", "(0.4 * 100 + 0.2 *1000)/1100"],
}
)

┌────────┬──────────────┬──────────────────────────────┐
│ entity ┆ total_amount ┆ pct_avg │
│ --- ┆ --- ┆ --- │
│ str ┆ i64 ┆ str │
╞════════╪══════════════╪══════════════════════════════╡
│ x ┆ 11 ┆ (0.5 * 1 + 0.4 * 10)/11 │
│ y ┆ 1100 ┆ (0.4 * 100 + 0.2 *1000)/1100 │
└────────┴──────────────┴──────────────────────────────┘


Подробнее здесь: https://stackoverflow.com/questions/798 ... ith-polars