polars_python/expr/
list.rs

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

polars_python/expr/
list.rs
