Как выполнить условную группировку с помощью поляров и Python ⇐ Python

[Anonymous]

Я хотел бы выполнить группировку в кадре данных Polars ниже:

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

import polars as pl

df = pl.DataFrame({
'A' : ['foo', 'bar', 'foo', 'bar', 'foo', 'bar'],
'B' : ['one', 'one', 'two', 'three', 'two', 'two'],
'C' : [1, 5, 5, 2, 5, 7],
'D' : [2.0, 5., 8., 1., 2., 9.]
})

Я попытался выполнить следующую операцию:

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

df.group_by('A').agg( pl.col('C').map_elements(lambda x: x >= x.mean()).alias('C') )

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

shape: (2, 2)
┌─────┬─────────────────────┐
│ A   ┆ C                   │
│ --- ┆ ---                 │
│ str ┆ list[bool]          │
╞═════╪═════════════════════╡
│ foo ┆ [false, true, true] │
│ bar ┆ [true, false, true] │
└─────┴─────────────────────┘

Однако результат не тот, который мне нужен: для каждой группы «A» должны быть возвращены значения C, большие или равные среднему значению в каждой группе A.


Подробнее здесь: https://stackoverflow.com/questions/770 ... and-python