polars_python/conversion/
chunked_array.rs

1use chrono::NaiveTime;
2use polars_core::utils::arrow::temporal_conversions::date32_to_date;
3use pyo3::prelude::*;
4use pyo3::types::{PyBytes, PyList, PyNone, PyTuple};
5use pyo3::{BoundObject, intern};
6
7use super::datetime::{
8    datetime_to_py_object, elapsed_offset_to_timedelta, nanos_since_midnight_to_naivetime,
9};
10use super::{decimal_to_digits, struct_dict};
11use crate::prelude::*;
12use crate::py_modules::pl_utils;
13
14impl<'py> IntoPyObject<'py> for &Wrap<&StringChunked> {
15    type Target = PyList;
16    type Output = Bound<'py, Self::Target>;
17    type Error = PyErr;
18
19    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
20        let iter = self.0.iter();
21        PyList::new(py, iter)
22    }
23}
24
25impl<'py> IntoPyObject<'py> for &Wrap<&BinaryChunked> {
26    type Target = PyList;
27    type Output = Bound<'py, Self::Target>;
28    type Error = PyErr;
29
30    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
31        let iter = self
32            .0
33            .iter()
34            .map(|opt_bytes| opt_bytes.map(|bytes| PyBytes::new(py, bytes)));
35        PyList::new(py, iter)
36    }
37}
38
39impl<'py> IntoPyObject<'py> for &Wrap<&StructChunked> {
40    type Target = PyList;
41    type Output = Bound<'py, Self::Target>;
42    type Error = PyErr;
43    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
44        let s = self.0.clone().into_series();
45        // todo! iterate its chunks and flatten.
46        // make series::iter() accept a chunk index.
47        let s = s.rechunk();
48        let iter = s.iter().map(|av| match av {
49            AnyValue::Struct(_, _, flds) => struct_dict(py, av._iter_struct_av(), flds)
50                .unwrap()
51                .into_any(),
52            AnyValue::Null => PyNone::get(py).into_bound().into_any(),
53            _ => unreachable!(),
54        });
55
56        PyList::new(py, iter)
57    }
58}
59
60impl<'py> IntoPyObject<'py> for &Wrap<&DurationChunked> {
61    type Target = PyList;
62    type Output = Bound<'py, Self::Target>;
63    type Error = PyErr;
64
65    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
66        let time_unit = self.0.time_unit();
67        let iter = self
68            .0
69            .physical()
70            .iter()
71            .map(|opt_v| opt_v.map(|v| elapsed_offset_to_timedelta(v, time_unit)));
72        PyList::new(py, iter)
73    }
74}
75
76impl<'py> IntoPyObject<'py> for &Wrap<&DatetimeChunked> {
77    type Target = PyList;
78    type Output = Bound<'py, Self::Target>;
79    type Error = PyErr;
80
81    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
82        let time_zone = self.0.time_zone().as_ref();
83        let time_unit = self.0.time_unit();
84        let iter = self.0.physical().iter().map(|opt_v| {
85            opt_v.map(|v| datetime_to_py_object(py, v, time_unit, time_zone).unwrap())
86        });
87        PyList::new(py, iter)
88    }
89}
90
91impl<'py> IntoPyObject<'py> for &Wrap<&TimeChunked> {
92    type Target = PyList;
93    type Output = Bound<'py, Self::Target>;
94    type Error = PyErr;
95
96    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
97        let iter = time_to_pyobject_iter(self.0);
98        PyList::new(py, iter)
99    }
100}
101
102pub(crate) fn time_to_pyobject_iter(
103    ca: &TimeChunked,
104) -> impl '_ + ExactSizeIterator<Item = Option<NaiveTime>> {
105    ca.phys
106        .iter()
107        .map(move |opt_v| opt_v.map(nanos_since_midnight_to_naivetime))
108}
109
110impl<'py> IntoPyObject<'py> for &Wrap<&DateChunked> {
111    type Target = PyList;
112    type Output = Bound<'py, Self::Target>;
113    type Error = PyErr;
114
115    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
116        let iter = self
117            .0
118            .physical()
119            .into_iter()
120            .map(|opt_v| opt_v.map(date32_to_date));
121        PyList::new(py, iter)
122    }
123}
124
125impl<'py> IntoPyObject<'py> for &Wrap<&DecimalChunked> {
126    type Target = PyList;
127    type Output = Bound<'py, Self::Target>;
128    type Error = PyErr;
129    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
130        let iter = decimal_to_pyobject_iter(py, self.0)?;
131        PyList::new(py, iter)
132    }
133}
134
135pub(crate) fn decimal_to_pyobject_iter<'py, 'a>(
136    py: Python<'py>,
137    ca: &'a DecimalChunked,
138) -> PyResult<impl ExactSizeIterator<Item = Option<Bound<'py, PyAny>>> + use<'py, 'a>> {
139    let utils = pl_utils(py).bind(py);
140    let convert = utils.getattr(intern!(py, "to_py_decimal"))?;
141    let py_scale = (-(ca.scale() as i32)).into_pyobject(py)?;
142    // if we don't know precision, the only safe bet is to set it to 39
143    let py_precision = ca.precision().unwrap_or(39).into_pyobject(py)?;
144    Ok(ca.physical().iter().map(move |opt_v| {
145        opt_v.map(|v| {
146            // TODO! use AnyValue so that we have a single impl.
147            const N: usize = 3;
148            let mut buf = [0_u128; N];
149            let n_digits = decimal_to_digits(v.abs(), &mut buf);
150            let buf = unsafe {
151                std::slice::from_raw_parts(
152                    buf.as_slice().as_ptr() as *const u8,
153                    N * size_of::<u128>(),
154                )
155            };
156            let digits = PyTuple::new(py, buf.iter().take(n_digits)).unwrap();
157            convert
158                .call1((v.is_negative() as u8, digits, &py_precision, &py_scale))
159                .unwrap()
160        })
161    }))
162}
polars_python/conversion/chunked_array.rs

polars_python/conversion/
chunked_array.rs
