插件

表达式插件是创建用户定义函数的首选方式。它们允许您编译一个Rust函数并将其注册为Polars库中的表达式。Polars引擎将在运行时动态链接您的函数，并且您的表达式将几乎与原生表达式一样快地运行。请注意，这不需要Python的任何干预，因此没有GIL争用。

他们将受益于默认表达式所具有的相同优势：

优化
并行性
Rust 原生性能

首先，我们将了解创建自定义表达式需要什么。

我们的第一个自定义表达式：Pig Latin

对于我们的第一个表达式，我们将创建一个猪拉丁语转换器。猪拉丁语是一种愚蠢的语言，其中每个单词的第一个字母被移除，添加到后面，最后加上“ay”。所以单词“pig”会转换为“igpay”。

我们当然已经可以通过表达式来实现这一点，例如 col("name").str.slice(1) + col("name").str.slice(0, 1) + "ay"，但一个专门用于此的函数会表现得更好，并且让我们能够了解插件。

设置

我们从一个新的库开始，如下所示的Cargo.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-py38"] }
pyo3-polars = { version = "*", features = ["derive"] }
serde = { version = "*", features = ["derive"] }

编写表达式

在这个库中，我们创建了一个辅助函数，将&str转换为猪拉丁语，并且我们创建了将作为表达式公开的函数。要公开一个函数，我们必须添加#[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输出，那么你可能需要查看polars::prelude::arity中的binary_elementwise_into_string_amortized实用函数。

这是Rust端所需的所有内容。在Python端，我们必须设置一个与Cargo.toml中定义的名称相同的文件夹，在这个例子中是"expression_lib"。我们将在与Rust src文件夹相同的目录下创建一个名为expression_lib的文件夹，并创建一个expression_lib/__init__.py文件。最终的文件结构应该如下所示：

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

然后我们创建一个新的类Language，它将保存我们新的expr.language命名空间的表达式。我们的表达式的函数名称可以被注册。请注意，这个名称必须正确，否则主Polars包无法解析函数名称。此外，我们可以设置额外的关键字参数，向Polars解释这个表达式的行为。在这种情况下，我们告诉Polars这个函数是逐元素的。这允许Polars以批处理方式运行这个表达式。而对于其他操作，这是不允许的，例如排序或切片。

# 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,
    )

然后我们可以在我们的环境中通过安装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"))

或者，你可以注册一个自定义命名空间，这使你可以编写：

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

接受kwargs

如果你想在polars表达式中接受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())
}

在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,
    )

输出数据类型

输出数据类型当然不必是固定的。它们通常取决于表达式的输入类型。为了适应这一点，你可以为#[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)
}

这就是你需要知道的入门知识。查看这个仓库来了解这一切是如何结合在一起的，并查看这个教程以获得更深入的理解。

社区插件

以下是社区实现的插件精选（非详尽）列表。

polars-xdt Polars 插件，提供额外的日期时间相关功能，这些功能不在主库的范围内
polars-distance 用于成对距离函数的Polars插件
polars-ds Polars 扩展旨在简化常见的数值/字符串数据分析过程
polars-hash 用于Polars的稳定非加密和加密哈希函数
polars-reverse-geocode 离线反向地理编码器，用于找到最接近给定（纬度，经度）对的城市

其他材料

Ritchie Vink - 关于Polars插件的主题演讲
Polars plugins tutorial 通过学习一些非常简单和最小的示例来了解如何编写插件
cookiecutter-polars-plugin Polars 插件项目模板