查找极坐标中的二阶聚合

Question

假设我有一个数据框：

df = pl.DataFrame({'a':[1,1,1,1,2,2,2,2],'b':[1,2,1,2,1,2,1,2],'c':[10,11,12,13,14,15,18,17]})
print(df)
output:
shape: (8, 3)
┌─────┬─────┬─────┐
│ a   ┆ b   ┆ c   │
│ --- ┆ --- ┆ --- │
│ i64 ┆ i64 ┆ i64 │
╞═════╪═════╪═════╡
│ 1   ┆ 1   ┆ 10  │
│ 1   ┆ 2   ┆ 11  │
│ 1   ┆ 1   ┆ 12  │
│ 1   ┆ 2   ┆ 13  │
│ 2   ┆ 1   ┆ 14  │
│ 2   ┆ 2   ┆ 15  │
│ 2   ┆ 1   ┆ 18  │
│ 2   ┆ 2   ┆ 17  │
└─────┴─────┴─────┘

现在，在根据

a

的每个组中，我想根据b

中的第二级组找到最后一个值的最大值

通过以下代码实现的东西

df.with_columns(pl.col("c").last().over(["a", "b"])).group_by(pl.col("a")).agg(pl.col("c").max())
output:
shape: (2, 2)
┌─────┬─────┐
│ a   ┆ c   │
│ --- ┆ --- │
│ i64 ┆ i64 │
╞═════╪═════╡
│ 2   ┆ 18  │
│ 1   ┆ 13  │
└─────┴─────┘

有没有办法用一个表达式来做到这一点？目前，我必须为这种情况编写一个单独的类，提供一个用于访问

with_columns

的表达式，以及用于

group_by

的另一个表达式。我可以做到这一点，但我也想在极地方面做得更好，并且想知道是否可以仅用一种表达来完成

Answer 1

我认为您正在要求链式

over

调用：https://github.com/pola-rs/polars/issues/14361

pl.col("c").last().over("a", "b").max().over("a")
# InvalidOperationError: window expression not allowed in aggregation

使用 map_batches 将复杂功能包装为

“单一表达式”

“可能”

df.with_columns(
    pl.struct("a", "c").last().over("a", "b").map_batches(lambda s:
        s.struct.unnest()
         .select(pl.col("c").max().over("a"))
         .to_series()
    )
    .alias("result")
)

shape: (8, 4)
┌─────┬─────┬─────┬────────┐
│ a   ┆ b   ┆ c   ┆ result │
│ --- ┆ --- ┆ --- ┆ ---    │
│ i64 ┆ i64 ┆ i64 ┆ i64    │
╞═════╪═════╪═════╪════════╡
│ 1   ┆ 1   ┆ 10  ┆ 13     │
│ 1   ┆ 2   ┆ 11  ┆ 13     │
│ 1   ┆ 1   ┆ 12  ┆ 13     │
│ 1   ┆ 2   ┆ 13  ┆ 13     │
│ 2   ┆ 1   ┆ 14  ┆ 18     │
│ 2   ┆ 2   ┆ 15  ┆ 18     │
│ 2   ┆ 1   ┆ 18  ┆ 18     │
│ 2   ┆ 2   ┆ 17  ┆ 18     │
└─────┴─────┴─────┴────────┘

但是要使其可重复使用可能相当棘手。

Answer 2

从逻辑上讲，我认为你甚至不需要使用

with_columns()

。

由于您只需要 Expr.last()

 内的 

GroupBy.last

 返回的值（您也可以在您的情况下使用 

a,b）来计算

Expr.max()

，我想说您只需要执行

 DataFrame.group_by()

两次，每次都缩小DataFrame的大小：

(
    df
    .group_by("a","b").last()
    .group_by("a").agg(pl.col("c").max())
)

┌─────┬─────┐
│ a   ┆ c   │
│ --- ┆ --- │
│ i64 ┆ i64 │
╞═════╪═════╡
│ 2   ┆ 18  │
│ 1   ┆ 13  │
└─────┴─────┘

当然，这不是一个表达方式，但我认为这已经足够清楚了，这是适合您的任务的适当分组流程。

如果您确实想将其作为单个操作运行，则可以使用

GroupBy.map_groups()

：

(
    df
    .group_by("a")
    .map_groups(lambda g: g.group_by("b").last().max())
    .drop("b")
)

┌─────┬─────┐
│ a   ┆ c   │
│ --- ┆ --- │
│ i64 ┆ i64 │
╞═════╪═════╡
│ 1   ┆ 13  │
│ 2   ┆ 18  │
└─────┴─────┘