Polars：嵌套“over”调用

Question

上下文。我编写了一个函数，用于计算列中除当前组中的元素之外的所有元素的平均值。

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

---

问题。现在，我面临的问题是我想在over上下文中运行此函数，以获取组中除当前子组中的元素之外的所有元素的平均值。

我希望以下内容能够发挥作用。

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

调用的问题？

Answer 1

为“所有行”创建一个组以使模式可重复。

def over_exclude(self, *over):
    if len(over) == 1:
        over = [True, *over]
    return self.over(over[0]) - self.over(over)
    
pl.Expr.over_exclude = over_exclude

df.with_columns(
   pl.col("value").sum().over_exclude("group")
      / pl.col("value").is_not_null().sum().over_exclude("group")
)

# shape: (6, 3)
# ┌───────┬──────────┬───────┐
# │ group ┆ subgroup ┆ value │
# │ ---   ┆ ---      ┆ ---   │
# │ str   ┆ str      ┆ f64   │
# ╞═══════╪══════════╪═══════╡
# │ A     ┆ a        ┆ 3.5   │
# │ A     ┆ b        ┆ 3.5   │
# │ B     ┆ c        ┆ 3.0   │
# │ B     ┆ d        ┆ 3.0   │
# │ C     ┆ e        ┆ 2.5   │
# │ C     ┆ f        ┆ 2.5   │
# └───────┴──────────┴───────┘

df.with_columns(
   pl.col("value").sum().over_exclude("group", "subgroup")
      / pl.col("value").is_not_null().sum().over_exclude("group", "subgroup")
)

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

也许有一些巧妙的方法可以避免重复

.over_exclude