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

polars_python/map/
series.rs
