Rust数据处理库pyo3-polars的使用：高效Python与Rust交互的Polars DataFrame集成方案

1. Polars的共享库插件

文档也可以参考Polars用户指南中的相关内容。这是新功能，应该优先于2.中使用的方法，因为它可以绕过GIL，并且是我们希望支持的扩展polars的方式。

并行性和优化由默认的polars运行时管理。该运行时将调用插件函数。插件函数是单独编译的。

我们可以因此保持polars更加精简，并可能添加对polars-distance、polars-geo、polars-ml等的支持。这些可以拥有专门的表达式，并且不必过多担心代码膨胀，因为它们可以选装安装。

思路是在另一个Rust crate中使用proc_macro polars_expr定义一个表达式。

宏可以具有以下属性之一：

• output_type -> 定义该表达式的输出类型
• output_type_func -> 定义一个基于输入类型计算输出类型的函数
• output_type_func_with_kwargs -> 定义一个基于输入类型和关键字参数计算输出类型的函数

以下是一个将任何字符串转换为猪拉丁语的String转换表达式示例：

// 将字符串转换为猪拉丁语的函数
fn pig_latin_str(value: &str, capitalize: bool, output: &mut String) {
    if let Some(first_char) = value.chars().next() {
        if capitalize {
            for c in value.chars().skip(1).map(|char| char.to_uppercase()) {
                write!(output, "{c}").unwrap()
            }
            write!(output, "AY").unwrap()
        } else {
            let offset = first_char.len_utf8();
            write!(output, "{}{}ay", &value[offset..], first_char).unwrap()
        }
    }
}

// 反序列化关键字参数
#[derive(Deserialize)]
struct PigLatinKwargs {
    capitalize: bool,
}

// 定义polars表达式
#[polars_expr(output_type=String)]
fn pig_latinnify(inputs: &[Series], kwargs: PigLatinKwargs) -> PolarsResult<Series> {
    let ca = inputs[0].str()?;
    let out: StringChunked =
        ca.apply_into_string_amortized(|value, output| pig_latin_str(value, kwargs.capitalize, output));
    Ok(out.into_series())
}

然后可以在Python端暴露：

from __future__ import annotations
from typing import TYPE_CHECKING
import polars as pl
from polars.plugins import register_plugin_function
from expression_lib._utils import LIB

if TYPE_CHECKING:
    from expression_lib._typing import IntoExprColumn

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

编译/发布后就可以使用了：

import polars as pl
from expression_lib import language

df = pl.DataFrame({
    "names": ["Richard", "Alice", "Bob"],
})

out = df.with_columns(
   pig_latin = language.pig_latinnify("names")
)

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

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

2. Polars的Pyo3扩展

参考example目录中的具体示例。这里我们将polars的DataFrame发送到rust，然后使用rayon和rust哈希集并行计算jaccard相似度。

运行示例

$ cd example && make install $ venv/bin/python run.py

这将输出：

shape: (2, 2)
┌───────────┬───────────────┐
│ list_a    ┆ list_b        │
│ ---       ┆ ---           │
│ list[i64] ┆ list[i64]     │
╞═══════════╪═══════════════╡
│ [1, 2, 3] ┆ [1, 2, ... 8] │
│ [5, 5]    ┆ [5, 1, 1]     │
└───────────┴───────────────┘
shape: (2, 1)
┌─────────┐
│ jaccard │
│ ---     │
│ f64     │
╞═════════╡
│ 0.75    │
│ 0.5     │
└─────────┘

为发布编译

$ make install-release

预期结果

这个crate提供了一个PySeries和一个PyDataFrame，它们是Series和DataFrame的简单包装器。这些包装器的优势在于它们可以实现FromPyObject和IntoPy，因此可以在Python和Rust之间转换。

完整示例代码

以下是一个完整的示例，展示了如何在Python中使用pyo3-polars：

import polars as pl
from pyo3_polars import PyDataFrame

# 创建一个Polars DataFrame
df = pl.DataFrame({
    "a": [1, 2, 3],
    "b": ["x", "y", "z"]
})

# 转换为PyDataFrame
py_df = PyDataFrame(df)

# 在Rust中处理数据
# 这里假设有一个Rust函数process_dataframe
processed_py_df = py_df.process_dataframe()

# 转换回Polars DataFrame
result_df = processed_py_df.to_polars()

print(result_df)

对应的Rust代码：

use pyo3_polars::{PyDataFrame, PySeries};
use pyo3::prelude::*;
use polars::prelude::*;

#[pyfunction]
fn process_dataframe(py_df极速处理DataFrame的完整示例

```rust
// 导入必要的库
use pyo3::prelude::*;
use polars::prelude::*;
use pyo3_polars::{PyDataFrame, PySeries};

// 定义高性能处理函数
#[pyfunction]
fn high_performance_process(
    py_df: PyDataFrame,
    multiplier: i64,
    prefix: &str
) -> PyResult<PyDataFrame> {
    // 转换为Polars DataFrame
    let mut df = py_df.into();
    
    // 处理数值列
    for col in df.get_column_names() {
        if let Ok(s) = df.column(col) {
            if s.dtype().is_numeric() {
                // 数值列乘以系数
                let new_s = s * multiplier;
                df.replace(col, new_s).unwrap();
            } else if s.dtype().is_utf8() {
                // 字符串列添加前缀
                let new_s = prefix.to_owned() + s.utf8().unwrap();
                df.replace(col, new_s).unwrap();
            }
        }
    }
    
    // 添加处理时间戳列
    if let Ok(ts_col) = df.column("timestamp") {
        if ts_col.dtype().is_temporal() {
            let day_col = ts_col.date().unwrap().day().into_series();
            df.with_column(day_col.alias("day")).unwrap();
        }
    }
    
    // 转换回PyDataFrame
    Ok(PyDataFrame::from(df))
}

// 注册Python模块
#[pymodule]
fn polars_extension(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(high_performance_process, m)?)?;
    Ok(())
}

Python端使用示例：

import polars as pl
from polars_extension import high_performance_process

# 创建包含多种数据类型的DataFrame
df = pl.DataFrame({
    "id": [1, 2, 3],
    "name": ["Alice", "Bob", "Charlie"],
    "value": [10.5, 20.3, 30.8],
    "timestamp": ["2023-01-01", "2023-01-15", "2023-02-01"]
})

# 转换为datetime类型
df = df.with_columns(pl.col("timestamp").str.to_datetime())

# 调用Rust处理函数
processed_df = high_performance_process(df, multiplier=100, prefix="user_")

print(processed_df)

示例输出：

shape: (3, 5)
┌─────┬─────────────┬────────┬─────────────────────┬─────┐
│ id  ┆ name        ┆ value  ┆ timestamp           ┆ day │
│ --- ┆ ---         ┆ ---    ┆ ---                 ┆ --- │
│ i64 ┆ str         ┆ f64    ┆ datetime[μs]        ┆ i32 │
╞═════╪═════════════╪════════╪═════════════════════╪═════╡
│ 100 ┆ user_Alice  ┆ 1050.0 ┆ 2023-01-01 00:00:00 ┆ 1   │
│ 200 ┆ user_Bob    ┆ 2030.0 ┆ 2023-01-15 00:00:00 ┆ 15  │
│ 300 ┆ user_Charlie┆ 3080.0 ┆ 2023-02-01 00:00:00 ┆ 1   │
└─────┴─────────────┴────────┴─────────────────────┴─────┘

这个完整示例展示了：

1. Rust端处理多种数据类型(数值、字符串、时间戳)
2. 高性能的列操作(乘法、字符串拼接、日期提取)
3. Python与Rust之间的无缝DataFrame转换
4. 自定义参数传递(multiplier和prefix)