1use std::sync::Arc;
2
3use polars_compute::rolling::QuantileMethod;
4use polars_core::error::PolarsResult;
5use polars_core::frame::DataFrame;
6use polars_core::prelude::{Column, GroupPositions};
7use polars_plan::dsl::{ColumnsUdf, SpecialEq};
8use polars_plan::plans::{IRBooleanFunction, IRFunctionExpr, IRPowFunction};
9use polars_utils::IdxSize;
10
11use crate::prelude::{AggregationContext, PhysicalExpr};
12use crate::state::ExecutionState;
13
14#[macro_export]
15macro_rules! wrap {
16 ($e:expr) => {
17 SpecialEq::new(Arc::new($e))
18 };
19
20 ($e:expr, $($args:expr),*) => {{
21 let f = move |s: &mut [::polars_core::prelude::Column]| {
22 $e(s, $($args),*)
23 };
24
25 SpecialEq::new(Arc::new(f))
26 }};
27}
28
29#[macro_export]
33macro_rules! map_as_slice {
34 ($func:path) => {{
35 let f = move |s: &mut [::polars_core::prelude::Column]| {
36 $func(s)
37 };
38
39 SpecialEq::new(Arc::new(f))
40 }};
41
42 ($func:path, $($args:expr),*) => {{
43 let f = move |s: &mut [::polars_core::prelude::Column]| {
44 $func(s, $($args),*)
45 };
46
47 SpecialEq::new(Arc::new(f))
48 }};
49}
50
51#[macro_export]
54macro_rules! map_owned {
55 ($func:path) => {{
56 let f = move |c: &mut [::polars_core::prelude::Column]| {
57 let c = std::mem::take(&mut c[0]);
58 $func(c)
59 };
60
61 SpecialEq::new(Arc::new(f))
62 }};
63
64 ($func:path, $($args:expr),*) => {{
65 let f = move |c: &mut [::polars_core::prelude::Column]| {
66 let c = std::mem::take(&mut c[0]);
67 $func(c, $($args),*)
68 };
69
70 SpecialEq::new(Arc::new(f))
71 }};
72}
73
74#[macro_export]
76macro_rules! map {
77 ($func:path) => {{
78 let f = move |c: &mut [::polars_core::prelude::Column]| {
79 let c = &c[0];
80 $func(c)
81 };
82
83 SpecialEq::new(Arc::new(f))
84 }};
85
86 ($func:path, $($args:expr),*) => {{
87 let f = move |c: &mut [::polars_core::prelude::Column]| {
88 let c = &c[0];
89 $func(c, $($args),*)
90 };
91
92 SpecialEq::new(Arc::new(f))
93 }};
94}
95
96#[cfg(feature = "dtype-array")]
97mod array;
98mod binary;
99#[cfg(feature = "bitwise")]
100mod bitwise;
101mod boolean;
102#[cfg(feature = "business")]
103mod business;
104#[cfg(feature = "dtype-categorical")]
105mod cat;
106#[cfg(feature = "cum_agg")]
107mod cum;
108#[cfg(feature = "temporal")]
109mod datetime;
110#[cfg(feature = "dtype-extension")]
111mod extension;
112mod groups_dispatch;
113mod horizontal;
114mod list;
115mod misc;
116mod pow;
117#[cfg(feature = "random")]
118mod random;
119#[cfg(feature = "range")]
120mod range;
121#[cfg(feature = "rolling_window")]
122mod rolling;
123#[cfg(feature = "rolling_window_by")]
124mod rolling_by;
125#[cfg(feature = "round_series")]
126mod round;
127mod shift_and_fill;
128#[cfg(feature = "strings")]
129mod strings;
130#[cfg(feature = "dtype-struct")]
131pub(crate) mod struct_;
132#[cfg(feature = "temporal")]
133mod temporal;
134#[cfg(feature = "trigonometry")]
135mod trigonometry;
136
137pub use groups_dispatch::drop_items;
138
139pub fn function_expr_to_udf(func: IRFunctionExpr) -> SpecialEq<Arc<dyn ColumnsUdf>> {
140 use IRFunctionExpr as F;
141 match func {
142 #[cfg(feature = "dtype-array")]
144 F::ArrayExpr(func) => array::function_expr_to_udf(func),
145 F::BinaryExpr(func) => binary::function_expr_to_udf(func),
146 #[cfg(feature = "dtype-categorical")]
147 F::Categorical(func) => cat::function_expr_to_udf(func),
148 #[cfg(feature = "dtype-extension")]
149 F::Extension(func) => extension::function_expr_to_udf(func),
150 F::ListExpr(func) => list::function_expr_to_udf(func),
151 #[cfg(feature = "strings")]
152 F::StringExpr(func) => strings::function_expr_to_udf(func),
153 #[cfg(feature = "dtype-struct")]
154 F::StructExpr(func) => struct_::function_expr_to_udf(func),
155 #[cfg(feature = "temporal")]
156 F::TemporalExpr(func) => temporal::temporal_func_to_udf(func),
157 #[cfg(feature = "bitwise")]
158 F::Bitwise(func) => bitwise::function_expr_to_udf(func),
159
160 F::Boolean(func) => boolean::function_expr_to_udf(func),
162 #[cfg(feature = "business")]
163 F::Business(func) => business::function_expr_to_udf(func),
164 #[cfg(feature = "abs")]
165 F::Abs => map!(misc::abs),
166 F::Negate => map!(misc::negate),
167 F::NullCount => {
168 let f = |s: &mut [Column]| {
169 let s = &s[0];
170 Ok(Column::new(s.name().clone(), [s.null_count() as IdxSize]))
171 };
172 wrap!(f)
173 },
174 F::Pow(func) => match func {
175 IRPowFunction::Generic => wrap!(pow::pow),
176 IRPowFunction::Sqrt => map!(pow::sqrt),
177 IRPowFunction::Cbrt => map!(pow::cbrt),
178 },
179 #[cfg(feature = "row_hash")]
180 F::Hash(k0, k1, k2, k3) => {
181 map!(misc::row_hash, k0, k1, k2, k3)
182 },
183 #[cfg(feature = "arg_where")]
184 F::ArgWhere => {
185 wrap!(misc::arg_where)
186 },
187 #[cfg(feature = "index_of")]
188 F::IndexOf => {
189 map_as_slice!(misc::index_of)
190 },
191 #[cfg(feature = "search_sorted")]
192 F::SearchSorted { side, descending } => {
193 map_as_slice!(misc::search_sorted_impl, side, descending)
194 },
195 #[cfg(feature = "range")]
196 F::Range(func) => range::function_expr_to_udf(func),
197
198 #[cfg(feature = "trigonometry")]
199 F::Trigonometry(trig_function) => {
200 map!(trigonometry::apply_trigonometric_function, trig_function)
201 },
202 #[cfg(feature = "trigonometry")]
203 F::Atan2 => {
204 wrap!(trigonometry::apply_arctan2)
205 },
206
207 #[cfg(feature = "sign")]
208 F::Sign => {
209 map!(misc::sign)
210 },
211 F::FillNull => {
212 map_as_slice!(misc::fill_null)
213 },
214 #[cfg(feature = "rolling_window")]
215 F::RollingExpr { function, options } => {
216 use IRRollingFunction::*;
217 use polars_plan::plans::IRRollingFunction;
218 match function {
219 Min => map!(rolling::rolling_min, options.clone()),
220 Max => map!(rolling::rolling_max, options.clone()),
221 Mean => map!(rolling::rolling_mean, options.clone()),
222 Sum => map!(rolling::rolling_sum, options.clone()),
223 Quantile => map!(rolling::rolling_quantile, options.clone()),
224 Var => map!(rolling::rolling_var, options.clone()),
225 Std => map!(rolling::rolling_std, options.clone()),
226 Rank => map!(rolling::rolling_rank, options.clone()),
227 #[cfg(feature = "moment")]
228 Skew => map!(rolling::rolling_skew, options.clone()),
229 #[cfg(feature = "moment")]
230 Kurtosis => map!(rolling::rolling_kurtosis, options.clone()),
231 #[cfg(feature = "cov")]
232 CorrCov {
233 corr_cov_options,
234 is_corr,
235 } => {
236 map_as_slice!(
237 rolling::rolling_corr_cov,
238 options.clone(),
239 corr_cov_options,
240 is_corr
241 )
242 },
243 Map(f) => {
244 map!(rolling::rolling_map, options.clone(), f.clone())
245 },
246 }
247 },
248 #[cfg(feature = "rolling_window_by")]
249 F::RollingExprBy {
250 function_by,
251 options,
252 } => {
253 use IRRollingFunctionBy::*;
254 use polars_plan::plans::IRRollingFunctionBy;
255 match function_by {
256 MinBy => map_as_slice!(rolling_by::rolling_min_by, options.clone()),
257 MaxBy => map_as_slice!(rolling_by::rolling_max_by, options.clone()),
258 MeanBy => map_as_slice!(rolling_by::rolling_mean_by, options.clone()),
259 SumBy => map_as_slice!(rolling_by::rolling_sum_by, options.clone()),
260 QuantileBy => {
261 map_as_slice!(rolling_by::rolling_quantile_by, options.clone())
262 },
263 VarBy => map_as_slice!(rolling_by::rolling_var_by, options.clone()),
264 StdBy => map_as_slice!(rolling_by::rolling_std_by, options.clone()),
265 RankBy => map_as_slice!(rolling_by::rolling_rank_by, options.clone()),
266 }
267 },
268 #[cfg(feature = "hist")]
269 F::Hist {
270 bin_count,
271 include_category,
272 include_breakpoint,
273 } => {
274 map_as_slice!(misc::hist, bin_count, include_category, include_breakpoint)
275 },
276 F::Rechunk => map!(misc::rechunk),
277 F::ShiftAndFill => {
278 map_as_slice!(shift_and_fill::shift_and_fill)
279 },
280 F::DropNans => map_owned!(misc::drop_nans),
281 F::DropNulls => map!(misc::drop_nulls),
282 #[cfg(feature = "round_series")]
283 F::Clip { has_min, has_max } => {
284 map_as_slice!(misc::clip, has_min, has_max)
285 },
286 F::Quantile { method } => map_as_slice!(misc::quantile, method),
287 #[cfg(feature = "mode")]
288 F::Mode { maintain_order } => map!(misc::mode, maintain_order),
289 #[cfg(feature = "moment")]
290 F::Skew(bias) => map!(misc::skew, bias),
291 #[cfg(feature = "moment")]
292 F::Kurtosis(fisher, bias) => map!(misc::kurtosis, fisher, bias),
293 F::ArgUnique => map!(misc::arg_unique),
294 F::ArgMin => map!(misc::arg_min),
295 F::ArgMax => map!(misc::arg_max),
296 F::ArgSort {
297 descending,
298 nulls_last,
299 } => map!(misc::arg_sort, descending, nulls_last),
300 F::MinBy => map_as_slice!(misc::min_by),
301 F::MaxBy => map_as_slice!(misc::max_by),
302 F::Product => map!(misc::product),
303 F::Repeat => map_as_slice!(misc::repeat),
304 #[cfg(feature = "rank")]
305 F::Rank { options, seed } => map!(misc::rank, options, seed),
306 F::AsList => map_as_slice!(misc::as_list),
307 #[cfg(feature = "dtype-struct")]
308 F::AsStruct => {
309 map_as_slice!(misc::as_struct)
310 },
311 #[cfg(feature = "top_k")]
312 F::TopK { descending } => {
313 map_as_slice!(polars_ops::prelude::top_k, descending)
314 },
315 #[cfg(feature = "top_k")]
316 F::TopKBy { descending } => {
317 map_as_slice!(polars_ops::prelude::top_k_by, descending.clone())
318 },
319 F::Shift => map_as_slice!(shift_and_fill::shift),
320 #[cfg(feature = "cum_agg")]
321 F::CumCount { reverse } => map!(cum::cum_count, reverse),
322 #[cfg(feature = "cum_agg")]
323 F::CumSum { reverse } => map!(cum::cum_sum, reverse),
324 #[cfg(feature = "cum_agg")]
325 F::CumProd { reverse } => map!(cum::cum_prod, reverse),
326 #[cfg(feature = "cum_agg")]
327 F::CumMin { reverse } => map!(cum::cum_min, reverse),
328 #[cfg(feature = "cum_agg")]
329 F::CumMax { reverse } => map!(cum::cum_max, reverse),
330 #[cfg(feature = "dtype-struct")]
331 F::ValueCounts {
332 sort,
333 parallel,
334 name,
335 normalize,
336 } => map!(misc::value_counts, sort, parallel, name.clone(), normalize),
337 #[cfg(feature = "unique_counts")]
338 F::UniqueCounts => map!(misc::unique_counts),
339 F::Reverse => map!(misc::reverse),
340 #[cfg(feature = "approx_unique")]
341 F::ApproxNUnique => map!(misc::approx_n_unique),
342 F::Coalesce => map_as_slice!(misc::coalesce),
343 #[cfg(feature = "diff")]
344 F::Diff(null_behavior) => map_as_slice!(misc::diff, null_behavior),
345 #[cfg(feature = "pct_change")]
346 F::PctChange => map_as_slice!(misc::pct_change),
347 #[cfg(feature = "interpolate")]
348 F::Interpolate(method) => {
349 map!(misc::interpolate, method)
350 },
351 #[cfg(feature = "interpolate_by")]
352 F::InterpolateBy => {
353 map_as_slice!(misc::interpolate_by)
354 },
355 #[cfg(feature = "log")]
356 F::Entropy { base, normalize } => map!(misc::entropy, base, normalize),
357 #[cfg(feature = "log")]
358 F::Log => map_as_slice!(misc::log),
359 #[cfg(feature = "log")]
360 F::Log1p => map!(misc::log1p),
361 #[cfg(feature = "log")]
362 F::Exp => map!(misc::exp),
363 F::Unique(stable) => map!(misc::unique, stable),
364 #[cfg(feature = "round_series")]
365 F::Round { decimals, mode } => map!(round::round, decimals, mode),
366 #[cfg(feature = "round_series")]
367 F::RoundSF { digits } => map!(round::round_sig_figs, digits),
368 #[cfg(feature = "round_series")]
369 F::Truncate { decimals } => map!(round::truncate, decimals),
370 #[cfg(feature = "round_series")]
371 F::Floor => map!(round::floor),
372 #[cfg(feature = "round_series")]
373 F::Ceil => map!(round::ceil),
374 #[cfg(feature = "fused")]
375 F::Fused(op) => map_as_slice!(misc::fused, op),
376 F::ConcatExpr { rechunk } => map_as_slice!(misc::concat_expr, rechunk),
377 #[cfg(feature = "cov")]
378 F::Correlation { method } => map_as_slice!(misc::corr, method),
379 #[cfg(feature = "peaks")]
380 F::PeakMin => map!(misc::peak_min),
381 #[cfg(feature = "peaks")]
382 F::PeakMax => map!(misc::peak_max),
383 #[cfg(feature = "repeat_by")]
384 F::RepeatBy => map_as_slice!(misc::repeat_by),
385 #[cfg(feature = "dtype-array")]
386 F::Reshape(dims) => map!(misc::reshape, &dims),
387 #[cfg(feature = "cutqcut")]
388 F::Cut {
389 breaks,
390 labels,
391 left_closed,
392 include_breaks,
393 } => map!(
394 misc::cut,
395 breaks.clone(),
396 labels.clone(),
397 left_closed,
398 include_breaks
399 ),
400 #[cfg(feature = "cutqcut")]
401 F::QCut {
402 probs,
403 labels,
404 left_closed,
405 allow_duplicates,
406 include_breaks,
407 } => map!(
408 misc::qcut,
409 probs.clone(),
410 labels.clone(),
411 left_closed,
412 allow_duplicates,
413 include_breaks
414 ),
415 #[cfg(feature = "rle")]
416 F::RLE => map!(polars_ops::series::rle),
417 #[cfg(feature = "rle")]
418 F::RLEID => map!(polars_ops::series::rle_id),
419 F::ToPhysical => map!(misc::to_physical),
420 #[cfg(feature = "random")]
421 F::Random { method, seed } => {
422 use IRRandomMethod::*;
423 use polars_plan::plans::IRRandomMethod;
424 match method {
425 Shuffle => map!(random::shuffle, seed),
426 Sample {
427 is_fraction,
428 with_replacement,
429 shuffle,
430 } => {
431 if is_fraction {
432 map_as_slice!(random::sample_frac, with_replacement, shuffle, seed)
433 } else {
434 map_as_slice!(random::sample_n, with_replacement, shuffle, seed)
435 }
436 },
437 }
438 },
439 F::SetSortedFlag(sortedness) => map!(misc::set_sorted_flag, sortedness),
440 #[cfg(feature = "ffi_plugin")]
441 F::FfiPlugin {
442 flags: _,
443 lib,
444 symbol,
445 kwargs,
446 } => unsafe {
447 map_as_slice!(
448 polars_plan::plans::plugin::call_plugin,
449 lib.as_ref(),
450 symbol.as_ref(),
451 kwargs.as_ref()
452 )
453 },
454
455 F::FoldHorizontal {
456 callback,
457 returns_scalar,
458 return_dtype,
459 } => map_as_slice!(
460 horizontal::fold,
461 &callback,
462 returns_scalar,
463 return_dtype.as_ref()
464 ),
465 F::ReduceHorizontal {
466 callback,
467 returns_scalar,
468 return_dtype,
469 } => map_as_slice!(
470 horizontal::reduce,
471 &callback,
472 returns_scalar,
473 return_dtype.as_ref()
474 ),
475 #[cfg(feature = "dtype-struct")]
476 F::CumReduceHorizontal {
477 callback,
478 returns_scalar,
479 return_dtype,
480 } => map_as_slice!(
481 horizontal::cum_reduce,
482 &callback,
483 returns_scalar,
484 return_dtype.as_ref()
485 ),
486 #[cfg(feature = "dtype-struct")]
487 F::CumFoldHorizontal {
488 callback,
489 returns_scalar,
490 return_dtype,
491 include_init,
492 } => map_as_slice!(
493 horizontal::cum_fold,
494 &callback,
495 returns_scalar,
496 return_dtype.as_ref(),
497 include_init
498 ),
499
500 F::MaxHorizontal => wrap!(misc::max_horizontal),
501 F::MinHorizontal => wrap!(misc::min_horizontal),
502 F::SumHorizontal { ignore_nulls } => wrap!(misc::sum_horizontal, ignore_nulls),
503 F::MeanHorizontal { ignore_nulls } => wrap!(misc::mean_horizontal, ignore_nulls),
504 #[cfg(feature = "ewma")]
505 F::EwmMean { options } => map!(misc::ewm_mean, options),
506 #[cfg(feature = "ewma_by")]
507 F::EwmMeanBy { half_life } => map_as_slice!(misc::ewm_mean_by, half_life),
508 #[cfg(feature = "ewma")]
509 F::EwmSum { options } => map!(misc::ewm_sum, options),
510 #[cfg(feature = "ewma_by")]
511 F::EwmSumBy { half_life } => map_as_slice!(misc::ewm_sum_by, half_life),
512 #[cfg(feature = "ewma")]
513 F::EwmStd { options } => map!(misc::ewm_std, options),
514 #[cfg(feature = "ewma")]
515 F::EwmVar { options } => map!(misc::ewm_var, options),
516 #[cfg(feature = "replace")]
517 F::Replace => {
518 map_as_slice!(misc::replace)
519 },
520 #[cfg(feature = "replace")]
521 F::ReplaceStrict { return_dtype } => {
522 map_as_slice!(misc::replace_strict, return_dtype.clone())
523 },
524
525 F::FillNullWithStrategy(strategy) => map!(misc::fill_null_with_strategy, strategy),
526 F::GatherEvery { n, offset } => map!(misc::gather_every, n, offset),
527 #[cfg(feature = "reinterpret")]
528 F::Reinterpret(dtype) => map!(misc::reinterpret, &dtype),
529 F::ExtendConstant => map_as_slice!(misc::extend_constant),
530
531 F::RowEncode(dts, variants) => {
532 map_as_slice!(misc::row_encode, dts.clone(), variants.clone())
533 },
534 #[cfg(feature = "dtype-struct")]
535 F::RowDecode(fs, variants) => {
536 map_as_slice!(misc::row_decode, fs.clone(), variants.clone())
537 },
538 F::DynamicPred { pred } => {
539 map_as_slice!(misc::dynamic_pred, &pred)
540 },
541 }
542}
543
544pub trait GroupsUdf: Send + Sync + 'static {
545 fn evaluate_on_groups<'a>(
546 &self,
547 inputs: &[Arc<dyn PhysicalExpr>],
548 df: &DataFrame,
549 groups: &'a GroupPositions,
550 state: &ExecutionState,
551 ) -> PolarsResult<AggregationContext<'a>>;
552}
553
554pub fn function_expr_to_groups_udf(func: &IRFunctionExpr) -> Option<SpecialEq<Arc<dyn GroupsUdf>>> {
555 macro_rules! wrap_groups {
556 ($f:expr$(, ($arg:expr, $n:ident:$ty:ty))*) => {{
557 struct Wrap($($ty),*);
558 impl GroupsUdf for Wrap {
559 fn evaluate_on_groups<'a>(
560 &self,
561 inputs: &[Arc<dyn PhysicalExpr>],
562 df: &DataFrame,
563 groups: &'a GroupPositions,
564 state: &ExecutionState,
565 ) -> PolarsResult<AggregationContext<'a>> {
566 let Wrap($($n),*) = self;
567 $f(inputs, df, groups, state$(, *$n)*)
568 }
569 }
570
571 SpecialEq::new(Arc::new(Wrap($($arg),*)) as Arc<dyn GroupsUdf>)
572 }};
573 }
574 use IRFunctionExpr as F;
575 Some(match func {
576 F::NullCount => wrap_groups!(groups_dispatch::null_count),
577 F::Reverse => wrap_groups!(groups_dispatch::reverse),
578 F::Boolean(IRBooleanFunction::HasNulls) => wrap_groups!(groups_dispatch::has_nulls),
579 F::Boolean(IRBooleanFunction::Any { ignore_nulls }) => {
580 let ignore_nulls = *ignore_nulls;
581 wrap_groups!(groups_dispatch::any, (ignore_nulls, v: bool))
582 },
583 F::Boolean(IRBooleanFunction::All { ignore_nulls }) => {
584 let ignore_nulls = *ignore_nulls;
585 wrap_groups!(groups_dispatch::all, (ignore_nulls, v: bool))
586 },
587 F::Boolean(IRBooleanFunction::IsEmpty { ignore_nulls }) => {
588 let ignore_nulls = *ignore_nulls;
589 wrap_groups!(groups_dispatch::is_empty, (ignore_nulls, v: bool))
590 },
591 #[cfg(feature = "bitwise")]
592 F::Bitwise(f) => {
593 use polars_plan::plans::IRBitwiseFunction as B;
594 match f {
595 B::And => wrap_groups!(groups_dispatch::bitwise_and),
596 B::Or => wrap_groups!(groups_dispatch::bitwise_or),
597 B::Xor => wrap_groups!(groups_dispatch::bitwise_xor),
598 _ => return None,
599 }
600 },
601 F::DropNans => wrap_groups!(groups_dispatch::drop_nans),
602 F::DropNulls => wrap_groups!(groups_dispatch::drop_nulls),
603
604 F::Quantile { method } => {
605 wrap_groups!(groups_dispatch::quantile, (*method, v: QuantileMethod))
606 },
607 #[cfg(feature = "moment")]
608 F::Skew(bias) => wrap_groups!(groups_dispatch::skew, (*bias, v: bool)),
609 #[cfg(feature = "moment")]
610 F::Kurtosis(fisher, bias) => {
611 wrap_groups!(groups_dispatch::kurtosis, (*fisher, v1: bool), (*bias, v2: bool))
612 },
613
614 F::Unique(stable) => wrap_groups!(groups_dispatch::unique, (*stable, v: bool)),
615 F::FillNullWithStrategy(polars_core::prelude::FillNullStrategy::Forward(limit)) => {
616 wrap_groups!(groups_dispatch::forward_fill_null, (*limit, v: Option<IdxSize>))
617 },
618 F::FillNullWithStrategy(polars_core::prelude::FillNullStrategy::Backward(limit)) => {
619 wrap_groups!(groups_dispatch::backward_fill_null, (*limit, v: Option<IdxSize>))
620 },
621
622 _ => return None,
623 })
624}