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, parallel).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(
60                Expr::Column(PlSmallStr::EMPTY).filter(predicate.inner),
61                false,
62            )
63            .into()
64    }
65
66    fn list_get(&self, index: PyExpr, null_on_oob: bool) -> Self {
67        self.inner
68            .clone()
69            .list()
70            .get(index.inner, null_on_oob)
71            .into()
72    }
73
74    fn list_join(&self, separator: PyExpr, ignore_nulls: bool) -> Self {
75        self.inner
76            .clone()
77            .list()
78            .join(separator.inner, ignore_nulls)
79            .into()
80    }
81
82    fn list_len(&self) -> Self {
83        self.inner.clone().list().len().into()
84    }
85
86    fn list_max(&self) -> Self {
87        self.inner.clone().list().max().into()
88    }
89
90    fn list_mean(&self) -> Self {
91        self.inner
92            .clone()
93            .list()
94            .mean()
95            .with_fmt("list.mean")
96            .into()
97    }
98
99    fn list_median(&self) -> Self {
100        self.inner
101            .clone()
102            .list()
103            .median()
104            .with_fmt("list.median")
105            .into()
106    }
107
108    fn list_std(&self, ddof: u8) -> Self {
109        self.inner
110            .clone()
111            .list()
112            .std(ddof)
113            .with_fmt("list.std")
114            .into()
115    }
116
117    fn list_var(&self, ddof: u8) -> Self {
118        self.inner
119            .clone()
120            .list()
121            .var(ddof)
122            .with_fmt("list.var")
123            .into()
124    }
125
126    fn list_min(&self) -> Self {
127        self.inner.clone().list().min().into()
128    }
129
130    fn list_reverse(&self) -> Self {
131        self.inner.clone().list().reverse().into()
132    }
133
134    fn list_shift(&self, periods: PyExpr) -> Self {
135        self.inner.clone().list().shift(periods.inner).into()
136    }
137
138    #[pyo3(signature = (offset, length=None))]
139    fn list_slice(&self, offset: PyExpr, length: Option<PyExpr>) -> Self {
140        let length = match length {
141            Some(i) => i.inner,
142            None => lit(i64::MAX),
143        };
144        self.inner.clone().list().slice(offset.inner, length).into()
145    }
146
147    fn list_tail(&self, n: PyExpr) -> Self {
148        self.inner.clone().list().tail(n.inner).into()
149    }
150
151    fn list_sort(&self, descending: bool, nulls_last: bool) -> Self {
152        self.inner
153            .clone()
154            .list()
155            .sort(
156                SortOptions::default()
157                    .with_order_descending(descending)
158                    .with_nulls_last(nulls_last),
159            )
160            .into()
161    }
162
163    fn list_sum(&self) -> Self {
164        self.inner.clone().list().sum().with_fmt("list.sum").into()
165    }
166
167    #[cfg(feature = "list_drop_nulls")]
168    fn list_drop_nulls(&self) -> Self {
169        self.inner.clone().list().drop_nulls().into()
170    }
171
172    #[cfg(feature = "list_sample")]
173    #[pyo3(signature = (n, with_replacement, shuffle, seed=None))]
174    fn list_sample_n(
175        &self,
176        n: PyExpr,
177        with_replacement: bool,
178        shuffle: bool,
179        seed: Option<u64>,
180    ) -> Self {
181        self.inner
182            .clone()
183            .list()
184            .sample_n(n.inner, with_replacement, shuffle, seed)
185            .into()
186    }
187
188    #[cfg(feature = "list_sample")]
189    #[pyo3(signature = (fraction, with_replacement, shuffle, seed=None))]
190    fn list_sample_fraction(
191        &self,
192        fraction: PyExpr,
193        with_replacement: bool,
194        shuffle: bool,
195        seed: Option<u64>,
196    ) -> Self {
197        self.inner
198            .clone()
199            .list()
200            .sample_fraction(fraction.inner, with_replacement, shuffle, seed)
201            .into()
202    }
203
204    #[cfg(feature = "list_gather")]
205    fn list_gather(&self, index: PyExpr, null_on_oob: bool) -> Self {
206        self.inner
207            .clone()
208            .list()
209            .gather(index.inner, null_on_oob)
210            .into()
211    }
212
213    #[cfg(feature = "list_gather")]
214    fn list_gather_every(&self, n: PyExpr, offset: PyExpr) -> Self {
215        self.inner
216            .clone()
217            .list()
218            .gather_every(n.inner, offset.inner)
219            .into()
220    }
221
222    fn list_to_array(&self, width: usize) -> Self {
223        self.inner.clone().list().to_array(width).into()
224    }
225
226    #[pyo3(signature = (width_strat, name_gen, upper_bound))]
227    fn list_to_struct(
228        &self,
229        width_strat: Wrap<ListToStructWidthStrategy>,
230        name_gen: Option<PyObject>,
231        upper_bound: Option<usize>,
232    ) -> PyResult<Self> {
233        let name_gen = name_gen.map(|lambda| {
234            NameGenerator::from_func(move |idx: usize| {
235                Python::with_gil(|py| {
236                    let out = lambda.call1(py, (idx,)).unwrap();
237                    let out: PlSmallStr = out.extract::<Cow<str>>(py).unwrap().as_ref().into();
238                    out
239                })
240            })
241        });
242
243        Ok(self
244            .inner
245            .clone()
246            .list()
247            .to_struct(ListToStructArgs::InferWidth {
248                infer_field_strategy: width_strat.0,
249                get_index_name: name_gen,
250                max_fields: upper_bound,
251            })
252            .into())
253    }
254
255    #[pyo3(signature = (names))]
256    fn list_to_struct_fixed_width(&self, names: Bound<'_, PySequence>) -> PyResult<Self> {
257        Ok(self
258            .inner
259            .clone()
260            .list()
261            .to_struct(ListToStructArgs::FixedWidth(
262                names
263                    .try_iter()?
264                    .map(|x| Ok(x?.extract::<Wrap<PlSmallStr>>()?.0))
265                    .collect::<PyResult<Arc<[_]>>>()?,
266            ))
267            .into())
268    }
269
270    fn list_n_unique(&self) -> Self {
271        self.inner.clone().list().n_unique().into()
272    }
273
274    fn list_unique(&self, maintain_order: bool) -> Self {
275        let e = self.inner.clone();
276
277        if maintain_order {
278            e.list().unique_stable().into()
279        } else {
280            e.list().unique().into()
281        }
282    }
283
284    #[cfg(feature = "list_sets")]
285    fn list_set_operation(&self, other: PyExpr, operation: Wrap<SetOperation>) -> Self {
286        let e = self.inner.clone().list();
287        match operation.0 {
288            SetOperation::Intersection => e.set_intersection(other.inner),
289            SetOperation::Difference => e.set_difference(other.inner),
290            SetOperation::Union => e.union(other.inner),
291            SetOperation::SymmetricDifference => e.set_symmetric_difference(other.inner),
292        }
293        .into()
294    }
295}
polars_python/expr/list.rs

polars_python/expr/
list.rs
