Полярные: вложенные вызовы «over» ⇐ Python

[Anonymous]

Контекст. Я написал функцию, которая вычисляет среднее значение всех элементов в столбце, кроме элементов в текущей группе.

Код: Выделить всё

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

def sum_excl_group(val_exp: pl.Expr, group_expr: pl.Expr) -> pl.Expr:
return val_exp.sum() - val_exp.sum().over(group_expr)

def count_non_null_excl_group(val_exp: pl.Expr, group_expr: pl.Expr) -> pl.Expr:
return val_exp.is_not_null().sum() - val_exp.is_not_null().sum().over(group_expr)

def mean_excl_group(val_exp: pl.Expr, group_expr: pl.Expr) -> pl.Expr:
return sum_excl_group(val_exp,group_expr) / count_non_null_excl_group(val_exp, group_expr)

(
df
.with_columns(
mean_excl_group(pl.col("value"), pl.col("group")).alias("mean_excl_group"),
)
)

Это дает ожидаемый результат.

Код: Выделить всё

shape: (6, 3)
┌───────┬───────┬─────────────────┐
│ group ┆ value ┆ mean_excl_group │
│ ---   ┆ ---   ┆ ---             │
│ str   ┆ i64   ┆ f64             │
╞═══════╪═══════╪═════════════════╡
│ A     ┆ 1     ┆ 3.5             │
│ A     ┆ 3     ┆ 3.5             │
│ B     ┆ 2     ┆ 3.0             │
│ B     ┆ 4     ┆ 3.0             │
│ C     ┆ 3     ┆ 2.5             │
│ C     ┆ 5     ┆ 2.5             │
└───────┴───────┴─────────────────┘

Проблема. Теперь я столкнулся с проблемой: мне хотелось бы запустить эту функцию внутри контекста над, чтобы получить среднее значение всех элементов в группе, кроме элементов в текущей подгруппе.
Я ожидал, что сработает следующее.

Код: Выделить всё

df = pl.DataFrame({
"group": ["A", "A", "B", "B", "C", "C"],
"subgroup": ["a", "b", "c", "d", "e", "f"],
"value": [1, 3, 2, 4, 3, 5],
})

(
df
.with_columns(
mean_excl_group(pl.col("value"), pl.col("subgroup")).over("group").alias("mean_excl_group"),
)
)

но получаю ошибку InvalidOperationError

Код: Выделить всё

InvalidOperationError: window expression not allowed in aggregation

Попытка. На данный момент я «решил» эту проблему, избегая вложенных вызовов over.

Код: Выделить всё

def sum_excl_group(val_exp: pl.Expr, coarse_group_expr: pl.Expr, fine_group_expr: pl.Expr) -> pl.Expr:
return val_exp.sum().over(coarse_group_expr) - val_exp.sum().over(fine_group_expr)

def count_non_null_excl_group(val_exp: pl.Expr, coarse_group_expr: pl.Expr, fine_group_expr: pl.Expr) -> pl.Expr:
return val_exp.is_not_null().sum().over(coarse_group_expr) - val_exp.is_not_null().sum().over(fine_group_expr)

def mean_excl_group(val_exp: pl.Expr, coarse_group_expr: pl.Expr, fine_group_expr: pl.Expr) -> pl.Expr:
return sum_excl_group(val_exp, coarse_group_expr, fine_group_expr) / count_non_null_excl_group(val_exp, coarse_group_expr, fine_group_expr)

(
df
.with_columns(
mean_excl_group(pl.col("value"), pl.col("group"), pl.col("subgroup")).alias("mean_excl_group"),
)
)

Это дает ожидаемый результат.

Код: Выделить всё

shape: (6, 4)
┌───────┬──────────┬───────┬─────────────────┐
│ group ┆ subgroup ┆ value ┆ mean_excl_group │
│ ---   ┆ ---      ┆ ---   ┆ ---             │
│ str   ┆ str      ┆ i64   ┆ f64             │
╞═══════╪══════════╪═══════╪═════════════════╡
│ A     ┆ a        ┆ 1     ┆ 3.0             │
│ A     ┆ b        ┆ 3     ┆ 1.0             │
│ B     ┆ c        ┆ 2     ┆ 4.0             │
│ B     ┆ d        ┆ 4     ┆ 2.0             │
│ C     ┆ e        ┆ 3     ┆ 5.0             │
│ C     ┆ f        ┆ 5     ┆ 3.0             │
└───────┴──────────┴───────┴─────────────────┘

Однако это требует от меня передачи обеих деталей всем функциям, что делает код более раздутым, чем (я надеюсь) он должен быть. Есть ли более чистый и полярный способ решить проблему вложенных вызовов over?

Подробнее здесь: https://stackoverflow.com/questions/774 ... over-calls