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