Skip to main content

polars_python/conversion/
chunked_array.rs

1use chrono::NaiveTime;
2use polars_compute::decimal::DecimalFmtBuffer;
3use polars_core::utils::arrow::temporal_conversions::date32_to_date;
4use pyo3::BoundObject;
5use pyo3::prelude::*;
6use pyo3::types::{PyBytes, PyList, PyNone};
7
8use super::datetime::{
9    datetime_to_py_object, elapsed_offset_to_timedelta, nanos_since_midnight_to_naivetime,
10};
11use super::struct_dict;
12use crate::interned;
13use crate::prelude::*;
14use crate::py_modules::pl_utils;
15
16impl<'py> IntoPyObject<'py> for &Wrap<&StringChunked> {
17    type Target = PyList;
18    type Output = Bound<'py, Self::Target>;
19    type Error = PyErr;
20
21    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
22        let iter = self.0.iter();
23        PyList::new(py, iter)
24    }
25}
26
27impl<'py> IntoPyObject<'py> for &Wrap<&BinaryChunked> {
28    type Target = PyList;
29    type Output = Bound<'py, Self::Target>;
30    type Error = PyErr;
31
32    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
33        let iter = self
34            .0
35            .iter()
36            .map(|opt_bytes| opt_bytes.map(|bytes| PyBytes::new(py, bytes)));
37        PyList::new(py, iter)
38    }
39}
40
41impl<'py> IntoPyObject<'py> for &Wrap<&StructChunked> {
42    type Target = PyList;
43    type Output = Bound<'py, Self::Target>;
44    type Error = PyErr;
45    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
46        let s = self.0.clone().into_series();
47        let iter = s.iter().map(|av| match av {
48            AnyValue::Struct(_, _, flds) => struct_dict(py, av._iter_struct_av(), flds)
49                .unwrap()
50                .into_any(),
51            AnyValue::Null => PyNone::get(py).into_bound().into_any(),
52            _ => unreachable!(),
53        });
54
55        PyList::new(py, iter)
56    }
57}
58
59impl<'py> IntoPyObject<'py> for &Wrap<&DurationChunked> {
60    type Target = PyList;
61    type Output = Bound<'py, Self::Target>;
62    type Error = PyErr;
63
64    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
65        let time_unit = self.0.time_unit();
66        let iter = self
67            .0
68            .physical()
69            .iter()
70            .map(|opt_v| opt_v.map(|v| elapsed_offset_to_timedelta(v, time_unit)));
71        PyList::new(py, iter)
72    }
73}
74
75impl<'py> IntoPyObject<'py> for &Wrap<&DatetimeChunked> {
76    type Target = PyList;
77    type Output = Bound<'py, Self::Target>;
78    type Error = PyErr;
79
80    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
81        let time_zone = self.0.time_zone().as_ref();
82        let time_unit = self.0.time_unit();
83        let iter = self.0.physical().iter().map(|opt_v| {
84            opt_v.map(|v| datetime_to_py_object(py, v, time_unit, time_zone).unwrap())
85        });
86        PyList::new(py, iter)
87    }
88}
89
90impl<'py> IntoPyObject<'py> for &Wrap<&TimeChunked> {
91    type Target = PyList;
92    type Output = Bound<'py, Self::Target>;
93    type Error = PyErr;
94
95    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
96        let iter = time_to_pyobject_iter(self.0);
97        PyList::new(py, iter)
98    }
99}
100
101pub(crate) fn time_to_pyobject_iter(
102    ca: &TimeChunked,
103) -> impl '_ + ExactSizeIterator<Item = Option<NaiveTime>> {
104    ca.phys
105        .iter()
106        .map(move |opt_v| opt_v.map(nanos_since_midnight_to_naivetime))
107}
108
109impl<'py> IntoPyObject<'py> for &Wrap<&DateChunked> {
110    type Target = PyList;
111    type Output = Bound<'py, Self::Target>;
112    type Error = PyErr;
113
114    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
115        let iter = self
116            .0
117            .physical()
118            .iter()
119            .map(|opt_v| opt_v.map(date32_to_date));
120        PyList::new(py, iter)
121    }
122}
123
124impl<'py> IntoPyObject<'py> for &Wrap<&DecimalChunked> {
125    type Target = PyList;
126    type Output = Bound<'py, Self::Target>;
127    type Error = PyErr;
128    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
129        let iter = decimal_to_pyobject_iter(py, self.0)?;
130        PyList::new(py, iter)
131    }
132}
133
134pub(crate) fn decimal_to_pyobject_iter<'py, 'a>(
135    py: Python<'py>,
136    ca: &'a DecimalChunked,
137) -> PyResult<impl ExactSizeIterator<Item = Option<Bound<'py, PyAny>>> + use<'py, 'a>> {
138    let utils = pl_utils(py).bind(py);
139    let convert = utils.getattr(interned::TO_PY_DECIMAL.get(py))?;
140    let py_precision = ca.precision().into_pyobject(py)?;
141    let mut buf = DecimalFmtBuffer::new();
142    Ok(ca.physical().iter().map(move |opt_v| {
143        opt_v.map(|v| {
144            let s = buf.format_dec128(v, ca.scale(), false, false);
145            convert.call1((&py_precision, s)).unwrap()
146        })
147    }))
148}