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