如何在Polars中编写可嵌套调用的自定义表达式

问题分析 #

之前的代码报错核心原因是：Polars的Expr对象本身没有挂载自定义函数，无法直接通过.运算符调用expr1/expr2这类自定义方法；另外apply()仅支持Python原生函数，不能传入Polars Expr对象，所以也会触发错误。

解决方案 #

方法一：让自定义函数接收Expr参数（推荐） #

把自定义函数改成接受Expr对象作为参数，而非硬编码固定列，这样就能灵活实现嵌套和链式调用：

import polars as pl

def expr2(expr: pl.Expr, method: str) -> pl.Expr:
    if method == 'ceil':
        return expr.ceil()
    elif method == 'floor':
        return expr.floor()
    else:
        raise ValueError("仅支持'ceil'或'floor'方法")

# 嵌套调用示例：先执行expr2逻辑，再取绝对值
def expr1(expr: pl.Expr, method: str) -> pl.Expr:
    return expr2(expr, method).abs()

# 测试数据
df = pl.Series('A', [-1.2, 3.7, -0.5]).to_frame()

# 调用方式
result = df.select(
    expr1(pl.col('A'), 'ceil')
)
print(result)

# 如果expr1不需要额外参数，更简洁的链式写法
def expr1(expr: pl.Expr) -> pl.Expr:
    return expr.abs()

result2 = df.select(
    expr1(expr2(pl.col('A'), 'floor'))
)
print(result2)

方法二：给Expr类打猴子补丁（不推荐） #

如果一定要用.链式调用的语法，可以通过猴子补丁给Expr类添加自定义方法，但这种方式可能在Polars版本更新时出现兼容性问题，需谨慎使用：

import polars as pl

# 定义自定义方法，第一个参数是self（即Expr对象）
def expr2(self: pl.Expr, method: str) -> pl.Expr:
    if method == 'ceil':
        return self.ceil()
    elif method == 'floor':
        return self.floor()
    else:
        raise ValueError("仅支持'ceil'或'floor'方法")

def expr1(self: pl.Expr) -> pl.Expr:
    return self.abs()

# 给Expr类挂载自定义方法
pl.Expr.expr2 = expr2
pl.Expr.expr1 = expr1

# 测试数据
df = pl.Series('A', [-1.2, 3.7, -0.5]).to_frame()

# 链式调用
result = df.select(
    pl.col('A').expr2('ceil').expr1()
)
print(result)

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