polars_python/series/
construction.rs

1use std::borrow::Cow;
2
3use arrow::array::Array;
4use arrow::bitmap::BitmapBuilder;
5use arrow::types::NativeType;
6use numpy::{Element, PyArray1, PyArrayMethods};
7use polars_core::prelude::*;
8use polars_core::utils::CustomIterTools;
9use pyo3::exceptions::{PyTypeError, PyValueError};
10use pyo3::prelude::*;
11
12use crate::PySeries;
13use crate::conversion::any_value::py_object_to_any_value;
14use crate::conversion::{Wrap, reinterpret_vec};
15use crate::error::PyPolarsErr;
16use crate::interop::arrow::to_rust::array_to_rust;
17use crate::prelude::ObjectValue;
18use crate::utils::EnterPolarsExt;
19
20// Init with numpy arrays.
21macro_rules! init_method {
22    ($name:ident, $type:ty) => {
23        #[pymethods]
24        impl PySeries {
25            #[staticmethod]
26            fn $name(name: &str, array: &Bound<PyArray1<$type>>, _strict: bool) -> Self {
27                mmap_numpy_array(name, array)
28            }
29        }
30    };
31}
32
33init_method!(new_i8, i8);
34init_method!(new_i16, i16);
35init_method!(new_i32, i32);
36init_method!(new_i64, i64);
37init_method!(new_u8, u8);
38init_method!(new_u16, u16);
39init_method!(new_u32, u32);
40init_method!(new_u64, u64);
41
42fn mmap_numpy_array<T: Element + NativeType>(name: &str, array: &Bound<PyArray1<T>>) -> PySeries {
43    let vals = unsafe { array.as_slice().unwrap() };
44
45    let arr = unsafe { arrow::ffi::mmap::slice_and_owner(vals, array.clone().unbind()) };
46    Series::from_arrow(name.into(), arr.to_boxed())
47        .unwrap()
48        .into()
49}
50
51#[pymethods]
52impl PySeries {
53    #[staticmethod]
54    fn new_bool(
55        py: Python,
56        name: &str,
57        array: &Bound<PyArray1<bool>>,
58        _strict: bool,
59    ) -> PyResult<Self> {
60        let array = array.readonly();
61        let vals = array.as_slice().unwrap();
62        py.enter_polars_series(|| Ok(Series::new(name.into(), vals)))
63    }
64
65    #[staticmethod]
66    fn new_f32(
67        py: Python,
68        name: &str,
69        array: &Bound<PyArray1<f32>>,
70        nan_is_null: bool,
71    ) -> PyResult<Self> {
72        if nan_is_null {
73            let array = array.readonly();
74            let vals = array.as_slice().unwrap();
75            py.enter_polars_series(|| {
76                let ca: Float32Chunked = vals
77                    .iter()
78                    .map(|&val| if f32::is_nan(val) { None } else { Some(val) })
79                    .collect_trusted();
80                Ok(ca.with_name(name.into()))
81            })
82        } else {
83            Ok(mmap_numpy_array(name, array))
84        }
85    }
86
87    #[staticmethod]
88    fn new_f64(
89        py: Python,
90        name: &str,
91        array: &Bound<PyArray1<f64>>,
92        nan_is_null: bool,
93    ) -> PyResult<Self> {
94        if nan_is_null {
95            let array = array.readonly();
96            let vals = array.as_slice().unwrap();
97            py.enter_polars_series(|| {
98                let ca: Float64Chunked = vals
99                    .iter()
100                    .map(|&val| if f64::is_nan(val) { None } else { Some(val) })
101                    .collect_trusted();
102                Ok(ca.with_name(name.into()))
103            })
104        } else {
105            Ok(mmap_numpy_array(name, array))
106        }
107    }
108}
109
110#[pymethods]
111impl PySeries {
112    #[staticmethod]
113    fn new_opt_bool(name: &str, values: &Bound<PyAny>, _strict: bool) -> PyResult<Self> {
114        let len = values.len()?;
115        let mut builder = BooleanChunkedBuilder::new(name.into(), len);
116
117        for res in values.try_iter()? {
118            let value = res?;
119            if value.is_none() {
120                builder.append_null()
121            } else {
122                let v = value.extract::<bool>()?;
123                builder.append_value(v)
124            }
125        }
126
127        let ca = builder.finish();
128        let s = ca.into_series();
129        Ok(s.into())
130    }
131}
132
133fn new_primitive<'a, T>(name: &str, values: &'a Bound<PyAny>, _strict: bool) -> PyResult<PySeries>
134where
135    T: PolarsNumericType,
136    ChunkedArray<T>: IntoSeries,
137    T::Native: FromPyObject<'a>,
138{
139    let len = values.len()?;
140    let mut builder = PrimitiveChunkedBuilder::<T>::new(name.into(), len);
141
142    for res in values.try_iter()? {
143        let value = res?;
144        if value.is_none() {
145            builder.append_null()
146        } else {
147            let v = value.extract::<T::Native>()?;
148            builder.append_value(v)
149        }
150    }
151
152    let ca = builder.finish();
153    let s = ca.into_series();
154    Ok(s.into())
155}
156
157// Init with lists that can contain Nones
158macro_rules! init_method_opt {
159    ($name:ident, $type:ty, $native: ty) => {
160        #[pymethods]
161        impl PySeries {
162            #[staticmethod]
163            fn $name(name: &str, obj: &Bound<PyAny>, strict: bool) -> PyResult<Self> {
164                new_primitive::<$type>(name, obj, strict)
165            }
166        }
167    };
168}
169
170init_method_opt!(new_opt_u8, UInt8Type, u8);
171init_method_opt!(new_opt_u16, UInt16Type, u16);
172init_method_opt!(new_opt_u32, UInt32Type, u32);
173init_method_opt!(new_opt_u64, UInt64Type, u64);
174init_method_opt!(new_opt_i8, Int8Type, i8);
175init_method_opt!(new_opt_i16, Int16Type, i16);
176init_method_opt!(new_opt_i32, Int32Type, i32);
177init_method_opt!(new_opt_i64, Int64Type, i64);
178init_method_opt!(new_opt_i128, Int128Type, i64);
179init_method_opt!(new_opt_f32, Float32Type, f32);
180init_method_opt!(new_opt_f64, Float64Type, f64);
181
182fn convert_to_avs<'a>(
183    values: &'a Bound<'a, PyAny>,
184    strict: bool,
185    allow_object: bool,
186) -> PyResult<Vec<AnyValue<'a>>> {
187    values
188        .try_iter()?
189        .map(|v| py_object_to_any_value(&(v?).as_borrowed(), strict, allow_object))
190        .collect()
191}
192
193#[pymethods]
194impl PySeries {
195    #[staticmethod]
196    fn new_from_any_values(name: &str, values: &Bound<PyAny>, strict: bool) -> PyResult<Self> {
197        let any_values_result = values
198            .try_iter()?
199            .map(|v| py_object_to_any_value(&(v?).as_borrowed(), strict, true))
200            .collect::<PyResult<Vec<AnyValue>>>();
201        let result = any_values_result.and_then(|avs| {
202            let s = Series::from_any_values(name.into(), avs.as_slice(), strict).map_err(|e| {
203                PyTypeError::new_err(format!(
204                    "{e}\n\nHint: Try setting `strict=False` to allow passing data with mixed types."
205                ))
206            })?;
207            Ok(s.into())
208        });
209
210        // Fall back to Object type for non-strict construction.
211        if !strict && result.is_err() {
212            return Python::with_gil(|py| {
213                let objects = values
214                    .try_iter()?
215                    .map(|v| v?.extract())
216                    .collect::<PyResult<Vec<ObjectValue>>>()?;
217                Ok(Self::new_object(py, name, objects, strict))
218            });
219        }
220
221        result
222    }
223
224    #[staticmethod]
225    fn new_from_any_values_and_dtype(
226        name: &str,
227        values: &Bound<PyAny>,
228        dtype: Wrap<DataType>,
229        strict: bool,
230    ) -> PyResult<Self> {
231        let avs = convert_to_avs(values, strict, false)?;
232        let s = Series::from_any_values_and_dtype(name.into(), avs.as_slice(), &dtype.0, strict)
233            .map_err(|e| {
234                PyTypeError::new_err(format!(
235                "{e}\n\nHint: Try setting `strict=False` to allow passing data with mixed types."
236            ))
237            })?;
238        Ok(s.into())
239    }
240
241    #[staticmethod]
242    fn new_str(name: &str, values: &Bound<PyAny>, _strict: bool) -> PyResult<Self> {
243        let len = values.len()?;
244        let mut builder = StringChunkedBuilder::new(name.into(), len);
245
246        for res in values.try_iter()? {
247            let value = res?;
248            if value.is_none() {
249                builder.append_null()
250            } else {
251                let v = value.extract::<Cow<str>>()?;
252                builder.append_value(v)
253            }
254        }
255
256        let ca = builder.finish();
257        let s = ca.into_series();
258        Ok(s.into())
259    }
260
261    #[staticmethod]
262    fn new_binary(name: &str, values: &Bound<PyAny>, _strict: bool) -> PyResult<Self> {
263        let len = values.len()?;
264        let mut builder = BinaryChunkedBuilder::new(name.into(), len);
265
266        for res in values.try_iter()? {
267            let value = res?;
268            if value.is_none() {
269                builder.append_null()
270            } else {
271                let v = value.extract::<&[u8]>()?;
272                builder.append_value(v)
273            }
274        }
275
276        let ca = builder.finish();
277        let s = ca.into_series();
278        Ok(s.into())
279    }
280
281    #[staticmethod]
282    fn new_decimal(name: &str, values: &Bound<PyAny>, strict: bool) -> PyResult<Self> {
283        Self::new_from_any_values(name, values, strict)
284    }
285
286    #[staticmethod]
287    fn new_series_list(name: &str, values: Vec<Option<PySeries>>, _strict: bool) -> PyResult<Self> {
288        let series = reinterpret_vec(values);
289        if let Some(s) = series.iter().flatten().next() {
290            if s.dtype().is_object() {
291                return Err(PyValueError::new_err(
292                    "list of objects isn't supported; try building a 'object' only series",
293                ));
294            }
295        }
296        Ok(Series::new(name.into(), series).into())
297    }
298
299    #[staticmethod]
300    #[pyo3(signature = (name, values, strict, dtype))]
301    fn new_array(
302        name: &str,
303        values: &Bound<PyAny>,
304        strict: bool,
305        dtype: Wrap<DataType>,
306    ) -> PyResult<Self> {
307        Self::new_from_any_values_and_dtype(name, values, dtype, strict)
308    }
309
310    #[staticmethod]
311    pub fn new_object(py: Python, name: &str, values: Vec<ObjectValue>, _strict: bool) -> Self {
312        #[cfg(feature = "object")]
313        {
314            let mut validity = BitmapBuilder::with_capacity(values.len());
315            values.iter().for_each(|v| {
316                let is_valid = !v.inner.is_none(py);
317                // SAFETY: we can ensure that validity has correct capacity.
318                unsafe { validity.push_unchecked(is_valid) };
319            });
320            // Object builder must be registered. This is done on import.
321            let ca = ObjectChunked::<ObjectValue>::new_from_vec_and_validity(
322                name.into(),
323                values,
324                validity.into_opt_validity(),
325            );
326            let s = ca.into_series();
327            s.into()
328        }
329        #[cfg(not(feature = "object"))]
330        panic!("activate 'object' feature")
331    }
332
333    #[staticmethod]
334    fn new_null(name: &str, values: &Bound<PyAny>, _strict: bool) -> PyResult<Self> {
335        let len = values.len()?;
336        Ok(Series::new_null(name.into(), len).into())
337    }
338
339    #[staticmethod]
340    fn from_arrow(name: &str, array: &Bound<PyAny>) -> PyResult<Self> {
341        let arr = array_to_rust(array)?;
342
343        match arr.dtype() {
344            ArrowDataType::LargeList(_) => {
345                let array = arr.as_any().downcast_ref::<LargeListArray>().unwrap();
346                let fast_explode = array.offsets().as_slice().windows(2).all(|w| w[0] != w[1]);
347
348                let mut out = ListChunked::with_chunk(name.into(), array.clone());
349                if fast_explode {
350                    out.set_fast_explode()
351                }
352                Ok(out.into_series().into())
353            },
354            _ => {
355                let series: Series =
356                    Series::try_new(name.into(), arr).map_err(PyPolarsErr::from)?;
357                Ok(series.into())
358            },
359        }
360    }
361}
polars_python/series/construction.rs

polars_python/series/
construction.rs
