Polars嵌套when-then-otherwise在group_by/窗口上下文行为异常

编写了一个自定义表达式封装Polars原生rank方法，目标是实现排序时自动忽略空值，空值位置最终返回null的效果：

import polars as pl
import numpy as np

def rank(_exp,method='average',descending=False):
    # 填充空值避免影响排名
    fill = -np.Inf if descending else np.Inf
    tmp = pl.when(_exp.is_not_null()).then(_exp).otherwise(fill).rank(descending=descending,method=method)
    # 回填空值
    exp = pl.when(_exp.is_not_null()).then(tmp).otherwise(_exp)
    return exp

对于包含空值的普通DataFrame，该函数返回结果符合预期：

df_with_nan = pl.DataFrame({'X': np.where(np.isin(np.arange(10),[1,3]),np.NaN,np.arange(10)), 'G1':np.array(['A']*5 + ['B']*5), 'G2':np.array(['C','D']*5)})
df_with_null = df_with_nan.fill_nan(None)

print(df_with_null.with_columns(rank(pl.col('X')).name.suffix('_rnk')))

shape: (10, 4)
┌──────┬─────┬─────┬───────┐
│ X    ┆ G1  ┆ G2  ┆ X_rnk │
│ ---  ┆ --- ┆ --- ┆ ---   │
│ f64  ┆ str ┆ str ┆ f64   │
╞══════╪═════╪═════╪═══════╡
│ 0.0  ┆ A   ┆ C   ┆ 1.0   │
│ null ┆ A   ┆ D   ┆ null  │
│ 2.0  ┆ A   ┆ C   ┆ 2.0   │
│ null ┆ A   ┆ D   ┆ null  │
│ 4.0  ┆ A   ┆ C   ┆ 3.0   │
│ 5.0  ┆ B   ┆ D   ┆ 4.0   │
│ 6.0  ┆ B   ┆ C   ┆ 5.0   │
│ 7.0  ┆ B   ┆ D   ┆ 6.0   │
│ 8.0  ┆ B   ┆ C   ┆ 7.0   │
│ 9.0  ┆ B   ┆ D   ┆ 8.0   │
└──────┴─────┴─────┴───────┘

但在两类场景下使用时，误以为结果出现错乱：

1. 窗口上下文（使用over()分区计算）：

print(df_with_null.with_columns(rank(pl.col('X')).over('G1','G2').name.suffix('_rnk')))

shape: (10, 4)
┌──────┬─────┬─────┬───────┐
│ X    ┆ G1  ┆ G2  ┆ X_rnk │
│ ---  ┆ --- ┆ --- ┆ ---   │
│ f64  ┆ str ┆ str ┆ f64   │
╞══════╪═════╪═════╪═══════╡
│ 0.0  ┆ A   ┆ C   ┆ 1.0   │
│ null ┆ A   ┆ D   ┆ null  │
│ 2.0  ┆ A   ┆ C   ┆ 2.0   │
│ null ┆ A   ┆ D   ┆ null  │
│ 4.0  ┆ A   ┆ C   ┆ 3.0   │
│ 5.0  ┆ B   ┆ D   ┆ 1.0   │
│ 6.0  ┆ B   ┆ C   ┆ 1.0   │
│ 7.0  ┆ B   ┆ D   ┆ 2.0   │
│ 8.0  ┆ B   ┆ C   ┆ 2.0   │
│ 9.0  ┆ B   ┆ D   ┆ 3.0   │
└──────┴─────┴─────┴───────┘

2. group_by分组聚合上下文：

print(df_with_null.group_by('G1','G2').agg(pl.col('X'),rank(pl.col('X')).name.suffix('_rnk')).explode(['X','X_rnk']))

shape: (10, 4)
┌─────┬─────┬──────┬───────┐
│ G1  ┆ G2  ┆ X    ┆ X_rnk │
│ --- ┆ --- ┆ ---  ┆ ---   │
│ str ┆ str ┆ f64  ┆ f64   │
╞═════╪═════╪══════╪═══════╡
│ B   ┆ C   ┆ 6.0  ┆ 1.0   │
│ B   ┆ C   ┆ 8.0  ┆ 2.0   │
│ B   ┆ D   ┆ 5.0  ┆ 1.0   │
│ B   ┆ D   ┆ 7.0  ┆ 2.0   │
│ B   ┆ D   ┆ 9.0  ┆ 3.0   │
│ A   ┆ D   ┆ null ┆ null  │
│ A   ┆ D   ┆ null ┆ null  │
│ A   ┆ C   ┆ 0.0  ┆ 1.0   │
│ A   ┆ C   ┆ 2.0  ┆ 2.0   │
│ A   ┆ C   ┆ 4.0  ┆ 3.0   │
└─────┴─────┴──────┴───────┘

排查后怀疑问题出在回填空值的when-then-otherwise语句，但不清楚触发原因，也需要可稳定实现忽略空值排序的替代方案。

1. 上述两个场景的结果本身不存在值错位问题：
   • over('G1','G2')是按G1+G2分区做组内排名，和全局排名逻辑不同，比如(B,D)组内的X值只有5、7、9，三个值的组内排名本来就是1、2、3，返回结果完全正确。
   • group_by默认使用哈希分组实现，不保留原表的行顺序，输出的分组顺序是不确定的，但每个分组内的X值和对应排名完全对齐，不存在错位。
2. 自定义填充+回填的写法存在兼容隐患：在部分Polars版本的窗口/分组表达式求值逻辑中，otherwise(_exp)会重复引用原表达式，在分组重排、分块计算时可能出现不可预期的对齐问题，不是官方推荐的写法。

直接使用Polars原生rank方法自带的空值处理参数即可，不需要自行实现填充、回填逻辑，全场景（全局计算、窗口over、分组agg）兼容，不会出现对齐问题：

def rank(_exp, method='average', descending=False):
    return _exp.rank(
        method=method,
        descending=descending,
        null_policy="ignore"  # 空值不参与排名，直接返回null
    )

如果需要group_by结果保留原表的行顺序，在调用group_by时加上maintain_order=True参数即可：

print(
    df_with_null.group_by('G1','G2', maintain_order=True)
    .agg(pl.col('X'),rank(pl.col('X')).name.suffix('_rnk'))
    .explode(['X','X_rnk'])
)

如果使用的是较早版本的Polars，没有null_policy="ignore"参数，可以把回填逻辑里的otherwise(_exp)改成otherwise(None)，避免重复引用原表达式导致的分组对齐问题：

def rank(_exp,method='average',descending=False):
    fill = -np.Inf if descending else np.Inf
    tmp = pl.when(_exp.is_not_null()).then(_exp).otherwise(fill).rank(descending=descending,method=method)
    exp = pl.when(_exp.is_not_null()).then(tmp).otherwise(None)
    return exp

内容的提问来源于stack exchange，提问作者lowmotion