[Python] How to perform a conditional groupby with polars and python

I'd like to perform a groupby on the polars dataframe below:

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.]
})


I tried to perform the following operation:

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


However, the result is not what I need, which would be, for each group “A” to return the values of C greater than or equal to the mean within each group A.

Continue reading...