polars_python/expr/
array.rs

1use polars::prelude::*;
2use polars_utils::python_function::PythonObject;
3use pyo3::prelude::*;
4use pyo3::pymethods;
5
6use crate::error::PyPolarsErr;
7use crate::expr::PyExpr;
8
9#[pymethods]
10impl PyExpr {
11    fn arr_len(&self) -> Self {
12        self.inner.clone().arr().len().into()
13    }
14
15    fn arr_max(&self) -> Self {
16        self.inner.clone().arr().max().into()
17    }
18
19    fn arr_min(&self) -> Self {
20        self.inner.clone().arr().min().into()
21    }
22
23    fn arr_sum(&self) -> Self {
24        self.inner.clone().arr().sum().into()
25    }
26
27    fn arr_std(&self, ddof: u8) -> Self {
28        self.inner.clone().arr().std(ddof).into()
29    }
30
31    fn arr_var(&self, ddof: u8) -> Self {
32        self.inner.clone().arr().var(ddof).into()
33    }
34
35    fn arr_mean(&self) -> Self {
36        self.inner.clone().arr().mean().into()
37    }
38
39    fn arr_median(&self) -> Self {
40        self.inner.clone().arr().median().into()
41    }
42
43    fn arr_unique(&self, maintain_order: bool) -> Self {
44        if maintain_order {
45            self.inner.clone().arr().unique_stable().into()
46        } else {
47            self.inner.clone().arr().unique().into()
48        }
49    }
50
51    fn arr_n_unique(&self) -> Self {
52        self.inner.clone().arr().n_unique().into()
53    }
54
55    fn arr_to_list(&self) -> Self {
56        self.inner.clone().arr().to_list().into()
57    }
58
59    fn arr_all(&self) -> Self {
60        self.inner.clone().arr().all().into()
61    }
62
63    fn arr_any(&self) -> Self {
64        self.inner.clone().arr().any().into()
65    }
66
67    fn arr_sort(&self, descending: bool, nulls_last: bool) -> Self {
68        self.inner
69            .clone()
70            .arr()
71            .sort(SortOptions {
72                descending,
73                nulls_last,
74                ..Default::default()
75            })
76            .into()
77    }
78
79    fn arr_reverse(&self) -> Self {
80        self.inner.clone().arr().reverse().into()
81    }
82
83    fn arr_arg_min(&self) -> Self {
84        self.inner.clone().arr().arg_min().into()
85    }
86
87    fn arr_arg_max(&self) -> Self {
88        self.inner.clone().arr().arg_max().into()
89    }
90
91    fn arr_get(&self, index: PyExpr, null_on_oob: bool) -> Self {
92        self.inner
93            .clone()
94            .arr()
95            .get(index.inner, null_on_oob)
96            .into()
97    }
98
99    fn arr_join(&self, separator: PyExpr, ignore_nulls: bool) -> Self {
100        self.inner
101            .clone()
102            .arr()
103            .join(separator.inner, ignore_nulls)
104            .into()
105    }
106
107    #[cfg(feature = "is_in")]
108    fn arr_contains(&self, other: PyExpr, nulls_equal: bool) -> Self {
109        self.inner
110            .clone()
111            .arr()
112            .contains(other.inner, nulls_equal)
113            .into()
114    }
115
116    #[cfg(feature = "array_count")]
117    fn arr_count_matches(&self, expr: PyExpr) -> Self {
118        self.inner.clone().arr().count_matches(expr.inner).into()
119    }
120
121    #[pyo3(signature = (name_gen))]
122    fn arr_to_struct(&self, name_gen: Option<PyObject>) -> Self {
123        let name_gen = name_gen.map(|o| PlanCallback::new_python(PythonObject(o)));
124        self.inner.clone().arr().to_struct(name_gen).into()
125    }
126
127    fn arr_slice(&self, offset: PyExpr, length: Option<PyExpr>, as_array: bool) -> PyResult<Self> {
128        let length = match length {
129            Some(i) => i.inner,
130            None => lit(i64::MAX),
131        };
132        Ok(self
133            .inner
134            .clone()
135            .arr()
136            .slice(offset.inner, length, as_array)
137            .map_err(PyPolarsErr::from)?
138            .into())
139    }
140
141    fn arr_tail(&self, n: PyExpr, as_array: bool) -> PyResult<Self> {
142        Ok(self
143            .inner
144            .clone()
145            .arr()
146            .tail(n.inner, as_array)
147            .map_err(PyPolarsErr::from)?
148            .into())
149    }
150
151    fn arr_shift(&self, n: PyExpr) -> Self {
152        self.inner.clone().arr().shift(n.inner).into()
153    }
154
155    fn arr_explode(&self) -> Self {
156        self.inner.clone().arr().explode().into()
157    }
158}