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        // todo! iterate its chunks and flatten.
47        // make series::iter() accept a chunk index.
48        let s = s.rechunk();
49        let iter = s.iter().map(|av| match av {
50            AnyValue::Struct(_, _, flds) => struct_dict(py, av._iter_struct_av(), flds)
51                .unwrap()
52                .into_any(),
53            AnyValue::Null => PyNone::get(py).into_bound().into_any(),
54            _ => unreachable!(),
55        });
56
57        PyList::new(py, iter)
58    }
59}
60
61impl<'py> IntoPyObject<'py> for &Wrap<&DurationChunked> {
62    type Target = PyList;
63    type Output = Bound<'py, Self::Target>;
64    type Error = PyErr;
65
66    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
67        let time_unit = self.0.time_unit();
68        let iter = self
69            .0
70            .physical()
71            .iter()
72            .map(|opt_v| opt_v.map(|v| elapsed_offset_to_timedelta(v, time_unit)));
73        PyList::new(py, iter)
74    }
75}
76
77impl<'py> IntoPyObject<'py> for &Wrap<&DatetimeChunked> {
78    type Target = PyList;
79    type Output = Bound<'py, Self::Target>;
80    type Error = PyErr;
81
82    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
83        let time_zone = self.0.time_zone().as_ref();
84        let time_unit = self.0.time_unit();
85        let iter = self.0.physical().iter().map(|opt_v| {
86            opt_v.map(|v| datetime_to_py_object(py, v, time_unit, time_zone).unwrap())
87        });
88        PyList::new(py, iter)
89    }
90}
91
92impl<'py> IntoPyObject<'py> for &Wrap<&TimeChunked> {
93    type Target = PyList;
94    type Output = Bound<'py, Self::Target>;
95    type Error = PyErr;
96
97    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
98        let iter = time_to_pyobject_iter(self.0);
99        PyList::new(py, iter)
100    }
101}
102
103pub(crate) fn time_to_pyobject_iter(
104    ca: &TimeChunked,
105) -> impl '_ + ExactSizeIterator<Item = Option<NaiveTime>> {
106    ca.phys
107        .iter()
108        .map(move |opt_v| opt_v.map(nanos_since_midnight_to_naivetime))
109}
110
111impl<'py> IntoPyObject<'py> for &Wrap<&DateChunked> {
112    type Target = PyList;
113    type Output = Bound<'py, Self::Target>;
114    type Error = PyErr;
115
116    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
117        let iter = self
118            .0
119            .physical()
120            .into_iter()
121            .map(|opt_v| opt_v.map(date32_to_date));
122        PyList::new(py, iter)
123    }
124}
125
126impl<'py> IntoPyObject<'py> for &Wrap<&DecimalChunked> {
127    type Target = PyList;
128    type Output = Bound<'py, Self::Target>;
129    type Error = PyErr;
130    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
131        let iter = decimal_to_pyobject_iter(py, self.0)?;
132        PyList::new(py, iter)
133    }
134}
135
136pub(crate) fn decimal_to_pyobject_iter<'py, 'a>(
137    py: Python<'py>,
138    ca: &'a DecimalChunked,
139) -> PyResult<impl ExactSizeIterator<Item = Option<Bound<'py, PyAny>>> + use<'py, 'a>> {
140    let utils = pl_utils(py).bind(py);
141    let convert = utils.getattr(intern!(py, "to_py_decimal"))?;
142    let py_precision = ca.precision().into_pyobject(py)?;
143    let mut buf = DecimalFmtBuffer::new();
144    Ok(ca.physical().iter().map(move |opt_v| {
145        opt_v.map(|v| {
146            let s = buf.format_dec128(v, ca.scale(), false, false);
147            convert.call1((&py_precision, s)).unwrap()
148        })
149    }))
150}