minor: implement with_new_expressions for AggregateFunctionExpr (#16897)

* minor

* Update aggregate.rs

Author: berkaysynnada

datafusion/expr/src/window_state.rs

@@ -34,7 +34,7 @@ use datafusion_common::{
 };
 
 /// Holds the state of evaluating a window function
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub struct WindowAggState {
     /// The range that we calculate the window function
     pub window_frame_range: Range<usize>,
@@ -112,7 +112,7 @@ impl WindowAggState {
 }
 
 /// This object stores the window frame state for use in incremental calculations.
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub enum WindowFrameContext {
     /// ROWS frames are inherently stateless.
     Rows(Arc<WindowFrame>),
@@ -240,7 +240,7 @@ impl WindowFrameContext {
 }
 
 /// State for each unique partition determined according to PARTITION BY column(s)
-#[derive(Debug)]
+#[derive(Debug, Clone, PartialEq)]
 pub struct PartitionBatchState {
     /// The record batch belonging to current partition
     pub record_batch: RecordBatch,
@@ -282,7 +282,7 @@ impl PartitionBatchState {
 /// ranges of data while processing RANGE frames.
 /// Attribute `sort_options` stores the column ordering specified by the ORDER
 /// BY clause. This information is used to calculate the range.
-#[derive(Debug, Default)]
+#[derive(Debug, Default, Clone)]
 pub struct WindowFrameStateRange {
     sort_options: Vec<SortOptions>,
 }
@@ -454,7 +454,7 @@ impl WindowFrameStateRange {
 
 /// This structure encapsulates all the state information we require as we
 /// scan groups of data while processing window frames.
-#[derive(Debug, Default)]
+#[derive(Debug, Default, Clone)]
 pub struct WindowFrameStateGroups {
     /// A tuple containing group values and the row index where the group ends.
     /// Example: [[1, 1], [1, 1], [2, 1], [2, 1], ...] would correspond to

datafusion/physical-expr/src/aggregate.rs

@@ -616,10 +616,42 @@ impl AggregateFunctionExpr {
     /// Returns `Some(Arc<dyn AggregateExpr>)` if re-write is supported, otherwise returns `None`.
     pub fn with_new_expressions(
         &self,
-        _args: Vec<Arc<dyn PhysicalExpr>>,
-        _order_by_exprs: Vec<Arc<dyn PhysicalExpr>>,
+        args: Vec<Arc<dyn PhysicalExpr>>,
+        order_by_exprs: Vec<Arc<dyn PhysicalExpr>>,
     ) -> Option<AggregateFunctionExpr> {
-        None
+        if args.len() != self.args.len()
+            || (self.order_sensitivity() != AggregateOrderSensitivity::Insensitive
+                && order_by_exprs.len() != self.order_bys.len())
+        {
+            return None;
+        }
+
+        let new_order_bys = self
+            .order_bys
+            .iter()
+            .zip(order_by_exprs)
+            .map(|(req, new_expr)| PhysicalSortExpr {
+                expr: new_expr,
+                options: req.options,
+            })
+            .collect();
+
+        Some(AggregateFunctionExpr {
+            fun: self.fun.clone(),
+            args,
+            return_field: Arc::clone(&self.return_field),
+            name: self.name.clone(),
+            // TODO: Human name should be updated after re-write to not mislead
+            human_display: self.human_display.clone(),
+            schema: self.schema.clone(),
+            order_bys: new_order_bys,
+            ignore_nulls: self.ignore_nulls,
+            ordering_fields: self.ordering_fields.clone(),
+            is_distinct: self.is_distinct,
+            is_reversed: false,
+            input_fields: self.input_fields.clone(),
+            is_nullable: self.is_nullable,
+        })
     }
 
     /// If this function is max, return (output_field, true)

datafusion/physical-expr/src/expressions/literal.rs

@@ -36,7 +36,7 @@ use datafusion_expr_common::interval_arithmetic::Interval;
 use datafusion_expr_common::sort_properties::{ExprProperties, SortProperties};
 
 /// Represents a literal value
-#[derive(Debug, PartialEq, Eq)]
+#[derive(Debug, PartialEq, Eq, Clone)]
 pub struct Literal {
     value: ScalarValue,
     field: FieldRef,

datafusion/physical-expr/src/window/aggregate.rs

@@ -23,7 +23,7 @@ use std::sync::Arc;
 
 use crate::aggregate::AggregateFunctionExpr;
 use crate::window::standard::add_new_ordering_expr_with_partition_by;
-use crate::window::window_expr::AggregateWindowExpr;
+use crate::window::window_expr::{AggregateWindowExpr, WindowFn};
 use crate::window::{
     PartitionBatches, PartitionWindowAggStates, SlidingAggregateWindowExpr, WindowExpr,
 };
@@ -211,6 +211,10 @@ impl WindowExpr for PlainAggregateWindowExpr {
     fn uses_bounded_memory(&self) -> bool {
         !self.window_frame.end_bound.is_unbounded()
     }
+
+    fn create_window_fn(&self) -> Result<WindowFn> {
+        Ok(WindowFn::Aggregate(self.get_accumulator()?))
+    }
 }
 
 impl AggregateWindowExpr for PlainAggregateWindowExpr {

datafusion/physical-expr/src/window/sliding_aggregate.rs

@@ -22,7 +22,7 @@ use std::ops::Range;
 use std::sync::Arc;
 
 use crate::aggregate::AggregateFunctionExpr;
-use crate::window::window_expr::AggregateWindowExpr;
+use crate::window::window_expr::{AggregateWindowExpr, WindowFn};
 use crate::window::{
     PartitionBatches, PartitionWindowAggStates, PlainAggregateWindowExpr, WindowExpr,
 };
@@ -175,6 +175,10 @@ impl WindowExpr for SlidingAggregateWindowExpr {
             window_frame: Arc::clone(&self.window_frame),
         }))
     }
+
+    fn create_window_fn(&self) -> Result<WindowFn> {
+        Ok(WindowFn::Aggregate(self.get_accumulator()?))
+    }
 }
 
 impl AggregateWindowExpr for SlidingAggregateWindowExpr {

datafusion/physical-expr/src/window/standard.rs

@@ -275,6 +275,10 @@ impl WindowExpr for StandardWindowExpr {
             false
         }
     }
+
+    fn create_window_fn(&self) -> Result<WindowFn> {
+        Ok(WindowFn::Builtin(self.expr.create_evaluator()?))
+    }
 }
 
 /// Adds a new ordering expression into existing ordering equivalence class(es) based on

datafusion/physical-expr/src/window/window_expr.rs

@@ -130,6 +130,12 @@ pub trait WindowExpr: Send + Sync + Debug {
     /// Get the reverse expression of this [WindowExpr].
     fn get_reverse_expr(&self) -> Option<Arc<dyn WindowExpr>>;
 
+    /// Creates a new instance of the window function evaluator.
+    ///
+    /// Returns `WindowFn::Builtin` for built-in window functions (e.g., ROW_NUMBER, RANK)
+    /// or `WindowFn::Aggregate` for aggregate window functions (e.g., SUM, AVG).
+    fn create_window_fn(&self) -> Result<WindowFn>;
+
     /// Returns all expressions used in the [`WindowExpr`].
     /// These expressions are (1) function arguments, (2) partition by expressions, (3) order by expressions.
     fn all_expressions(&self) -> WindowPhysicalExpressions {

datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes.rs

@@ -15,11 +15,14 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use std::mem::size_of;
+
 use crate::aggregates::group_values::GroupValues;
+
 use arrow::array::{Array, ArrayRef, OffsetSizeTrait, RecordBatch};
+use datafusion_common::Result;
 use datafusion_expr::EmitTo;
 use datafusion_physical_expr_common::binary_map::{ArrowBytesMap, OutputType};
-use std::mem::size_of;
 
 /// A [`GroupValues`] storing single column of Utf8/LargeUtf8/Binary/LargeBinary values
 ///
@@ -42,11 +45,7 @@ impl<O: OffsetSizeTrait> GroupValuesByes<O> {
 }
 
 impl<O: OffsetSizeTrait> GroupValues for GroupValuesByes<O> {
-    fn intern(
-        &mut self,
-        cols: &[ArrayRef],
-        groups: &mut Vec<usize>,
-    ) -> datafusion_common::Result<()> {
+    fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec<usize>) -> Result<()> {
         assert_eq!(cols.len(), 1);
 
         // look up / add entries in the table
@@ -85,7 +84,7 @@ impl<O: OffsetSizeTrait> GroupValues for GroupValuesByes<O> {
         self.num_groups
     }
 
-    fn emit(&mut self, emit_to: EmitTo) -> datafusion_common::Result<Vec<ArrayRef>> {
+    fn emit(&mut self, emit_to: EmitTo) -> Result<Vec<ArrayRef>> {
         // Reset the map to default, and convert it into a single array
         let map_contents = self.map.take().into_state();
 

datafusion/physical-plan/src/test.rs

@@ -131,7 +131,7 @@ impl ExecutionPlan for TestMemoryExec {
     }
 
     fn as_any(&self) -> &dyn Any {
-        unimplemented!()
+        self
     }
 
     fn properties(&self) -> &PlanProperties {