polars_python/series/
import.rs

1use arrow::array::{Array, PrimitiveArray};
2use arrow::ffi;
3use arrow::ffi::{ArrowArray, ArrowArrayStream, ArrowArrayStreamReader, ArrowSchema};
4use polars::prelude::*;
5use polars_ffi::version_0::SeriesExport;
6use pyo3::exceptions::{PyTypeError, PyValueError};
7use pyo3::prelude::*;
8use pyo3::pybacked::PyBackedBytes;
9use pyo3::types::{PyCapsule, PyTuple, PyType};
10
11use super::PySeries;
12use crate::error::PyPolarsErr;
13
14/// Import `__arrow_c_array__` across Python boundary
15pub(crate) fn call_arrow_c_array<'py>(
16    ob: &Bound<'py, PyAny>,
17) -> PyResult<(Bound<'py, PyCapsule>, Bound<'py, PyCapsule>)> {
18    if !ob.hasattr("__arrow_c_array__")? {
19        return Err(PyValueError::new_err(
20            "Expected an object with dunder __arrow_c_array__",
21        ));
22    }
23
24    let tuple = ob.getattr("__arrow_c_array__")?.call0()?;
25    if !tuple.is_instance_of::<PyTuple>() {
26        return Err(PyTypeError::new_err(
27            "Expected __arrow_c_array__ to return a tuple.",
28        ));
29    }
30
31    let schema_capsule = tuple.get_item(0)?.cast_into()?;
32    let array_capsule = tuple.get_item(1)?.cast_into()?;
33    Ok((schema_capsule, array_capsule))
34}
35
36pub(crate) fn import_array_pycapsules(
37    schema_capsule: &Bound<PyCapsule>,
38    array_capsule: &Bound<PyCapsule>,
39) -> PyResult<(arrow::datatypes::Field, Box<dyn Array>)> {
40    let field = import_schema_pycapsule(schema_capsule)?;
41
42    // # Safety
43    // array_capsule holds a valid C ArrowArray pointer, as defined by the Arrow PyCapsule
44    // Interface
45    unsafe {
46        let array_ptr = std::ptr::replace(
47            array_capsule
48                .pointer_checked(Some(c"arrow_array"))?
49                .as_ptr() as _,
50            ArrowArray::empty(),
51        );
52        let array = ffi::import_array_from_c(array_ptr, field.dtype().clone()).unwrap();
53
54        Ok((field, array))
55    }
56}
57
58pub(crate) fn import_schema_pycapsule(
59    schema_capsule: &Bound<PyCapsule>,
60) -> PyResult<arrow::datatypes::Field> {
61    // # Safety
62    // schema_capsule holds a valid C ArrowSchema pointer, as defined by the Arrow PyCapsule
63    // Interface
64    unsafe {
65        let schema_ptr = schema_capsule
66            .pointer_checked(Some(c"arrow_schema"))?
67            .cast::<ArrowSchema>()
68            .as_ref();
69        let field = ffi::import_field_from_c(schema_ptr).unwrap();
70
71        Ok(field)
72    }
73}
74
75/// Import `__arrow_c_stream__` across Python boundary.
76pub(crate) fn call_arrow_c_stream<'py>(ob: &Bound<'py, PyAny>) -> PyResult<Bound<'py, PyCapsule>> {
77    if !ob.hasattr("__arrow_c_stream__")? {
78        return Err(PyValueError::new_err(
79            "Expected an object with dunder __arrow_c_stream__",
80        ));
81    }
82
83    let capsule = ob.getattr("__arrow_c_stream__")?.call0()?.cast_into()?;
84    Ok(capsule)
85}
86
87/// Takes ownership of the `ArrowArrayStream` behind a stream capsule and wraps it
88/// for iteration.
89///
90/// # Safety
91/// `capsule` must hold a valid C `ArrowArrayStream` pointer, as defined by the Arrow
92/// PyCapsule Interface.
93pub(crate) fn open_stream_capsule(
94    capsule: &Bound<PyCapsule>,
95) -> PyResult<ArrowArrayStreamReader<Box<ArrowArrayStream>>> {
96    unsafe {
97        let stream_ptr = Box::new(std::ptr::replace(
98            capsule
99                .pointer_checked(Some(c"arrow_array_stream"))?
100                .as_ptr() as _,
101            ArrowArrayStream::empty(),
102        ));
103        ArrowArrayStreamReader::try_new(stream_ptr)
104            .map_err(|err| PyValueError::new_err(err.to_string()))
105    }
106}
107
108pub(crate) fn import_stream_pycapsule(capsule: &Bound<PyCapsule>) -> PyResult<PySeries> {
109    let mut stream = open_stream_capsule(capsule)?;
110
111    let mut produced_arrays: Vec<Box<dyn Array>> = vec![];
112    while let Some(array) = unsafe { stream.next() } {
113        produced_arrays.push(array.map_err(PyPolarsErr::from)?);
114    }
115
116    // Series::try_from fails for an empty vec of chunks
117    let s = if produced_arrays.is_empty() {
118        let polars_dt = DataType::from_arrow_field(stream.field());
119        Series::new_empty(stream.field().name.clone(), &polars_dt)
120    } else {
121        Series::try_from((stream.field(), produced_arrays)).map_err(PyPolarsErr::from)?
122    };
123    Ok(PySeries::new(s))
124}
125#[pymethods]
126impl PySeries {
127    #[classmethod]
128    pub fn from_arrow_c_array(_cls: &Bound<PyType>, ob: &Bound<'_, PyAny>) -> PyResult<Self> {
129        let (schema_capsule, array_capsule) = call_arrow_c_array(ob)?;
130        let (field, array) = import_array_pycapsules(&schema_capsule, &array_capsule)?;
131        let s = Series::try_from((&field, array)).unwrap();
132        Ok(PySeries::new(s))
133    }
134
135    #[classmethod]
136    pub fn from_arrow_c_stream(_cls: &Bound<PyType>, ob: &Bound<'_, PyAny>) -> PyResult<Self> {
137        let capsule = call_arrow_c_stream(ob)?;
138        import_stream_pycapsule(&capsule)
139    }
140
141    #[classmethod]
142    /// Import a series via polars-ffi
143    /// Takes ownership of the [`SeriesExport`] at [`location`]
144    /// # Safety
145    /// [`location`] should be the address of an allocated and initialized [`SeriesExport`]
146    pub unsafe fn _import(_cls: &Bound<PyType>, location: usize) -> PyResult<Self> {
147        let location = location as *mut SeriesExport;
148
149        // # Safety
150        // `location` should be valid for reading
151        let series = unsafe {
152            let export = location.read();
153            polars_ffi::version_0::import_series(export).map_err(PyPolarsErr::from)?
154        };
155        Ok(PySeries::from(series))
156    }
157
158    #[staticmethod]
159    pub fn _import_decimal_from_iceberg_binary_repr(
160        bytes_list: &Bound<PyAny>, // list[bytes | None]
161        precision: usize,
162        scale: usize,
163    ) -> PyResult<Self> {
164        // From iceberg spec:
165        // * Decimal(P, S): Stores unscaled value as two’s-complement
166        //   big-endian binary, using the minimum number of bytes for the
167        //   value.
168        let max_abs_decimal_value = 10_i128.pow(u32::try_from(precision).unwrap()) - 1;
169
170        let out: Vec<i128> = bytes_list
171            .try_iter()?
172            .map(|bytes| {
173                let be_bytes: Option<PyBackedBytes> = bytes?.extract()?;
174
175                let mut le_bytes: [u8; 16] = [0; _];
176
177                if let Some(be_bytes) = be_bytes.as_deref() {
178                    if be_bytes.len() > le_bytes.len() {
179                        return Err(PyValueError::new_err(format!(
180                            "iceberg binary data for decimal exceeded 16 bytes: {}",
181                            be_bytes.len()
182                        )));
183                    }
184
185                    for (i, byte) in be_bytes.iter().rev().enumerate() {
186                        le_bytes[i] = *byte;
187                    }
188                }
189
190                let value = i128::from_le_bytes(le_bytes);
191
192                if value.abs() > max_abs_decimal_value {
193                    return Err(PyValueError::new_err(format!(
194                        "iceberg decoded value for decimal exceeded precision: \
195                        value: {value}, precision: {precision}",
196                    )));
197                }
198
199                Ok(value)
200            })
201            .collect::<PyResult<_>>()?;
202
203        Ok(PySeries::from(unsafe {
204            Series::from_chunks_and_dtype_unchecked(
205                PlSmallStr::EMPTY,
206                vec![PrimitiveArray::<i128>::from_vec(out).boxed()],
207                &DataType::Decimal(precision, scale),
208            )
209        }))
210    }
211}
polars_python/series/import.rs

polars_python/series/
import.rs
