move comparators initialization

Author: akoshchiy

datafusion/core/tests/physical_optimizer/window_optimize.rs

@@ -50,7 +50,6 @@ mod test {
             &[],
             Arc::new(frame),
             None,
-            Arc::clone(&schema),
         );
 
         let bounded_agg_exec = BoundedWindowAggExec::try_new(

datafusion/expr/src/window_frame.rs

@@ -57,25 +57,26 @@ impl WindowFrameBoundsComparators {
     /// - `N FOLLOWING` (asc)  → right_delta = `current + N`  (AddWrapping)
     /// - `N FOLLOWING` (desc) → right_delta = `current - N`  (SubWrapping)
     /// - `CURRENT ROW`        → no delta
-    pub fn new(frame: &WindowFrame, columns: &[(DataType, SortOptions)]) -> Option<Self> {
+    pub fn new(
+        frame: &WindowFrame,
+        columns: &[(DataType, SortOptions)],
+    ) -> Result<Option<Self>> {
         // Comparators are only needed for RANGE frames.
         if frame.units != WindowFrameUnits::Range {
-            return None;
+            return Ok(None);
         }
 
         let sort_descending = columns.first().map(|(_, o)| o.descending).unwrap_or(false);
 
         let start_bound_comparators =
-            Self::build_bound_comparators(&frame.start_bound, columns, sort_descending)
-                .ok()?;
+            Self::build_bound_comparators(&frame.start_bound, columns, sort_descending)?;
         let end_bound_comparators =
-            Self::build_bound_comparators(&frame.end_bound, columns, sort_descending)
-                .ok()?;
+            Self::build_bound_comparators(&frame.end_bound, columns, sort_descending)?;
 
-        Some(Self {
+        Ok(Some(Self {
             start_bound_comparators,
             end_bound_comparators,
-        })
+        }))
     }
 
     /// Builds per-column `ArrayComparator`s for a single bound.

datafusion/expr/src/window_state.rs

@@ -689,7 +689,7 @@ mod tests {
         expected_results: Vec<Range<usize>>,
     ) -> Result<()> {
         let (range_columns, _) = get_test_data();
-        let frame_comparators = WindowFrameBoundsComparators::new(window_frame, &[]);
+        let frame_comparators = WindowFrameBoundsComparators::new(window_frame, &[])?;
         let mut window_frame_context =
             WindowFrameContext::new(Arc::clone(window_frame), frame_comparators);
         let n_row = range_columns[0].len();

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

@@ -34,7 +34,6 @@ use arrow::array::BooleanArray;
 use arrow::datatypes::{FieldRef, SchemaRef};
 use arrow::record_batch::RecordBatch;
 use datafusion_common::{Result, ScalarValue, exec_datafusion_err};
-use datafusion_expr::window_frame::WindowFrameBoundsComparators;
 use datafusion_expr::{Accumulator, WindowFrame, WindowFrameBound, WindowFrameUnits};
 use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr;
 
@@ -49,8 +48,6 @@ pub struct PlainAggregateWindowExpr {
     window_frame: Arc<WindowFrame>,
     is_constant_in_partition: bool,
     filter: Option<Arc<dyn PhysicalExpr>>,
-    schema: SchemaRef,
-    frame_comparators: Option<WindowFrameBoundsComparators>,
 }
 
 impl PlainAggregateWindowExpr {
@@ -61,26 +58,16 @@ impl PlainAggregateWindowExpr {
         order_by: &[PhysicalSortExpr],
         window_frame: Arc<WindowFrame>,
         filter: Option<Arc<dyn PhysicalExpr>>,
-        schema: SchemaRef,
     ) -> Self {
         let is_constant_in_partition =
             Self::is_window_constant_in_partition(order_by, &window_frame);
-        let frame_comparators = {
-            let cols: Option<Vec<_>> = order_by
-                .iter()
-                .map(|o| o.expr.data_type(&schema).ok().map(|dt| (dt, o.options)))
-                .collect();
-            cols.and_then(|c| WindowFrameBoundsComparators::new(&window_frame, &c))
-        };
         Self {
             aggregate,
             partition_by: partition_by.to_vec(),
             order_by: order_by.to_vec(),
             window_frame,
             is_constant_in_partition,
             filter,
-            schema,
-            frame_comparators,
         }
     }
 
@@ -159,10 +146,15 @@ impl WindowExpr for PlainAggregateWindowExpr {
 
     fn evaluate_stateful(
         &self,
+        input_schema: &SchemaRef,
         partition_batches: &PartitionBatches,
         window_agg_state: &mut PartitionWindowAggStates,
     ) -> Result<()> {
-        self.aggregate_evaluate_stateful(partition_batches, window_agg_state)?;
+        self.aggregate_evaluate_stateful(
+            input_schema,
+            partition_batches,
+            window_agg_state,
+        )?;
 
         // Update window frame range for each partition. As we know that
         // non-sliding aggregations will never call `retract_batch`, this value
@@ -208,7 +200,6 @@ impl WindowExpr for PlainAggregateWindowExpr {
                         .collect::<Vec<_>>(),
                     Arc::new(self.window_frame.reverse()),
                     self.filter.clone(),
-                    Arc::clone(&self.schema),
                 )) as _
             } else {
                 Arc::new(SlidingAggregateWindowExpr::new(
@@ -221,7 +212,6 @@ impl WindowExpr for PlainAggregateWindowExpr {
                         .collect::<Vec<_>>(),
                     Arc::new(self.window_frame.reverse()),
                     self.filter.clone(),
-                    Arc::clone(&self.schema),
                 )) as _
             }
         })
@@ -234,10 +224,6 @@ impl WindowExpr for PlainAggregateWindowExpr {
     fn create_window_fn(&self) -> Result<WindowFn> {
         Ok(WindowFn::Aggregate(self.get_accumulator()?))
     }
-
-    fn frame_bounds_comparators(&self) -> Option<WindowFrameBoundsComparators> {
-        self.frame_comparators.clone()
-    }
 }
 
 impl AggregateWindowExpr for PlainAggregateWindowExpr {

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

@@ -33,7 +33,6 @@ use arrow::datatypes::{FieldRef, SchemaRef};
 use arrow::record_batch::RecordBatch;
 
 use datafusion_common::{Result, ScalarValue};
-use datafusion_expr::window_frame::WindowFrameBoundsComparators;
 use datafusion_expr::{Accumulator, WindowFrame};
 
 /// A window expr that takes the form of an aggregate function that
@@ -47,8 +46,6 @@ pub struct SlidingAggregateWindowExpr {
     order_by: Vec<PhysicalSortExpr>,
     window_frame: Arc<WindowFrame>,
     filter: Option<Arc<dyn PhysicalExpr>>,
-    schema: SchemaRef,
-    frame_comparators: Option<WindowFrameBoundsComparators>,
 }
 
 impl SlidingAggregateWindowExpr {
@@ -59,23 +56,13 @@ impl SlidingAggregateWindowExpr {
         order_by: &[PhysicalSortExpr],
         window_frame: Arc<WindowFrame>,
         filter: Option<Arc<dyn PhysicalExpr>>,
-        schema: SchemaRef,
     ) -> Self {
-        let frame_comparators = {
-            let cols: Option<Vec<_>> = order_by
-                .iter()
-                .map(|o| o.expr.data_type(&schema).ok().map(|dt| (dt, o.options)))
-                .collect();
-            cols.and_then(|c| WindowFrameBoundsComparators::new(&window_frame, &c))
-        };
         Self {
             aggregate,
             partition_by: partition_by.to_vec(),
             order_by: order_by.to_vec(),
             window_frame,
             filter,
-            schema,
-            frame_comparators,
         }
     }
 
@@ -114,10 +101,15 @@ impl WindowExpr for SlidingAggregateWindowExpr {
 
     fn evaluate_stateful(
         &self,
+        input_schema: &SchemaRef,
         partition_batches: &PartitionBatches,
         window_agg_state: &mut PartitionWindowAggStates,
     ) -> Result<()> {
-        self.aggregate_evaluate_stateful(partition_batches, window_agg_state)
+        self.aggregate_evaluate_stateful(
+            input_schema,
+            partition_batches,
+            window_agg_state,
+        )
     }
 
     fn partition_by(&self) -> &[Arc<dyn PhysicalExpr>] {
@@ -144,7 +136,6 @@ impl WindowExpr for SlidingAggregateWindowExpr {
                     &reversed_order_by,
                     Arc::new(reverse_window_frame),
                     self.filter.clone(),
-                    self.schema.clone(),
                 )) as _
             } else {
                 Arc::new(SlidingAggregateWindowExpr::new(
@@ -157,7 +148,6 @@ impl WindowExpr for SlidingAggregateWindowExpr {
                         .collect::<Vec<_>>(),
                     Arc::new(self.window_frame.reverse()),
                     self.filter.clone(),
-                    self.schema.clone(),
                 )) as _
             }
         })
@@ -193,17 +183,12 @@ impl WindowExpr for SlidingAggregateWindowExpr {
             &new_order_by,
             Arc::clone(&self.window_frame),
             self.filter.clone(),
-            self.schema.clone(),
         )))
     }
 
     fn create_window_fn(&self) -> Result<WindowFn> {
         Ok(WindowFn::Aggregate(self.get_accumulator()?))
     }
-
-    fn frame_bounds_comparators(&self) -> Option<WindowFrameBoundsComparators> {
-        self.frame_comparators.clone()
-    }
 }
 
 impl AggregateWindowExpr for SlidingAggregateWindowExpr {

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

@@ -32,7 +32,6 @@ use arrow::record_batch::RecordBatch;
 use datafusion_common::utils::evaluate_partition_ranges;
 use datafusion_common::{Result, ScalarValue};
 use datafusion_expr::WindowFrame;
-use datafusion_expr::window_frame::WindowFrameBoundsComparators;
 use datafusion_expr::window_state::{WindowAggState, WindowFrameContext};
 use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr;
 
@@ -43,8 +42,6 @@ pub struct StandardWindowExpr {
     partition_by: Vec<Arc<dyn PhysicalExpr>>,
     order_by: Vec<PhysicalSortExpr>,
     window_frame: Arc<WindowFrame>,
-    schema: SchemaRef,
-    frame_comparators: Option<WindowFrameBoundsComparators>,
 }
 
 impl StandardWindowExpr {
@@ -54,22 +51,12 @@ impl StandardWindowExpr {
         partition_by: &[Arc<dyn PhysicalExpr>],
         order_by: &[PhysicalSortExpr],
         window_frame: Arc<WindowFrame>,
-        schema: SchemaRef,
     ) -> Self {
-        let frame_comparators = {
-            let cols: Option<Vec<_>> = order_by
-                .iter()
-                .map(|o| o.expr.data_type(&schema).ok().map(|dt| (dt, o.options)))
-                .collect();
-            cols.and_then(|c| WindowFrameBoundsComparators::new(&window_frame, &c))
-        };
         Self {
             expr,
             partition_by: partition_by.to_vec(),
             order_by: order_by.to_vec(),
             window_frame,
-            schema,
-            frame_comparators,
         }
     }
 
@@ -139,7 +126,7 @@ impl WindowExpr for StandardWindowExpr {
 
             let mut window_frame_ctx = WindowFrameContext::new(
                 Arc::clone(&self.window_frame),
-                self.frame_comparators.clone(),
+                self.build_frame_bounds_comparators(batch.schema_ref())?,
             );
             let mut last_range = Range { start: 0, end: 0 };
 
@@ -170,6 +157,7 @@ impl WindowExpr for StandardWindowExpr {
     /// stateful, bounded-memory implementations.
     fn evaluate_stateful(
         &self,
+        input_schema: &SchemaRef,
         partition_batches: &PartitionBatches,
         window_agg_state: &mut PartitionWindowAggStates,
     ) -> Result<()> {
@@ -178,6 +166,10 @@ impl WindowExpr for StandardWindowExpr {
         // create a WindowAggState to clone when `window_agg_state` does not contain the respective
         // group, which is faster than potentially creating a new one at every iteration
         let new_state = WindowAggState::new(out_type)?;
+
+        let frame_bounds_comparators =
+            self.build_frame_bounds_comparators(input_schema)?;
+
         for (partition_row, partition_batch_state) in partition_batches.iter() {
             let window_state =
                 if let Some(window_state) = window_agg_state.get_mut(partition_row) {
@@ -225,7 +217,7 @@ impl WindowExpr for StandardWindowExpr {
                         .get_or_insert_with(|| {
                             WindowFrameContext::new(
                                 Arc::clone(&self.window_frame),
-                                self.frame_comparators.clone(),
+                                frame_bounds_comparators.clone(),
                             )
                         })
                         .calculate_range(
@@ -283,7 +275,6 @@ impl WindowExpr for StandardWindowExpr {
                     .map(|e| e.reverse())
                     .collect::<Vec<_>>(),
                 Arc::new(self.window_frame.reverse()),
-                Arc::clone(&self.schema),
             )) as _
         })
     }
@@ -301,10 +292,6 @@ impl WindowExpr for StandardWindowExpr {
     fn create_window_fn(&self) -> Result<WindowFn> {
         Ok(WindowFn::Builtin(self.expr.create_evaluator()?))
     }
-
-    fn frame_bounds_comparators(&self) -> Option<WindowFrameBoundsComparators> {
-        self.frame_comparators.clone()
-    }
 }
 
 /// Adds a new ordering expression into existing ordering equivalence class(es) based on

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

@@ -27,7 +27,7 @@ use arrow::array::{Array, ArrayRef, new_empty_array};
 use arrow::compute::SortOptions;
 use arrow::compute::filter as arrow_filter;
 use arrow::compute::kernels::sort::SortColumn;
-use arrow::datatypes::FieldRef;
+use arrow::datatypes::{FieldRef, SchemaRef};
 use arrow::record_batch::RecordBatch;
 use datafusion_common::cast::as_boolean_array;
 
@@ -71,7 +71,17 @@ use indexmap::IndexMap;
 /// [`PlainAggregateWindowExpr`]: crate::window::PlainAggregateWindowExpr
 /// [`SlidingAggregateWindowExpr`]: crate::window::SlidingAggregateWindowExpr
 pub trait WindowExpr: Send + Sync + Debug {
-    fn frame_bounds_comparators(&self) -> Option<WindowFrameBoundsComparators>;
+    fn build_frame_bounds_comparators(
+        &self,
+        input_schema: &SchemaRef,
+    ) -> Result<Option<WindowFrameBoundsComparators>> {
+        let cols: Result<Vec<_>> = self
+            .order_by()
+            .iter()
+            .map(|o| o.expr.data_type(input_schema).map(|dt| (dt, o.options)))
+            .collect();
+        cols.and_then(|c| WindowFrameBoundsComparators::new(self.get_window_frame(), &c))
+    }
 
     /// Returns the window expression as [`Any`] so that it can be
     /// downcast to a specific implementation.
@@ -104,6 +114,7 @@ pub trait WindowExpr: Send + Sync + Debug {
     /// stateful, bounded-memory implementations.
     fn evaluate_stateful(
         &self,
+        _input_schema: &SchemaRef,
         _partition_batches: &PartitionBatches,
         _window_agg_state: &mut PartitionWindowAggStates,
     ) -> Result<()> {
@@ -211,7 +222,7 @@ pub trait AggregateWindowExpr: WindowExpr {
         let mut last_range = Range { start: 0, end: 0 };
         let mut window_frame_ctx = WindowFrameContext::new(
             Arc::clone(self.get_window_frame()),
-            self.frame_bounds_comparators(),
+            self.build_frame_bounds_comparators(batch.schema_ref())?,
         );
         self.get_result_column(
             &mut accumulator,
@@ -228,11 +239,16 @@ pub trait AggregateWindowExpr: WindowExpr {
     /// state so that it can work incrementally over multiple chunks.
     fn aggregate_evaluate_stateful(
         &self,
+        input_schema: &SchemaRef,
         partition_batches: &PartitionBatches,
         window_agg_state: &mut PartitionWindowAggStates,
     ) -> Result<()> {
         let field = self.field()?;
         let out_type = field.data_type();
+
+        let frame_bounds_comparators =
+            self.build_frame_bounds_comparators(input_schema)?;
+
         for (partition_row, partition_batch_state) in partition_batches.iter() {
             if !window_agg_state.contains_key(partition_row) {
                 let accumulator = self.get_accumulator()?;
@@ -259,7 +275,7 @@ pub trait AggregateWindowExpr: WindowExpr {
             let window_frame_ctx = state.window_frame_ctx.get_or_insert_with(|| {
                 WindowFrameContext::new(
                     Arc::clone(self.get_window_frame()),
-                    self.frame_bounds_comparators(),
+                    frame_bounds_comparators.clone(),
                 )
             });
             let out_col = self.get_result_column(