Как сделать if и else в контексте Polars group_by ⇐ Python

[Anonymous]

Для кадра данных цель состоит в том, чтобы получить среднее значение столбца - a group_by другого столбца - b, учитывая, что первое значение a в группе не равно нулю. Если это так, просто верните значение null.
Пример кадра данных

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

df = pl.DataFrame({"a": [None, 1, 2, 3, 4], "b": [1, 1, 2, 2, 2]})

Я пробовал что-то вроде

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

df.group_by("b").agg(
pl.when(pl.col("a").first().is_null()).then(None).otherwise(pl.mean("a"))
)

Результаты соответствуют ожиданиям, но появляется предупреждение о том, что не может гарантировать выполнение своей работы в контексте group_by.

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

The predicate 'col("a").first().is_null()' in 'when->then->otherwise' is not a valid aggregation and might produce a different number of rows than the groupby operation would. This behavior is experimental and may be subject to change
shape: (2, 2)
┌─────┬─────────┐
│ b   ┆ literal │
│ --- ┆ ---     │
│ i64 ┆ f64     │
╞═════╪═════════╡
│ 1   ┆ null    │
│ 2   ┆ 3.0     │
└─────┴─────────┘

Могу ли я узнать, почему и что может быть лучшим альтернативным способом сделать if-else в group_by?

Подробнее здесь: https://stackoverflow.com/questions/758 ... by-context