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

polars_python/expr/
array.rs
