Двухуровневая группа_by для расчета среднего значения максимальной агрегации в Polars ⇐ Python

[Anonymous]

Пример:

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

import polars as pl
df = pl.DataFrame(
{
'NE' : ["P1", "P1", "P1", "P1", "P2", "P2", "P2", "P2", "P3", "P3", "P3", "P3" ],
'DT' : ["D1", "D1", "D2", "D2", "D1", "D1", "D2", "D2", "D1", "D1", "D2", "D2" ],
'TM' : ["H1", "H2", "H1", "H2", "H1", "H2", "H1", "H2", "H1", "H2", "H1", "H2" ],
'UT' : [ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12 ]
}
)
df
shape: (12, 4)
┌─────┬─────┬─────┬─────┐
│ NE  ┆ DT  ┆ TM  ┆ UT  │
│ --- ┆ --- ┆ --- ┆ --- │
│ str ┆ str ┆ str ┆ i64 │
╞═════╪═════╪═════╪═════╡
│ P1  ┆ D1  ┆ H1  ┆ 1   │
│ P1  ┆ D1  ┆ H2  ┆ 2   │
│ P1  ┆ D2  ┆ H1  ┆ 3   │
│ P1  ┆ D2  ┆ H2  ┆ 4   │
│ …   ┆ …   ┆ …   ┆ …   │
│ P3  ┆ D1  ┆ H1  ┆ 9   │
│ P3  ┆ D1  ┆ H2  ┆ 10  │
│ P3  ┆ D2  ┆ H1  ┆ 11  │
│ P3  ┆ D2  ┆ H2  ┆ 12  │
└─────┴─────┴─────┴─────┘

Мне нужно выполнить два уровня агрегирования в DF со многими столбцами и несколькими миллионами строк

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

df.group_by('NE','DT').agg(pl.max('UT')).group_by('NE').agg(pl.max('UT').alias('UT_max'), pl.mean('UT').alias("UT_avg"))

shape: (3, 3)
┌─────┬────────┬────────┐
│ NE  ┆ UT_max ┆ UT_avg │
│ --- ┆ ---    ┆ ---    │
│ str ┆ i64    ┆ f64    │
╞═════╪════════╪════════╡
│ P2  ┆ 8      ┆ 7.0    │
│ P3  ┆ 12     ┆ 11.0   │
│ P1  ┆ 4      ┆ 3.0    │
└─────┴────────┴────────┘

Есть ли эффективный способ сделать это в одной group_by?
Заранее спасибо.

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