refactor[rust]: remove explicit shift expression node (#4667)

polars/polars-lazy/src/dsl/expr.rs

@@ -314,10 +314,6 @@ pub enum Expr {
         function: FunctionExpr,
         options: FunctionOptions,
     },
-    Shift {
-        input: Box<Expr>,
-        periods: i64,
-    },
     Reverse(Box<Expr>),
     Duplicated(Box<Expr>),
     IsUnique(Box<Expr>),

polars/polars-lazy/src/dsl/function_expr/mod.rs

@@ -16,6 +16,7 @@ mod row_hash;
 mod schema;
 #[cfg(feature = "search_sorted")]
 mod search_sorted;
+mod shift;
 mod shift_and_fill;
 #[cfg(feature = "sign")]
 mod sign;
@@ -92,6 +93,7 @@ pub enum FunctionExpr {
         k: usize,
         reverse: bool,
     },
+    Shift(i64),
 }
 
 #[cfg(feature = "trigonometry")]
@@ -131,19 +133,6 @@ macro_rules! map_as_slice {
     }};
 }
 
-// Fn(&Series)
-#[macro_export(super)]
-macro_rules! map_without_args {
-    ($func:path) => {{
-        let f = move |s: &mut [Series]| {
-            let s = &s[0];
-            $func(s)
-        };
-
-        SpecialEq::new(Arc::new(f))
-    }};
-}
-
 // FnOnce(Series)
 // FnOnce(Series, args)
 #[macro_export(super)]
@@ -168,7 +157,17 @@ macro_rules! map_owned {
 }
 
 // Fn(&Series, args)
-macro_rules! map_with_args {
+#[macro_export(super)]
+macro_rules! map {
+    ($func:path) => {{
+        let f = move |s: &mut [Series]| {
+            let s = &s[0];
+            $func(s)
+        };
+
+        SpecialEq::new(Arc::new(f))
+    }};
+
     ($func:path, $($args:expr),*) => {{
         let f = move |s: &mut [Series]| {
             let s = &s[0];
@@ -195,7 +194,7 @@ impl From<FunctionExpr> for SpecialEq<Arc<dyn SeriesUdf>> {
             }
             #[cfg(feature = "row_hash")]
             Hash(k0, k1, k2, k3) => {
-                map_with_args!(row_hash::row_hash, k0, k1, k2, k3)
+                map!(row_hash::row_hash, k0, k1, k2, k3)
             }
             #[cfg(feature = "is_in")]
             IsIn => {
@@ -218,11 +217,11 @@ impl From<FunctionExpr> for SpecialEq<Arc<dyn SeriesUdf>> {
             }
             #[cfg(feature = "trigonometry")]
             Trigonometry(trig_function) => {
-                map_with_args!(trigonometry::apply_trigonometric_function, trig_function)
+                map!(trigonometry::apply_trigonometric_function, trig_function)
             }
             #[cfg(feature = "sign")]
             Sign => {
-                map_without_args!(sign::sign)
+                map!(sign::sign)
             }
             FillNull { super_type } => {
                 map_as_slice!(fill_null::fill_null, &super_type)
@@ -234,7 +233,7 @@ impl From<FunctionExpr> for SpecialEq<Arc<dyn SeriesUdf>> {
             }
             #[cfg(all(feature = "rolling_window", feature = "moment"))]
             RollingSkew { window_size, bias } => {
-                map_with_args!(rolling::rolling_skew, window_size, bias)
+                map!(rolling::rolling_skew, window_size, bias)
             }
             ShiftAndFill { periods } => {
                 map_as_slice!(shift_and_fill::shift_and_fill, periods)
@@ -255,14 +254,15 @@ impl From<FunctionExpr> for SpecialEq<Arc<dyn SeriesUdf>> {
             StructExpr(sf) => {
                 use StructFunction::*;
                 match sf {
-                    FieldByIndex(index) => map_with_args!(struct_::get_by_index, index),
-                    FieldByName(name) => map_with_args!(struct_::get_by_name, name.clone()),
+                    FieldByIndex(index) => map!(struct_::get_by_index, index),
+                    FieldByName(name) => map!(struct_::get_by_name, name.clone()),
                 }
             }
             #[cfg(feature = "top_k")]
             TopK { k, reverse } => {
-                map_with_args!(top_k, k, reverse)
+                map!(top_k, k, reverse)
             }
+            Shift(periods) => map!(shift::shift, periods),
         }
     }
 }
@@ -273,45 +273,45 @@ impl From<StringFunction> for SpecialEq<Arc<dyn SeriesUdf>> {
         use StringFunction::*;
         match func {
             Contains { pat, literal } => {
-                map_with_args!(strings::contains, &pat, literal)
+                map!(strings::contains, &pat, literal)
             }
             EndsWith(sub) => {
-                map_with_args!(strings::ends_with, &sub)
+                map!(strings::ends_with, &sub)
             }
             StartsWith(sub) => {
-                map_with_args!(strings::starts_with, &sub)
+                map!(strings::starts_with, &sub)
             }
             Extract { pat, group_index } => {
-                map_with_args!(strings::extract, &pat, group_index)
+                map!(strings::extract, &pat, group_index)
             }
             ExtractAll(pat) => {
-                map_with_args!(strings::extract_all, &pat)
+                map!(strings::extract_all, &pat)
             }
             CountMatch(pat) => {
-                map_with_args!(strings::count_match, &pat)
+                map!(strings::count_match, &pat)
             }
             #[cfg(feature = "string_justify")]
             Zfill(alignment) => {
-                map_with_args!(strings::zfill, alignment)
+                map!(strings::zfill, alignment)
             }
             #[cfg(feature = "string_justify")]
             LJust { width, fillchar } => {
-                map_with_args!(strings::ljust, width, fillchar)
+                map!(strings::ljust, width, fillchar)
             }
             #[cfg(feature = "string_justify")]
             RJust { width, fillchar } => {
-                map_with_args!(strings::rjust, width, fillchar)
+                map!(strings::rjust, width, fillchar)
             }
             #[cfg(feature = "temporal")]
             Strptime(options) => {
-                map_with_args!(strings::strptime, &options)
+                map!(strings::strptime, &options)
             }
             #[cfg(feature = "concat_str")]
-            Concat(delimiter) => map_with_args!(strings::concat, &delimiter),
+            Concat(delimiter) => map!(strings::concat, &delimiter),
             #[cfg(feature = "regex")]
             Replace { all, literal } => map_as_slice!(strings::replace, literal, all),
-            Uppercase => map_without_args!(strings::uppercase),
-            Lowercase => map_without_args!(strings::lowercase),
+            Uppercase => map!(strings::uppercase),
+            Lowercase => map!(strings::lowercase),
         }
     }
 }

polars/polars-lazy/src/dsl/function_expr/nan.rs

@@ -2,7 +2,7 @@
 use serde::{Deserialize, Serialize};
 
 use super::*;
-use crate::{map_owned, map_without_args};
+use crate::{map, map_owned};
 
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
 #[derive(Clone, PartialEq, Debug, Eq, Hash)]
@@ -56,8 +56,8 @@ impl NanFunction {
 impl From<NanFunction> for SpecialEq<Arc<dyn SeriesUdf>> {
     fn from(nan_function: NanFunction) -> Self {
         match nan_function {
-            NanFunction::IsNan => map_without_args!(is_nan),
-            NanFunction::IsNotNan => map_without_args!(is_not_nan),
+            NanFunction::IsNan => map!(is_nan),
+            NanFunction::IsNotNan => map!(is_not_nan),
             NanFunction::DropNans => map_owned!(drop_nans),
         }
     }

polars/polars-lazy/src/dsl/function_expr/schema.rs

@@ -154,6 +154,7 @@ impl FunctionExpr {
             }
             #[cfg(feature = "top_k")]
             TopK { .. } => same_type(),
+            Shift(..) => same_type(),
         }
     }
 }

polars/polars-lazy/src/dsl/function_expr/shift.rs

@@ -0,0 +1,5 @@
+use super::*;
+
+pub(super) fn shift(s: &Series, periods: i64) -> Result<Series> {
+    Ok(s.shift(periods))
+}

polars/polars-lazy/src/dsl/mod.rs

@@ -847,10 +847,7 @@ impl Expr {
 
     /// Shift the values in the array by some period. See [the eager implementation](polars_core::series::SeriesTrait::shift).
     pub fn shift(self, periods: i64) -> Self {
-        Expr::Shift {
-            input: Box::new(self),
-            periods,
-        }
+        self.apply_private(FunctionExpr::Shift(periods), "shift")
     }
 
     /// Shift the values in the array by some period and fill the resulting empty values.

polars/polars-lazy/src/logical_plan/aexpr.rs

@@ -91,10 +91,6 @@ pub enum AExpr {
         function: FunctionExpr,
         options: FunctionOptions,
     },
-    Shift {
-        input: Node,
-        periods: i64,
-    },
     Window {
         function: Node,
         partition_by: Vec<Node>,
@@ -132,7 +128,6 @@ impl AExpr {
             Sort { .. }
             | SortBy { .. }
             | Agg { .. }
-            | Shift { .. }
             | Window { .. }
             | Count
             | Slice { .. }
@@ -393,7 +388,6 @@ impl AExpr {
                     .collect::<Result<Vec<_>>>()?;
                 function.get_field(schema, ctxt, &fields)
             }
-            Shift { input, .. } => arena.get(*input).to_field(schema, ctxt, arena),
             Slice { input, .. } => arena.get(*input).to_field(schema, ctxt, arena),
             Wildcard => panic!("should be no wildcard at this point"),
             Nth(_) => panic!("should be no nth at this point"),

polars/polars-lazy/src/logical_plan/conversion.rs

@@ -116,10 +116,6 @@ pub(crate) fn to_aexpr(expr: Expr, arena: &mut Arena<AExpr>) -> Node {
             function,
             options,
         },
-        Expr::Shift { input, periods } => AExpr::Shift {
-            input: to_aexpr(*input, arena),
-            periods,
-        },
         Expr::Window {
             function,
             partition_by,
@@ -598,13 +594,6 @@ pub(crate) fn node_to_expr(node: Node, expr_arena: &Arena<AExpr>) -> Expr {
                 AggExpr::Count(Box::new(exp)).into()
             }
         },
-        AExpr::Shift { input, periods } => {
-            let e = node_to_expr(input, expr_arena);
-            Expr::Shift {
-                input: Box::new(e),
-                periods,
-            }
-        }
         AExpr::Ternary {
             predicate,
             truthy,

polars/polars-lazy/src/logical_plan/format.rs

@@ -284,7 +284,6 @@ impl fmt::Debug for Expr {
                     write!(f, "{:?}.{}()", input[0], options.fmt_str)
                 }
             }
-            Shift { input, periods, .. } => write!(f, "SHIFT {:?} by {}", input, periods),
             Slice {
                 input,
                 offset,

polars/polars-lazy/src/logical_plan/iterator.rs

@@ -64,7 +64,6 @@ macro_rules! push_expr {
             // as the root columns/ input columns by `_suffix` and `_keep_name` etc.
             AnonymousFunction { input, .. } => input.$iter().rev().for_each(|e| $push(e)),
             Function { input, .. } => input.$iter().rev().for_each(|e| $push(e)),
-            Shift { input, .. } => $push(input),
             Reverse(e) => $push(e),
             Duplicated(e) => $push(e),
             IsUnique(e) => $push(e),
@@ -235,7 +234,6 @@ impl AExpr {
             {
                 input.iter().rev().for_each(push)
             }
-            Shift { input, .. } => push(input),
             Reverse(e) => push(e),
             Duplicated(e) => push(e),
             IsUnique(e) => push(e),

polars/polars-lazy/src/logical_plan/optimizer/predicate_pushdown/utils.rs

@@ -114,7 +114,7 @@ pub(super) fn predicate_is_pushdown_boundary(node: Node, expr_arena: &Arena<AExp
     let matches = |e: &AExpr| {
         matches!(
             e,
-            AExpr::Shift { .. } | AExpr::Sort { .. } | AExpr::SortBy { .. }
+            AExpr::Sort { .. } | AExpr::SortBy { .. }
             | AExpr::Agg(_) // an aggregation needs all rows
             | AExpr::Reverse(_)
             // Apply groups can be something like shift, sort, or an aggregation like skew
@@ -142,7 +142,7 @@ pub(super) fn project_other_column_is_predicate_pushdown_boundary(
     let matches = |e: &AExpr| {
         matches!(
             e,
-            AExpr::Shift { .. } | AExpr::Sort { .. } | AExpr::SortBy { .. }
+            AExpr::Sort { .. } | AExpr::SortBy { .. }
             | AExpr::Agg(_) // an aggregation needs all rows
             | AExpr::Reverse(_)
             // Apply groups can be something like shift, sort, or an aggregation like skew
@@ -175,7 +175,7 @@ pub(super) fn projection_column_is_predicate_pushdown_boundary(
     let matches = |e: &AExpr| {
         matches!(
             e,
-            AExpr::Shift { .. } | AExpr::Sort { .. } | AExpr::SortBy { .. }
+            AExpr::Sort { .. } | AExpr::SortBy { .. }
             | AExpr::Agg(_) // an aggregation needs all rows
             | AExpr::Reverse(_)
             // everything that works on groups likely changes to order of elements w/r/t the other columns

polars/polars-lazy/src/physical_plan/expressions/apply.rs

@@ -154,7 +154,16 @@ impl PhysicalExpr for ApplyExpr {
                 }
                 ApplyOptions::ApplyFlat => {
                     // make sure the groups are updated because we are about to throw away
-                    // the series length information
+                    // the series' length information
+                    let set_update_groups = match ac.update_groups {
+                        UpdateGroups::WithSeriesLen => {
+                            ac.groups();
+                            true
+                        }
+                        UpdateGroups::WithSeriesLenOwned(_) => false,
+                        UpdateGroups::No | UpdateGroups::WithGroupsLen => false,
+                    };
+
                     if let UpdateGroups::WithSeriesLen = ac.update_groups {
                         ac.groups();
                     }
@@ -164,6 +173,13 @@ impl PhysicalExpr for ApplyExpr {
 
                     check_map_output_len(input_len, s.len())?;
                     ac.with_series(s, false);
+
+                    if set_update_groups {
+                        // The flat_naive orders by groups, so we must create new groups
+                        // not by series length as we don't have an agg_list, but by original
+                        // groups length
+                        ac.update_groups = UpdateGroups::WithGroupsLen;
+                    }
                     Ok(ac)
                 }
                 ApplyOptions::ApplyList => {

polars/polars-lazy/src/physical_plan/expressions/mod.rs

@@ -11,7 +11,6 @@ mod is_not_null;
 mod is_null;
 mod literal;
 mod not;
-mod shift;
 mod slice;
 mod sort;
 mod sortby;
@@ -29,8 +28,8 @@ use polars_core::prelude::*;
 use polars_io::predicates::PhysicalIoExpr;
 pub(crate) use {
     aggregation::*, alias::*, apply::*, binary::*, cast::*, column::*, count::*, filter::*,
-    is_not_null::*, is_null::*, literal::*, not::*, shift::*, slice::*, sort::*, sortby::*,
-    take::*, ternary::*, window::*,
+    is_not_null::*, is_null::*, literal::*, not::*, slice::*, sort::*, sortby::*, take::*,
+    ternary::*, window::*,
 };
 
 use crate::physical_plan::state::ExecutionState;