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