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

polars_python/map/
series.rs
