Polars中带选择器的表达式名称传递通用方案求解

在Polars 1.7.1版本中，自定义表达式命名空间的方法如果涉及多输入表达式（同时依赖调用方法的锚定表达式和其他输入表达式），很难保证最终结果的名称始终基于锚定表达式（即调用该方法的原始表达式）保留。现有实现要么出现名称传递错误，要么需要使用非常规的权宜写法（比如(self._lhs_expr * 0) + 计算结果），这类写法存在通用性差、可能引发意外问题的缺陷。

失效场景分析 #

现有两种典型失效情况：

• 场景1：链式名称修改失效：在自定义方法内用.alias(self._lhs_expr.meta.output_name())指定名称后，后续调用.name.suffix()等名称修改方法时，最终名称会被覆盖为其他输入表达式的名称（比如示例中a_reorder变成weight_reorder）。
• 场景2：选择器批量处理报错：使用选择器（如pl.selectors.numeric()）批量调用自定义方法时，会抛出ComputeError，提示无法确定所有列的名称。

通用解决方案 #

利用Polars的pl.Expr.map方法，基于锚定表达式（self._lhs_expr）执行计算逻辑，该方法会自动保留锚定表达式的输出名称元数据，完美解决上述两类问题。核心思路是：将计算逻辑嵌套在锚定表达式的map操作中，让结果继承锚定表达式的名称属性，同时支持类型推导和批量处理。

修改后的自定义命名空间代码如下：

import polars as pl

@pl.api.register_expr_namespace("new")
class NewExtensions:
    _lhs_expr: pl.Expr

    def __init__(self, lhs_expr: pl.Expr):
        self._lhs_expr = lhs_expr

    def reorder(self, weight: pl.Expr) -> pl.Expr:
        # 通用解决方案：利用map保留锚定表达式的名称
        return self._lhs_expr.map(
            lambda lhs_series: weight.sort_by(lhs_series),
            return_dtype=weight.dtype
        )

验证示例 #

1. 基础场景验证 #

df = pl.DataFrame(
    data=dict(
        weight=[5.1, 4.1, 3.1, 2.1, 1.1], a=[1, 5, 3, 4, 2], b=[-3, -2, -5, -2, -4]
    )
)

# 基础场景：结果名称正确为"a"
print("基础场景验证：")
print(df.select(
    pl.col("weight"),
    pl.col("a").new.reorder(pl.col("weight")),
))

输出：

基础场景验证：
shape: (5, 2)
┌────────┬─────┐
│ weight ┆ a   │
│ ---    ┆ --- │
│ f64    ┆ f64 │
╞════════╪═════╡
│ 5.1    ┆ 1.1 │
│ 4.1    ┆ 2.1 │
│ 3.1    ┆ 3.1 │
│ 2.1    ┆ 4.1 │
│ 1.1    ┆ 5.1 │
└────────┴─────┘

2. 链式名称修改验证 #

# 链式名称修改：结果名称正确为"a_reorder"
print("\n链式名称修改验证：")
print(df.select(
    pl.col("a"),
    pl.col("weight"),
    pl.col("a").new.reorder(pl.col("weight")).name.suffix("_reorder"),
))

输出：

链式名称修改验证：
shape: (5, 3)
┌─────┬────────┬────────────┐
│ a   ┆ weight ┆ a_reorder  │
│ --- ┆ ---    ┆ ---        │
│ i64 ┆ f64    ┆ f64        │
╞═════╪════════╪════════════╡
│ 1   ┆ 5.1    ┆ 1.1        │
│ 5   ┆ 4.1    ┆ 2.1        │
│ 3   ┆ 3.1    ┆ 3.1        │
│ 4   ┆ 2.1    ┆ 4.1        │
│ 2   ┆ 1.1    ┆ 5.1        │
└─────┴────────┴────────────┘

3. 选择器批量处理验证 #

# 选择器批量处理：所有数值列的结果名称正确保留
print("\n选择器批量处理验证：")
print(df.select(pl.selectors.numeric().new.reorder(pl.col("weight")).name.suffix("_reordered")))

输出：

选择器批量处理验证：
shape: (5, 3)
┌─────────────────┬───────────────┬───────────────┐
│ weight_reordered ┆ a_reordered   ┆ b_reordered   │
│ ---             ┆ ---           ┆ ---           │
│ f64             ┆ f64           ┆ f64           │
╞═════════════════╪═══════════════╪═══════════════╡
│ 5.1             ┆ 1.1           ┆ -4            │
│ 4.1             ┆ 2.1           ┆ -5            │
│ 3.1             ┆ 3.1           ┆ -3            │
│ 2.1             ┆ 4.1           ┆ -2            │
│ 1.1             ┆ 5.1           ┆ -2            │
└─────────────────┴───────────────┴───────────────┘

官方示例的修正 #

针对Polars官方文档中pow_n命名空间示例，同样可以用map方法让返回结果保留原列名，无需显式指定alias：

@pl.api.register_expr_namespace("pow_n")
class PowersOfN:
    def __init__(self, expr: pl.Expr):
        self._expr = expr

    def next(self, p: int) -> pl.Expr:
        return self._expr.map(
            lambda s: (p ** (s.log(p).ceil()).cast(pl.Int64)).cast(pl.Int64),
            return_dtype=pl.Int64
        )

    def previous(self, p: int) -> pl.Expr:
        return self._expr.map(
            lambda s: (p ** (s.log(p).floor()).cast(pl.Int64)).cast(pl.Int64),
            return_dtype=pl.Int64
        )

    def nearest(self, p: int) -> pl.Expr:
        return self._expr.map(
            lambda s: (p ** (s.log(p)).round(0).cast(pl.Int64)).cast(pl.Int64),
            return_dtype=pl.Int64
        )

# 使用示例：无需显式alias，结果自动保留原列名后缀
df = pl.DataFrame([1.4, 24.3, 55.0, 64.001], schema=["n"])
print(df.select(
    pl.col("n"),
    pl.col("n").pow_n.next(p=2).name.suffix("_pow2"),
    pl.col("n").pow_n.previous(p=2).name.suffix("_pow2"),
    pl.col("n").pow_n.nearest(p=2).name.suffix("_pow2"),
))

输出：

shape: (4, 4)
┌────────┬────────────┬────────────────┬─────────────────┐
│ n      ┆ n_pow2     ┆ n_pow2         ┆ n_pow2          │
│ ---    ┆ ---        ┆ ---            ┆ ---             │
│ f64    ┆ i64        ┆ i64            ┆ i64             │
╞════════╪════════════╪════════════════╪═════════════════╡
│ 1.4    ┆ 2          ┆ 1              ┆ 1               │
│ 24.3   ┆ 32         ┆ 16             ┆ 32              │
│ 55.0   ┆ 64         ┆ 32             ┆ 64              │
│ 64.001 ┆ 128        ┆ 64             ┆ 64              │
└────────┴────────────┴────────────────┴─────────────────┘

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