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

polars_python/map/
series.rs
