polars_python/interop/arrow/
to_py.rs

1use arrow::datatypes::ArrowDataType;
2use arrow::ffi;
3use arrow::record_batch::RecordBatch;
4use polars::datatypes::CompatLevel;
5use polars::frame::DataFrame;
6use polars::prelude::{ArrayRef, ArrowField, PlSmallStr, SchemaExt};
7use polars::series::Series;
8use polars_core::utils::arrow;
9use polars_error::PolarsResult;
10use pyo3::ffi::Py_uintptr_t;
11use pyo3::prelude::*;
12use pyo3::types::PyCapsule;
13
14/// Arrow array to Python.
15pub(crate) fn to_py_array(
16    array: ArrayRef,
17    field: &ArrowField,
18    pyarrow: &Bound<PyModule>,
19) -> PyResult<Py<PyAny>> {
20    let schema = Box::new(ffi::export_field_to_c(field));
21    let array = Box::new(ffi::export_array_to_c(array));
22
23    let schema_ptr: *const ffi::ArrowSchema = &*schema;
24    let array_ptr: *const ffi::ArrowArray = &*array;
25
26    let array = pyarrow.getattr("Array")?.call_method1(
27        "_import_from_c",
28        (array_ptr as Py_uintptr_t, schema_ptr as Py_uintptr_t),
29    )?;
30
31    Ok(array.unbind())
32}
33
34/// RecordBatch to Python.
35pub(crate) fn to_py_rb(
36    rb: &RecordBatch,
37    py: Python<'_>,
38    pyarrow: &Bound<PyModule>,
39) -> PyResult<Py<PyAny>> {
40    let mut arrays = Vec::with_capacity(rb.width());
41
42    for (array, field) in rb.columns().iter().zip(rb.schema().iter_values()) {
43        let array_object = to_py_array(array.clone(), field, pyarrow)?;
44        arrays.push(array_object);
45    }
46
47    let schema = Box::new(ffi::export_field_to_c(&ArrowField {
48        name: PlSmallStr::EMPTY,
49        dtype: ArrowDataType::Struct(rb.schema().iter_values().cloned().collect()),
50        is_nullable: false,
51        metadata: None,
52    }));
53    let schema_ptr: *const ffi::ArrowSchema = &*schema;
54
55    let schema = pyarrow
56        .getattr("Schema")?
57        .call_method1("_import_from_c", (schema_ptr as Py_uintptr_t,))?;
58    let record = pyarrow
59        .getattr("RecordBatch")?
60        .call_method1("from_arrays", (arrays, py.None(), schema))?;
61
62    Ok(record.unbind())
63}
64
65/// Export a series to a C stream via a PyCapsule according to the Arrow PyCapsule Interface
66/// https://arrow.apache.org/docs/dev/format/CDataInterface/PyCapsuleInterface.html
67pub(crate) fn series_to_stream<'py>(
68    series: &Series,
69    py: Python<'py>,
70) -> PyResult<Bound<'py, PyCapsule>> {
71    let field = series.field().to_arrow(CompatLevel::newest());
72    let series = series.clone();
73    let iter = Box::new(
74        (0..series.n_chunks()).map(move |i| Ok(series.to_arrow(i, CompatLevel::newest()))),
75    ) as _;
76
77    let stream = ffi::export_iterator(iter, field);
78    PyCapsule::new_with_value(py, stream, c"arrow_array_stream")
79}
80
81pub(crate) fn dataframe_to_stream<'py>(
82    df: &DataFrame,
83    py: Python<'py>,
84) -> PyResult<Bound<'py, PyCapsule>> {
85    let iter = Box::new(DataFrameStreamIterator::new(df));
86    let field = iter.field();
87    let stream = ffi::export_iterator(iter, field);
88    PyCapsule::new_with_value(py, stream, c"arrow_array_stream")
89}
90
91#[cfg(feature = "c_api")]
92#[pyfunction]
93pub(crate) fn polars_schema_to_pycapsule<'py>(
94    py: Python<'py>,
95    schema: crate::prelude::Wrap<polars::prelude::Schema>,
96    compat_level: crate::prelude::PyCompatLevel,
97) -> PyResult<Bound<'py, PyCapsule>> {
98    let schema: arrow::ffi::ArrowSchema = arrow::ffi::export_field_to_c(&ArrowField::new(
99        PlSmallStr::EMPTY,
100        ArrowDataType::Struct(
101            schema
102                .0
103                .iter_fields()
104                .map(|x| x.to_arrow(compat_level.0))
105                .collect(),
106        ),
107        false,
108    ));
109
110    PyCapsule::new_with_value(py, schema, c"arrow_schema")
111}
112
113pub struct DataFrameStreamIterator {
114    columns: Vec<Series>,
115    dtype: ArrowDataType,
116    idx: usize,
117    n_chunks: usize,
118    height: usize,
119}
120
121impl DataFrameStreamIterator {
122    fn new(df: &DataFrame) -> Self {
123        let schema = df.schema().to_arrow(CompatLevel::newest());
124        let dtype = ArrowDataType::Struct(schema.into_iter_values().collect());
125        let n_chunks = if df.width() == 0 {
126            usize::from(df.height() > 0)
127        } else {
128            df.first_col_n_chunks()
129        };
130
131        Self {
132            columns: df
133                .columns()
134                .iter()
135                .map(|v| v.as_materialized_series().clone())
136                .collect(),
137            dtype,
138            idx: 0,
139            n_chunks,
140            height: df.height(),
141        }
142    }
143
144    fn field(&self) -> ArrowField {
145        ArrowField::new(PlSmallStr::EMPTY, self.dtype.clone(), false)
146    }
147}
148
149impl Iterator for DataFrameStreamIterator {
150    type Item = PolarsResult<ArrayRef>;
151
152    fn next(&mut self) -> Option<Self::Item> {
153        if self.idx >= self.n_chunks {
154            None
155        } else {
156            // create a batch of the columns with the same chunk no.
157            let batch_cols = self
158                .columns
159                .iter()
160                .map(|s| s.to_arrow(self.idx, CompatLevel::newest()))
161                .collect::<Vec<_>>();
162            self.idx += 1;
163
164            let col_len = batch_cols.first().map_or(self.height, |c| c.len());
165            let array =
166                arrow::array::StructArray::new(self.dtype.clone(), col_len, batch_cols, None);
167            Some(Ok(Box::new(array)))
168        }
169    }
170}
polars_python/interop/arrow/to_py.rs

polars_python/interop/arrow/
to_py.rs
