polars_python/expr/
array.rs

1use polars::prelude::*;
2use polars_ops::prelude::array::ArrToStructNameGenerator;
3use polars_utils::pl_str::PlSmallStr;
4use pyo3::prelude::*;
5use pyo3::pybacked::PyBackedStr;
6use pyo3::pymethods;
7
8use crate::error::PyPolarsErr;
9use crate::expr::PyExpr;
10
11#[pymethods]
12impl PyExpr {
13    fn arr_max(&self) -> Self {
14        self.inner.clone().arr().max().into()
15    }
16
17    fn arr_min(&self) -> Self {
18        self.inner.clone().arr().min().into()
19    }
20
21    fn arr_sum(&self) -> Self {
22        self.inner.clone().arr().sum().into()
23    }
24
25    fn arr_std(&self, ddof: u8) -> Self {
26        self.inner.clone().arr().std(ddof).into()
27    }
28
29    fn arr_var(&self, ddof: u8) -> Self {
30        self.inner.clone().arr().var(ddof).into()
31    }
32
33    fn arr_median(&self) -> Self {
34        self.inner.clone().arr().median().into()
35    }
36
37    fn arr_unique(&self, maintain_order: bool) -> Self {
38        if maintain_order {
39            self.inner.clone().arr().unique_stable().into()
40        } else {
41            self.inner.clone().arr().unique().into()
42        }
43    }
44
45    fn arr_n_unique(&self) -> Self {
46        self.inner.clone().arr().n_unique().into()
47    }
48
49    fn arr_to_list(&self) -> Self {
50        self.inner.clone().arr().to_list().into()
51    }
52
53    fn arr_all(&self) -> Self {
54        self.inner.clone().arr().all().into()
55    }
56
57    fn arr_any(&self) -> Self {
58        self.inner.clone().arr().any().into()
59    }
60
61    fn arr_sort(&self, descending: bool, nulls_last: bool) -> Self {
62        self.inner
63            .clone()
64            .arr()
65            .sort(SortOptions {
66                descending,
67                nulls_last,
68                ..Default::default()
69            })
70            .into()
71    }
72
73    fn arr_reverse(&self) -> Self {
74        self.inner.clone().arr().reverse().into()
75    }
76
77    fn arr_arg_min(&self) -> Self {
78        self.inner.clone().arr().arg_min().into()
79    }
80
81    fn arr_arg_max(&self) -> Self {
82        self.inner.clone().arr().arg_max().into()
83    }
84
85    fn arr_get(&self, index: PyExpr, null_on_oob: bool) -> Self {
86        self.inner
87            .clone()
88            .arr()
89            .get(index.inner, null_on_oob)
90            .into()
91    }
92
93    fn arr_join(&self, separator: PyExpr, ignore_nulls: bool) -> Self {
94        self.inner
95            .clone()
96            .arr()
97            .join(separator.inner, ignore_nulls)
98            .into()
99    }
100
101    #[cfg(feature = "is_in")]
102    fn arr_contains(&self, other: PyExpr) -> Self {
103        self.inner.clone().arr().contains(other.inner).into()
104    }
105
106    #[cfg(feature = "array_count")]
107    fn arr_count_matches(&self, expr: PyExpr) -> Self {
108        self.inner.clone().arr().count_matches(expr.inner).into()
109    }
110
111    #[pyo3(signature = (name_gen))]
112    fn arr_to_struct(&self, name_gen: Option<PyObject>) -> PyResult<Self> {
113        let name_gen = name_gen.map(|lambda| {
114            Arc::new(move |idx: usize| {
115                Python::with_gil(|py| {
116                    let out = lambda.call1(py, (idx,)).unwrap();
117                    let out: PlSmallStr = (&*out.extract::<PyBackedStr>(py).unwrap()).into();
118                    out
119                })
120            }) as ArrToStructNameGenerator
121        });
122
123        Ok(self
124            .inner
125            .clone()
126            .arr()
127            .to_struct(name_gen)
128            .map_err(PyPolarsErr::from)?
129            .into())
130    }
131
132    fn arr_shift(&self, n: PyExpr) -> Self {
133        self.inner.clone().arr().shift(n.inner).into()
134    }
135
136    fn arr_explode(&self) -> Self {
137        self.inner.clone().arr().explode().into()
138    }
139}
polars_python/expr/array.rs

polars_python/expr/
array.rs
