Polars 插件（Plugins）扩展机制全指南：表达式插件与 IO 插件的原理与实战

Polars 插件（Plugins）扩展机制全指南：表达式插件与 IO 插件的原理与实战

【免费下载链接】polarsExtremely fast Query Engine for DataFrames, written in Rust项目地址: https://gitcode.com/GitHub_Trending/po/polars

Polars 内置了大量表达式与数据源能力，但真实业务中总有引擎覆盖不到的自定义算法与私有文件格式。为此，Polars 在官方用户手册中专门开辟了plugins一节，系统性地给出「表达式插件（Expression plugins）」与「IO 插件（IO plugins）」两条官方扩展路线，并整理了社区插件与配套学习资料。本文以 plugins/index.md 为主线，结合仓库内pyo3-polars派生宏、Python 侧注册入口与可运行示例的源码实现，完整讲解如何在 Polars 中编写、注册并分发自己的插件，读完即可照例复刻出一个编译期链接、性能接近原生表达式的自定义函数，或一个接入懒执行/流式引擎的私有数据源。

插件体系的整体蓝图：一份「入口索引」读懂两条扩展路线

Polars 官方手册的 plugins 一节本质是一张扩展能力地图，它把插件分为互补的两类：

• 表达式插件（Expression plugins）：把一段编译好的 Rust 函数注册成 Polars 表达式，适合扩展计算内核；
• IO 插件（IO plugins）：把自定义的读取源注册进查询引擎，适合接入 Polars 尚不支持的输入格式。

两者详细指南分别见 Expression plugins 与 IO plugins，它们共同构成plugins/index.md这一入口页的全部主线内容。理解这个「入口页 → 两篇子指南」的层次，也就理解了插件系统的设计分界：凡是「一行输入产生一行输出（或一组聚合）的计算」，走表达式插件；凡是「从某种介质把数据灌进引擎」的读取，走 IO 插件。

社区插件生态与配套学习资料（原索引页内容）

入口页还维护了一份精选的、非穷举（curated, non-exhaustive）的社区插件清单，可作为扩展方向的参考（以下插件均为仓库外项目，具体可用性以各项目为准）：

• 综合类：polars-xdt（补充主库暂未纳入的日期时间类功能）、polars-hash（稳定的非加密与加密哈希函数）；
• 数据科学类：polars-distance（成对距离函数）、polars-ds（面向常规数值/字符串分析流程的简化扩展）；
• 生物信息类：polars-bio（构建在 Polars、Arrow、DataFusion 之上的基因组学 Python 库）；
• 地理空间类：polars-st（在 DataFrame/Series/Expr 上提供类似 Shapely、Geopandas 的空间运算）、polars-reverse-geocode（离线逆地理编码）、polars-h3（接入 H3 离散全球网格系统，把点与几何直接索引为六边形）。

入口页同时给出了三类补充材料：官方作者的插件主题演讲（Ritchie Vink Keynote on Polars Plugins）、polars-plugins-tutorial极简教程，以及cookiecutter-polars-plugin项目脚手架模板。若想在本仓库内直接看到「完整可编译」的对应样例，推荐优先研读两个官方示例目录：pyo3-polars/example/derive_expression（表达式插件完整工程）与 pyo3-polars/example/io_plugin（Rust 实现的 IO 源），后文所有代码讲解都能在这两处找到一一对应的落地实现。

表达式插件：接近原生性能的自定义函数

表达式插件是 Polars 实现用户自定义函数（UDF）的首选方式。它把一个编译好的 Rust 函数注册为 Polars 表达式，查询引擎在运行时动态链接该函数，执行速度几乎与原生表达式相当；整个过程完全不经过 Python、不产生 GIL 竞争。作为一等公民，自定义表达式自动继承原生表达式的三类收益：

• 查询优化（Optimization）：可被谓词下推、投影下推、常量折叠等优化规则感知与重排；
• 并行执行（Parallelism）：由引擎调度到多线程；
• Rust 原生性能（Rust native performance）：无解释器开销。

文档以「Pig Latin 转换」作为第一个教学案例——把每个单词的首字母挪到词尾并追加ay，如pig→igpay。该功能当然能用纯表达式拼接col("name").str.slice(1) + col("name").str.slice(0, 1) + "ay"实现，但专函数化既更高效，也是学习插件机制的最小切口。

工程搭建：Cargo.toml与编译目标

先创建一个新的 Rust library，其Cargo.toml关键点在于：crate-type = ["cdylib"]（最终产物是供动态链接的共享库），并依赖polars、pyo3（启用extension-module与abi3-py310）以及带derivefeature 的pyo3-polars（提供#[polars_expr]过程宏），示例如下：

[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"] }

仓库内的真实样例 pyo3-polars/example/derive_expression/expression_lib/Cargo.toml 结构与此一致，并且其lib.rs还会安装一个全局内存分配器：

use pyo3_polars::PolarsAllocator; mod distances; mod expressions; #[global_allocator] static ALLOC: PolarsAllocator = PolarsAllocator::new();

如 pyo3-polars/example/derive_expression/expression_lib/src/lib.rs 所示，安装PolarsAllocator是为保证插件内 Rust 分配的缓冲区与 Polars 的分配机制保持一致，避免跨库free引发内存错误——编写真实插件时这一步值得保留。

Rust 侧：#[polars_expr]宏与 Series 处理

定义插件函数的核心约束有两条：必须给函数标注#[polars_expr(output_type=DataType)]；函数签名必须把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_type=String)] fn pig_latinnify(inputs: &[Series]) -> PolarsResult<Series> { 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。若插件接受多个输入、按元素操作并产出String，文档建议关注polars::prelude::arity下提供的binary_elementwise_into_string_amortized等工具函数，它们按同样的零额外分配思路处理二元输入。

#[polars_expr]宏在编译期替你生成了跨 FFI 的接线代码。从派生宏实现 pyo3-polars/pyo3-polars-derive/src/lib.rs 可以看清其底层契约：宏会把你的函数调用包装成fn(inputs, [kwargs] [, context]) -> PolarsResult<Series>，成功时通过polars_ffi::version_0::export_series(&out)把Series导出给 Polars，失败时则调用_update_last_error记录错误并返回空值——这也就是文档所说的「引擎在运行时动态链接」的真实载体。

Python 侧：包结构、注册与调用

Rust 侧到此即全部完成。Python 侧需要一个与Cargo.toml中[lib] name完全同名的文件夹（本例为expression_lib），与 Rust 的src平级放置，里面至少有一个__init__.py。最终目录结构如下：

├── 📁 expression_lib/ # name must match "lib.name" in Cargo.toml │ └── __init__.py │ ├── 📁src/ │ ├── lib.rs │ └── expressions.rs │ ├── Cargo.toml └── pyproject.toml

注册函数时，function_name必须与 Rust 侧函数名完全一致，否则主 Polars 包无法解析；同时可通过关键字参数向引擎声明函数行为。下面声明is_elementwise=True，表示函数是逐元素操作——只有这样才能允许 Polars 按批次执行；像sort、slice这类会改变数据位置或长度的操作则绝不能如此标记。

# 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_path=PLUGIN_PATH, function_name="pig_latinnify", args=expr, is_elementwise=True, )

注意这里PLUGIN_PATH = Path(__file__).parent指向插件包目录本身，Python 侧注册入口在 py-polars/src/polars/plugins.py 中定义。仓库里真实示例的注册方式如 expression_lib/language.py 所示——把「指向.so的位置」与「函数名」显式分离：

def pig_latinnify(expr: IntoExprColumn, capitalize: bool = False) -> pl.Expr: return register_plugin_function( plugin_path=LIB, args=[expr], function_name="pig_latinnify", is_elementwise=True, kwargs={"capitalize": capitalize}, )

register_plugin_function的行为声明参数

Python 注册函数的行为声明直接决定 Polars 引擎如何处理你的函数，官方实现（py-polars/src/polars/plugins.py）支持以下参数，它们的含义即为引擎优化与调度所依赖的「事实」：

源码 docstring 同时给出严厉警告：注册参数若与真实行为不符，会导致错误结果，因此这些开关必须与 Rust 实现语义严格对应。

编译与首次运行

在当前环境安装maturin后执行maturin develop --release即可把共享库构建并安装进当前虚拟环境。随后插件立即可用：

import polars as pl from expression_lib import pig_latinnify df = pl.DataFrame( { "convert": ["pig", "latin", "is", "silly"], } ) out = df.with_columns(pig_latin=pig_latinnify("convert"))

如果你希望把它收进一个更像原生语法的命名空间，也可以注册自定义 namespace（register_expr_namespace），让用户写出：

out = df.with_columns( pig_latin=pl.col("convert").language.pig_latinnify(), )

即注册一个Expr.language命名空间，把多个相关插件函数统一收纳。

进阶一：通过 kwargs 传递自定义参数

若插件需要接收普通标量参数，只需在 Rust 侧定义一个struct并派生serde::Deserialize：

/// Provide your own kwargs struct with the proper schema and accept that type /// in your plugin expression. #[derive(Deserialize)] pub struct MyKwargs { float_arg: f64, integer_arg: i64, string_arg: String, boolean_arg: bool, } /// If you want to accept `kwargs`. You define a `kwargs` argument /// on the second position in you plugin. You can provide any custom struct that is deserializable /// with the pickle protocol (on the Rust side). #[polars_expr(output_type=String)] fn append_kwargs(input: &[Series], kwargs: MyKwargs) -> PolarsResult<Series> { 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()) }

关键点有二：kwargs必须位于函数签名第二个位置；struct 的字段类型会决定 Python 端必须传什么。Python 侧在注册时把同名参数放入kwargs字典即可：

def append_args( expr: IntoExpr, float_arg: float, integer_arg: int, string_arg: str, boolean_arg: bool, ) -> pl.Expr: """ This example shows how arguments other than `Series` can be used. """ return register_plugin_function( plugin_path=PLUGIN_PATH, function_name="append_kwargs", args=expr, kwargs={ "float_arg": float_arg, "integer_arg": integer_arg, "string_arg": string_arg, "boolean_arg": boolean_arg, }, is_elementwise=True, )

仓库示例 expression_lib/language.py 中的append_args与 Rust 侧append_kwargs一一对应，是一份可编译对照的现成参考。派生宏内部对 kwargs 的处理也可见于 pyo3-polars-derive/src/lib.rs：它通过_parse_kwargs完成反序列化，失败时返回带提示语的InvalidOperation错误。

进阶二：根据输入类型动态决定输出类型

输出类型并不总是固定的，常常取决于输入表达式类型。此时可给#[polars_expr()]提供output_type_func参数，指向一个把输入字段&[Field]映射为输出Field（含列名与数据类型）的函数。典型做法是借助工具结构FieldsMapper：

use polars_plan::dsl::FieldsMapper; fn haversine_output(input_fields: &[Field]) -> PolarsResult<Field> { FieldsMapper::new(input_fields).map_to_float_dtype() } #[polars_expr(output_type_func=haversine_output)] fn haversine(inputs: &[Series]) -> PolarsResult<Series> { let out = match inputs[0].dtype() { DataType::Float32 => { let start_lat = inputs[0].f32().unwrap(); let start_long = inputs[1].f32().unwrap(); let end_lat = inputs[2].f32().unwrap(); let end_long = inputs[3].f32().unwrap(); crate::distances::naive_haversine(start_lat, start_long, end_lat, end_long)? .into_series() } DataType::Float64 => { let start_lat = inputs[0].f64().unwrap(); let start_long = inputs[1].f64().unwrap(); let end_lat = inputs[2].f64().unwrap(); let end_long = inputs[3].f64().unwrap(); crate::distances::naive_haversine(start_lat, start_long, end_lat, end_long)? .into_series() } _ => polars_bail!(InvalidOperation: "only supported for float types"), }; Ok(out) }

上例让输出类型跟随输入在Float32/Float64间变化，遇到其他类型则以polars_bail!报错。其底层实现（haversine 距离计算等）在仓库示例 pyo3-polars/example/derive_expression/expression_lib/src/distances.rs 与expressions.rs中均有完整实现，读者可对照output_type_func与函数体的「类型分派」写法理解这套双轨机制。

IO 插件：为懒执行引擎接入自定义数据源

除表达式插件外，Polars 还支持IO 插件——把自定义文件格式注册为查询引擎的数据源（source）。IO 源可以通过 Arrow FFI 以零拷贝方式把大块数据交给引擎，因此官方选择现阶段用 Python 作为 IO 插件的接口层：引擎认为 GIL 仅在短暂的数据交接点被持有，不会成为瓶颈。例如一个 IO 源可以在 Rust 中解析自己的 DataFrame，仅在 rendez-vous 时刻零拷贝移交数据、短时占用 GIL。

适用场景：当你有 Polars 原生不支持的源文件，却想继续享受投影下推、谓词下推、提前终止（early stopping）与流式引擎（streaming engine）等优化时，就该考虑 IO 插件。

教学示例：一个刻意写得很慢的自定义 CSV 源

文档强调以下示例仅为教学目的，性能刻意做差。首先准备导入，其中关键的是register_io_source（每次实例化时注册一个新的生成器）：

# Use python for csv parsing. import csv import polars as pl # Used to register a new generator on every instantiation. from polars.io.plugins import register_io_source from typing import Iterator import io

第一步：解析 schema

Polars 中每个scan函数都必须能提供其读取数据的 schema。这个简化版 CSV 解析器把所有列一律读成pl.String，差异只在字段名与字段个数：

def parse_schema(csv_str: str) -> pl.Schema: first_line = csv_str.split("\n")[0] return pl.Schema({k: pl.String for k in first_line.split(",")})

对小型 CSV"a,b,c\n1,2,3"运行将得到Schema([('a', String), ('b', String), ('c', String)])。

第二步：写数据源（外层包装 + 内层生成器）

IO 源采用「外函数 + 内函数」双层结构。外层my_scan_csv是面向用户的入口，接收文件名等自定义参数（CSV 场景下如delimiter、quote_char）；内层才是数据源实现，其签名是预定义、必须接受的，共四个参数：

• with_columns：被投影的列。若应用了投影下推，读取端应尽量只产出这些列；
• predicate：Polars 表达式。读取端应据此过滤行；
• n_rows：只需物化前 n 行，读取端读完即可停止；
• batch_size：期望的批大小提示，读取端生成器应尽量按此产块。

def my_scan_csv(csv_str: str) -> pl.LazyFrame: schema = parse_schema(csv_str) def source_generator( with_columns: list[str] | None, predicate: pl.Expr | None, n_rows: int | None, batch_size: int | None, ) -> Iterator[pl.DataFrame]: """ Generator function that creates the source. This function will be registered as IO source. """ if batch_size is None: batch_size = 100 # Initialize the reader. reader = csv.reader(io.StringIO(csv_str), delimiter=',') # Skip the header. _ = next(reader) # Ensure we don't read more rows than requested from the engine while n_rows is None or n_rows > 0: if n_rows is not None: batch_size = min(batch_size, n_rows) rows = [] for _ in range(batch_size): try: row = next(reader) except StopIteration: n_rows = 0 break rows.append(row) df = pl.from_records(rows, schema=schema, orient="row") n_rows -= df.height # If we would make a performant reader, we would not read these # columns at all. if with_columns is not None: df = df.select(with_columns) # If the source supports predicate pushdown, the expression can be parsed # to skip rows/groups. if predicate is not None: df = df.filter(predicate) yield df return register_io_source(io_source=source_generator, schema=schema)

最外层调用register_io_source，返回一个真正的pl.LazyFrame——于是它自动被并入懒执行计划，享受下推与流式执行。

第三步：跑一个（很慢的）实测

用两组字符串模拟「文件内容」，验证完整 collect 与「头部截断 + 谓词下推」两种路径：

csv_str1 = """a,b,c,d 1,2,3,4 9,10,11,2 1,2,3,4 1,122,3,4""" print(my_scan_csv(csv_str1).collect()) csv_str2 = """a,b 1,2 9,10 1,2 1,122""" print(my_scan_csv(csv_str2).head(2).collect())

运行输出分别如下：第一个完整打印 4 行 4 列；第二个受.head(2)驱动，内部以n_rows截断只产出了前 2 行（这正是文档所述「reader can stop when n_rows are read」的体现）：

shape: (4, 4) ┌─────┬─────┬─────┬─────┐ │ a ┆ b ┆ c ┆ d │ │ --- ┆ --- ┆ --- ┆ --- │ │ str ┆ str ┆ str ┆ str │ ╞═════╪═════╪═════╪═════╡ │ 1 ┆ 2 ┆ 3 ┆ 4 │ │ 9 ┆ 10 ┆ 11 ┆ 2 │ │ 1 ┆ 2 ┆ 3 ┆ 4 │ │ 1 ┆ 122 ┆ 3 ┆ 4 │ └─────┴─────┴─────┴─────┘ shape: (2, 2) ┌─────┬─────┐ │ a ┆ b │ │ --- ┆ --- │ │ str ┆ str │ ╞═════╪═════╡ │ 1 ┆ 2 │ │ 9 ┌10 ┐ ... └─────┴─────┘

（注：第二个 DataFrame 仅两行两列，示意投影与行数限制均已生效。）

源码级细节：register_io_source的完整签名与引擎约定

Python 侧实现位于 py-polars/src/polars/io/plugins.py，从它的签名与 docstring 可以提炼出官方对 IO 插件作者的约定：

• 接口本身被标记为@unstable()：即 API 可能在不视为破坏性变更的前提下随时调整，生产接入需留意版本锁定；
• io_source：一个签名为(with_columns, predicate, n_rows, batch_size) -> Iterator[DataFrame]的生成器函数；
• schema：接受一个SchemaDict或返回 schema 的可调用对象，描述投影下推之前的完整输出 schema；
• validate_schema：是否让引擎校验每个批次与给定 schema 匹配——不匹配属于实现错误，校验能避免难以排查的隐性 bug；
• is_pure：若源是纯的（同参数同结果），优化器可对计划中重复出现的同一 IO 源做去重（de-duplication）；
• explain_name/explain_detail：自定义explain查询计划中该 scan 节点的标签与细节说明（渲染形如PYTHON[<name>] SCAN ...）。

引擎传入的predicate会被序列化为字节并尝试用pl.Expr.deserialize还原；若还原失败，插件会收到None，由 Polars 引擎自行在解析后过滤（见 io/plugins.py 中对parsed_predicate_success的处理）。

仓库佐证：一个真正在 Rust 中实现的 IO 源

若想看到「底层用 Rust 实现 IO 源」的完整形态，仓库提供了官方示例 pyo3-polars/example/io_plugin/io_plugin/src/lib.rs。它定义了一个RandomSource（一个#[pyclass]），其接口恰好映射引擎的四项约定：

• schema()返回PySchema，让引擎拿到完整输出 schema；
• try_set_predicate(predicate: PyExpr)与set_with_columns(columns: Vec<String>)分别承接谓词下推与投影下推；
• next()每次按size_hint/n_rows产出一个PyDataFrame，内部用注释明确区分「投影下推/切片下推应在采样前做」「谓词下推可以在更低层完成、也可事后过滤」——与 Python 教学版的两阶段处理一一对应。

对应 Python 侧封装在 pyo3-polars/example/io_plugin/io_plugin/io_plugin/init.py：scan_random先创建采样器，再在生成器中把with_columns、predicate透传给 Rust 对象；若 Rust 端因无法反序列化谓词而返回失败标记（predicate_set = False），则退回在 Python 端out.filter(predicate)。这套「先尝试下推、失败则本地过滤」的兜底逻辑，是编写高性能 IO 源时值得复用的模式。

如何选型：表达式插件还是 IO 插件

综合两条路线，可归纳出清晰的选型建议：

两条路线的权威说明分别见 expr_plugins.md 与 io_plugins.md；而 plugins/index.md 作为总入口，负责把读者导流到对应方向并呈现社区生态。需要特别提醒的事实边界是：入口页所列举的社区插件均为仓库外独立项目，具体能力与维护状态以其各自仓库为准，本仓库内可直接验证的仅是官方示例工程与 Python/Rust 侧的注册、宏展开与 FFI 实现。

对想动手写第一个插件的开发者，最经济的路径是：先跑通derive_expression示例的make/run.py，再对照本文把 Pig Latin 与 kwargs 两个案例亲手改一遍；当需要处理私有格式时，再基于io_plugin示例把RandomSource替换为真实解析器即可。

【免费下载链接】polarsExtremely fast Query Engine for DataFrames, written in Rust项目地址: https://gitcode.com/GitHub_Trending/po/polars