polars_python/expr/
list.rs

1use polars::prelude::*;
2use polars::series::ops::NullBehavior;
3use polars_utils::pl_str::PlSmallStr;
4use polars_utils::python_function::PythonObject;
5use pyo3::prelude::*;
6use pyo3::types::PySequence;
7
8use crate::PyExpr;
9use crate::conversion::Wrap;
10
11#[pymethods]
12impl PyExpr {
13    #[cfg(feature = "list_any_all")]
14    fn list_all(&self) -> Self {
15        self.inner.clone().list().all().into()
16    }
17
18    #[cfg(feature = "list_any_all")]
19    fn list_any(&self) -> Self {
20        self.inner.clone().list().any().into()
21    }
22
23    fn list_arg_max(&self) -> Self {
24        self.inner.clone().list().arg_max().into()
25    }
26
27    fn list_arg_min(&self) -> Self {
28        self.inner.clone().list().arg_min().into()
29    }
30
31    #[cfg(feature = "is_in")]
32    fn list_contains(&self, other: PyExpr, nulls_equal: bool) -> Self {
33        self.inner
34            .clone()
35            .list()
36            .contains(other.inner, nulls_equal)
37            .into()
38    }
39
40    #[cfg(feature = "list_count")]
41    fn list_count_matches(&self, expr: PyExpr) -> Self {
42        self.inner.clone().list().count_matches(expr.inner).into()
43    }
44
45    fn list_diff(&self, n: i64, null_behavior: Wrap<NullBehavior>) -> PyResult<Self> {
46        Ok(self.inner.clone().list().diff(n, null_behavior.0).into())
47    }
48
49    fn list_eval(&self, expr: PyExpr, _parallel: bool) -> Self {
50        self.inner.clone().list().eval(expr.inner).into()
51    }
52
53    #[cfg(feature = "list_filter")]
54    fn list_filter(&self, predicate: PyExpr) -> Self {
55        self.inner
56            .clone()
57            .list()
58            .eval(Expr::Column(PlSmallStr::EMPTY).filter(predicate.inner))
59            .into()
60    }
61
62    fn list_get(&self, index: PyExpr, null_on_oob: bool) -> Self {
63        self.inner
64            .clone()
65            .list()
66            .get(index.inner, null_on_oob)
67            .into()
68    }
69
70    fn list_join(&self, separator: PyExpr, ignore_nulls: bool) -> Self {
71        self.inner
72            .clone()
73            .list()
74            .join(separator.inner, ignore_nulls)
75            .into()
76    }
77
78    fn list_len(&self) -> Self {
79        self.inner.clone().list().len().into()
80    }
81
82    fn list_max(&self) -> Self {
83        self.inner.clone().list().max().into()
84    }
85
86    fn list_mean(&self) -> Self {
87        self.inner.clone().list().mean().into()
88    }
89
90    fn list_median(&self) -> Self {
91        self.inner.clone().list().median().into()
92    }
93
94    fn list_std(&self, ddof: u8) -> Self {
95        self.inner.clone().list().std(ddof).into()
96    }
97
98    fn list_var(&self, ddof: u8) -> Self {
99        self.inner.clone().list().var(ddof).into()
100    }
101
102    fn list_min(&self) -> Self {
103        self.inner.clone().list().min().into()
104    }
105
106    fn list_reverse(&self) -> Self {
107        self.inner.clone().list().reverse().into()
108    }
109
110    fn list_shift(&self, periods: PyExpr) -> Self {
111        self.inner.clone().list().shift(periods.inner).into()
112    }
113
114    #[pyo3(signature = (offset, length=None))]
115    fn list_slice(&self, offset: PyExpr, length: Option<PyExpr>) -> Self {
116        let length = match length {
117            Some(i) => i.inner,
118            None => lit(i64::MAX),
119        };
120        self.inner.clone().list().slice(offset.inner, length).into()
121    }
122
123    fn list_tail(&self, n: PyExpr) -> Self {
124        self.inner.clone().list().tail(n.inner).into()
125    }
126
127    fn list_sort(&self, descending: bool, nulls_last: bool) -> Self {
128        self.inner
129            .clone()
130            .list()
131            .sort(
132                SortOptions::default()
133                    .with_order_descending(descending)
134                    .with_nulls_last(nulls_last),
135            )
136            .into()
137    }
138
139    fn list_sum(&self) -> Self {
140        self.inner.clone().list().sum().into()
141    }
142
143    #[cfg(feature = "list_drop_nulls")]
144    fn list_drop_nulls(&self) -> Self {
145        self.inner.clone().list().drop_nulls().into()
146    }
147
148    #[cfg(feature = "list_sample")]
149    #[pyo3(signature = (n, with_replacement, shuffle, seed=None))]
150    fn list_sample_n(
151        &self,
152        n: PyExpr,
153        with_replacement: bool,
154        shuffle: bool,
155        seed: Option<u64>,
156    ) -> Self {
157        self.inner
158            .clone()
159            .list()
160            .sample_n(n.inner, with_replacement, shuffle, seed)
161            .into()
162    }
163
164    #[cfg(feature = "list_sample")]
165    #[pyo3(signature = (fraction, with_replacement, shuffle, seed=None))]
166    fn list_sample_fraction(
167        &self,
168        fraction: PyExpr,
169        with_replacement: bool,
170        shuffle: bool,
171        seed: Option<u64>,
172    ) -> Self {
173        self.inner
174            .clone()
175            .list()
176            .sample_fraction(fraction.inner, with_replacement, shuffle, seed)
177            .into()
178    }
179
180    #[cfg(feature = "list_gather")]
181    fn list_gather(&self, index: PyExpr, null_on_oob: bool) -> Self {
182        self.inner
183            .clone()
184            .list()
185            .gather(index.inner, null_on_oob)
186            .into()
187    }
188
189    #[cfg(feature = "list_gather")]
190    fn list_gather_every(&self, n: PyExpr, offset: PyExpr) -> Self {
191        self.inner
192            .clone()
193            .list()
194            .gather_every(n.inner, offset.inner)
195            .into()
196    }
197
198    fn list_to_array(&self, width: usize) -> Self {
199        self.inner.clone().list().to_array(width).into()
200    }
201
202    #[pyo3(signature = (width_strat, name_gen, upper_bound))]
203    fn list_to_struct(
204        &self,
205        width_strat: Wrap<ListToStructWidthStrategy>,
206        name_gen: Option<PyObject>,
207        upper_bound: Option<usize>,
208    ) -> PyResult<Self> {
209        let name_gen = name_gen.map(|lambda| PlanCallback::new_python(PythonObject(lambda)));
210
211        Ok(self
212            .inner
213            .clone()
214            .list()
215            .to_struct(ListToStruct::InferWidth {
216                infer_field_strategy: width_strat.0,
217                get_index_name: name_gen,
218                max_fields: upper_bound,
219            })
220            .into())
221    }
222
223    #[pyo3(signature = (names))]
224    fn list_to_struct_fixed_width(&self, names: Bound<'_, PySequence>) -> PyResult<Self> {
225        Ok(self
226            .inner
227            .clone()
228            .list()
229            .to_struct(ListToStruct::FixedWidth(
230                names
231                    .try_iter()?
232                    .map(|x| Ok(x?.extract::<Wrap<PlSmallStr>>()?.0))
233                    .collect::<PyResult<Arc<[_]>>>()?,
234            ))
235            .into())
236    }
237
238    fn list_n_unique(&self) -> Self {
239        self.inner.clone().list().n_unique().into()
240    }
241
242    fn list_unique(&self, maintain_order: bool) -> Self {
243        let e = self.inner.clone();
244
245        if maintain_order {
246            e.list().unique_stable().into()
247        } else {
248            e.list().unique().into()
249        }
250    }
251
252    #[cfg(feature = "list_sets")]
253    fn list_set_operation(&self, other: PyExpr, operation: Wrap<SetOperation>) -> Self {
254        let e = self.inner.clone().list();
255        match operation.0 {
256            SetOperation::Intersection => e.set_intersection(other.inner),
257            SetOperation::Difference => e.set_difference(other.inner),
258            SetOperation::Union => e.union(other.inner),
259            SetOperation::SymmetricDifference => e.set_symmetric_difference(other.inner),
260        }
261        .into()
262    }
263}
polars_python/expr/list.rs

polars_python/expr/
list.rs
