polars_python/series/
map.rs

1use pyo3::Python;
2use pyo3::prelude::*;
3use pyo3::types::PyCFunction;
4
5use super::PySeries;
6use crate::error::PyPolarsErr;
7use crate::map::check_nested_object;
8use crate::map::series::{ApplyLambda, call_lambda_and_extract};
9use crate::prelude::*;
10use crate::py_modules::pl_series;
11use crate::{apply_method_all_arrow_series2, raise_err};
12
13#[pymethods]
14impl PySeries {
15    #[pyo3(signature = (function, return_dtype, skip_nulls))]
16    fn map_elements(
17        &self,
18        function: &Bound<PyAny>,
19        return_dtype: Option<Wrap<DataType>>,
20        skip_nulls: bool,
21    ) -> PyResult<PySeries> {
22        let series = &self.series.read().clone(); // Clone so we don't deadlock on re-entrance.
23
24        if skip_nulls && (series.null_count() == series.len()) {
25            if let Some(return_dtype) = return_dtype {
26                return Ok(
27                    Series::full_null(series.name().clone(), series.len(), &return_dtype.0).into(),
28                );
29            }
30            let msg = "The output type of the 'map_elements' function cannot be determined.\n\
31            The function was never called because 'skip_nulls=True' and all values are null.\n\
32            Consider setting 'skip_nulls=False' or setting the 'return_dtype'.";
33            raise_err!(msg, ComputeError)
34        }
35
36        let return_dtype = return_dtype.map(|dt| dt.0);
37
38        macro_rules! dispatch_apply {
39            ($self:expr, $method:ident, $($args:expr),*) => {
40                match $self.dtype() {
41                    #[cfg(feature = "object")]
42                    DataType::Object(_) => {
43                        let ca = $self.0.unpack::<ObjectType<ObjectValue>>().unwrap();
44                        ca.$method($($args),*)
45                    },
46                    _ => {
47                        apply_method_all_arrow_series2!(
48                            $self,
49                            $method,
50                            $($args),*
51                        )
52                    }
53
54                }
55            }
56
57        }
58
59        Python::attach(|py| {
60            if matches!(
61                series.dtype(),
62                DataType::Datetime(_, _)
63                    | DataType::Date
64                    | DataType::Duration(_)
65                    | DataType::Categorical(_, _)
66                    | DataType::Enum(_, _)
67                    | DataType::Binary
68                    | DataType::Array(_, _)
69                    | DataType::Time
70                    | DataType::Decimal(_, _)
71            ) || !skip_nulls
72            {
73                let mut avs = Vec::with_capacity(series.len());
74                let s = series.rechunk();
75
76                for av in s.iter() {
77                    let out = match (skip_nulls, av) {
78                        (true, AnyValue::Null) => AnyValue::Null,
79                        (_, av) => {
80                            let av: Option<Wrap<AnyValue>> =
81                                call_lambda_and_extract(py, function, Wrap(av))?;
82                            match av {
83                                None => AnyValue::Null,
84                                Some(av) => av.0,
85                            }
86                        },
87                    };
88                    avs.push(out)
89                }
90                let out = Series::new(series.name().clone(), &avs);
91                let dtype = out.dtype();
92                if dtype.is_nested() {
93                    check_nested_object(dtype)?;
94                }
95
96                return Ok(out.into());
97            }
98
99            let out = match return_dtype {
100                Some(DataType::Int8) => {
101                    let ca: Int8Chunked = dispatch_apply!(
102                        series,
103                        apply_lambda_with_primitive_out_type,
104                        py,
105                        function,
106                        0,
107                        None
108                    )?;
109                    ca.into_series()
110                },
111                Some(DataType::Int16) => {
112                    let ca: Int16Chunked = dispatch_apply!(
113                        series,
114                        apply_lambda_with_primitive_out_type,
115                        py,
116                        function,
117                        0,
118                        None
119                    )?;
120                    ca.into_series()
121                },
122                Some(DataType::Int32) => {
123                    let ca: Int32Chunked = dispatch_apply!(
124                        series,
125                        apply_lambda_with_primitive_out_type,
126                        py,
127                        function,
128                        0,
129                        None
130                    )?;
131                    ca.into_series()
132                },
133                Some(DataType::Int64) => {
134                    let ca: Int64Chunked = dispatch_apply!(
135                        series,
136                        apply_lambda_with_primitive_out_type,
137                        py,
138                        function,
139                        0,
140                        None
141                    )?;
142                    ca.into_series()
143                },
144                Some(DataType::Int128) => {
145                    let ca: Int128Chunked = dispatch_apply!(
146                        series,
147                        apply_lambda_with_primitive_out_type,
148                        py,
149                        function,
150                        0,
151                        None
152                    )?;
153                    ca.into_series()
154                },
155                Some(DataType::UInt8) => {
156                    let ca: UInt8Chunked = dispatch_apply!(
157                        series,
158                        apply_lambda_with_primitive_out_type,
159                        py,
160                        function,
161                        0,
162                        None
163                    )?;
164                    ca.into_series()
165                },
166                Some(DataType::UInt16) => {
167                    let ca: UInt16Chunked = dispatch_apply!(
168                        series,
169                        apply_lambda_with_primitive_out_type,
170                        py,
171                        function,
172                        0,
173                        None
174                    )?;
175                    ca.into_series()
176                },
177                Some(DataType::UInt32) => {
178                    let ca: UInt32Chunked = dispatch_apply!(
179                        series,
180                        apply_lambda_with_primitive_out_type,
181                        py,
182                        function,
183                        0,
184                        None
185                    )?;
186                    ca.into_series()
187                },
188                Some(DataType::UInt64) => {
189                    let ca: UInt64Chunked = dispatch_apply!(
190                        series,
191                        apply_lambda_with_primitive_out_type,
192                        py,
193                        function,
194                        0,
195                        None
196                    )?;
197                    ca.into_series()
198                },
199                Some(DataType::UInt128) => {
200                    let ca: UInt128Chunked = dispatch_apply!(
201                        series,
202                        apply_lambda_with_primitive_out_type,
203                        py,
204                        function,
205                        0,
206                        None
207                    )?;
208                    ca.into_series()
209                },
210                Some(DataType::Float32) => {
211                    let ca: Float32Chunked = dispatch_apply!(
212                        series,
213                        apply_lambda_with_primitive_out_type,
214                        py,
215                        function,
216                        0,
217                        None
218                    )?;
219                    ca.into_series()
220                },
221                Some(DataType::Float64) => {
222                    let ca: Float64Chunked = dispatch_apply!(
223                        series,
224                        apply_lambda_with_primitive_out_type,
225                        py,
226                        function,
227                        0,
228                        None
229                    )?;
230                    ca.into_series()
231                },
232                Some(DataType::Boolean) => {
233                    let ca: BooleanChunked = dispatch_apply!(
234                        series,
235                        apply_lambda_with_bool_out_type,
236                        py,
237                        function,
238                        0,
239                        None
240                    )?;
241                    ca.into_series()
242                },
243                Some(DataType::String) => {
244                    let ca = dispatch_apply!(
245                        series,
246                        apply_lambda_with_string_out_type,
247                        py,
248                        function,
249                        0,
250                        None
251                    )?;
252
253                    ca.into_series()
254                },
255                Some(DataType::List(inner)) => {
256                    check_nested_object(&inner)?;
257                    // Make sure the function returns a Series of the correct data type.
258                    let function_owned = function.clone().unbind();
259                    let dtype_py = Wrap((*inner).clone());
260                    let function_wrapped =
261                        PyCFunction::new_closure(py, None, None, move |args, _kwargs| {
262                            Python::attach(|py| {
263                                let out = function_owned.call1(py, args)?;
264                                if out.is_none(py) {
265                                    Ok(py.None())
266                                } else {
267                                    pl_series(py).call1(py, ("", out, &dtype_py))
268                                }
269                            })
270                        })?
271                        .into_any()
272                        .unbind();
273
274                    let ca = dispatch_apply!(
275                        series,
276                        apply_lambda_with_list_out_type,
277                        py,
278                        function_wrapped,
279                        0,
280                        None,
281                        inner.as_ref()
282                    )?;
283
284                    ca.into_series()
285                },
286                #[cfg(feature = "object")]
287                Some(DataType::Object(_)) => {
288                    let ca = dispatch_apply!(
289                        series,
290                        apply_lambda_with_object_out_type,
291                        py,
292                        function,
293                        0,
294                        None
295                    )?;
296                    ca.into_series()
297                },
298                None => return dispatch_apply!(series, apply_lambda_unknown, py, function),
299
300                _ => return dispatch_apply!(series, apply_lambda_unknown, py, function),
301            };
302
303            assert!(out.dtype().is_known());
304            Ok(out.into())
305        })
306    }
307}
polars_python/series/map.rs

polars_python/series/
map.rs
