polars_python/series/
c_interface.rs

1#![allow(unsafe_op_in_unsafe_fn)]
2use polars::prelude::*;
3use pyo3::ffi::Py_uintptr_t;
4use pyo3::prelude::*;
5
6use super::PySeries;
7use crate::error::PyPolarsErr;
8
9// Import arrow data directly without requiring pyarrow (used in pyo3-polars)
10#[pymethods]
11impl PySeries {
12    #[staticmethod]
13    unsafe fn _import_arrow_from_c(
14        name: &str,
15        chunks: Vec<(Py_uintptr_t, Py_uintptr_t)>,
16    ) -> PyResult<Self> {
17        let chunks = chunks
18            .into_iter()
19            .map(|(schema_ptr, array_ptr)| {
20                let schema_ptr = schema_ptr as *mut arrow::ffi::ArrowSchema;
21                let array_ptr = array_ptr as *mut arrow::ffi::ArrowArray;
22
23                // Don't take the box from raw as the other process must deallocate that memory.
24                let array = std::ptr::read_unaligned(array_ptr);
25                let schema = &*schema_ptr;
26
27                let field = arrow::ffi::import_field_from_c(schema).unwrap();
28                arrow::ffi::import_array_from_c(array, field.dtype).unwrap()
29            })
30            .collect::<Vec<_>>();
31
32        let s = Series::try_new(name.into(), chunks).map_err(PyPolarsErr::from)?;
33        Ok(s.into())
34    }
35
36    unsafe fn _export_arrow_to_c(
37        &self,
38        out_ptr: Py_uintptr_t,
39        out_schema_ptr: Py_uintptr_t,
40    ) -> PyResult<()> {
41        export_chunk(&self.series, out_ptr, out_schema_ptr).map_err(PyPolarsErr::from)?;
42        Ok(())
43    }
44}
45
46unsafe fn export_chunk(
47    s: &Series,
48    out_ptr: Py_uintptr_t,
49    out_schema_ptr: Py_uintptr_t,
50) -> PolarsResult<()> {
51    polars_ensure!(s.chunks().len() == 1, InvalidOperation: "expect a single chunk");
52
53    let c_array = arrow::ffi::export_array_to_c(s.chunks()[0].clone());
54    let out_ptr = out_ptr as *mut arrow::ffi::ArrowArray;
55    *out_ptr = c_array;
56
57    let field = ArrowField::new(
58        s.name().clone(),
59        s.dtype().to_arrow(CompatLevel::newest()),
60        true,
61    );
62    let c_schema = arrow::ffi::export_field_to_c(&field);
63
64    let out_schema_ptr = out_schema_ptr as *mut arrow::ffi::ArrowSchema;
65    *out_schema_ptr = c_schema;
66    Ok(())
67}