[MINOR] Remove update state api from PartitionEvaluator (#6966)

mustafasrepo · ozankabak · web-flow · commit afc9c9daecb3 · 2023-07-15T06:27:50.000-04:00
* remove update_state api from partition_evaluator

* Resolve linter errors

* Simplifications

* remove row_idx argument from evaluate

* Simplifications

* Update datafusion/expr/src/partition_evaluator.rs

Co-authored-by: Mehmet Ozan Kabak &lt;ozankabak@gmail.com&gt;

* Update comment

Co-authored-by: Mehmet Ozan Kabak &lt;ozankabak@gmail.com&gt;

* Update document

* Use boolean operator instead of bitwise

---------

Co-authored-by: Mehmet Ozan Kabak &lt;ozankabak@gmail.com&gt;
diff --git a/datafusion/core/tests/fuzz_cases/window_fuzz.rs b/datafusion/core/tests/fuzz_cases/window_fuzz.rs
@@ -208,6 +208,13 @@ fn get_random_function(
                 vec![],
             ),
         );
+        window_fn_map.insert(
+            "dense_rank",
+            (
+                WindowFunction::BuiltInWindowFunction(BuiltInWindowFunction::DenseRank),
+                vec![],
+            ),
+        );
         window_fn_map.insert(
             "lead",
             (
diff --git a/datafusion/core/tests/user_defined/user_defined_window_functions.rs b/datafusion/core/tests/user_defined/user_defined_window_functions.rs
@@ -32,8 +32,8 @@ use arrow_schema::DataType;
 use datafusion::{assert_batches_eq, prelude::SessionContext};
 use datafusion_common::{Result, ScalarValue};
 use datafusion_expr::{
-    function::PartitionEvaluatorFactory, window_state::WindowAggState,
-    PartitionEvaluator, ReturnTypeFunction, Signature, Volatility, WindowUDF,
+    function::PartitionEvaluatorFactory, PartitionEvaluator, ReturnTypeFunction,
+    Signature, Volatility, WindowUDF,
 };
 
 /// A query with a window function evaluated over the entire partition
@@ -195,7 +195,6 @@ async fn test_stateful_udwf() {
         &execute(&ctx, UNBOUNDED_WINDOW_QUERY).await.unwrap()
     );
     assert_eq!(test_state.evaluate_called(), 10);
-    assert_eq!(test_state.update_state_called(), 10);
     assert_eq!(test_state.evaluate_all_called(), 0);
 }
 
@@ -229,7 +228,6 @@ async fn test_stateful_udwf_bounded_window() {
     );
     // Evaluate and update_state is called for each input row
     assert_eq!(test_state.evaluate_called(), 10);
-    assert_eq!(test_state.update_state_called(), 10);
     assert_eq!(test_state.evaluate_all_called(), 0);
 }
 
@@ -388,8 +386,6 @@ struct TestState {
     evaluate_all_called: AtomicUsize,
     /// How many times was `evaluate` called?
     evaluate_called: AtomicUsize,
-    /// How many times was `update_state` called?
-    update_state_called: AtomicUsize,
     /// How many times was `evaluate_all_with_rank` called?
     evaluate_all_with_rank_called: AtomicUsize,
     /// should the functions say they use the window frame?
@@ -451,16 +447,6 @@ impl TestState {
         self.evaluate_called.fetch_add(1, Ordering::SeqCst);
     }
 
-    /// return the update_state_called counter
-    fn update_state_called(&self) -> usize {
-        self.update_state_called.load(Ordering::SeqCst)
-    }
-
-    /// update the update_state_called counter
-    fn inc_update_state_called(&self) {
-        self.update_state_called.fetch_add(1, Ordering::SeqCst);
-    }
-
     /// return the evaluate_all_with_rank_called counter
     fn evaluate_all_with_rank_called(&self) -> usize {
         self.evaluate_all_with_rank_called.load(Ordering::SeqCst)
@@ -555,17 +541,6 @@ impl PartitionEvaluator for OddCounter {
         Ok(Arc::new(array))
     }
 
-    fn update_state(
-        &mut self,
-        _state: &WindowAggState,
-        _idx: usize,
-        _range_columns: &[ArrayRef],
-        _sort_partition_points: &[Range<usize>],
-    ) -> Result<()> {
-        self.test_state.inc_update_state_called();
-        Ok(())
-    }
-
     fn supports_bounded_execution(&self) -> bool {
         self.test_state.supports_bounded_execution
     }
diff --git a/datafusion/expr/src/partition_evaluator.rs b/datafusion/expr/src/partition_evaluator.rs
@@ -69,27 +69,10 @@ use crate::window_state::WindowAggState;
 /// capabilities described by [`supports_bounded_execution`],
 /// [`uses_window_frame`], and [`include_rank`],
 ///
-/// # Stateless `PartitionEvaluator`s
-///
-/// In this case, `PartitionEvaluator` holds no state, and either
-/// [`evaluate_all`] or [`evaluate_all_with_rank`] is called with
-/// values for the entire partition.
-///
-/// # Stateful `PartitionEvaluator`s
-///
-/// In this case, [`Self::evaluate`] is called to calculate the window
-/// function incrementally for each new batch.
-///
-/// For example, when computing `ROW_NUMBER` incrementally,
-/// [`Self::evaluate`] will be called multiple times with
-/// different batches. For all batches after the first, the output
-/// `row_number` must start from last `row_number` produced for the
-/// previous batch. The previous row number is saved and restored as
-/// the state.
-///
 /// When implementing a new `PartitionEvaluator`, implement
 /// corresponding evaluator according to table below.
 ///
+/// # Implementation Table
 ///
 /// |[`uses_window_frame`]|[`supports_bounded_execution`]|[`include_rank`]|function_to_implement|
 /// |---|---|----|----|
@@ -105,25 +88,6 @@ use crate::window_state::WindowAggState;
 /// [`include_rank`]: Self::include_rank
 /// [`supports_bounded_execution`]: Self::supports_bounded_execution
 pub trait PartitionEvaluator: Debug + Send {
-    /// Updates the internal state for window function
-    ///
-    /// Only used for stateful evaluation
-    ///
-    /// `state`: is useful to update internal state for window function.
-    /// `idx`: is the index of last row for which result is calculated.
-    /// `range_columns`: is the result of order by column values. It is used to calculate rank boundaries
-    /// `sort_partition_points`: is the boundaries of each rank in the range_column. It is used to update rank.
-    fn update_state(
-        &mut self,
-        _state: &WindowAggState,
-        _idx: usize,
-        _range_columns: &[ArrayRef],
-        _sort_partition_points: &[Range<usize>],
-    ) -> Result<()> {
-        // If we do not use state, update_state does nothing
-        Ok(())
-    }
-
     /// When the window frame has a fixed beginning (e.g UNBOUNDED
     /// PRECEDING), some functions such as FIRST_VALUE, LAST_VALUE and
     /// NTH_VALUE do not need the (unbounded) input once they have
@@ -220,7 +184,10 @@ pub trait PartitionEvaluator: Debug + Send {
     /// trait.
     ///
     /// Returns a [`ScalarValue`] that is the value of the window
-    /// function within `range` for the entire partition
+    /// function within `range` for the entire partition. Argument
+    /// `values` contains the evaluation result of function arguments
+    /// and evaluation results of ORDER BY expressions. If function has a
+    /// single argument, `values[1..]` will contain ORDER BY expression results.
     fn evaluate(
         &mut self,
         _values: &[ArrayRef],
diff --git a/datafusion/physical-expr/src/window/built_in.rs b/datafusion/physical-expr/src/window/built_in.rs
@@ -23,7 +23,7 @@ use std::sync::Arc;
 
 use super::BuiltInWindowFunctionExpr;
 use super::WindowExpr;
-use crate::window::window_expr::WindowFn;
+use crate::window::window_expr::{get_orderby_values, WindowFn};
 use crate::window::{PartitionBatches, PartitionWindowAggStates, WindowState};
 use crate::{expressions::PhysicalSortExpr, reverse_order_bys, PhysicalExpr};
 use arrow::array::{new_empty_array, ArrayRef};
@@ -101,14 +101,19 @@ impl WindowExpr for BuiltInWindowExpr {
                 self.order_by.iter().map(|o| o.options).collect();
             let mut row_wise_results = vec![];
 
-            let (values, order_bys) = self.get_values_orderbys(batch)?;
+            let mut values = self.evaluate_args(batch)?;
+            let order_bys = get_orderby_values(self.order_by_columns(batch)?);
+            let n_args = values.len();
+            values.extend(order_bys);
+            let order_bys_ref = &values[n_args..];
+
             let mut window_frame_ctx =
                 WindowFrameContext::new(self.window_frame.clone(), sort_options);
             let mut last_range = Range { start: 0, end: 0 };
             // We iterate on each row to calculate window frame range and and window function result
             for idx in 0..num_rows {
                 let range = window_frame_ctx.calculate_range(
-                    &order_bys,
+                    order_bys_ref,
                     &last_range,
                     num_rows,
                     idx,
@@ -119,11 +124,11 @@ impl WindowExpr for BuiltInWindowExpr {
             }
             ScalarValue::iter_to_array(row_wise_results.into_iter())
         } else if evaluator.include_rank() {
-            let columns = self.sort_columns(batch)?;
+            let columns = self.order_by_columns(batch)?;
             let sort_partition_points = evaluate_partition_ranges(num_rows, &columns)?;
             evaluator.evaluate_all_with_rank(num_rows, &sort_partition_points)
         } else {
-            let (values, _) = self.get_values_orderbys(batch)?;
+            let values = self.evaluate_args(batch)?;
             evaluator.evaluate_all(&values, num_rows)
         }
     }
@@ -157,18 +162,20 @@ impl WindowExpr for BuiltInWindowExpr {
             };
             let state = &mut window_state.state;
 
-            let (values, order_bys) =
-                self.get_values_orderbys(&partition_batch_state.record_batch)?;
+            let batch_ref = &partition_batch_state.record_batch;
+            let mut values = self.evaluate_args(batch_ref)?;
+            let order_bys = if evaluator.uses_window_frame() || evaluator.include_rank() {
+                get_orderby_values(self.order_by_columns(batch_ref)?)
+            } else {
+                vec![]
+            };
+            let n_args = values.len();
+            values.extend(order_bys);
+            let order_bys_ref = &values[n_args..];
 
             // We iterate on each row to perform a running calculation.
             let record_batch = &partition_batch_state.record_batch;
             let num_rows = record_batch.num_rows();
-            let sort_partition_points = if evaluator.include_rank() {
-                let columns = self.sort_columns(record_batch)?;
-                evaluate_partition_ranges(num_rows, &columns)?
-            } else {
-                vec![]
-            };
             let mut row_wise_results: Vec<ScalarValue> = vec![];
             for idx in state.last_calculated_index..num_rows {
                 let frame_range = if evaluator.uses_window_frame() {
@@ -181,7 +188,7 @@ impl WindowExpr for BuiltInWindowExpr {
                             )
                         })
                         .calculate_range(
-                            &order_bys,
+                            order_bys_ref,
                             // Start search from the last range
                             &state.window_frame_range,
                             num_rows,
@@ -197,7 +204,6 @@ impl WindowExpr for BuiltInWindowExpr {
                 }
                 // Update last range
                 state.window_frame_range = frame_range;
-                evaluator.update_state(state, idx, &order_bys, &sort_partition_points)?;
                 row_wise_results
                     .push(evaluator.evaluate(&values, &state.window_frame_range)?);
             }
diff --git a/datafusion/physical-expr/src/window/lead_lag.rs b/datafusion/physical-expr/src/window/lead_lag.rs
@@ -18,15 +18,13 @@
 //! Defines physical expression for `lead` and `lag` that can evaluated
 //! at runtime during query execution
 
-use crate::window::window_expr::LeadLagState;
 use crate::window::BuiltInWindowFunctionExpr;
 use crate::PhysicalExpr;
 use arrow::array::ArrayRef;
 use arrow::compute::cast;
 use arrow::datatypes::{DataType, Field};
 use datafusion_common::ScalarValue;
 use datafusion_common::{DataFusionError, Result};
-use datafusion_expr::window_state::WindowAggState;
 use datafusion_expr::PartitionEvaluator;
 use std::any::Any;
 use std::cmp::min;
@@ -105,7 +103,6 @@ impl BuiltInWindowFunctionExpr for WindowShift {
 
     fn create_evaluator(&self) -> Result<Box<dyn PartitionEvaluator>> {
         Ok(Box::new(WindowShiftEvaluator {
-            state: LeadLagState { idx: 0 },
             shift_offset: self.shift_offset,
             default_value: self.default_value.clone(),
         }))
@@ -124,7 +121,6 @@ impl BuiltInWindowFunctionExpr for WindowShift {
 
 #[derive(Debug)]
 pub(crate) struct WindowShiftEvaluator {
-    state: LeadLagState,
     shift_offset: i64,
     default_value: Option<ScalarValue>,
 }
@@ -179,17 +175,6 @@ fn shift_with_default_value(
 }
 
 impl PartitionEvaluator for WindowShiftEvaluator {
-    fn update_state(
-        &mut self,
-        _state: &WindowAggState,
-        idx: usize,
-        _range_columns: &[ArrayRef],
-        _sort_partition_points: &[Range<usize>],
-    ) -> Result<()> {
-        self.state.idx = idx;
-        Ok(())
-    }
-
     fn get_range(&self, idx: usize, n_rows: usize) -> Result<Range<usize>> {
         if self.shift_offset > 0 {
             let offset = self.shift_offset as usize;
@@ -206,11 +191,18 @@ impl PartitionEvaluator for WindowShiftEvaluator {
     fn evaluate(
         &mut self,
         values: &[ArrayRef],
-        _range: &Range<usize>,
+        range: &Range<usize>,
     ) -> Result<ScalarValue> {
         let array = &values[0];
         let dtype = array.data_type();
-        let idx = self.state.idx as i64 - self.shift_offset;
+        // LAG mode
+        let idx = if self.shift_offset > 0 {
+            range.end as i64 - self.shift_offset - 1
+        } else {
+            // LEAD mode
+            range.start as i64 - self.shift_offset
+        };
+
         if idx < 0 || idx as usize >= array.len() {
             get_default_value(self.default_value.as_ref(), dtype)
         } else {
diff --git a/datafusion/physical-expr/src/window/nth_value.rs b/datafusion/physical-expr/src/window/nth_value.rs
@@ -145,18 +145,6 @@ pub(crate) struct NthValueEvaluator {
 }
 
 impl PartitionEvaluator for NthValueEvaluator {
-    fn update_state(
-        &mut self,
-        state: &WindowAggState,
-        _idx: usize,
-        _range_columns: &[ArrayRef],
-        _sort_partition_points: &[Range<usize>],
-    ) -> Result<()> {
-        // If we do not use state, update_state does nothing
-        self.state.range.clone_from(&state.window_frame_range);
-        Ok(())
-    }
-
     /// When the window frame has a fixed beginning (e.g UNBOUNDED
     /// PRECEDING), for some functions such as FIRST_VALUE, LAST_VALUE and
     /// NTH_VALUE we can memoize result.  Once result is calculated it
diff --git a/datafusion/physical-expr/src/window/rank.rs b/datafusion/physical-expr/src/window/rank.rs
@@ -26,7 +26,6 @@ use arrow::array::{Float64Array, UInt64Array};
 use arrow::datatypes::{DataType, Field};
 use datafusion_common::utils::get_row_at_idx;
 use datafusion_common::{DataFusionError, Result, ScalarValue};
-use datafusion_expr::window_state::WindowAggState;
 use datafusion_expr::PartitionEvaluator;
 use std::any::Any;
 use std::iter;
@@ -116,39 +115,26 @@ pub(crate) struct RankEvaluator {
 }
 
 impl PartitionEvaluator for RankEvaluator {
-    fn update_state(
+    /// Evaluates the window function inside the given range.
+    fn evaluate(
         &mut self,
-        state: &WindowAggState,
-        idx: usize,
-        range_columns: &[ArrayRef],
-        sort_partition_points: &[Range<usize>],
-    ) -> Result<()> {
-        // find range inside `sort_partition_points` containing `idx`
-        let chunk_idx = sort_partition_points
-            .iter()
-            .position(|elem| elem.start <= idx && idx < elem.end)
-            .ok_or_else(|| {
-                DataFusionError::Execution(
-                    "Expects sort_partition_points to contain idx".to_string(),
-                )
-            })?;
-        let chunk = &sort_partition_points[chunk_idx];
-        let last_rank_data = get_row_at_idx(range_columns, chunk.end - 1)?;
+        values: &[ArrayRef],
+        range: &Range<usize>,
+    ) -> Result<ScalarValue> {
+        let row_idx = range.start;
+        // There is no argument, values are order by column values (where rank is calculated)
+        let range_columns = values;
+        let last_rank_data = get_row_at_idx(range_columns, row_idx)?;
         let empty = self.state.last_rank_data.is_empty();
         if empty || self.state.last_rank_data != last_rank_data {
             self.state.last_rank_data = last_rank_data;
-            self.state.last_rank_boundary = state.offset_pruned_rows + chunk.start;
-            self.state.n_rank = 1 + if empty { chunk_idx } else { self.state.n_rank };
+            self.state.last_rank_boundary += self.state.current_group_count;
+            self.state.current_group_count = 1;
+            self.state.n_rank += 1;
+        } else {
+            // data is still in the same rank
+            self.state.current_group_count += 1;
         }
-        Ok(())
-    }
-
-    /// evaluate window function result inside given range
-    fn evaluate(
-        &mut self,
-        _values: &[ArrayRef],
-        _range: &Range<usize>,
-    ) -> Result<ScalarValue> {
         match self.rank_type {
             RankType::Basic => Ok(ScalarValue::UInt64(Some(
                 self.state.last_rank_boundary as u64 + 1,
diff --git a/datafusion/physical-expr/src/window/window_expr.rs b/datafusion/physical-expr/src/window/window_expr.rs
