Skip to main content

polars_python/series/
export.rs

1use polars_core::prelude::*;
2use polars_ffi::version_0::SeriesExport;
3use pyo3::IntoPyObjectExt;
4use pyo3::prelude::*;
5use pyo3::types::{PyCapsule, PyList};
6
7use super::PySeries;
8use crate::error::PyPolarsErr;
9use crate::interop;
10use crate::interop::arrow::to_py::series_to_stream;
11use crate::prelude::*;
12
13#[pymethods]
14impl PySeries {
15    /// Convert this Series to a Python list.
16    /// This operation copies data.
17    pub fn to_list<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
18        let series = &self.series.read();
19
20        fn to_list_recursive<'py>(py: Python<'py>, series: &Series) -> PyResult<Bound<'py, PyAny>> {
21            let pylist = match series.dtype() {
22                DataType::Boolean => PyList::new(py, series.bool().map_err(PyPolarsErr::from)?)?,
23                DataType::UInt8 => PyList::new(py, series.u8().map_err(PyPolarsErr::from)?)?,
24                DataType::UInt16 => PyList::new(py, series.u16().map_err(PyPolarsErr::from)?)?,
25                DataType::UInt32 => PyList::new(py, series.u32().map_err(PyPolarsErr::from)?)?,
26                DataType::UInt64 => PyList::new(py, series.u64().map_err(PyPolarsErr::from)?)?,
27                DataType::UInt128 => PyList::new(py, series.u128().map_err(PyPolarsErr::from)?)?,
28                DataType::Int8 => PyList::new(py, series.i8().map_err(PyPolarsErr::from)?)?,
29                DataType::Int16 => PyList::new(py, series.i16().map_err(PyPolarsErr::from)?)?,
30                DataType::Int32 => PyList::new(py, series.i32().map_err(PyPolarsErr::from)?)?,
31                DataType::Int64 => PyList::new(py, series.i64().map_err(PyPolarsErr::from)?)?,
32                DataType::Int128 => PyList::new(py, series.i128().map_err(PyPolarsErr::from)?)?,
33                DataType::Float16 => PyList::new(py, series.f16().map_err(PyPolarsErr::from)?)?,
34                DataType::Float32 => PyList::new(py, series.f32().map_err(PyPolarsErr::from)?)?,
35                DataType::Float64 => PyList::new(py, series.f64().map_err(PyPolarsErr::from)?)?,
36                DataType::Categorical(_, _) | DataType::Enum(_, _) => {
37                    with_match_categorical_physical_type!(series.dtype().cat_physical().unwrap(), |$C| {
38                        PyList::new(py, series.cat::<$C>().unwrap().iter_str())?
39                    })
40                },
41                #[cfg(feature = "object")]
42                DataType::Object(_) => {
43                    let v = PyList::empty(py);
44                    for i in 0..series.len() {
45                        let obj: Option<&ObjectValue> = series.get_object(i).map(|any| any.into());
46                        v.append(obj)?;
47                    }
48                    v
49                },
50                DataType::List(_) => {
51                    let v = PyList::empty(py);
52                    let ca = series.list().map_err(PyPolarsErr::from)?;
53                    for opt_s in ca.amortized_iter() {
54                        match opt_s {
55                            None => {
56                                v.append(py.None())?;
57                            },
58                            Some(s) => {
59                                let pylst = to_list_recursive(py, s.as_ref())?;
60                                v.append(pylst)?;
61                            },
62                        }
63                    }
64                    v
65                },
66                DataType::Array(_, _) => {
67                    let v = PyList::empty(py);
68                    let ca = series.array().map_err(PyPolarsErr::from)?;
69                    for opt_s in ca.amortized_iter() {
70                        match opt_s {
71                            None => {
72                                v.append(py.None())?;
73                            },
74                            Some(s) => {
75                                let pylst = to_list_recursive(py, s.as_ref())?;
76                                v.append(pylst)?;
77                            },
78                        }
79                    }
80                    v
81                },
82                DataType::Date => {
83                    let ca = series.date().map_err(PyPolarsErr::from)?;
84                    return Wrap(ca).into_bound_py_any(py);
85                },
86                DataType::Time => {
87                    let ca = series.time().map_err(PyPolarsErr::from)?;
88                    return Wrap(ca).into_bound_py_any(py);
89                },
90                DataType::Datetime(_, _) => {
91                    let ca = series.datetime().map_err(PyPolarsErr::from)?;
92                    return Wrap(ca).into_bound_py_any(py);
93                },
94                DataType::Decimal(_, _) => {
95                    let ca = series.decimal().map_err(PyPolarsErr::from)?;
96                    return Wrap(ca).into_bound_py_any(py);
97                },
98                DataType::String => {
99                    let ca = series.str().map_err(PyPolarsErr::from)?;
100                    return Wrap(ca).into_bound_py_any(py);
101                },
102                DataType::Struct(_) => {
103                    let ca = series.struct_().map_err(PyPolarsErr::from)?;
104                    return Wrap(ca).into_bound_py_any(py);
105                },
106                DataType::Duration(_) => {
107                    let ca = series.duration().map_err(PyPolarsErr::from)?;
108                    return Wrap(ca).into_bound_py_any(py);
109                },
110                DataType::Binary => {
111                    let ca = series.binary().map_err(PyPolarsErr::from)?;
112                    return Wrap(ca).into_bound_py_any(py);
113                },
114                DataType::Null => {
115                    let null: Option<u8> = None;
116                    let n = series.len();
117                    let iter = std::iter::repeat_n(null, n);
118                    use std::iter::RepeatN;
119                    struct NullIter {
120                        iter: RepeatN<Option<u8>>,
121                        n: usize,
122                    }
123                    impl Iterator for NullIter {
124                        type Item = Option<u8>;
125
126                        fn next(&mut self) -> Option<Self::Item> {
127                            self.iter.next()
128                        }
129                        fn size_hint(&self) -> (usize, Option<usize>) {
130                            (self.n, Some(self.n))
131                        }
132                    }
133                    impl ExactSizeIterator for NullIter {}
134
135                    PyList::new(py, NullIter { iter, n })?
136                },
137                DataType::Unknown(_) => {
138                    panic!("to_list not implemented for unknown")
139                },
140                DataType::BinaryOffset => {
141                    unreachable!()
142                },
143                DataType::Extension(_, _) => {
144                    return to_list_recursive(py, series.ext().unwrap().storage());
145                },
146            };
147            Ok(pylist.into_any())
148        }
149
150        to_list_recursive(py, series)
151    }
152
153    /// Return the underlying Arrow array.
154    #[allow(clippy::wrong_self_convention)]
155    fn to_arrow(&self, py: Python<'_>, compat_level: PyCompatLevel) -> PyResult<Py<PyAny>> {
156        self.rechunk(py, true)?;
157        let pyarrow = py.import("pyarrow")?;
158
159        let s = self.series.read();
160        interop::arrow::to_py::to_py_array(
161            s.to_arrow(0, compat_level.0),
162            &s.field().to_arrow(compat_level.0),
163            &pyarrow,
164        )
165    }
166
167    #[allow(unused_variables)]
168    #[pyo3(signature = (requested_schema=None))]
169    fn __arrow_c_stream__<'py>(
170        &self,
171        py: Python<'py>,
172        requested_schema: Option<Py<PyAny>>,
173    ) -> PyResult<Bound<'py, PyCapsule>> {
174        series_to_stream(&self.series.read(), py)
175    }
176
177    pub fn _export(&self, _py: Python<'_>, location: usize) {
178        let export = polars_ffi::version_0::export_series(&self.series.read());
179        unsafe {
180            (location as *mut SeriesExport).write(export);
181        }
182    }
183}