polars_python/map/
series.rs

1use polars::prelude::*;
2use pyo3::prelude::*;
3use pyo3::types::{PyNone, PyTuple};
4
5use super::*;
6use crate::error::PyPolarsErr;
7use crate::prelude::ObjectValue;
8use crate::{PySeries, Wrap};
9
10pub trait ApplyLambdaGeneric<'py> {
11    fn apply_generic(
12        &self,
13        py: Python<'py>,
14        lambda: &Bound<'py, PyAny>,
15        skip_nulls: bool,
16    ) -> PyResult<Series>;
17
18    fn apply_generic_with_dtype(
19        &self,
20        py: Python<'py>,
21        lambda: &Bound<'py, PyAny>,
22        datatype: &DataType,
23        skip_nulls: bool,
24    ) -> PyResult<Series>;
25}
26
27fn call_and_collect_anyvalues<'py, T, I>(
28    py: Python<'py>,
29    lambda: &Bound<'py, PyAny>,
30    len: usize,
31    iter: I,
32    skip_nulls: bool,
33) -> PyResult<Vec<AnyValue<'static>>>
34where
35    T: IntoPyObject<'py>,
36    I: Iterator<Item = Option<T>>,
37{
38    let mut avs = Vec::with_capacity(len);
39    for opt_val in iter {
40        let arg = match opt_val {
41            None if skip_nulls => {
42                avs.push(AnyValue::Null);
43                continue;
44            },
45            None => PyTuple::new(py, [PyNone::get(py)])?,
46            Some(val) => PyTuple::new(py, [val])?,
47        };
48        let out = lambda.call1(arg)?;
49        let av: Option<Wrap<AnyValue>> = if out.is_none() {
50            Ok(None)
51        } else {
52            out.extract().map(Some)
53        }?;
54        avs.push(av.map(|w| w.0).unwrap_or(AnyValue::Null));
55    }
56    Ok(avs)
57}
58
59impl<'py> ApplyLambdaGeneric<'py> for BooleanChunked {
60    fn apply_generic(
61        &self,
62        py: Python<'py>,
63        lambda: &Bound<'py, PyAny>,
64        skip_nulls: bool,
65    ) -> PyResult<Series> {
66        let avs = call_and_collect_anyvalues(py, lambda, self.len(), self.into_iter(), skip_nulls)?;
67        Ok(Series::from_any_values(self.name().clone(), &avs, true).map_err(PyPolarsErr::from)?)
68    }
69
70    fn apply_generic_with_dtype(
71        &self,
72        py: Python<'py>,
73        lambda: &Bound<'py, PyAny>,
74        datatype: &DataType,
75        skip_nulls: bool,
76    ) -> PyResult<Series> {
77        let avs = call_and_collect_anyvalues(py, lambda, self.len(), self.into_iter(), skip_nulls)?;
78        Ok(
79            Series::from_any_values_and_dtype(self.name().clone(), &avs, datatype, true)
80                .map_err(PyPolarsErr::from)?,
81        )
82    }
83}
84
85impl<'py, T> ApplyLambdaGeneric<'py> for ChunkedArray<T>
86where
87    T: PyPolarsNumericType,
88    T::Native: IntoPyObject<'py> + for<'a> FromPyObject<'a, 'py>,
89{
90    fn apply_generic(
91        &self,
92        py: Python<'py>,
93        lambda: &Bound<'py, PyAny>,
94        skip_nulls: bool,
95    ) -> PyResult<Series> {
96        let avs = call_and_collect_anyvalues(py, lambda, self.len(), self.into_iter(), skip_nulls)?;
97        Ok(Series::from_any_values(self.name().clone(), &avs, true).map_err(PyPolarsErr::from)?)
98    }
99
100    fn apply_generic_with_dtype(
101        &self,
102        py: Python<'py>,
103        lambda: &Bound<'py, PyAny>,
104        datatype: &DataType,
105        skip_nulls: bool,
106    ) -> PyResult<Series> {
107        let avs = call_and_collect_anyvalues(py, lambda, self.len(), self.into_iter(), skip_nulls)?;
108        Ok(
109            Series::from_any_values_and_dtype(self.name().clone(), &avs, datatype, true)
110                .map_err(PyPolarsErr::from)?,
111        )
112    }
113}
114
115impl<'py> ApplyLambdaGeneric<'py> for StringChunked {
116    fn apply_generic(
117        &self,
118        py: Python<'py>,
119        lambda: &Bound<'py, PyAny>,
120        skip_nulls: bool,
121    ) -> PyResult<Series> {
122        let avs = call_and_collect_anyvalues(py, lambda, self.len(), self.into_iter(), skip_nulls)?;
123        Ok(Series::from_any_values(self.name().clone(), &avs, true).map_err(PyPolarsErr::from)?)
124    }
125
126    fn apply_generic_with_dtype(
127        &self,
128        py: Python<'py>,
129        lambda: &Bound<'py, PyAny>,
130        datatype: &DataType,
131        skip_nulls: bool,
132    ) -> PyResult<Series> {
133        let avs = call_and_collect_anyvalues(py, lambda, self.len(), self.into_iter(), skip_nulls)?;
134        Ok(
135            Series::from_any_values_and_dtype(self.name().clone(), &avs, datatype, true)
136                .map_err(PyPolarsErr::from)?,
137        )
138    }
139}
140
141impl<'py> ApplyLambdaGeneric<'py> for ListChunked {
142    fn apply_generic(
143        &self,
144        py: Python<'py>,
145        lambda: &Bound<'py, PyAny>,
146        skip_nulls: bool,
147    ) -> PyResult<Series> {
148        let it = self.into_iter().map(|opt_s| opt_s.map(Wrap));
149        let avs = call_and_collect_anyvalues(py, lambda, self.len(), it, skip_nulls)?;
150        Ok(Series::from_any_values(self.name().clone(), &avs, true).map_err(PyPolarsErr::from)?)
151    }
152
153    fn apply_generic_with_dtype(
154        &self,
155        py: Python<'py>,
156        lambda: &Bound<'py, PyAny>,
157        datatype: &DataType,
158        skip_nulls: bool,
159    ) -> PyResult<Series> {
160        let it = self.into_iter().map(|opt_s| opt_s.map(Wrap));
161        let avs = call_and_collect_anyvalues(py, lambda, self.len(), it, skip_nulls)?;
162        Ok(
163            Series::from_any_values_and_dtype(self.name().clone(), &avs, datatype, true)
164                .map_err(PyPolarsErr::from)?,
165        )
166    }
167}
168
169#[cfg(feature = "dtype-array")]
170impl<'py> ApplyLambdaGeneric<'py> for ArrayChunked {
171    fn apply_generic(
172        &self,
173        py: Python<'py>,
174        lambda: &Bound<'py, PyAny>,
175        skip_nulls: bool,
176    ) -> PyResult<Series> {
177        let it = self.into_iter().map(|opt_s| Some(PySeries::new(opt_s?)));
178        let avs = call_and_collect_anyvalues(py, lambda, self.len(), it, skip_nulls)?;
179        Ok(Series::from_any_values(self.name().clone(), &avs, true).map_err(PyPolarsErr::from)?)
180    }
181
182    fn apply_generic_with_dtype(
183        &self,
184        py: Python<'py>,
185        lambda: &Bound<'py, PyAny>,
186        datatype: &DataType,
187        skip_nulls: bool,
188    ) -> PyResult<Series> {
189        let it = self.into_iter().map(|opt_s| Some(PySeries::new(opt_s?)));
190        let avs = call_and_collect_anyvalues(py, lambda, self.len(), it, skip_nulls)?;
191        Ok(
192            Series::from_any_values_and_dtype(self.name().clone(), &avs, datatype, true)
193                .map_err(PyPolarsErr::from)?,
194        )
195    }
196}
197
198#[cfg(feature = "object")]
199impl<'py> ApplyLambdaGeneric<'py> for ObjectChunked<ObjectValue> {
200    fn apply_generic(
201        &self,
202        py: Python<'py>,
203        lambda: &Bound<'py, PyAny>,
204        skip_nulls: bool,
205    ) -> PyResult<Series> {
206        let avs = call_and_collect_anyvalues(py, lambda, self.len(), self.into_iter(), skip_nulls)?;
207        Ok(Series::from_any_values(self.name().clone(), &avs, true).map_err(PyPolarsErr::from)?)
208    }
209
210    fn apply_generic_with_dtype(
211        &self,
212        py: Python<'py>,
213        lambda: &Bound<'py, PyAny>,
214        datatype: &DataType,
215        skip_nulls: bool,
216    ) -> PyResult<Series> {
217        let avs = call_and_collect_anyvalues(py, lambda, self.len(), self.into_iter(), skip_nulls)?;
218        Ok(
219            Series::from_any_values_and_dtype(self.name().clone(), &avs, datatype, true)
220                .map_err(PyPolarsErr::from)?,
221        )
222    }
223}
224
225impl<'py> ApplyLambdaGeneric<'py> for StructChunked {
226    fn apply_generic(
227        &self,
228        py: Python<'py>,
229        lambda: &Bound<'py, PyAny>,
230        skip_nulls: bool,
231    ) -> PyResult<Series> {
232        let it = (0..self.len())
233            .map(|i| unsafe { self.get_any_value_unchecked(i).null_to_none().map(Wrap) });
234        let avs = call_and_collect_anyvalues(py, lambda, self.len(), it, skip_nulls)?;
235        Ok(Series::from_any_values(self.name().clone(), &avs, true).map_err(PyPolarsErr::from)?)
236    }
237
238    fn apply_generic_with_dtype(
239        &self,
240        py: Python<'py>,
241        lambda: &Bound<'py, PyAny>,
242        datatype: &DataType,
243        skip_nulls: bool,
244    ) -> PyResult<Series> {
245        let it = (0..self.len())
246            .map(|i| unsafe { self.get_any_value_unchecked(i).null_to_none().map(Wrap) });
247        let avs = call_and_collect_anyvalues(py, lambda, self.len(), it, skip_nulls)?;
248        Ok(
249            Series::from_any_values_and_dtype(self.name().clone(), &avs, datatype, true)
250                .map_err(PyPolarsErr::from)?,
251        )
252    }
253}
254
255impl<'py> ApplyLambdaGeneric<'py> for BinaryChunked {
256    fn apply_generic(
257        &self,
258        py: Python<'py>,
259        lambda: &Bound<'py, PyAny>,
260        skip_nulls: bool,
261    ) -> PyResult<Series> {
262        let avs = call_and_collect_anyvalues(py, lambda, self.len(), self.into_iter(), skip_nulls)?;
263        Ok(Series::from_any_values(self.name().clone(), &avs, true).map_err(PyPolarsErr::from)?)
264    }
265
266    fn apply_generic_with_dtype(
267        &self,
268        py: Python<'py>,
269        lambda: &Bound<'py, PyAny>,
270        datatype: &DataType,
271        skip_nulls: bool,
272    ) -> PyResult<Series> {
273        let avs = call_and_collect_anyvalues(py, lambda, self.len(), self.into_iter(), skip_nulls)?;
274        Ok(
275            Series::from_any_values_and_dtype(self.name().clone(), &avs, datatype, true)
276                .map_err(PyPolarsErr::from)?,
277        )
278    }
279}
280
281impl<'py, L, P> ApplyLambdaGeneric<'py> for Logical<L, P>
282where
283    L: PolarsDataType,
284    P: PolarsDataType,
285    Logical<L, P>: LogicalType,
286{
287    fn apply_generic(
288        &self,
289        py: Python<'py>,
290        lambda: &Bound<'py, PyAny>,
291        skip_nulls: bool,
292    ) -> PyResult<Series> {
293        let it = (0..self.len())
294            .map(|i| unsafe { self.get_any_value_unchecked(i).null_to_none().map(Wrap) });
295        let avs = call_and_collect_anyvalues(py, lambda, self.len(), it, skip_nulls)?;
296        Ok(Series::from_any_values(self.name().clone(), &avs, true).map_err(PyPolarsErr::from)?)
297    }
298
299    fn apply_generic_with_dtype(
300        &self,
301        py: Python<'py>,
302        lambda: &Bound<'py, PyAny>,
303        datatype: &DataType,
304        skip_nulls: bool,
305    ) -> PyResult<Series> {
306        let it = (0..self.len())
307            .map(|i| unsafe { self.get_any_value_unchecked(i).null_to_none().map(Wrap) });
308        let avs = call_and_collect_anyvalues(py, lambda, self.len(), it, skip_nulls)?;
309        Ok(
310            Series::from_any_values_and_dtype(self.name().clone(), &avs, datatype, true)
311                .map_err(PyPolarsErr::from)?,
312        )
313    }
314}
315
316impl<'py> ApplyLambdaGeneric<'py> for NullChunked {
317    fn apply_generic(
318        &self,
319        py: Python<'py>,
320        lambda: &Bound<'py, PyAny>,
321        skip_nulls: bool,
322    ) -> PyResult<Series> {
323        let it = (0..self.len()).map(|_| None::<Wrap<AnyValue<'static>>>);
324        let avs = call_and_collect_anyvalues(py, lambda, self.len(), it, skip_nulls)?;
325        Ok(Series::from_any_values(self.name().clone(), &avs, true).map_err(PyPolarsErr::from)?)
326    }
327
328    fn apply_generic_with_dtype(
329        &self,
330        py: Python<'py>,
331        lambda: &Bound<'py, PyAny>,
332        datatype: &DataType,
333        skip_nulls: bool,
334    ) -> PyResult<Series> {
335        let it = (0..self.len()).map(|_| None::<Wrap<AnyValue<'static>>>);
336        let avs = call_and_collect_anyvalues(py, lambda, self.len(), it, skip_nulls)?;
337        Ok(
338            Series::from_any_values_and_dtype(self.name().clone(), &avs, datatype, true)
339                .map_err(PyPolarsErr::from)?,
340        )
341    }
342}
343
344impl<'py> ApplyLambdaGeneric<'py> for ExtensionChunked {
345    fn apply_generic(
346        &self,
347        _py: Python<'py>,
348        _lambda: &Bound<'py, PyAny>,
349        _skip_nulls: bool,
350    ) -> PyResult<Series> {
351        unreachable!()
352    }
353
354    fn apply_generic_with_dtype(
355        &self,
356        _py: Python<'py>,
357        _lambda: &Bound<'py, PyAny>,
358        _datatype: &DataType,
359        _skip_nulls: bool,
360    ) -> PyResult<Series> {
361        unreachable!()
362    }
363}
polars_python/map/series.rs

polars_python/map/
series.rs
