Как установить замаскированные значения внутри каждой группы в контексте group_by, используя py-polars ⇐ Python

[Anonymous]

Поскольку Rank не обрабатывает нулевые значения, я хочу написать функцию Rank, которая может обрабатывать нулевые значения.
import numpy as np
import polars as pl

df = pl.DataFrame({
'group': ['a'] * 3 + ['b'] * 3,
'value': [2, 1, None, 4, 5, 6],
})
df

shape: (6, 2)
┌───────┬───────┐
│ group ┆ value │
│ --- ┆ --- │
│ str ┆ i64 │
╞═══════╪═══════╡
│ a ┆ 2 │
│ a ┆ 1 │
│ a ┆ null │
│ b ┆ 4 │
│ b ┆ 5 │
│ b ┆ 6 │
└───────┴───────┘

Это работает хорошо, если я не использовал group_by, поскольку я могу использовать «когда-то-иначе» для установки значений.
def valid_rank(expr: pl.Expr, descending=False):
"""handle null values when rank"""
FLOAT_MAX, FLOAT_MIN = np.finfo(float).max, np.finfo(float).min
mask = expr.is_null()
expr = expr.fill_null(FLOAT_MIN) if descending else expr.fill_null(FLOAT_MAX)
return pl.when(~mask).then(expr.rank(descending=descending)).otherwise(None)

df.with_columns(valid_rank(pl.col('value')))

shape: (6, 2)
┌───────┬───────┐
│ group ┆ value │
│ --- ┆ --- │
│ str ┆ f32 │
╞═══════╪═══════╡
│ a ┆ 2.0 │
│ a ┆ 1.0 │
│ a ┆ null │
│ b ┆ 3.0 │
│ b ┆ 4.0 │
│ b ┆ 5.0 │
└───────┴───────┘

Однако в контексте group_by предикат col("value").is_not_null() в if->then->иначе не является агрегацией, поэтому я получу
ComputeError: the predicate 'not(col("value").is_null())' in 'when->then->otherwise' is not a valid aggregation and might produce a different number of rows than the group_by operation would

Обычно мне приходится выполнять некоторые вычисления внутри каждой группы после ранга, и меня беспокоит производительность, если я использую part_by для разделения DataFrame. Поэтому я надеюсь, что в Polars могут быть такие выражения, как np.putmask или подобные функции, которые смогут устанавливать значения внутри каждой группы.
def valid_rank(expr: pl.Expr, descending=False):
"""handle null values when rank"""
FLOAT_MAX, FLOAT_MIN = np.finfo(float).max, np.finfo(float).min
mask = expr.is_null()
expr = expr.fill_null(FLOAT_MIN) if descending else expr.fill_null(FLOAT_MAX)
# return pl.putmask(expr.rank(descending=descending), mask, None) # hope
# return expr.rank(descending).set(mask, None) # hope


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