polars_python/map/
series.rs

1use pyo3::prelude::*;
2use pyo3::pybacked::PyBackedStr;
3use pyo3::types::{PyBool, PyCFunction, PyFloat, PyList, PyString, PyTuple};
4
5use super::*;
6use crate::py_modules::{pl_series, polars};
7
8/// Find the output type and dispatch to that implementation.
9fn infer_and_finish<'py, A: ApplyLambda<'py>>(
10    applyer: &A,
11    py: Python<'py>,
12    lambda: &Bound<'py, PyAny>,
13    out: &Bound<'py, PyAny>,
14    null_count: usize,
15) -> PyResult<PySeries> {
16    if out.is_instance_of::<PyBool>() {
17        let first_value = out.extract::<bool>().unwrap();
18        applyer
19            .apply_lambda_with_bool_out_type(py, lambda, null_count, Some(first_value))
20            .map(|ca| ca.into_series().into())
21    } else if out.is_instance_of::<PyFloat>() {
22        let first_value = out.extract::<f64>().unwrap();
23        applyer
24            .apply_lambda_with_primitive_out_type::<Float64Type>(
25                py,
26                lambda,
27                null_count,
28                Some(first_value),
29            )
30            .map(|ca| ca.into_series().into())
31    } else if out.is_instance_of::<PyString>() {
32        let first_value = out.extract::<PyBackedStr>().unwrap();
33        applyer
34            .apply_lambda_with_string_out_type(py, lambda, null_count, Some(first_value))
35            .map(|ca| ca.into_series().into())
36    } else if out.hasattr("_s")? {
37        let py_pyseries = out.getattr("_s").unwrap();
38        let series = py_pyseries.extract::<PySeries>().unwrap().series;
39        let dt = series.dtype();
40        applyer
41            .apply_lambda_with_list_out_type(
42                py,
43                lambda.to_owned().unbind(),
44                null_count,
45                Some(&series),
46                dt,
47            )
48            .map(|ca| ca.into_series().into())
49    } else if out.is_instance_of::<PyList>() || out.is_instance_of::<PyTuple>() {
50        let series = pl_series(py).call1(py, (out,))?;
51        let py_pyseries = series.getattr(py, "_s").unwrap();
52        let series = py_pyseries.extract::<PySeries>(py).unwrap().series;
53
54        let dt = series.dtype();
55
56        // Null dtype may be incorrect, fall back to AnyValues logic.
57        if dt.is_nested_null() {
58            let av = out.extract::<Wrap<AnyValue>>()?;
59            return applyer
60                .apply_extract_any_values(py, lambda, null_count, av.0)
61                .map(|s| s.into());
62        }
63
64        // make a new python function that is:
65        // def new_lambda(lambda: Callable):
66        //     pl.Series(lambda(value))
67        let lambda_owned = lambda.to_owned().unbind();
68        let new_lambda = PyCFunction::new_closure(py, None, None, move |args, _kwargs| {
69            Python::with_gil(|py| {
70                let out = lambda_owned.call1(py, args)?;
71                // check if Series, if not, call series constructor on it
72                pl_series(py).call1(py, (out,))
73            })
74        })?
75        .into_any()
76        .unbind();
77
78        let result = applyer
79            .apply_lambda_with_list_out_type(py, new_lambda, null_count, Some(&series), dt)
80            .map(|ca| ca.into_series().into());
81        match result {
82            Ok(out) => Ok(out),
83            // Try AnyValue
84            Err(_) => {
85                let av = out.extract::<Wrap<AnyValue>>()?;
86                applyer
87                    .apply_extract_any_values(py, lambda, null_count, av.0)
88                    .map(|s| s.into())
89            },
90        }
91    } else if out.is_instance_of::<PyDict>() {
92        let first = out.extract::<Wrap<AnyValue<'_>>>()?;
93        applyer.apply_into_struct(py, lambda, null_count, first.0)
94    }
95    // this succeeds for numpy ints as well, where checking if it is pyint fails
96    // we do this later in the chain so that we don't extract integers from string chars.
97    else if out.extract::<i64>().is_ok() {
98        let first_value = out.extract::<i64>().unwrap();
99        applyer
100            .apply_lambda_with_primitive_out_type::<Int64Type>(
101                py,
102                lambda,
103                null_count,
104                Some(first_value),
105            )
106            .map(|ca| ca.into_series().into())
107    } else if let Ok(av) = out.extract::<Wrap<AnyValue>>() {
108        applyer
109            .apply_extract_any_values(py, lambda, null_count, av.0)
110            .map(|s| s.into())
111    } else {
112        #[cfg(feature = "object")]
113        {
114            applyer
115                .apply_lambda_with_object_out_type(
116                    py,
117                    lambda,
118                    null_count,
119                    Some(out.to_owned().unbind().into()),
120                )
121                .map(|ca| ca.into_series().into())
122        }
123        #[cfg(not(feature = "object"))]
124        {
125            todo!()
126        }
127    }
128}
129
130pub trait ApplyLambda<'py> {
131    fn apply_lambda_unknown(
132        &self,
133        _py: Python<'py>,
134        _lambda: &Bound<'py, PyAny>,
135    ) -> PyResult<PySeries>;
136
137    // Used to store a struct type
138    fn apply_into_struct(
139        &self,
140        py: Python<'py>,
141        lambda: &Bound<'py, PyAny>,
142        init_null_count: usize,
143        first_value: AnyValue<'static>,
144    ) -> PyResult<PySeries>;
145
146    /// Apply a lambda with a primitive output type
147    fn apply_lambda_with_primitive_out_type<D>(
148        &self,
149        py: Python<'py>,
150        lambda: &Bound<'py, PyAny>,
151        init_null_count: usize,
152        first_value: Option<D::Native>,
153    ) -> PyResult<ChunkedArray<D>>
154    where
155        D: PyPolarsNumericType,
156        D::Native: IntoPyObject<'py> + FromPyObject<'py>;
157
158    /// Apply a lambda with a boolean output type
159    fn apply_lambda_with_bool_out_type(
160        &self,
161        py: Python<'py>,
162        lambda: &Bound<'py, PyAny>,
163        init_null_count: usize,
164        first_value: Option<bool>,
165    ) -> PyResult<ChunkedArray<BooleanType>>;
166
167    /// Apply a lambda with string output type
168    fn apply_lambda_with_string_out_type(
169        &self,
170        py: Python<'py>,
171        lambda: &Bound<'py, PyAny>,
172        init_null_count: usize,
173        first_value: Option<PyBackedStr>,
174    ) -> PyResult<StringChunked>;
175
176    /// Apply a lambda with list output type
177    fn apply_lambda_with_list_out_type(
178        &self,
179        py: Python<'py>,
180        lambda: PyObject,
181        init_null_count: usize,
182        first_value: Option<&Series>,
183        dt: &DataType,
184    ) -> PyResult<ListChunked>;
185
186    fn apply_extract_any_values(
187        &self,
188        py: Python<'py>,
189        lambda: &Bound<'py, PyAny>,
190        init_null_count: usize,
191        first_value: AnyValue<'static>,
192    ) -> PyResult<Series>;
193
194    /// Apply a lambda with list output type
195    #[cfg(feature = "object")]
196    fn apply_lambda_with_object_out_type(
197        &self,
198        py: Python<'py>,
199        lambda: &Bound<'py, PyAny>,
200        init_null_count: usize,
201        first_value: Option<ObjectValue>,
202    ) -> PyResult<ObjectChunked<ObjectValue>>;
203}
204
205pub fn call_lambda<'py, T>(
206    py: Python<'py>,
207    lambda: &Bound<'py, PyAny>,
208    in_val: T,
209) -> PyResult<Bound<'py, PyAny>>
210where
211    T: IntoPyObject<'py>,
212{
213    let arg = PyTuple::new(py, [in_val])?;
214    lambda.call1(arg)
215}
216
217pub(crate) fn call_lambda_and_extract<'py, T, S>(
218    py: Python<'py>,
219    lambda: &Bound<'py, PyAny>,
220    in_val: T,
221) -> PyResult<Option<S>>
222where
223    T: IntoPyObject<'py>,
224    S: FromPyObject<'py>,
225{
226    let out = call_lambda(py, lambda, in_val)?;
227    if out.is_none() {
228        Ok(None)
229    } else {
230        out.extract::<S>().map(Some)
231    }
232}
233
234fn call_lambda_series_out<'py, T>(
235    py: Python<'py>,
236    lambda: &Bound<'py, PyAny>,
237    in_val: T,
238) -> PyResult<Series>
239where
240    T: IntoPyObject<'py>,
241{
242    let arg = PyTuple::new(py, [in_val])?;
243    let out = lambda.call1(arg)?;
244    let py_series = out.getattr("_s")?;
245    py_series.extract::<PySeries>().map(|s| s.series)
246}
247
248fn extract_anyvalues<'py, T, I>(
249    py: Python<'py>,
250    lambda: &Bound<'py, PyAny>,
251    len: usize,
252    init_null_count: usize,
253    iter: I,
254    first_value: AnyValue<'static>,
255) -> PyResult<Vec<AnyValue<'static>>>
256where
257    T: IntoPyObject<'py>,
258    I: Iterator<Item = Option<T>>,
259{
260    let mut avs = Vec::with_capacity(len);
261    avs.extend(std::iter::repeat_n(AnyValue::Null, init_null_count));
262    avs.push(first_value);
263
264    for opt_val in iter {
265        let av = match opt_val {
266            None => AnyValue::Null,
267            Some(val) => {
268                let val: Option<Wrap<AnyValue>> = call_lambda_and_extract(py, lambda, val)?;
269                match val {
270                    None => AnyValue::Null,
271                    Some(av) => av.0,
272                }
273            },
274        };
275        avs.push(av)
276    }
277    Ok(avs)
278}
279
280impl<'py> ApplyLambda<'py> for BooleanChunked {
281    fn apply_lambda_unknown(
282        &self,
283        py: Python<'py>,
284        lambda: &Bound<'py, PyAny>,
285    ) -> PyResult<PySeries> {
286        let mut null_count = 0;
287        for opt_v in self.into_iter() {
288            if let Some(v) = opt_v {
289                let arg = PyTuple::new(py, [v])?;
290                let out = lambda.call1(arg)?;
291                if out.is_none() {
292                    null_count += 1;
293                    continue;
294                }
295                return infer_and_finish(self, py, lambda, &out, null_count);
296            } else {
297                null_count += 1
298            }
299        }
300        Ok(Self::full_null(self.name().clone(), self.len())
301            .into_series()
302            .into())
303    }
304
305    fn apply_into_struct(
306        &self,
307        py: Python<'py>,
308        lambda: &Bound<'py, PyAny>,
309        init_null_count: usize,
310        first_value: AnyValue<'static>,
311    ) -> PyResult<PySeries> {
312        let skip = 1;
313        let it = self
314            .into_iter()
315            .skip(init_null_count + skip)
316            .map(|opt_val| opt_val.map(|val| call_lambda(py, lambda, val)).transpose());
317        iterator_to_struct(
318            py,
319            it,
320            init_null_count,
321            first_value,
322            self.name().clone(),
323            self.len(),
324        )
325    }
326
327    fn apply_lambda_with_primitive_out_type<D>(
328        &self,
329        py: Python<'py>,
330        lambda: &Bound<'py, PyAny>,
331        init_null_count: usize,
332        first_value: Option<D::Native>,
333    ) -> PyResult<ChunkedArray<D>>
334    where
335        D: PyPolarsNumericType,
336        D::Native: IntoPyObject<'py> + FromPyObject<'py>,
337    {
338        let skip = usize::from(first_value.is_some());
339        if init_null_count == self.len() {
340            Ok(ChunkedArray::full_null(self.name().clone(), self.len()))
341        } else {
342            let it = self
343                .into_iter()
344                .skip(init_null_count + skip)
345                .map(|opt_val| {
346                    opt_val
347                        .and_then(|val| call_lambda_and_extract(py, lambda, val).transpose())
348                        .transpose()
349                });
350            iterator_to_primitive(
351                it,
352                init_null_count,
353                first_value,
354                self.name().clone(),
355                self.len(),
356            )
357        }
358    }
359
360    fn apply_lambda_with_bool_out_type(
361        &self,
362        py: Python<'py>,
363        lambda: &Bound<'py, PyAny>,
364        init_null_count: usize,
365        first_value: Option<bool>,
366    ) -> PyResult<BooleanChunked> {
367        let skip = usize::from(first_value.is_some());
368        if init_null_count == self.len() {
369            Ok(ChunkedArray::full_null(self.name().clone(), self.len()))
370        } else {
371            let it = self
372                .into_iter()
373                .skip(init_null_count + skip)
374                .map(|opt_val| {
375                    opt_val
376                        .and_then(|val| call_lambda_and_extract(py, lambda, val).transpose())
377                        .transpose()
378                });
379            iterator_to_bool(
380                it,
381                init_null_count,
382                first_value,
383                self.name().clone(),
384                self.len(),
385            )
386        }
387    }
388
389    fn apply_lambda_with_string_out_type(
390        &self,
391        py: Python<'py>,
392        lambda: &Bound<'py, PyAny>,
393        init_null_count: usize,
394        first_value: Option<PyBackedStr>,
395    ) -> PyResult<StringChunked> {
396        let skip = usize::from(first_value.is_some());
397        if init_null_count == self.len() {
398            Ok(ChunkedArray::full_null(self.name().clone(), self.len()))
399        } else {
400            let it = self
401                .into_iter()
402                .skip(init_null_count + skip)
403                .map(|opt_val| {
404                    opt_val
405                        .and_then(|val| call_lambda_and_extract(py, lambda, val).transpose())
406                        .transpose()
407                });
408            iterator_to_string(
409                it,
410                init_null_count,
411                first_value,
412                self.name().clone(),
413                self.len(),
414            )
415        }
416    }
417
418    fn apply_lambda_with_list_out_type(
419        &self,
420        py: Python<'py>,
421        lambda: PyObject,
422        init_null_count: usize,
423        first_value: Option<&Series>,
424        dt: &DataType,
425    ) -> PyResult<ListChunked> {
426        let skip = usize::from(first_value.is_some());
427        let lambda = lambda.bind(py);
428        if init_null_count == self.len() {
429            Ok(ChunkedArray::full_null(self.name().clone(), self.len()))
430        } else {
431            let it = self
432                .into_iter()
433                .skip(init_null_count + skip)
434                .map(|opt_val| {
435                    opt_val
436                        .map(|val| call_lambda_series_out(py, lambda, val))
437                        .transpose()
438                });
439            iterator_to_list(
440                dt,
441                it,
442                init_null_count,
443                first_value,
444                self.name().clone(),
445                self.len(),
446            )
447        }
448    }
449
450    fn apply_extract_any_values(
451        &self,
452        py: Python<'py>,
453        lambda: &Bound<'py, PyAny>,
454        init_null_count: usize,
455        first_value: AnyValue<'static>,
456    ) -> PyResult<Series> {
457        let iter = self.into_iter().skip(init_null_count + 1);
458        let avs = extract_anyvalues(py, lambda, self.len(), init_null_count, iter, first_value)?;
459
460        Ok(Series::new(self.name().clone(), &avs))
461    }
462
463    #[cfg(feature = "object")]
464    fn apply_lambda_with_object_out_type(
465        &self,
466        py: Python<'py>,
467        lambda: &Bound<'py, PyAny>,
468        init_null_count: usize,
469        first_value: Option<ObjectValue>,
470    ) -> PyResult<ObjectChunked<ObjectValue>> {
471        let skip = usize::from(first_value.is_some());
472        if init_null_count == self.len() {
473            Ok(ChunkedArray::full_null(self.name().clone(), self.len()))
474        } else {
475            let it = self
476                .into_iter()
477                .skip(init_null_count + skip)
478                .map(|opt_val| {
479                    opt_val
480                        .and_then(|val| call_lambda_and_extract(py, lambda, val).transpose())
481                        .transpose()
482                });
483            iterator_to_object(
484                it,
485                init_null_count,
486                first_value,
487                self.name().clone(),
488                self.len(),
489            )
490        }
491    }
492}
493
494impl<'py, T> ApplyLambda<'py> for ChunkedArray<T>
495where
496    T: PyPolarsNumericType,
497    T::Native: IntoPyObject<'py> + FromPyObject<'py>,
498    ChunkedArray<T>: IntoSeries,
499{
500    fn apply_lambda_unknown(
501        &self,
502        py: Python<'py>,
503        lambda: &Bound<'py, PyAny>,
504    ) -> PyResult<PySeries> {
505        let mut null_count = 0;
506        for opt_v in self.into_iter() {
507            if let Some(v) = opt_v {
508                let arg = PyTuple::new(py, [v])?;
509                let out = lambda.call1(arg)?;
510                if out.is_none() {
511                    null_count += 1;
512                    continue;
513                }
514                return infer_and_finish(self, py, lambda, &out, null_count);
515            } else {
516                null_count += 1
517            }
518        }
519        Ok(Self::full_null(self.name().clone(), self.len())
520            .into_series()
521            .into())
522    }
523
524    fn apply_into_struct(
525        &self,
526        py: Python<'py>,
527        lambda: &Bound<'py, PyAny>,
528        init_null_count: usize,
529        first_value: AnyValue<'static>,
530    ) -> PyResult<PySeries> {
531        let skip = 1;
532
533        let it = self
534            .into_iter()
535            .skip(init_null_count + skip)
536            .map(|opt_val| opt_val.map(|val| call_lambda(py, lambda, val)).transpose());
537        iterator_to_struct(
538            py,
539            it,
540            init_null_count,
541            first_value,
542            self.name().clone(),
543            self.len(),
544        )
545    }
546
547    fn apply_lambda_with_primitive_out_type<D>(
548        &self,
549        py: Python<'py>,
550        lambda: &Bound<'py, PyAny>,
551        init_null_count: usize,
552        first_value: Option<D::Native>,
553    ) -> PyResult<ChunkedArray<D>>
554    where
555        D: PyPolarsNumericType,
556        D::Native: IntoPyObject<'py> + FromPyObject<'py>,
557    {
558        let skip = usize::from(first_value.is_some());
559        if init_null_count == self.len() {
560            Ok(ChunkedArray::full_null(self.name().clone(), self.len()))
561        } else {
562            let it = self
563                .into_iter()
564                .skip(init_null_count + skip)
565                .map(|opt_val| {
566                    opt_val
567                        .and_then(|val| call_lambda_and_extract(py, lambda, val).transpose())
568                        .transpose()
569                });
570            iterator_to_primitive(
571                it,
572                init_null_count,
573                first_value,
574                self.name().clone(),
575                self.len(),
576            )
577        }
578    }
579
580    fn apply_lambda_with_bool_out_type(
581        &self,
582        py: Python<'py>,
583        lambda: &Bound<'py, PyAny>,
584        init_null_count: usize,
585        first_value: Option<bool>,
586    ) -> PyResult<BooleanChunked> {
587        let skip = usize::from(first_value.is_some());
588        if init_null_count == self.len() {
589            Ok(ChunkedArray::full_null(self.name().clone(), self.len()))
590        } else {
591            let it = self
592                .into_iter()
593                .skip(init_null_count + skip)
594                .map(|opt_val| {
595                    opt_val
596                        .and_then(|val| call_lambda_and_extract(py, lambda, val).transpose())
597                        .transpose()
598                });
599            iterator_to_bool(
600                it,
601                init_null_count,
602                first_value,
603                self.name().clone(),
604                self.len(),
605            )
606        }
607    }
608
609    fn apply_lambda_with_string_out_type(
610        &self,
611        py: Python<'py>,
612        lambda: &Bound<'py, PyAny>,
613        init_null_count: usize,
614        first_value: Option<PyBackedStr>,
615    ) -> PyResult<StringChunked> {
616        let skip = usize::from(first_value.is_some());
617        if init_null_count == self.len() {
618            Ok(ChunkedArray::full_null(self.name().clone(), self.len()))
619        } else {
620            let it = self
621                .into_iter()
622                .skip(init_null_count + skip)
623                .map(|opt_val| {
624                    opt_val
625                        .and_then(|val| call_lambda_and_extract(py, lambda, val).transpose())
626                        .transpose()
627                });
628            iterator_to_string(
629                it,
630                init_null_count,
631                first_value,
632                self.name().clone(),
633                self.len(),
634            )
635        }
636    }
637
638    fn apply_lambda_with_list_out_type(
639        &self,
640        py: Python<'py>,
641        lambda: PyObject,
642        init_null_count: usize,
643        first_value: Option<&Series>,
644        dt: &DataType,
645    ) -> PyResult<ListChunked> {
646        let skip = usize::from(first_value.is_some());
647        let lambda = lambda.bind(py);
648        if init_null_count == self.len() {
649            Ok(ChunkedArray::full_null(self.name().clone(), self.len()))
650        } else {
651            let it = self
652                .into_iter()
653                .skip(init_null_count + skip)
654                .map(|opt_val| {
655                    opt_val
656                        .map(|val| call_lambda_series_out(py, lambda, val))
657                        .transpose()
658                });
659            iterator_to_list(
660                dt,
661                it,
662                init_null_count,
663                first_value,
664                self.name().clone(),
665                self.len(),
666            )
667        }
668    }
669
670    fn apply_extract_any_values(
671        &self,
672        py: Python<'py>,
673        lambda: &Bound<'py, PyAny>,
674        init_null_count: usize,
675        first_value: AnyValue<'static>,
676    ) -> PyResult<Series> {
677        let iter = self.into_iter().skip(init_null_count + 1);
678        let avs = extract_anyvalues(py, lambda, self.len(), init_null_count, iter, first_value)?;
679
680        Ok(Series::new(self.name().clone(), &avs))
681    }
682
683    #[cfg(feature = "object")]
684    fn apply_lambda_with_object_out_type(
685        &self,
686        py: Python<'py>,
687        lambda: &Bound<'py, PyAny>,
688        init_null_count: usize,
689        first_value: Option<ObjectValue>,
690    ) -> PyResult<ObjectChunked<ObjectValue>> {
691        let skip = usize::from(first_value.is_some());
692        if init_null_count == self.len() {
693            Ok(ChunkedArray::full_null(self.name().clone(), self.len()))
694        } else {
695            let it = self
696                .into_iter()
697                .skip(init_null_count + skip)
698                .map(|opt_val| {
699                    opt_val
700                        .and_then(|val| call_lambda_and_extract(py, lambda, val).transpose())
701                        .transpose()
702                });
703            iterator_to_object(
704                it,
705                init_null_count,
706                first_value,
707                self.name().clone(),
708                self.len(),
709            )
710        }
711    }
712}
713
714impl<'py> ApplyLambda<'py> for StringChunked {
715    fn apply_lambda_unknown(
716        &self,
717        py: Python<'py>,
718        lambda: &Bound<'py, PyAny>,
719    ) -> PyResult<PySeries> {
720        let mut null_count = 0;
721        for opt_v in self.into_iter() {
722            if let Some(v) = opt_v {
723                let arg = PyTuple::new(py, [v])?;
724                let out = lambda.call1(arg)?;
725                if out.is_none() {
726                    null_count += 1;
727                    continue;
728                }
729                return infer_and_finish(self, py, lambda, &out, null_count);
730            } else {
731                null_count += 1
732            }
733        }
734        Ok(Self::full_null(self.name().clone(), self.len())
735            .into_series()
736            .into())
737    }
738
739    fn apply_into_struct(
740        &self,
741        py: Python<'py>,
742        lambda: &Bound<'py, PyAny>,
743        init_null_count: usize,
744        first_value: AnyValue<'static>,
745    ) -> PyResult<PySeries> {
746        let skip = 1;
747
748        let it = self
749            .into_iter()
750            .skip(init_null_count + skip)
751            .map(|opt_val| opt_val.map(|val| call_lambda(py, lambda, val)).transpose());
752        iterator_to_struct(
753            py,
754            it,
755            init_null_count,
756            first_value,
757            self.name().clone(),
758            self.len(),
759        )
760    }
761
762    fn apply_lambda_with_primitive_out_type<D>(
763        &self,
764        py: Python<'py>,
765        lambda: &Bound<'py, PyAny>,
766        init_null_count: usize,
767        first_value: Option<D::Native>,
768    ) -> PyResult<ChunkedArray<D>>
769    where
770        D: PyPolarsNumericType,
771        D::Native: IntoPyObject<'py> + FromPyObject<'py>,
772    {
773        let skip = usize::from(first_value.is_some());
774        if init_null_count == self.len() {
775            Ok(ChunkedArray::full_null(self.name().clone(), self.len()))
776        } else {
777            let it = self
778                .into_iter()
779                .skip(init_null_count + skip)
780                .map(|opt_val| {
781                    opt_val
782                        .and_then(|val| call_lambda_and_extract(py, lambda, val).transpose())
783                        .transpose()
784                });
785            iterator_to_primitive(
786                it,
787                init_null_count,
788                first_value,
789                self.name().clone(),
790                self.len(),
791            )
792        }
793    }
794
795    fn apply_lambda_with_bool_out_type(
796        &self,
797        py: Python<'py>,
798        lambda: &Bound<'py, PyAny>,
799        init_null_count: usize,
800        first_value: Option<bool>,
801    ) -> PyResult<BooleanChunked> {
802        let skip = usize::from(first_value.is_some());
803        if init_null_count == self.len() {
804            Ok(ChunkedArray::full_null(self.name().clone(), self.len()))
805        } else {
806            let it = self
807                .into_iter()
808                .skip(init_null_count + skip)
809                .map(|opt_val| {
810                    opt_val
811                        .and_then(|val| call_lambda_and_extract(py, lambda, val).transpose())
812                        .transpose()
813                });
814            iterator_to_bool(
815                it,
816                init_null_count,
817                first_value,
818                self.name().clone(),
819                self.len(),
820            )
821        }
822    }
823
824    fn apply_lambda_with_string_out_type(
825        &self,
826        py: Python<'py>,
827        lambda: &Bound<'py, PyAny>,
828        init_null_count: usize,
829        first_value: Option<PyBackedStr>,
830    ) -> PyResult<StringChunked> {
831        let skip = usize::from(first_value.is_some());
832        if init_null_count == self.len() {
833            Ok(ChunkedArray::full_null(self.name().clone(), self.len()))
834        } else {
835            let it = self
836                .into_iter()
837                .skip(init_null_count + skip)
838                .map(|opt_val| {
839                    opt_val
840                        .and_then(|val| call_lambda_and_extract(py, lambda, val).transpose())
841                        .transpose()
842                });
843            iterator_to_string(
844                it,
845                init_null_count,
846                first_value,
847                self.name().clone(),
848                self.len(),
849            )
850        }
851    }
852    fn apply_lambda_with_list_out_type(
853        &self,
854        py: Python<'py>,
855        lambda: PyObject,
856        init_null_count: usize,
857        first_value: Option<&Series>,
858        dt: &DataType,
859    ) -> PyResult<ListChunked> {
860        let skip = usize::from(first_value.is_some());
861        let lambda = lambda.bind(py);
862        if init_null_count == self.len() {
863            Ok(ChunkedArray::full_null(self.name().clone(), self.len()))
864        } else {
865            let it = self
866                .into_iter()
867                .skip(init_null_count + skip)
868                .map(|opt_val| {
869                    opt_val
870                        .map(|val| call_lambda_series_out(py, lambda, val))
871                        .transpose()
872                });
873            iterator_to_list(
874                dt,
875                it,
876                init_null_count,
877                first_value,
878                self.name().clone(),
879                self.len(),
880            )
881        }
882    }
883
884    fn apply_extract_any_values(
885        &self,
886        py: Python<'py>,
887        lambda: &Bound<'py, PyAny>,
888        init_null_count: usize,
889        first_value: AnyValue<'static>,
890    ) -> PyResult<Series> {
891        let iter = self.into_iter().skip(init_null_count + 1);
892        let avs = extract_anyvalues(py, lambda, self.len(), init_null_count, iter, first_value)?;
893        Ok(Series::new(self.name().clone(), &avs))
894    }
895
896    #[cfg(feature = "object")]
897    fn apply_lambda_with_object_out_type(
898        &self,
899        py: Python<'py>,
900        lambda: &Bound<'py, PyAny>,
901        init_null_count: usize,
902        first_value: Option<ObjectValue>,
903    ) -> PyResult<ObjectChunked<ObjectValue>> {
904        let skip = usize::from(first_value.is_some());
905        if init_null_count == self.len() {
906            Ok(ChunkedArray::full_null(self.name().clone(), self.len()))
907        } else {
908            let it = self
909                .into_iter()
910                .skip(init_null_count + skip)
911                .map(|opt_val| {
912                    opt_val
913                        .and_then(|val| call_lambda_and_extract(py, lambda, val).transpose())
914                        .transpose()
915                });
916            iterator_to_object(
917                it,
918                init_null_count,
919                first_value,
920                self.name().clone(),
921                self.len(),
922            )
923        }
924    }
925}
926
927fn call_series_lambda(
928    pypolars: &Bound<PyModule>,
929    lambda: &Bound<PyAny>,
930    series: Series,
931) -> PyResult<Option<Series>> {
932    // create a PySeries struct/object for Python
933    let pyseries = PySeries::new(series);
934    // Wrap this PySeries object in the python side Series wrapper
935    let python_series_wrapper = pypolars
936        .getattr("wrap_s")
937        .unwrap()
938        .call1((pyseries,))
939        .unwrap();
940
941    // call the lambda en get a python side Series wrapper
942    let out = lambda.call1((python_series_wrapper,))?;
943    // unpack the wrapper in a PySeries
944    let py_pyseries = out
945        .getattr("_s")
946        .expect("could not get Series attribute '_s'");
947    Ok(py_pyseries.extract::<PySeries>().ok().map(|s| s.series))
948}
949
950impl<'py> ApplyLambda<'py> for ListChunked {
951    fn apply_lambda_unknown(
952        &self,
953        py: Python<'py>,
954        lambda: &Bound<'py, PyAny>,
955    ) -> PyResult<PySeries> {
956        let pypolars = polars(py).bind(py);
957        let mut null_count = 0;
958        for opt_v in self.into_iter() {
959            if let Some(v) = opt_v {
960                // create a PySeries struct/object for Python
961                let pyseries = PySeries::new(v);
962                // Wrap this PySeries object in the python side Series wrapper
963                let python_series_wrapper = pypolars
964                    .getattr("wrap_s")
965                    .unwrap()
966                    .call1((pyseries,))
967                    .unwrap();
968
969                let out = lambda.call1((python_series_wrapper,))?;
970                if out.is_none() {
971                    null_count += 1;
972                    continue;
973                }
974                return infer_and_finish(self, py, lambda, &out, null_count);
975            } else {
976                null_count += 1
977            }
978        }
979        Ok(Self::full_null(self.name().clone(), self.len())
980            .into_series()
981            .into())
982    }
983
984    fn apply_into_struct(
985        &self,
986        py: Python<'py>,
987        lambda: &Bound<'py, PyAny>,
988        init_null_count: usize,
989        first_value: AnyValue<'static>,
990    ) -> PyResult<PySeries> {
991        let skip = 1;
992        // get the pypolars module
993        let pypolars = polars(py).bind(py);
994
995        let it = self
996            .into_iter()
997            .skip(init_null_count + skip)
998            .map(|opt_val| {
999                opt_val
1000                    .map(|val| {
1001                        // create a PySeries struct/object for Python
1002                        let pyseries = PySeries::new(val);
1003                        // Wrap this PySeries object in the python side Series wrapper
1004                        let python_series_wrapper = pypolars
1005                            .getattr("wrap_s")
1006                            .unwrap()
1007                            .call1((pyseries,))
1008                            .unwrap();
1009                        call_lambda(py, lambda, python_series_wrapper)
1010                    })
1011                    .transpose()
1012            });
1013        iterator_to_struct(
1014            py,
1015            it,
1016            init_null_count,
1017            first_value,
1018            self.name().clone(),
1019            self.len(),
1020        )
1021    }
1022
1023    fn apply_lambda_with_primitive_out_type<D>(
1024        &self,
1025        py: Python<'py>,
1026        lambda: &Bound<'py, PyAny>,
1027        init_null_count: usize,
1028        first_value: Option<D::Native>,
1029    ) -> PyResult<ChunkedArray<D>>
1030    where
1031        D: PyPolarsNumericType,
1032        D::Native: IntoPyObject<'py> + FromPyObject<'py>,
1033    {
1034        let skip = usize::from(first_value.is_some());
1035        let pypolars = polars(py).bind(py);
1036        if init_null_count == self.len() {
1037            Ok(ChunkedArray::full_null(self.name().clone(), self.len()))
1038        } else {
1039            let it = self
1040                .into_iter()
1041                .skip(init_null_count + skip)
1042                .map(|opt_val| {
1043                    opt_val
1044                        .and_then(|val| {
1045                            // create a PySeries struct/object for Python
1046                            let pyseries = PySeries::new(val);
1047                            // Wrap this PySeries object in the python side Series wrapper
1048                            let python_series_wrapper = pypolars
1049                                .getattr("wrap_s")
1050                                .unwrap()
1051                                .call1((pyseries,))
1052                                .unwrap();
1053                            call_lambda_and_extract(py, lambda, python_series_wrapper).transpose()
1054                        })
1055                        .transpose()
1056                });
1057            iterator_to_primitive(
1058                it,
1059                init_null_count,
1060                first_value,
1061                self.name().clone(),
1062                self.len(),
1063            )
1064        }
1065    }
1066
1067    fn apply_lambda_with_bool_out_type(
1068        &self,
1069        py: Python<'py>,
1070        lambda: &Bound<'py, PyAny>,
1071        init_null_count: usize,
1072        first_value: Option<bool>,
1073    ) -> PyResult<BooleanChunked> {
1074        let skip = usize::from(first_value.is_some());
1075        let pypolars = polars(py).bind(py);
1076        if init_null_count == self.len() {
1077            Ok(ChunkedArray::full_null(self.name().clone(), self.len()))
1078        } else {
1079            let it = self
1080                .into_iter()
1081                .skip(init_null_count + skip)
1082                .map(|opt_val| {
1083                    opt_val
1084                        .and_then(|val| {
1085                            // create a PySeries struct/object for Python
1086                            let pyseries = PySeries::new(val);
1087                            // Wrap this PySeries object in the python side Series wrapper
1088                            let python_series_wrapper = pypolars
1089                                .getattr("wrap_s")
1090                                .unwrap()
1091                                .call1((pyseries,))
1092                                .unwrap();
1093                            call_lambda_and_extract(py, lambda, python_series_wrapper).transpose()
1094                        })
1095                        .transpose()
1096                });
1097            iterator_to_bool(
1098                it,
1099                init_null_count,
1100                first_value,
1101                self.name().clone(),
1102                self.len(),
1103            )
1104        }
1105    }
1106
1107    fn apply_lambda_with_string_out_type(
1108        &self,
1109        py: Python<'py>,
1110        lambda: &Bound<'py, PyAny>,
1111        init_null_count: usize,
1112        first_value: Option<PyBackedStr>,
1113    ) -> PyResult<StringChunked> {
1114        let skip = usize::from(first_value.is_some());
1115        // get the pypolars module
1116        let pypolars = polars(py).bind(py);
1117
1118        if init_null_count == self.len() {
1119            Ok(ChunkedArray::full_null(self.name().clone(), self.len()))
1120        } else {
1121            let it = self
1122                .into_iter()
1123                .skip(init_null_count + skip)
1124                .map(|opt_val| {
1125                    opt_val
1126                        .and_then(|val| {
1127                            // create a PySeries struct/object for Python
1128                            let pyseries = PySeries::new(val);
1129                            // Wrap this PySeries object in the python side Series wrapper
1130                            let python_series_wrapper = pypolars
1131                                .getattr("wrap_s")
1132                                .unwrap()
1133                                .call1((pyseries,))
1134                                .unwrap();
1135                            call_lambda_and_extract(py, lambda, python_series_wrapper).transpose()
1136                        })
1137                        .transpose()
1138                });
1139            iterator_to_string(
1140                it,
1141                init_null_count,
1142                first_value,
1143                self.name().clone(),
1144                self.len(),
1145            )
1146        }
1147    }
1148    fn apply_lambda_with_list_out_type(
1149        &self,
1150        py: Python<'py>,
1151        lambda: PyObject,
1152        init_null_count: usize,
1153        first_value: Option<&Series>,
1154        dt: &DataType,
1155    ) -> PyResult<ListChunked> {
1156        let skip = usize::from(first_value.is_some());
1157        let pypolars = polars(py).bind(py);
1158        let lambda = lambda.bind(py);
1159        if init_null_count == self.len() {
1160            Ok(ChunkedArray::full_null(self.name().clone(), self.len()))
1161        } else {
1162            let it = self
1163                .into_iter()
1164                .skip(init_null_count + skip)
1165                .map(|opt_val| {
1166                    opt_val
1167                        .map(|val| call_series_lambda(pypolars, lambda, val))
1168                        .transpose()
1169                        .map(|v| v.flatten())
1170                });
1171            iterator_to_list(
1172                dt,
1173                it,
1174                init_null_count,
1175                first_value,
1176                self.name().clone(),
1177                self.len(),
1178            )
1179        }
1180    }
1181
1182    fn apply_extract_any_values(
1183        &self,
1184        py: Python<'py>,
1185        lambda: &Bound<'py, PyAny>,
1186        init_null_count: usize,
1187        first_value: AnyValue<'static>,
1188    ) -> PyResult<Series> {
1189        let pypolars = polars(py).bind(py);
1190        let mut avs = Vec::with_capacity(self.len());
1191        avs.extend(std::iter::repeat_n(AnyValue::Null, init_null_count));
1192        avs.push(first_value);
1193
1194        let call_with_value = |val: Series| {
1195            // create a PySeries struct/object for Python
1196            let pyseries = PySeries::new(val);
1197            // Wrap this PySeries object in the python side Series wrapper
1198            let python_series_wrapper = pypolars
1199                .getattr("wrap_s")
1200                .unwrap()
1201                .call1((pyseries,))
1202                .unwrap();
1203            call_lambda_and_extract::<_, Wrap<AnyValue>>(py, lambda, python_series_wrapper).map(
1204                |opt_wrap| match opt_wrap {
1205                    None => AnyValue::Null,
1206                    Some(w) => w.0,
1207                },
1208            )
1209        };
1210
1211        for opt_val in self.into_iter().skip(init_null_count + 1) {
1212            if let Some(s) = opt_val {
1213                let av = call_with_value(s)?;
1214                avs.push(av);
1215            } else {
1216                avs.push(AnyValue::Null);
1217            }
1218        }
1219        Ok(Series::new(self.name().clone(), &avs))
1220    }
1221
1222    #[cfg(feature = "object")]
1223    fn apply_lambda_with_object_out_type(
1224        &self,
1225        py: Python<'py>,
1226        lambda: &Bound<'py, PyAny>,
1227        init_null_count: usize,
1228        first_value: Option<ObjectValue>,
1229    ) -> PyResult<ObjectChunked<ObjectValue>> {
1230        let skip = usize::from(first_value.is_some());
1231        let pypolars = polars(py).bind(py);
1232        if init_null_count == self.len() {
1233            Ok(ChunkedArray::full_null(self.name().clone(), self.len()))
1234        } else {
1235            let it = self
1236                .into_iter()
1237                .skip(init_null_count + skip)
1238                .map(|opt_val| {
1239                    opt_val
1240                        .and_then(|val| {
1241                            // create a PySeries struct/object for Python
1242                            let pyseries = PySeries::new(val);
1243                            // Wrap this PySeries object in the python side Series wrapper
1244                            let python_series_wrapper = pypolars
1245                                .getattr("wrap_s")
1246                                .unwrap()
1247                                .call1((pyseries,))
1248                                .unwrap();
1249                            call_lambda_and_extract(py, lambda, python_series_wrapper).transpose()
1250                        })
1251                        .transpose()
1252                });
1253            iterator_to_object(
1254                it,
1255                init_null_count,
1256                first_value,
1257                self.name().clone(),
1258                self.len(),
1259            )
1260        }
1261    }
1262}
1263
1264#[cfg(feature = "dtype-array")]
1265impl<'py> ApplyLambda<'py> for ArrayChunked {
1266    fn apply_lambda_unknown(
1267        &self,
1268        py: Python<'py>,
1269        lambda: &Bound<'py, PyAny>,
1270    ) -> PyResult<PySeries> {
1271        let pypolars = polars(py).bind(py);
1272        let mut null_count = 0;
1273        for opt_v in self.into_iter() {
1274            if let Some(v) = opt_v {
1275                // create a PySeries struct/object for Python
1276                let pyseries = PySeries::new(v);
1277                // Wrap this PySeries object in the python side Series wrapper
1278                let python_series_wrapper = pypolars
1279                    .getattr("wrap_s")
1280                    .unwrap()
1281                    .call1((pyseries,))
1282                    .unwrap();
1283
1284                let out = lambda.call1((python_series_wrapper,))?;
1285                if out.is_none() {
1286                    null_count += 1;
1287                    continue;
1288                }
1289                return infer_and_finish(self, py, lambda, &out, null_count);
1290            } else {
1291                null_count += 1
1292            }
1293        }
1294        Ok(Self::full_null(self.name().clone(), self.len())
1295            .into_series()
1296            .into())
1297    }
1298
1299    fn apply_into_struct(
1300        &self,
1301        py: Python<'py>,
1302        lambda: &Bound<'py, PyAny>,
1303        init_null_count: usize,
1304        first_value: AnyValue<'static>,
1305    ) -> PyResult<PySeries> {
1306        let skip = 1;
1307        // get the pypolars module
1308        let pypolars = polars(py).bind(py);
1309
1310        let it = self
1311            .into_iter()
1312            .skip(init_null_count + skip)
1313            .map(|opt_val| {
1314                opt_val
1315                    .map(|val| {
1316                        // create a PySeries struct/object for Python
1317                        let pyseries = PySeries::new(val);
1318                        // Wrap this PySeries object in the python side Series wrapper
1319                        let python_series_wrapper = pypolars
1320                            .getattr("wrap_s")
1321                            .unwrap()
1322                            .call1((pyseries,))
1323                            .unwrap();
1324                        call_lambda(py, lambda, python_series_wrapper)
1325                    })
1326                    .transpose()
1327            });
1328        iterator_to_struct(
1329            py,
1330            it,
1331            init_null_count,
1332            first_value,
1333            self.name().clone(),
1334            self.len(),
1335        )
1336    }
1337
1338    fn apply_lambda_with_primitive_out_type<D>(
1339        &self,
1340        py: Python<'py>,
1341        lambda: &Bound<'py, PyAny>,
1342        init_null_count: usize,
1343        first_value: Option<D::Native>,
1344    ) -> PyResult<ChunkedArray<D>>
1345    where
1346        D: PyPolarsNumericType,
1347        D::Native: IntoPyObject<'py> + FromPyObject<'py>,
1348    {
1349        let skip = usize::from(first_value.is_some());
1350        let pypolars = polars(py).bind(py);
1351        if init_null_count == self.len() {
1352            Ok(ChunkedArray::full_null(self.name().clone(), self.len()))
1353        } else {
1354            let it = self
1355                .into_iter()
1356                .skip(init_null_count + skip)
1357                .map(|opt_val| {
1358                    opt_val
1359                        .and_then(|val| {
1360                            // create a PySeries struct/object for Python
1361                            let pyseries = PySeries::new(val);
1362                            // Wrap this PySeries object in the python side Series wrapper
1363                            let python_series_wrapper = pypolars
1364                                .getattr("wrap_s")
1365                                .unwrap()
1366                                .call1((pyseries,))
1367                                .unwrap();
1368                            call_lambda_and_extract(py, lambda, python_series_wrapper).transpose()
1369                        })
1370                        .transpose()
1371                });
1372            iterator_to_primitive(
1373                it,
1374                init_null_count,
1375                first_value,
1376                self.name().clone(),
1377                self.len(),
1378            )
1379        }
1380    }
1381
1382    fn apply_lambda_with_bool_out_type(
1383        &self,
1384        py: Python<'py>,
1385        lambda: &Bound<'py, PyAny>,
1386        init_null_count: usize,
1387        first_value: Option<bool>,
1388    ) -> PyResult<BooleanChunked> {
1389        let skip = usize::from(first_value.is_some());
1390        let pypolars = polars(py).bind(py);
1391        if init_null_count == self.len() {
1392            Ok(ChunkedArray::full_null(self.name().clone(), self.len()))
1393        } else {
1394            let it = self
1395                .into_iter()
1396                .skip(init_null_count + skip)
1397                .map(|opt_val| {
1398                    opt_val
1399                        .and_then(|val| {
1400                            // create a PySeries struct/object for Python
1401                            let pyseries = PySeries::new(val);
1402                            // Wrap this PySeries object in the python side Series wrapper
1403                            let python_series_wrapper = pypolars
1404                                .getattr("wrap_s")
1405                                .unwrap()
1406                                .call1((pyseries,))
1407                                .unwrap();
1408                            call_lambda_and_extract(py, lambda, python_series_wrapper).transpose()
1409                        })
1410                        .transpose()
1411                });
1412            iterator_to_bool(
1413                it,
1414                init_null_count,
1415                first_value,
1416                self.name().clone(),
1417                self.len(),
1418            )
1419        }
1420    }
1421
1422    fn apply_lambda_with_string_out_type(
1423        &self,
1424        py: Python<'py>,
1425        lambda: &Bound<'py, PyAny>,
1426        init_null_count: usize,
1427        first_value: Option<PyBackedStr>,
1428    ) -> PyResult<StringChunked> {
1429        let skip = usize::from(first_value.is_some());
1430        // get the pypolars module
1431        let pypolars = polars(py).bind(py);
1432
1433        if init_null_count == self.len() {
1434            Ok(ChunkedArray::full_null(self.name().clone(), self.len()))
1435        } else {
1436            let it = self
1437                .into_iter()
1438                .skip(init_null_count + skip)
1439                .map(|opt_val| {
1440                    opt_val
1441                        .and_then(|val| {
1442                            // create a PySeries struct/object for Python
1443                            let pyseries = PySeries::new(val);
1444                            // Wrap this PySeries object in the python side Series wrapper
1445                            let python_series_wrapper = pypolars
1446                                .getattr("wrap_s")
1447                                .unwrap()
1448                                .call1((pyseries,))
1449                                .unwrap();
1450                            call_lambda_and_extract(py, lambda, python_series_wrapper).transpose()
1451                        })
1452                        .transpose()
1453                });
1454            iterator_to_string(
1455                it,
1456                init_null_count,
1457                first_value,
1458                self.name().clone(),
1459                self.len(),
1460            )
1461        }
1462    }
1463    fn apply_lambda_with_list_out_type(
1464        &self,
1465        py: Python<'py>,
1466        lambda: PyObject,
1467        init_null_count: usize,
1468        first_value: Option<&Series>,
1469        dt: &DataType,
1470    ) -> PyResult<ListChunked> {
1471        let skip = usize::from(first_value.is_some());
1472        let pypolars = polars(py).bind(py);
1473        let lambda = lambda.bind(py);
1474        if init_null_count == self.len() {
1475            Ok(ChunkedArray::full_null(self.name().clone(), self.len()))
1476        } else {
1477            let it = self
1478                .into_iter()
1479                .skip(init_null_count + skip)
1480                .map(|opt_val| {
1481                    opt_val
1482                        .map(|val| call_series_lambda(pypolars, lambda, val))
1483                        .transpose()
1484                        .map(|v| v.flatten())
1485                });
1486            iterator_to_list(
1487                dt,
1488                it,
1489                init_null_count,
1490                first_value,
1491                self.name().clone(),
1492                self.len(),
1493            )
1494        }
1495    }
1496
1497    fn apply_extract_any_values(
1498        &self,
1499        py: Python<'py>,
1500        lambda: &Bound<'py, PyAny>,
1501        init_null_count: usize,
1502        first_value: AnyValue<'static>,
1503    ) -> PyResult<Series> {
1504        let pypolars = polars(py).bind(py);
1505        let mut avs = Vec::with_capacity(self.len());
1506        avs.extend(std::iter::repeat_n(AnyValue::Null, init_null_count));
1507        avs.push(first_value);
1508
1509        let call_with_value = |val: Series| {
1510            // create a PySeries struct/object for Python
1511            let pyseries = PySeries::new(val);
1512            // Wrap this PySeries object in the python side Series wrapper
1513            let python_series_wrapper = pypolars
1514                .getattr("wrap_s")
1515                .unwrap()
1516                .call1((pyseries,))
1517                .unwrap();
1518            call_lambda_and_extract::<_, Wrap<AnyValue>>(py, lambda, python_series_wrapper).map(
1519                |opt_wrap| match opt_wrap {
1520                    None => AnyValue::Null,
1521                    Some(w) => w.0,
1522                },
1523            )
1524        };
1525
1526        for opt_val in self.into_iter().skip(init_null_count + 1) {
1527            if let Some(s) = opt_val {
1528                let av = call_with_value(s)?;
1529                avs.push(av);
1530            } else {
1531                avs.push(AnyValue::Null);
1532            }
1533        }
1534
1535        Ok(Series::new(self.name().clone(), &avs))
1536    }
1537
1538    #[cfg(feature = "object")]
1539    fn apply_lambda_with_object_out_type(
1540        &self,
1541        py: Python<'py>,
1542        lambda: &Bound<'py, PyAny>,
1543        init_null_count: usize,
1544        first_value: Option<ObjectValue>,
1545    ) -> PyResult<ObjectChunked<ObjectValue>> {
1546        let skip = usize::from(first_value.is_some());
1547        let pypolars = polars(py).bind(py);
1548        if init_null_count == self.len() {
1549            Ok(ChunkedArray::full_null(self.name().clone(), self.len()))
1550        } else {
1551            let it = self
1552                .into_iter()
1553                .skip(init_null_count + skip)
1554                .map(|opt_val| {
1555                    opt_val
1556                        .and_then(|val| {
1557                            // create a PySeries struct/object for Python
1558                            let pyseries = PySeries::new(val);
1559                            // Wrap this PySeries object in the python side Series wrapper
1560                            let python_series_wrapper = pypolars
1561                                .getattr("wrap_s")
1562                                .unwrap()
1563                                .call1((pyseries,))
1564                                .unwrap();
1565                            call_lambda_and_extract(py, lambda, python_series_wrapper).transpose()
1566                        })
1567                        .transpose()
1568                });
1569            iterator_to_object(
1570                it,
1571                init_null_count,
1572                first_value,
1573                self.name().clone(),
1574                self.len(),
1575            )
1576        }
1577    }
1578}
1579
1580#[cfg(feature = "object")]
1581impl<'py> ApplyLambda<'py> for ObjectChunked<ObjectValue> {
1582    fn apply_lambda_unknown(
1583        &self,
1584        py: Python<'py>,
1585        lambda: &Bound<'py, PyAny>,
1586    ) -> PyResult<PySeries> {
1587        let mut null_count = 0;
1588        for opt_v in self.into_iter() {
1589            if let Some(v) = opt_v {
1590                let arg = PyTuple::new(py, [v])?;
1591                let out = lambda.call1(arg)?;
1592                if out.is_none() {
1593                    null_count += 1;
1594                    continue;
1595                }
1596                return infer_and_finish(self, py, lambda, &out, null_count);
1597            } else {
1598                null_count += 1
1599            }
1600        }
1601        Ok(Self::full_null(self.name().clone(), self.len())
1602            .into_series()
1603            .into())
1604    }
1605
1606    fn apply_into_struct(
1607        &self,
1608        py: Python<'py>,
1609        lambda: &Bound<'py, PyAny>,
1610        init_null_count: usize,
1611        first_value: AnyValue<'static>,
1612    ) -> PyResult<PySeries> {
1613        let skip = 1;
1614        let it = self
1615            .into_iter()
1616            .skip(init_null_count + skip)
1617            .map(|object_value| lambda.call1((object_value.map(|v| &v.inner),)).map(Some));
1618        iterator_to_struct(
1619            py,
1620            it,
1621            init_null_count,
1622            first_value,
1623            self.name().clone(),
1624            self.len(),
1625        )
1626    }
1627
1628    fn apply_lambda_with_primitive_out_type<D>(
1629        &self,
1630        py: Python<'py>,
1631        lambda: &Bound<'py, PyAny>,
1632        init_null_count: usize,
1633        first_value: Option<D::Native>,
1634    ) -> PyResult<ChunkedArray<D>>
1635    where
1636        D: PyPolarsNumericType,
1637        D::Native: IntoPyObject<'py> + FromPyObject<'py>,
1638    {
1639        let skip = usize::from(first_value.is_some());
1640        if init_null_count == self.len() {
1641            Ok(ChunkedArray::full_null(self.name().clone(), self.len()))
1642        } else {
1643            let it = self
1644                .into_iter()
1645                .skip(init_null_count + skip)
1646                .map(|opt_val| {
1647                    opt_val
1648                        .and_then(|val| call_lambda_and_extract(py, lambda, val).transpose())
1649                        .transpose()
1650                });
1651            iterator_to_primitive(
1652                it,
1653                init_null_count,
1654                first_value,
1655                self.name().clone(),
1656                self.len(),
1657            )
1658        }
1659    }
1660
1661    fn apply_lambda_with_bool_out_type(
1662        &self,
1663        py: Python<'py>,
1664        lambda: &Bound<'py, PyAny>,
1665        init_null_count: usize,
1666        first_value: Option<bool>,
1667    ) -> PyResult<BooleanChunked> {
1668        let skip = usize::from(first_value.is_some());
1669        if init_null_count == self.len() {
1670            Ok(ChunkedArray::full_null(self.name().clone(), self.len()))
1671        } else {
1672            let it = self
1673                .into_iter()
1674                .skip(init_null_count + skip)
1675                .map(|opt_val| {
1676                    opt_val
1677                        .and_then(|val| call_lambda_and_extract(py, lambda, val).transpose())
1678                        .transpose()
1679                });
1680            iterator_to_bool(
1681                it,
1682                init_null_count,
1683                first_value,
1684                self.name().clone(),
1685                self.len(),
1686            )
1687        }
1688    }
1689
1690    fn apply_lambda_with_string_out_type(
1691        &self,
1692        py: Python<'py>,
1693        lambda: &Bound<'py, PyAny>,
1694        init_null_count: usize,
1695        first_value: Option<PyBackedStr>,
1696    ) -> PyResult<StringChunked> {
1697        let skip = usize::from(first_value.is_some());
1698        if init_null_count == self.len() {
1699            Ok(ChunkedArray::full_null(self.name().clone(), self.len()))
1700        } else {
1701            let it = self
1702                .into_iter()
1703                .skip(init_null_count + skip)
1704                .map(|opt_val| {
1705                    opt_val
1706                        .and_then(|val| call_lambda_and_extract(py, lambda, val).transpose())
1707                        .transpose()
1708                });
1709            iterator_to_string(
1710                it,
1711                init_null_count,
1712                first_value,
1713                self.name().clone(),
1714                self.len(),
1715            )
1716        }
1717    }
1718
1719    fn apply_lambda_with_list_out_type(
1720        &self,
1721        py: Python<'py>,
1722        lambda: PyObject,
1723        init_null_count: usize,
1724        first_value: Option<&Series>,
1725        dt: &DataType,
1726    ) -> PyResult<ListChunked> {
1727        let skip = usize::from(first_value.is_some());
1728        let lambda = lambda.bind(py);
1729        if init_null_count == self.len() {
1730            Ok(ChunkedArray::full_null(self.name().clone(), self.len()))
1731        } else {
1732            let it = self
1733                .into_iter()
1734                .skip(init_null_count + skip)
1735                .map(|opt_val| {
1736                    opt_val
1737                        .map(|val| call_lambda_series_out(py, lambda, val))
1738                        .transpose()
1739                });
1740            iterator_to_list(
1741                dt,
1742                it,
1743                init_null_count,
1744                first_value,
1745                self.name().clone(),
1746                self.len(),
1747            )
1748        }
1749    }
1750
1751    fn apply_extract_any_values(
1752        &self,
1753        py: Python<'py>,
1754        lambda: &Bound<'py, PyAny>,
1755        init_null_count: usize,
1756        first_value: AnyValue<'static>,
1757    ) -> PyResult<Series> {
1758        let iter = self.into_iter().skip(init_null_count + 1);
1759        let avs = extract_anyvalues(py, lambda, self.len(), init_null_count, iter, first_value)?;
1760        Ok(Series::new(self.name().clone(), &avs))
1761    }
1762
1763    #[cfg(feature = "object")]
1764    fn apply_lambda_with_object_out_type(
1765        &self,
1766        py: Python<'py>,
1767        lambda: &Bound<'py, PyAny>,
1768        init_null_count: usize,
1769        first_value: Option<ObjectValue>,
1770    ) -> PyResult<ObjectChunked<ObjectValue>> {
1771        let skip = usize::from(first_value.is_some());
1772        if init_null_count == self.len() {
1773            Ok(ChunkedArray::full_null(self.name().clone(), self.len()))
1774        } else {
1775            let it = self
1776                .into_iter()
1777                .skip(init_null_count + skip)
1778                .map(|opt_val| {
1779                    opt_val
1780                        .and_then(|val| call_lambda_and_extract(py, lambda, val).transpose())
1781                        .transpose()
1782                });
1783            iterator_to_object(
1784                it,
1785                init_null_count,
1786                first_value,
1787                self.name().clone(),
1788                self.len(),
1789            )
1790        }
1791    }
1792}
1793
1794fn iter_struct(ca: &StructChunked) -> impl Iterator<Item = AnyValue> {
1795    (0..ca.len()).map(|i| unsafe { ca.get_any_value_unchecked(i) })
1796}
1797
1798impl<'py> ApplyLambda<'py> for StructChunked {
1799    fn apply_lambda_unknown(
1800        &self,
1801        py: Python<'py>,
1802        lambda: &Bound<'py, PyAny>,
1803    ) -> PyResult<PySeries> {
1804        let mut null_count = 0;
1805
1806        for val in iter_struct(self) {
1807            match val {
1808                AnyValue::Null => null_count += 1,
1809                _ => {
1810                    let out = lambda.call1((Wrap(val),))?;
1811                    if out.is_none() {
1812                        null_count += 1;
1813                        continue;
1814                    }
1815                    return infer_and_finish(self, py, lambda, &out, null_count);
1816                },
1817            }
1818        }
1819
1820        Ok(Self::full_null(self.name().clone(), self.len())
1821            .into_series()
1822            .into())
1823    }
1824
1825    fn apply_into_struct(
1826        &self,
1827        py: Python<'py>,
1828        lambda: &Bound<'py, PyAny>,
1829        init_null_count: usize,
1830        first_value: AnyValue<'static>,
1831    ) -> PyResult<PySeries> {
1832        let skip = 1;
1833        let it = iter_struct(self)
1834            .skip(init_null_count + skip)
1835            .map(|val| match val {
1836                AnyValue::Null => Ok(None),
1837                _ => lambda.call1((Wrap(val),)).map(Some),
1838            });
1839        iterator_to_struct(
1840            py,
1841            it,
1842            init_null_count,
1843            first_value,
1844            self.name().clone(),
1845            self.len(),
1846        )
1847    }
1848
1849    fn apply_lambda_with_primitive_out_type<D>(
1850        &self,
1851        py: Python<'py>,
1852        lambda: &Bound<'py, PyAny>,
1853        init_null_count: usize,
1854        first_value: Option<D::Native>,
1855    ) -> PyResult<ChunkedArray<D>>
1856    where
1857        D: PyPolarsNumericType,
1858        D::Native: IntoPyObject<'py> + FromPyObject<'py>,
1859    {
1860        let skip = usize::from(first_value.is_some());
1861        let it = iter_struct(self)
1862            .skip(init_null_count + skip)
1863            .map(|val| match val {
1864                AnyValue::Null => Ok(None),
1865                _ => call_lambda_and_extract(py, lambda, Wrap(val)),
1866            });
1867
1868        iterator_to_primitive(
1869            it,
1870            init_null_count,
1871            first_value,
1872            self.name().clone(),
1873            self.len(),
1874        )
1875    }
1876
1877    fn apply_lambda_with_bool_out_type(
1878        &self,
1879        py: Python<'py>,
1880        lambda: &Bound<'py, PyAny>,
1881        init_null_count: usize,
1882        first_value: Option<bool>,
1883    ) -> PyResult<BooleanChunked> {
1884        let skip = usize::from(first_value.is_some());
1885        let it = iter_struct(self)
1886            .skip(init_null_count + skip)
1887            .map(|val| match val {
1888                AnyValue::Null => Ok(None),
1889                _ => call_lambda_and_extract(py, lambda, Wrap(val)),
1890            });
1891
1892        iterator_to_bool(
1893            it,
1894            init_null_count,
1895            first_value,
1896            self.name().clone(),
1897            self.len(),
1898        )
1899    }
1900
1901    fn apply_lambda_with_string_out_type(
1902        &self,
1903        py: Python<'py>,
1904        lambda: &Bound<'py, PyAny>,
1905        init_null_count: usize,
1906        first_value: Option<PyBackedStr>,
1907    ) -> PyResult<StringChunked> {
1908        let skip = usize::from(first_value.is_some());
1909        let it = iter_struct(self)
1910            .skip(init_null_count + skip)
1911            .map(|val| match val {
1912                AnyValue::Null => Ok(None),
1913                _ => call_lambda_and_extract(py, lambda, Wrap(val)),
1914            });
1915
1916        iterator_to_string(
1917            it,
1918            init_null_count,
1919            first_value,
1920            self.name().clone(),
1921            self.len(),
1922        )
1923    }
1924    fn apply_lambda_with_list_out_type(
1925        &self,
1926        py: Python<'py>,
1927        lambda: PyObject,
1928        init_null_count: usize,
1929        first_value: Option<&Series>,
1930        dt: &DataType,
1931    ) -> PyResult<ListChunked> {
1932        let skip = usize::from(first_value.is_some());
1933        let lambda = lambda.bind(py);
1934        let it = iter_struct(self)
1935            .skip(init_null_count + skip)
1936            .map(|val| match val {
1937                AnyValue::Null => Ok(None),
1938                _ => call_lambda_series_out(py, lambda, Wrap(val)).map(Some),
1939            });
1940        iterator_to_list(
1941            dt,
1942            it,
1943            init_null_count,
1944            first_value,
1945            self.name().clone(),
1946            self.len(),
1947        )
1948    }
1949
1950    fn apply_extract_any_values(
1951        &self,
1952        py: Python<'py>,
1953        lambda: &Bound<'py, PyAny>,
1954        init_null_count: usize,
1955        first_value: AnyValue<'static>,
1956    ) -> PyResult<Series> {
1957        let mut avs = Vec::with_capacity(self.len());
1958        avs.extend(std::iter::repeat_n(AnyValue::Null, init_null_count));
1959        avs.push(first_value);
1960
1961        for val in iter_struct(self).skip(init_null_count + 1) {
1962            let av: Option<Wrap<AnyValue>> = call_lambda_and_extract(py, lambda, Wrap(val))?;
1963            let out = match av {
1964                None => AnyValue::Null,
1965                Some(av) => av.0,
1966            };
1967            avs.push(out)
1968        }
1969
1970        Ok(Series::new(self.name().clone(), &avs))
1971    }
1972
1973    #[cfg(feature = "object")]
1974    fn apply_lambda_with_object_out_type(
1975        &self,
1976        py: Python<'py>,
1977        lambda: &Bound<'py, PyAny>,
1978        init_null_count: usize,
1979        first_value: Option<ObjectValue>,
1980    ) -> PyResult<ObjectChunked<ObjectValue>> {
1981        let skip = usize::from(first_value.is_some());
1982        let it = iter_struct(self)
1983            .skip(init_null_count + skip)
1984            .map(|val| match val {
1985                AnyValue::Null => Ok(None),
1986                _ => call_lambda_and_extract(py, lambda, Wrap(val)),
1987            });
1988
1989        iterator_to_object(
1990            it,
1991            init_null_count,
1992            first_value,
1993            self.name().clone(),
1994            self.len(),
1995        )
1996    }
1997}
polars_python/map/series.rs

polars_python/map/
series.rs
