如何修改PyO3扩展的ffi.rs以从Rust返回Polars Dataframe到Python

需求背景 #

已通过PyO3+maturin实现Rust处理Polars Series并返回Series至Python，现需扩展为支持返回Polars Dataframe。

修改后的ffi.rs代码 #

use arrow::ffi;
use polars::prelude::*;
use pyo3::exceptions::PyValueError;
use pyo3::ffi::Py_uintptr_t;
use pyo3::prelude::*;
use pyo3::{PyAny, PyObject, PyResult};

/// 将Python中的Arrow数组转换为Rust端的Arrow数组（零拷贝）
fn array_to_rust(arrow_array: &PyAny) -> PyResult<ArrayRef> {
    let array = Box::new(ffi::ArrowArray::empty());
    let schema = Box::new(ffi::ArrowSchema::empty());

    let array_ptr = &*array as *const ffi::ArrowArray;
    let schema_ptr = &*schema as *const ffi::ArrowSchema;

    // 调用PyArrow私有API完成转换，操作内存需标记unsafe
    arrow_array.call_method1(
        "_export_to_c",
        (array_ptr as Py_uintptr_t, schema_ptr as Py_uintptr_t),
    )?;

    unsafe {
        let field = ffi::import_field_from_c(schema.as_ref()).unwrap();
        let array = ffi::import_array_from_c(*array, field.data_type).unwrap();
        Ok(array)
    }
}

/// 将Rust端的Arrow数组转换为Python端的Arrow数组（零拷贝）
pub(crate) fn to_py_array(py: Python, pyarrow: &PyModule, array: ArrayRef) -> PyResult<PyObject> {
    let schema = Box::new(ffi::export_field_to_c(&ArrowField::new(
        "",
        array.data_type().clone(),
        true,
    )));
    let array = Box::new(ffi::export_array_to_c(array));

    let schema_ptr: *const ffi::ArrowSchema = &*schema;
    let array_ptr: *const ffi::ArrowArray = &*array;

    let array = pyarrow.getattr("Array")?.call_method1(
        "_import_from_c",
        (array_ptr as Py_uintptr_t, schema_ptr as Py_uintptr_t),
    )?;

    Ok(array.to_object(py))
}

/// 将Python Polars Series转换为Rust Polars Series
pub fn py_series_to_rust_series(series: &PyAny) -> PyResult<Series> {
    // 先合并为单chunk，确保Arrow数组结构统一
    let series = series.call_method0("rechunk")?;

    let name = series.getattr("name")?.extract::<String>()?;
    let pyarrow_array = series.call_method0("to_arrow")?;
    let rust_array = array_to_rust(pyarrow_array)?;

    Series::try_from((name.as_str(), rust_array))
        .map_err(|e| PyValueError::new_err(format!("转换Series失败: {}", e)))
}

/// 将Rust Polars Series转换为Python Polars Series
pub fn rust_series_to_py_series(series: &Series) -> PyResult<PyObject> {
    let rechunked_series = series.rechunk();
    let arrow_array = rechunked_series.to_arrow(0);

    Python::with_gil(|py| {
        let pyarrow = py.import("pyarrow")?;
        let pyarrow_array = to_py_array(py, pyarrow, arrow_array)?;

        let polars = py.import("polars")?;
        let py_series = polars.call_method1("from_arrow", (pyarrow_array,))?;
        Ok(py_series.to_object(py))
    })
}

/// 将Python Polars Dataframe转换为Rust Polars Dataframe
pub fn py_dataframe_to_rust_dataframe(df: &PyAny) -> PyResult<DataFrame> {
    // 先合并所有列的chunk
    let df = df.call_method0("rechunk")?;
    let py_columns = df.call_method0("columns")?.extract::<Vec<String>>()?;

    let mut rust_series = Vec::with_capacity(py_columns.len());
    for col_name in py_columns {
        let py_series = df.call_method1("select", (col_name,))?.call_method0("to_series")?;
        rust_series.push(py_series_to_rust_series(&py_series)?);
    }

    DataFrame::new(rust_series)
        .map_err(|e| PyValueError::new_err(format!("转换Dataframe失败: {}", e)))
}

/// 将Rust Polars Dataframe转换为Python Polars Dataframe
pub fn rust_dataframe_to_py_dataframe(df: &DataFrame) -> PyResult<PyObject> {
    // 确保所有列都是单chunk
    let rechunked_df = df.rechunk();

    Python::with_gil(|py| {
        let pyarrow = py.import("pyarrow")?;
        let polars = py.import("polars")?;

        // 将每一列转换为Python端的Arrow数组，收集为列表
        let mut pyarrow_arrays = Vec::with_capacity(rechunked_df.width());
        for series in rechunked_df.get_columns() {
            let arrow_array = series.to_arrow(0);
            pyarrow_arrays.push(to_py_array(py, pyarrow, arrow_array)?);
        }

        // 创建PyArrow Table
        let pyarrow_table = pyarrow.getattr("Table")?.call1((pyarrow_arrays, rechunked_df.schema().names()))?;
        // 从PyArrow Table转换为Polars Dataframe
        let py_df = polars.call_method1("from_arrow", (pyarrow_table,))?;

        Ok(py_df.to_object(py))
    })
}

关键修改说明 #

• 新增py_dataframe_to_rust_dataframe函数：遍历Python Dataframe的所有列，逐个转换为Rust Series后，再构建成Rust Dataframe。
• 新增rust_dataframe_to_py_dataframe函数：
  1. 先将Rust Dataframe的每一列转为Arrow数组，再转换为Python端的Arrow数组；
  2. 用这些Arrow数组创建PyArrow Table；
  3. 最后通过Polars的from_arrow方法将Table转为Python端的Polars Dataframe。
• 所有转换基于Arrow的C接口实现零拷贝，保障性能。

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