在Maturin项目中实现Polars DataFrame到Rust结构体向量的转换

一、基础场景：结构体字段为字符串类型 #

1. Rust依赖配置 #

首先在Cargo.toml中添加必要依赖，确保支持Polars和Python交互：

[dependencies]
pyo3 = { version = "0.21", features = ["extension-module"] }
polars = { version = "0.38", features = ["strings"] }

2. Rust侧转换实现 #

通过PyO3接收Python传入的Polars DataFrame，转换为Rust的Polars DataFrame后逐行映射为MyStruct：

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

#[derive(Debug)]
struct MyStruct {
    a: i64,
    b: String,
}

#[pyfunction]
fn df_to_my_structs(py: Python, df: &PyAny) -> PyResult<Vec<MyStruct>> {
    // 将Python侧Polars DataFrame转为Rust侧实例
    let rust_df: DataFrame = df.extract(py)?;

    // 获取指定列并转换为对应类型的列视图
    let col_a = rust_df.column("a")?.i64()?;
    let col_b = rust_df.column("b")?.str()?;

    // 预分配内存提升性能
    let mut result = Vec::with_capacity(rust_df.height());
    for (a_val, b_val) in col_a.iter().zip(col_b.iter()) {
        match (a_val, b_val) {
            (Some(a), Some(b)) => {
                result.push(MyStruct {
                    a,
                    b: b.to_string(),
                });
            }
            _ => return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
                "DataFrame包含空值，无法转换为MyStruct",
            )),
        }
    }

    Ok(result)
}

#[pymodule]
fn polars_rust_converter(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(df_to_my_structs, m)?)?;
    Ok(())
}

3. Python侧调用示例 #

编译Maturin模块后，直接调用Rust暴露的函数：

import polars as pl
import polars_rust_converter

df = pl.DataFrame({'a': [1, 2], 'b': ['foo', 'bar']})
structs = polars_rust_converter.df_to_my_structs(df)
print(structs)

二、进阶场景：结构体字段为枚举类型 #

1. Rust侧枚举与转换逻辑 #

定义枚举MyEnum及字符串转换函数，修改结构体后调整转换逻辑：

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

#[derive(Debug)]
enum MyEnum {
    Foo,
    Bar,
}

#[derive(Debug)]
struct MyStruct {
    a: i64,
    b: MyEnum,
}

// 字符串转枚举的工具函数
fn string_to_myenum(s: &str) -> Result<MyEnum, String> {
    match s.to_lowercase().as_str() {
        "foo" => Ok(MyEnum::Foo),
        "bar" => Ok(MyEnum::Bar),
        other => Err(format!("无法识别的枚举值: {}", other)),
    }
}

#[pyfunction]
fn df_to_my_structs_with_enum(py: Python, df: &PyAny) -> PyResult<Vec<MyStruct>> {
    let rust_df: DataFrame = df.extract(py)?;

    let col_a = rust_df.column("a")?.i64()?;
    let col_b = rust_df.column("b")?.str()?;

    let mut result = Vec::with_capacity(rust_df.height());
    for (a_val, b_val) in col_a.iter().zip(col_b.iter()) {
        // 处理空值场景，转为Python可捕获的错误
        let a = a_val.ok_or_else(|| PyErr::new::<pyo3::exceptions::PyValueError, _>("列a包含空值"))?;
        let b_str = b_val.ok_or_else(|| PyErr::new::<pyo3::exceptions::PyValueError, _>("列b包含空值"))?;
        
        // 转换字符串为枚举，将Rust错误转为Python错误
        let b_enum = string_to_myenum(b_str)
            .map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e))?;

        result.push(MyStruct { a, b: b_enum });
    }

    Ok(result)
}

#[pymodule]
fn polars_rust_converter(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(df_to_my_structs_with_enum, m)?)?;
    Ok(())
}

2. Python侧调用示例 #

import polars as pl
import polars_rust_converter

df = pl.DataFrame({'a': [1, 2], 'b': ['foo', 'bar']})
structs = polars_rust_converter.df_to_my_structs_with_enum(df)
print(structs)

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