Polars中能否创建自定义表达式？Rust实现后供Python调用可行吗？

完全可以用Rust编写Polars自定义表达式并在Python端调用，而且相比Python的apply/map方法，Rust实现的表达式能直接利用Polars的内部执行机制，避开GIL开销，性能提升显著。

实现步骤 #

1. Rust端开发 #

首先创建Rust项目，配置依赖并编写自定义表达式逻辑：

Cargo.toml 配置

[package]
name = "polars_custom_expr"
version = "0.1.0"
edition = "2021"

[dependencies]
pyo3 = { version = "0.20", features = ["extension-module"] }
polars = { version = "0.32", features = ["full"] }
polars-core = "0.32"

lib.rs 代码示例

这里实现一个计算数值平方的自定义表达式：

use pyo3::prelude::*;
use polars::prelude::*;

// 定义自定义表达式函数：接收Expr，返回处理后的Expr
#[pyfunction]
fn custom_square(expr: Expr) -> Expr {
    expr.map(
        |series| {
            // 将系列转换为i64类型，计算每个值的平方
            series.i64()?
                .map_values(|val| val * val)
                .into_series()
        },
        // 指定输出数据类型
        GetOutput::from_type(DataType::Int64)
    )
}

// 暴露模块给Python
#[pymodule]
fn polars_custom_expr(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(custom_square, m)?)?;
    Ok(())
}

2. 编译为Python扩展 #

使用maturin工具编译并安装到Python环境：

1. 安装maturin：pip install maturin
2. 在Rust项目根目录执行：maturin develop

3. Python端调用 #

编译完成后，即可在Python中导入自定义模块，和Polars原生表达式链式使用：

import polars as pl
from polars_custom_expr import custom_square

# 测试数据
df = pl.DataFrame({"nums": [1, 2, 3, 4, 5]})

# 使用自定义表达式
result = df.select(
    pl.col("nums"),
    custom_square(pl.col("nums")).alias("nums_squared")
)

print(result)

输出结果：

shape: (5, 2)
┌──────┬─────────────┐
│ nums ┆ nums_squared│
│ ---  ┆ ---         │
│ i64  ┆ i64         │
╞══════╪═════════════╡
│ 1    ┆ 1           │
│ 2    ┆ 4           │
│ 3    ┆ 9           │
│ 4    ┆ 16          │
│ 5    ┆ 25          │
└──────┴─────────────┘

关键注意事项 #

• 确保Rust依赖的Polars版本与Python端安装的Polars版本一致，避免兼容性问题。
• 处理数据类型时要严谨，比如上述示例中默认输入为i64类型，若需支持多类型，可在Rust代码中添加类型判断逻辑。
• 自定义表达式支持Polars的懒执行模式，能和其他原生表达式无缝组合，充分利用Polars的查询优化能力。

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