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

polars_python/expr/
list.rs
