1use std::ops::Neg;
2
3use polars::lazy::dsl;
4use polars::prelude::*;
5use polars::series::ops::NullBehavior;
6use polars_core::chunked_array::cast::CastOptions;
7use polars_core::series::IsSorted;
8use polars_plan::plans::predicates::aexpr_to_skip_batch_predicate;
9use polars_plan::plans::{ExprToIRContext, RowEncodingVariant, node_to_expr, to_expr_ir};
10use polars_utils::arena::Arena;
11use pyo3::class::basic::CompareOp;
12use pyo3::prelude::*;
13
14use super::datatype::PyDataTypeExpr;
15use super::selector::PySelector;
16use crate::PyExpr;
17use crate::conversion::{Wrap, parse_fill_null_strategy};
18use crate::error::PyPolarsErr;
19use crate::utils::EnterPolarsExt;
20
21#[pymethods]
22impl PyExpr {
23 fn __richcmp__(&self, other: Self, op: CompareOp) -> Self {
24 match op {
25 CompareOp::Eq => self.eq(other),
26 CompareOp::Ne => self.neq(other),
27 CompareOp::Gt => self.gt(other),
28 CompareOp::Lt => self.lt(other),
29 CompareOp::Ge => self.gt_eq(other),
30 CompareOp::Le => self.lt_eq(other),
31 }
32 }
33
34 fn __add__(&self, rhs: Self) -> PyResult<Self> {
35 Ok(dsl::binary_expr(self.inner.clone(), Operator::Plus, rhs.inner).into())
36 }
37 fn __sub__(&self, rhs: Self) -> PyResult<Self> {
38 Ok(dsl::binary_expr(self.inner.clone(), Operator::Minus, rhs.inner).into())
39 }
40 fn __mul__(&self, rhs: Self) -> PyResult<Self> {
41 Ok(dsl::binary_expr(self.inner.clone(), Operator::Multiply, rhs.inner).into())
42 }
43 fn __truediv__(&self, rhs: Self) -> PyResult<Self> {
44 Ok(dsl::binary_expr(self.inner.clone(), Operator::TrueDivide, rhs.inner).into())
45 }
46 fn __mod__(&self, rhs: Self) -> PyResult<Self> {
47 Ok(dsl::binary_expr(self.inner.clone(), Operator::Modulus, rhs.inner).into())
48 }
49 fn __floordiv__(&self, rhs: Self) -> PyResult<Self> {
50 Ok(dsl::binary_expr(self.inner.clone(), Operator::FloorDivide, rhs.inner).into())
51 }
52 fn __neg__(&self) -> PyResult<Self> {
53 Ok(self.inner.clone().neg().into())
54 }
55
56 fn to_str(&self) -> String {
57 format!("{:?}", self.inner)
58 }
59 fn eq(&self, other: Self) -> Self {
60 self.inner.clone().eq(other.inner).into()
61 }
62
63 fn eq_missing(&self, other: Self) -> Self {
64 self.inner.clone().eq_missing(other.inner).into()
65 }
66 fn neq(&self, other: Self) -> Self {
67 self.inner.clone().neq(other.inner).into()
68 }
69 fn neq_missing(&self, other: Self) -> Self {
70 self.inner.clone().neq_missing(other.inner).into()
71 }
72 fn gt(&self, other: Self) -> Self {
73 self.inner.clone().gt(other.inner).into()
74 }
75 fn gt_eq(&self, other: Self) -> Self {
76 self.inner.clone().gt_eq(other.inner).into()
77 }
78 fn lt_eq(&self, other: Self) -> Self {
79 self.inner.clone().lt_eq(other.inner).into()
80 }
81 fn lt(&self, other: Self) -> Self {
82 self.inner.clone().lt(other.inner).into()
83 }
84
85 fn alias(&self, name: &str) -> Self {
86 self.inner.clone().alias(name).into()
87 }
88 fn not_(&self) -> Self {
89 self.inner.clone().not().into()
90 }
91 fn is_null(&self) -> Self {
92 self.inner.clone().is_null().into()
93 }
94 fn is_not_null(&self) -> Self {
95 self.inner.clone().is_not_null().into()
96 }
97
98 fn is_infinite(&self) -> Self {
99 self.inner.clone().is_infinite().into()
100 }
101
102 fn is_finite(&self) -> Self {
103 self.inner.clone().is_finite().into()
104 }
105
106 fn is_nan(&self) -> Self {
107 self.inner.clone().is_nan().into()
108 }
109
110 fn is_not_nan(&self) -> Self {
111 self.inner.clone().is_not_nan().into()
112 }
113
114 fn min(&self) -> Self {
115 self.inner.clone().min().into()
116 }
117 fn max(&self) -> Self {
118 self.inner.clone().max().into()
119 }
120 #[cfg(feature = "propagate_nans")]
121 fn nan_max(&self) -> Self {
122 self.inner.clone().nan_max().into()
123 }
124 #[cfg(feature = "propagate_nans")]
125 fn nan_min(&self) -> Self {
126 self.inner.clone().nan_min().into()
127 }
128 fn mean(&self) -> Self {
129 self.inner.clone().mean().into()
130 }
131 fn median(&self) -> Self {
132 self.inner.clone().median().into()
133 }
134 fn sum(&self) -> Self {
135 self.inner.clone().sum().into()
136 }
137 fn n_unique(&self) -> Self {
138 self.inner.clone().n_unique().into()
139 }
140 fn arg_unique(&self) -> Self {
141 self.inner.clone().arg_unique().into()
142 }
143 fn unique(&self) -> Self {
144 self.inner.clone().unique().into()
145 }
146 fn unique_stable(&self) -> Self {
147 self.inner.clone().unique_stable().into()
148 }
149 fn first(&self) -> Self {
150 self.inner.clone().first().into()
151 }
152 fn last(&self) -> Self {
153 self.inner.clone().last().into()
154 }
155 fn item(&self, allow_empty: bool) -> Self {
156 self.inner.clone().item(allow_empty).into()
157 }
158 fn implode(&self) -> Self {
159 self.inner.clone().implode().into()
160 }
161 fn quantile(&self, quantile: Self, interpolation: Wrap<QuantileMethod>) -> Self {
162 self.inner
163 .clone()
164 .quantile(quantile.inner, interpolation.0)
165 .into()
166 }
167
168 #[pyo3(signature = (breaks, labels, left_closed, include_breaks))]
169 #[cfg(feature = "cutqcut")]
170 fn cut(
171 &self,
172 breaks: Vec<f64>,
173 labels: Option<Vec<String>>,
174 left_closed: bool,
175 include_breaks: bool,
176 ) -> Self {
177 self.inner
178 .clone()
179 .cut(breaks, labels, left_closed, include_breaks)
180 .into()
181 }
182 #[pyo3(signature = (probs, labels, left_closed, allow_duplicates, include_breaks))]
183 #[cfg(feature = "cutqcut")]
184 fn qcut(
185 &self,
186 probs: Vec<f64>,
187 labels: Option<Vec<String>>,
188 left_closed: bool,
189 allow_duplicates: bool,
190 include_breaks: bool,
191 ) -> Self {
192 self.inner
193 .clone()
194 .qcut(probs, labels, left_closed, allow_duplicates, include_breaks)
195 .into()
196 }
197 #[pyo3(signature = (n_bins, labels, left_closed, allow_duplicates, include_breaks))]
198 #[cfg(feature = "cutqcut")]
199 fn qcut_uniform(
200 &self,
201 n_bins: usize,
202 labels: Option<Vec<String>>,
203 left_closed: bool,
204 allow_duplicates: bool,
205 include_breaks: bool,
206 ) -> Self {
207 self.inner
208 .clone()
209 .qcut_uniform(
210 n_bins,
211 labels,
212 left_closed,
213 allow_duplicates,
214 include_breaks,
215 )
216 .into()
217 }
218
219 #[cfg(feature = "rle")]
220 fn rle(&self) -> Self {
221 self.inner.clone().rle().into()
222 }
223 #[cfg(feature = "rle")]
224 fn rle_id(&self) -> Self {
225 self.inner.clone().rle_id().into()
226 }
227
228 fn agg_groups(&self) -> Self {
229 self.inner.clone().agg_groups().into()
230 }
231 fn count(&self) -> Self {
232 self.inner.clone().count().into()
233 }
234 fn len(&self) -> Self {
235 self.inner.clone().len().into()
236 }
237 fn value_counts(&self, sort: bool, parallel: bool, name: String, normalize: bool) -> Self {
238 self.inner
239 .clone()
240 .value_counts(sort, parallel, name.as_str(), normalize)
241 .into()
242 }
243 fn unique_counts(&self) -> Self {
244 self.inner.clone().unique_counts().into()
245 }
246 fn null_count(&self) -> Self {
247 self.inner.clone().null_count().into()
248 }
249 fn cast(&self, dtype: PyDataTypeExpr, strict: bool, wrap_numerical: bool) -> Self {
250 let options = if wrap_numerical {
251 CastOptions::Overflowing
252 } else if strict {
253 CastOptions::Strict
254 } else {
255 CastOptions::NonStrict
256 };
257
258 let expr = self.inner.clone().cast_with_options(dtype.inner, options);
259 expr.into()
260 }
261 fn sort_with(&self, descending: bool, nulls_last: bool) -> Self {
262 self.inner
263 .clone()
264 .sort(SortOptions {
265 descending,
266 nulls_last,
267 multithreaded: true,
268 maintain_order: false,
269 limit: None,
270 })
271 .into()
272 }
273
274 fn arg_sort(&self, descending: bool, nulls_last: bool) -> Self {
275 self.inner.clone().arg_sort(descending, nulls_last).into()
276 }
277
278 #[cfg(feature = "top_k")]
279 fn top_k(&self, k: Self) -> Self {
280 self.inner.clone().top_k(k.inner).into()
281 }
282
283 #[cfg(feature = "top_k")]
284 fn top_k_by(&self, by: Vec<Self>, k: Self, reverse: Vec<bool>) -> Self {
285 let by = by.into_iter().map(|e| e.inner).collect::<Vec<_>>();
286 self.inner.clone().top_k_by(k.inner, by, reverse).into()
287 }
288
289 #[cfg(feature = "top_k")]
290 fn bottom_k(&self, k: Self) -> Self {
291 self.inner.clone().bottom_k(k.inner).into()
292 }
293
294 #[cfg(feature = "top_k")]
295 fn bottom_k_by(&self, by: Vec<Self>, k: Self, reverse: Vec<bool>) -> Self {
296 let by = by.into_iter().map(|e| e.inner).collect::<Vec<_>>();
297 self.inner.clone().bottom_k_by(k.inner, by, reverse).into()
298 }
299
300 #[cfg(feature = "peaks")]
301 fn peak_min(&self) -> Self {
302 self.inner.clone().peak_min().into()
303 }
304
305 #[cfg(feature = "peaks")]
306 fn peak_max(&self) -> Self {
307 self.inner.clone().peak_max().into()
308 }
309
310 fn arg_max(&self) -> Self {
311 self.inner.clone().arg_max().into()
312 }
313
314 fn arg_min(&self) -> Self {
315 self.inner.clone().arg_min().into()
316 }
317
318 #[cfg(feature = "index_of")]
319 fn index_of(&self, element: Self) -> Self {
320 self.inner.clone().index_of(element.inner).into()
321 }
322
323 #[cfg(feature = "search_sorted")]
324 #[pyo3(signature = (element, side, descending))]
325 fn search_sorted(&self, element: Self, side: Wrap<SearchSortedSide>, descending: bool) -> Self {
326 self.inner
327 .clone()
328 .search_sorted(element.inner, side.0, descending)
329 .into()
330 }
331
332 fn gather(&self, idx: Self) -> Self {
333 self.inner.clone().gather(idx.inner).into()
334 }
335
336 fn get(&self, idx: Self) -> Self {
337 self.inner.clone().get(idx.inner).into()
338 }
339
340 fn sort_by(
341 &self,
342 by: Vec<Self>,
343 descending: Vec<bool>,
344 nulls_last: Vec<bool>,
345 multithreaded: bool,
346 maintain_order: bool,
347 ) -> Self {
348 let by = by.into_iter().map(|e| e.inner).collect::<Vec<_>>();
349 self.inner
350 .clone()
351 .sort_by(
352 by,
353 SortMultipleOptions {
354 descending,
355 nulls_last,
356 multithreaded,
357 maintain_order,
358 limit: None,
359 },
360 )
361 .into()
362 }
363
364 #[pyo3(signature = (n, fill_value))]
365 fn shift(&self, n: Self, fill_value: Option<Self>) -> Self {
366 let expr = self.inner.clone();
367 let out = match fill_value {
368 Some(v) => expr.shift_and_fill(n.inner, v.inner),
369 None => expr.shift(n.inner),
370 };
371 out.into()
372 }
373
374 fn fill_null(&self, expr: Self) -> Self {
375 self.inner.clone().fill_null(expr.inner).into()
376 }
377
378 fn fill_null_with_strategy(&self, strategy: &str, limit: FillNullLimit) -> PyResult<Self> {
379 let strategy = parse_fill_null_strategy(strategy, limit)?;
380 Ok(self.inner.clone().fill_null_with_strategy(strategy).into())
381 }
382
383 fn fill_nan(&self, expr: Self) -> Self {
384 self.inner.clone().fill_nan(expr.inner).into()
385 }
386
387 fn drop_nulls(&self) -> Self {
388 self.inner.clone().drop_nulls().into()
389 }
390
391 fn drop_nans(&self) -> Self {
392 self.inner.clone().drop_nans().into()
393 }
394
395 fn filter(&self, predicate: Self) -> Self {
396 self.inner.clone().filter(predicate.inner).into()
397 }
398
399 fn reverse(&self) -> Self {
400 self.inner.clone().reverse().into()
401 }
402
403 fn std(&self, ddof: u8) -> Self {
404 self.inner.clone().std(ddof).into()
405 }
406
407 fn var(&self, ddof: u8) -> Self {
408 self.inner.clone().var(ddof).into()
409 }
410
411 fn is_unique(&self) -> Self {
412 self.inner.clone().is_unique().into()
413 }
414
415 fn is_between(&self, lower: Self, upper: Self, closed: Wrap<ClosedInterval>) -> Self {
416 self.inner
417 .clone()
418 .is_between(lower.inner, upper.inner, closed.0)
419 .into()
420 }
421
422 fn is_close(&self, other: Self, abs_tol: f64, rel_tol: f64, nans_equal: bool) -> Self {
423 self.inner
424 .clone()
425 .is_close(other.inner, abs_tol, rel_tol, nans_equal)
426 .into()
427 }
428
429 #[cfg(feature = "approx_unique")]
430 fn approx_n_unique(&self) -> Self {
431 self.inner.clone().approx_n_unique().into()
432 }
433
434 fn is_first_distinct(&self) -> Self {
435 self.inner.clone().is_first_distinct().into()
436 }
437
438 fn is_last_distinct(&self) -> Self {
439 self.inner.clone().is_last_distinct().into()
440 }
441
442 fn explode(&self) -> Self {
443 self.inner.clone().explode().into()
444 }
445
446 fn gather_every(&self, n: usize, offset: usize) -> Self {
447 self.inner.clone().gather_every(n, offset).into()
448 }
449
450 fn slice(&self, offset: Self, length: Self) -> Self {
451 self.inner.clone().slice(offset.inner, length.inner).into()
452 }
453
454 fn append(&self, other: Self, upcast: bool) -> Self {
455 self.inner.clone().append(other.inner, upcast).into()
456 }
457
458 fn rechunk(&self) -> Self {
459 self.inner.clone().rechunk().into()
460 }
461
462 fn round(&self, decimals: u32, mode: Wrap<RoundMode>) -> Self {
463 self.inner.clone().round(decimals, mode.0).into()
464 }
465
466 fn round_sig_figs(&self, digits: i32) -> Self {
467 self.clone().inner.round_sig_figs(digits).into()
468 }
469
470 fn floor(&self) -> Self {
471 self.inner.clone().floor().into()
472 }
473
474 fn ceil(&self) -> Self {
475 self.inner.clone().ceil().into()
476 }
477
478 #[pyo3(signature = (min, max))]
479 fn clip(&self, min: Option<Self>, max: Option<Self>) -> Self {
480 let expr = self.inner.clone();
481 let out = match (min, max) {
482 (Some(min), Some(max)) => expr.clip(min.inner, max.inner),
483 (Some(min), None) => expr.clip_min(min.inner),
484 (None, Some(max)) => expr.clip_max(max.inner),
485 (None, None) => expr,
486 };
487 out.into()
488 }
489
490 fn abs(&self) -> Self {
491 self.inner.clone().abs().into()
492 }
493
494 #[cfg(feature = "trigonometry")]
495 fn sin(&self) -> Self {
496 self.inner.clone().sin().into()
497 }
498
499 #[cfg(feature = "trigonometry")]
500 fn cos(&self) -> Self {
501 self.inner.clone().cos().into()
502 }
503
504 #[cfg(feature = "trigonometry")]
505 fn tan(&self) -> Self {
506 self.inner.clone().tan().into()
507 }
508
509 #[cfg(feature = "trigonometry")]
510 fn cot(&self) -> Self {
511 self.inner.clone().cot().into()
512 }
513
514 #[cfg(feature = "trigonometry")]
515 fn arcsin(&self) -> Self {
516 self.inner.clone().arcsin().into()
517 }
518
519 #[cfg(feature = "trigonometry")]
520 fn arccos(&self) -> Self {
521 self.inner.clone().arccos().into()
522 }
523
524 #[cfg(feature = "trigonometry")]
525 fn arctan(&self) -> Self {
526 self.inner.clone().arctan().into()
527 }
528
529 #[cfg(feature = "trigonometry")]
530 fn arctan2(&self, y: Self) -> Self {
531 self.inner.clone().arctan2(y.inner).into()
532 }
533
534 #[cfg(feature = "trigonometry")]
535 fn sinh(&self) -> Self {
536 self.inner.clone().sinh().into()
537 }
538
539 #[cfg(feature = "trigonometry")]
540 fn cosh(&self) -> Self {
541 self.inner.clone().cosh().into()
542 }
543
544 #[cfg(feature = "trigonometry")]
545 fn tanh(&self) -> Self {
546 self.inner.clone().tanh().into()
547 }
548
549 #[cfg(feature = "trigonometry")]
550 fn arcsinh(&self) -> Self {
551 self.inner.clone().arcsinh().into()
552 }
553
554 #[cfg(feature = "trigonometry")]
555 fn arccosh(&self) -> Self {
556 self.inner.clone().arccosh().into()
557 }
558
559 #[cfg(feature = "trigonometry")]
560 fn arctanh(&self) -> Self {
561 self.inner.clone().arctanh().into()
562 }
563
564 #[cfg(feature = "trigonometry")]
565 pub fn degrees(&self) -> Self {
566 self.inner.clone().degrees().into()
567 }
568
569 #[cfg(feature = "trigonometry")]
570 pub fn radians(&self) -> Self {
571 self.inner.clone().radians().into()
572 }
573
574 #[cfg(feature = "sign")]
575 fn sign(&self) -> Self {
576 self.inner.clone().sign().into()
577 }
578
579 fn is_duplicated(&self) -> Self {
580 self.inner.clone().is_duplicated().into()
581 }
582
583 #[pyo3(signature = (partition_by, order_by, order_by_descending, order_by_nulls_last, mapping_strategy))]
584 fn over(
585 &self,
586 partition_by: Option<Vec<Self>>,
587 order_by: Option<Vec<Self>>,
588 order_by_descending: bool,
589 order_by_nulls_last: bool,
590 mapping_strategy: Wrap<WindowMapping>,
591 ) -> PyResult<Self> {
592 let partition_by = partition_by.map(|partition_by| {
593 partition_by
594 .into_iter()
595 .map(|e| e.inner)
596 .collect::<Vec<Expr>>()
597 });
598
599 let order_by = order_by.map(|order_by| {
600 (
601 order_by.into_iter().map(|e| e.inner).collect::<Vec<Expr>>(),
602 SortOptions {
603 descending: order_by_descending,
604 nulls_last: order_by_nulls_last,
605 maintain_order: false,
606 ..Default::default()
607 },
608 )
609 });
610
611 Ok(self
612 .inner
613 .clone()
614 .over_with_options(partition_by, order_by, mapping_strategy.0)
615 .map_err(PyPolarsErr::from)?
616 .into())
617 }
618
619 fn rolling(
620 &self,
621 index_column: &str,
622 period: &str,
623 offset: &str,
624 closed: Wrap<ClosedWindow>,
625 ) -> PyResult<Self> {
626 let options = RollingGroupOptions {
627 index_column: index_column.into(),
628 period: Duration::try_parse(period).map_err(PyPolarsErr::from)?,
629 offset: Duration::try_parse(offset).map_err(PyPolarsErr::from)?,
630 closed_window: closed.0,
631 };
632
633 Ok(self.inner.clone().rolling(options).into())
634 }
635
636 fn and_(&self, expr: Self) -> Self {
637 self.inner.clone().and(expr.inner).into()
638 }
639
640 fn or_(&self, expr: Self) -> Self {
641 self.inner.clone().or(expr.inner).into()
642 }
643
644 fn xor_(&self, expr: Self) -> Self {
645 self.inner.clone().xor(expr.inner).into()
646 }
647
648 #[cfg(feature = "is_in")]
649 fn is_in(&self, expr: Self, nulls_equal: bool) -> Self {
650 self.inner.clone().is_in(expr.inner, nulls_equal).into()
651 }
652
653 #[cfg(feature = "repeat_by")]
654 fn repeat_by(&self, by: Self) -> Self {
655 self.inner.clone().repeat_by(by.inner).into()
656 }
657
658 fn pow(&self, exponent: Self) -> Self {
659 self.inner.clone().pow(exponent.inner).into()
660 }
661
662 fn sqrt(&self) -> Self {
663 self.inner.clone().sqrt().into()
664 }
665
666 fn cbrt(&self) -> Self {
667 self.inner.clone().cbrt().into()
668 }
669
670 fn cum_sum(&self, reverse: bool) -> Self {
671 self.inner.clone().cum_sum(reverse).into()
672 }
673 fn cum_max(&self, reverse: bool) -> Self {
674 self.inner.clone().cum_max(reverse).into()
675 }
676 fn cum_min(&self, reverse: bool) -> Self {
677 self.inner.clone().cum_min(reverse).into()
678 }
679 fn cum_prod(&self, reverse: bool) -> Self {
680 self.inner.clone().cum_prod(reverse).into()
681 }
682 fn cum_count(&self, reverse: bool) -> Self {
683 self.inner.clone().cum_count(reverse).into()
684 }
685
686 fn cumulative_eval(&self, expr: Self, min_samples: usize) -> Self {
687 self.inner
688 .clone()
689 .cumulative_eval(expr.inner, min_samples)
690 .into()
691 }
692
693 fn product(&self) -> Self {
694 self.inner.clone().product().into()
695 }
696
697 fn dot(&self, other: Self) -> Self {
698 self.inner.clone().dot(other.inner).into()
699 }
700
701 fn reinterpret(&self, signed: bool) -> Self {
702 self.inner.clone().reinterpret(signed).into()
703 }
704 fn mode(&self) -> Self {
705 self.inner.clone().mode().into()
706 }
707 fn interpolate(&self, method: Wrap<InterpolationMethod>) -> Self {
708 self.inner.clone().interpolate(method.0).into()
709 }
710 fn interpolate_by(&self, by: PyExpr) -> Self {
711 self.inner.clone().interpolate_by(by.inner).into()
712 }
713
714 fn lower_bound(&self) -> Self {
715 self.inner.clone().lower_bound().into()
716 }
717
718 fn upper_bound(&self) -> Self {
719 self.inner.clone().upper_bound().into()
720 }
721
722 #[pyo3(signature = (method, descending, seed))]
723 fn rank(&self, method: Wrap<RankMethod>, descending: bool, seed: Option<u64>) -> Self {
724 let options = RankOptions {
725 method: method.0,
726 descending,
727 };
728 self.inner.clone().rank(options, seed).into()
729 }
730
731 fn diff(&self, n: PyExpr, null_behavior: Wrap<NullBehavior>) -> Self {
732 self.inner.clone().diff(n.inner, null_behavior.0).into()
733 }
734
735 #[cfg(feature = "pct_change")]
736 fn pct_change(&self, n: Self) -> Self {
737 self.inner.clone().pct_change(n.inner).into()
738 }
739
740 fn skew(&self, bias: bool) -> Self {
741 self.inner.clone().skew(bias).into()
742 }
743 fn kurtosis(&self, fisher: bool, bias: bool) -> Self {
744 self.inner.clone().kurtosis(fisher, bias).into()
745 }
746
747 #[cfg(feature = "dtype-array")]
748 fn reshape(&self, dims: Vec<i64>) -> Self {
749 self.inner.clone().reshape(&dims).into()
750 }
751
752 fn to_physical(&self) -> Self {
753 self.inner.clone().to_physical().into()
754 }
755
756 #[pyo3(signature = (seed))]
757 fn shuffle(&self, seed: Option<u64>) -> Self {
758 self.inner.clone().shuffle(seed).into()
759 }
760
761 #[pyo3(signature = (n, with_replacement, shuffle, seed))]
762 fn sample_n(&self, n: Self, with_replacement: bool, shuffle: bool, seed: Option<u64>) -> Self {
763 self.inner
764 .clone()
765 .sample_n(n.inner, with_replacement, shuffle, seed)
766 .into()
767 }
768
769 #[pyo3(signature = (frac, with_replacement, shuffle, seed))]
770 fn sample_frac(
771 &self,
772 frac: Self,
773 with_replacement: bool,
774 shuffle: bool,
775 seed: Option<u64>,
776 ) -> Self {
777 self.inner
778 .clone()
779 .sample_frac(frac.inner, with_replacement, shuffle, seed)
780 .into()
781 }
782
783 fn ewm_mean(&self, alpha: f64, adjust: bool, min_periods: usize, ignore_nulls: bool) -> Self {
784 let options = EWMOptions {
785 alpha,
786 adjust,
787 bias: false,
788 min_periods,
789 ignore_nulls,
790 };
791 self.inner.clone().ewm_mean(options).into()
792 }
793 fn ewm_mean_by(&self, times: PyExpr, half_life: &str) -> PyResult<Self> {
794 let half_life = Duration::try_parse(half_life).map_err(PyPolarsErr::from)?;
795 Ok(self
796 .inner
797 .clone()
798 .ewm_mean_by(times.inner, half_life)
799 .into())
800 }
801
802 fn ewm_std(
803 &self,
804 alpha: f64,
805 adjust: bool,
806 bias: bool,
807 min_periods: usize,
808 ignore_nulls: bool,
809 ) -> Self {
810 let options = EWMOptions {
811 alpha,
812 adjust,
813 bias,
814 min_periods,
815 ignore_nulls,
816 };
817 self.inner.clone().ewm_std(options).into()
818 }
819 fn ewm_var(
820 &self,
821 alpha: f64,
822 adjust: bool,
823 bias: bool,
824 min_periods: usize,
825 ignore_nulls: bool,
826 ) -> Self {
827 let options = EWMOptions {
828 alpha,
829 adjust,
830 bias,
831 min_periods,
832 ignore_nulls,
833 };
834 self.inner.clone().ewm_var(options).into()
835 }
836 fn extend_constant(&self, value: PyExpr, n: PyExpr) -> Self {
837 self.inner
838 .clone()
839 .extend_constant(value.inner, n.inner)
840 .into()
841 }
842
843 fn any(&self, ignore_nulls: bool) -> Self {
844 self.inner.clone().any(ignore_nulls).into()
845 }
846 fn all(&self, ignore_nulls: bool) -> Self {
847 self.inner.clone().all(ignore_nulls).into()
848 }
849
850 fn log(&self, base: PyExpr) -> Self {
851 self.inner.clone().log(base.inner).into()
852 }
853
854 fn log1p(&self) -> Self {
855 self.inner.clone().log1p().into()
856 }
857
858 fn exp(&self) -> Self {
859 self.inner.clone().exp().into()
860 }
861
862 fn entropy(&self, base: f64, normalize: bool) -> Self {
863 self.inner.clone().entropy(base, normalize).into()
864 }
865 fn hash(&self, seed: u64, seed_1: u64, seed_2: u64, seed_3: u64) -> Self {
866 self.inner.clone().hash(seed, seed_1, seed_2, seed_3).into()
867 }
868 fn set_sorted_flag(&self, descending: bool) -> Self {
869 let is_sorted = if descending {
870 IsSorted::Descending
871 } else {
872 IsSorted::Ascending
873 };
874 self.inner.clone().set_sorted_flag(is_sorted).into()
875 }
876
877 fn replace(&self, old: PyExpr, new: PyExpr) -> Self {
878 self.inner.clone().replace(old.inner, new.inner).into()
879 }
880
881 #[pyo3(signature = (old, new, default, return_dtype))]
882 fn replace_strict(
883 &self,
884 old: PyExpr,
885 new: PyExpr,
886 default: Option<PyExpr>,
887 return_dtype: Option<PyDataTypeExpr>,
888 ) -> Self {
889 self.inner
890 .clone()
891 .replace_strict(
892 old.inner,
893 new.inner,
894 default.map(|e| e.inner),
895 return_dtype.map(|dt| dt.inner),
896 )
897 .into()
898 }
899
900 #[cfg(feature = "hist")]
901 #[pyo3(signature = (bins, bin_count, include_category, include_breakpoint))]
902 fn hist(
903 &self,
904 bins: Option<PyExpr>,
905 bin_count: Option<usize>,
906 include_category: bool,
907 include_breakpoint: bool,
908 ) -> Self {
909 let bins = bins.map(|e| e.inner);
910 self.inner
911 .clone()
912 .hist(bins, bin_count, include_category, include_breakpoint)
913 .into()
914 }
915
916 #[pyo3(signature = (schema))]
917 fn skip_batch_predicate(&self, py: Python<'_>, schema: Wrap<Schema>) -> PyResult<Option<Self>> {
918 let mut aexpr_arena = Arena::new();
919 py.enter_polars(|| {
920 let mut ctx = ExprToIRContext::new(&mut aexpr_arena, &schema.0);
921 ctx.allow_unknown = true;
922 let node = to_expr_ir(self.inner.clone(), &mut ctx)?.node();
923 let Some(node) = aexpr_to_skip_batch_predicate(node, &mut aexpr_arena, &schema.0)
924 else {
925 return Ok(None);
926 };
927 let skip_batch_predicate = node_to_expr(node, &aexpr_arena);
928 PolarsResult::Ok(Some(Self {
929 inner: skip_batch_predicate,
930 }))
931 })
932 }
933
934 #[staticmethod]
935 fn row_encode_unordered(exprs: Vec<Self>) -> Self {
936 Expr::n_ary(
937 FunctionExpr::RowEncode(RowEncodingVariant::Unordered),
938 exprs.into_iter().map(|e| e.inner.clone()).collect(),
939 )
940 .into()
941 }
942
943 #[staticmethod]
944 fn row_encode_ordered(
945 exprs: Vec<Self>,
946 descending: Option<Vec<bool>>,
947 nulls_last: Option<Vec<bool>>,
948 ) -> Self {
949 Expr::n_ary(
950 FunctionExpr::RowEncode(RowEncodingVariant::Ordered {
951 descending,
952 nulls_last,
953 }),
954 exprs.into_iter().map(|e| e.inner.clone()).collect(),
955 )
956 .into()
957 }
958
959 fn row_decode_unordered(&self, names: Vec<String>, datatypes: Vec<PyDataTypeExpr>) -> Self {
960 let fields = names
961 .into_iter()
962 .zip(datatypes)
963 .map(|(name, dtype)| (PlSmallStr::from_string(name), dtype.inner))
964 .collect();
965 self.inner
966 .clone()
967 .map_unary(FunctionExpr::RowDecode(
968 fields,
969 RowEncodingVariant::Unordered,
970 ))
971 .into()
972 }
973
974 fn row_decode_ordered(
975 &self,
976 names: Vec<String>,
977 datatypes: Vec<PyDataTypeExpr>,
978 descending: Option<Vec<bool>>,
979 nulls_last: Option<Vec<bool>>,
980 ) -> Self {
981 let fields = names
982 .into_iter()
983 .zip(datatypes)
984 .map(|(name, dtype)| (PlSmallStr::from_string(name), dtype.inner))
985 .collect::<Vec<_>>();
986 self.inner
987 .clone()
988 .map_unary(FunctionExpr::RowDecode(
989 fields,
990 RowEncodingVariant::Ordered {
991 descending,
992 nulls_last,
993 },
994 ))
995 .into()
996 }
997
998 #[allow(clippy::wrong_self_convention)]
999 fn into_selector(&self) -> PyResult<PySelector> {
1000 Ok(self
1001 .inner
1002 .clone()
1003 .into_selector()
1004 .ok_or_else(
1005 || polars_err!(InvalidOperation: "expr `{}` is not a selector", &self.inner),
1006 )
1007 .map_err(PyPolarsErr::from)?
1008 .into())
1009 }
1010
1011 #[staticmethod]
1012 fn new_selector(selector: PySelector) -> Self {
1013 Expr::Selector(selector.inner).into()
1014 }
1015}