Polars 如何编写并注册一个 Rust 表达式插件（pyo3-polars）？

Polars 如何编写并注册一个 Rust 表达式插件pyo3-polars【免费下载链接】polarsExtremely fast Query Engine for DataFrames, written in Rust项目地址: https://gitcode.com/GitHub_Trending/po/polars内置表达式覆盖不了的自定义计算比如按特定规则批量改写一列字符串Polars 文档推荐的做法是用 Rust 写一个函数编译成动态链接库再通过register_plugin_function在 Python 侧注册成表达式。Polars 引擎会在运行时动态链接你的函数整个过程不经过 Python因此没有 GIL 竞争并且表达式享受与原生表达式相同的优化和并行调度。文档以 Pig Latin 转换器为例把每个单词的首字母移到词尾并追加 aypig 按文档定义应变为 igpay。完整路径是新建 Rust 库并声明cdylib→ 用#[polars_expr]宏暴露表达式函数 → 建立同名 Python 包注册表达式 → 用maturin编译安装 → 在DataFrame.with_columns中调用验证。准备环境Python 环境与polars包Rust 工具链maturin会通过 cargo 编译 Rust 代码。示例项目 requirements.txt 中固定maturin1.12.4并安装polars可参考此版本约束。构建后端声明在pyproject.toml中示例如下来自 示例项目 pyproject.toml[build-system] requires [maturin1.0,2.0] build-backend maturin [project] name expression_lib version 0.0.1 requires-python 3.8第一步创建 Rust 库文档要求的文件结构是Python 包目录与 Cargo.toml 中[lib].name同名和src/平级├── expression_lib/ │ └── __init__.py ├── src/ │ ├── lib.rs │ └── expressions.rs ├── Cargo.toml └── pyproject.tomlCargo.toml来自用户指南[package] name expression_lib version 0.1.0 edition 2021 [lib] name expression_lib crate-type [cdylib] [dependencies] polars { version * } pyo3 { version *, features [extension-module, abi3-py310] } pyo3-polars { version *, features [derive] } serde { version *, features [derive] }要点crate-type [cdylib]保证编译产物是动态链接库pyo3-polars需要启用derivefeature 才能使用#[polars_expr]宏。入口src/lib.rs来自 示例项目 lib.rsuse pyo3_polars::PolarsAllocator; mod expressions; #[global_allocator] static ALLOC: PolarsAllocator PolarsAllocator::new();表达式本身写在src/expressions.rs。暴露一个函数的规则是加#[polars_expr(output_typeDataType)]属性且函数第一个参数必须是inputs: [Series]// src/expressions.rs use polars::prelude::*; use pyo3_polars::derive::polars_expr; use std::fmt::Write; fn pig_latin_str(value: str, output: mut String) { if let Some(first_char) value.chars().next() { write!(output, {}{}ay, value[1..], first_char).unwrap() } } #[polars_expr(output_typeString)] fn pig_latinnify(inputs: [Series]) - PolarsResultSeries { let ca inputs[0].str()?; let out: StringChunked ca.apply_into_string_amortized(pig_latin_str); Ok(out.into_series()) }这里使用apply_into_string_amortized而不是apply_values是为了避免每行都分配一个新字符串如果插件接收多个输入、逐元素计算且输出String文档建议改用polars::prelude::arity中的binary_elementwise_into_string_amortized。第二步Python 侧注册表达式创建与[lib].name同名的 Python 包本例为expression_lib在__init__.py中调用register_plugin_function# expression_lib/__init__.py from pathlib import Path from typing import TYPE_CHECKING import polars as pl from polars.plugins import register_plugin_function from polars._typing import IntoExpr PLUGIN_PATH Path(__file__).parent def pig_latinnify(expr: IntoExpr) - pl.Expr: Pig-latinnify expression. return register_plugin_function( plugin_pathPLUGIN_PATH, function_namepig_latinnify, argsexpr, is_elementwiseTrue, )两处关键约束function_name必须与 Rust 函数名完全一致否则主 Polars 包无法解析该函数名即注册后调用失败。is_elementwiseTrue声明该表达式是逐元素的Polars 允许它以批量方式运行像排序、切片这类非逐元素操作不能这样声明。仓库中更完整的注册写法可参考 language.py它把plugin_path抽到expression_lib/_utils.py的LIB常量里并把注册函数挂到language命名空间下。第三步编译并安装用户指南给出的主路径是安装maturin后执行maturin develop --release另一种走仓库自带示例工程的构建方式Makefile。先说明副作用make install会在pyo3-polars/example/derive_expression/目录下创建.venv虚拟环境并安装requirements.txt中的依赖然后编译安装插件库cd pyo3-polars/example/derive_expression make install # 内部执行 maturin develop -m expression_lib/Cargo.toml --uv # 需要 release 构建时 make install-release # 内部执行 maturin develop --release -m expression_lib/Cargo.toml --uv第四步验证表达式可用按用户指南的最小用法调用插件编译安装后即可直接使用import polars as pl from expression_lib import pig_latinnify df pl.DataFrame( { convert: [pig, latin, is, silly], } ) out df.with_columns(pig_latinpig_latinnify(convert))成功条件out新增pig_latin列每个单词按文档定义的 Pig Latin 规则转换例如 pig → igpay这是文档给出的规则示例不是固定日志。也可以直接运行示例脚本做完整验证cd pyo3-polars/example/derive_expression make runrun.py 会构建一个含字符串、列表、日期、坐标等多列的 DataFrame用language命名空间调用pig_latinnify、hamming_distance、haversine、is_leap_year等插件函数并打印结果最后打印finished。脚本后半部分还覆盖了两条错误路径可作为注册失败时的对照传入错误类型的 kwargs如float_argTrue会抛出pl.exceptions.ComputeError消息包含the plugin failed with message触发panic的插件会抛出ComputeError消息包含the plugin panicked。如果你希望以命名空间方式书写表达式文档说明可以注册自定义 namespaceregister_expr_namespace之后即可写pl.col(convert).language.pig_latinnify()language.py 给出了这种组织方式。可选扩展kwargs 与动态输出类型接收 kwargsRust 侧定义一个派生serde::Deserialize的 struct并把它作为插件函数的第二个参数#[derive(Deserialize)] pub struct MyKwargs { float_arg: f64, integer_arg: i64, string_arg: String, boolean_arg: bool, } #[polars_expr(output_typeString)] fn append_kwargs(input: [Series], kwargs: MyKwargs) - PolarsResultSeries { let input input[0]; let input input.cast(DataType::String)?; let ca input.str().unwrap(); Ok(ca .apply_into_string_amortized(|val, buf| { write!( buf, {}-{}-{}-{}-{}, val, kwargs.float_arg, kwargs.integer_arg, kwargs.string_arg, kwargs.boolean_arg ) .unwrap() }) .into_series()) }Python 侧在注册时通过kwargs字典传入参数名与 struct 字段一一对应def append_args(expr, float_arg, integer_arg, string_arg, boolean_arg) - pl.Expr: return register_plugin_function( plugin_pathPLUGIN_PATH, function_nameappend_kwargs, argsexpr, kwargs{ float_arg: float_arg, integer_arg: integer_arg, string_arg: string_arg, boolean_arg: boolean_arg, }, is_elementwiseTrue, )输出类型依赖输入类型#[polars_expr]宏还支持三个属性output_type固定输出类型、output_type_func按输入字段计算输出类型、output_type_func_with_kwargs按输入字段和 kwargs 计算输出类型。后两者指向一个把[Field]映射为输出Field的函数文档示例用FieldsMapper工具完成映射use polars_plan::dsl::FieldsMapper; fn haversine_output(input_fields: [Field]) - PolarsResultField { FieldsMapper::new(input_fields).map_to_float_dtype() } #[polars_expr(output_type_funchaversine_output)] fn haversine(inputs: [Series]) - PolarsResultSeries { // 完整实现见示例工程 expressions.rs todo!() }haversine 完整实现按输入是 Float32/Float64 分支计算以及output_type_func_with_kwargs的change_time_zone用法都在 expressions.rs 中。限制与注意事项用户指南 文件自身带有一条 TODO 注释文中步骤arent clear and may not work that well。如果遇到与文档不符的问题优先对照仓库内可运行的 derive_expression 示例工程 排查。插件函数是单独编译的 crate并行与优化由默认 Polars 运行时管理见 pyo3-polars README不要与 PyO3 扩展类型PySeries/PyDataFrame包装器混用后者只是便捷包装除最新定义适配最新 Polars 版本外没有稳定性保证。is_elementwise只适用于逐元素表达式聚合、排序、切片类逻辑不能声明为 elementwise。function_name与 Rust 函数名不一致时主包解析不到该函数调用会直接失败。【免费下载链接】polarsExtremely fast Query Engine for DataFrames, written in Rust项目地址: https://gitcode.com/GitHub_Trending/po/polars创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考