Simplify and encapsulate window function state management #6621

alamb · 2023-06-09T20:58:22Z

Which issue does this PR close?

Related to #5781

Rationale for this change

PartitionEvaluator currently has some methods that expose an enum BuiltinWindowState specific to the built in functions. If we choose to expose PartitionEvaluator externally it needs to not include details specifically about the built i window functions

It turns out there is exactly one use for this state, which is memoization (for nth value calculation). We can move that functionality into the trait and avoid the need for BuiltinWindowState entirely

What changes are included in this PR?

Move the memoization code to the window function to the window function that uses it (nth_value)
Remove PartitionEvaluator::get_state, PartitionEvaluator::set_state() and BuiltinWindowState

Are these changes tested?

existing tests

Are there any user-facing changes?

No these are all (currently) internal APIs

alamb · 2023-06-09T20:59:24Z

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

-                    memoize_nth_value(state, nth_value_state)?;
-                    evaluator.set_state(&evaluator_state)?;
-                }
+                evaluator.memoize(state)?;


I think this is actually more readable as well as cleaning up the trait

alamb · 2023-06-09T21:00:04Z

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

@@ -100,14 +99,6 @@ pub trait PartitionEvaluator: Debug + Send {
        false
    }

-    /// Returns the internal state of the window function


The point of the PR is to get of the functions on this trait that serialized / set state (which has a single use)

alamb · 2023-06-09T21:04:58Z

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

@@ -327,16 +327,7 @@ pub struct LeadLagState {
    pub idx: usize,
 }

-#[derive(Debug, Clone, Default)]


It turns out this structure is not used anywhere else

alamb · 2023-06-09T21:06:19Z

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

@@ -127,13 +117,16 @@ pub trait PartitionEvaluator: Debug + Send {
        Ok(())
    }

-    /// Sets the internal state for window function
+    /// When the window frame has a fixed beginning (e.g UNBOUNDED


While writing this explanation it occurred to me that the term memoize, while technically accurate is not the entire story here. I wonder if there is a better name like adjust_window_state 🤔

I think adjust_window_state feels like a more general functionality. However, this function is a special case for nth_value when window frame boundary has UNBOUNDED PRECEDING. Hence I think memoize is better for this case.

I will leave it named memoize

mustafasrepo · 2023-06-12T07:03:09Z

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

-        if let BuiltinWindowState::NthValue(nth_value_state) = state {
-            self.state = nth_value_state.clone()
+    /// When the window frame has a fixed beginning (e.g UNBOUNDED
+    /// PRECEDING), some functions such as FIRST_VALUE, LAST_VALUE and


I think we can replace this line with PRECEDING), for some functions such as FIRST_VALUE, LAST_VALUE and

Done (BTW if you do "CTRL-G" you can get inline suggestions, to get the github UI to show the suggestion inline):

Suggested change

/// PRECEDING), some functions such as FIRST_VALUE, LAST_VALUE and

/// PRECEDING), for some functions such as FIRST_VALUE, LAST_VALUE and

Thanks for the tip

mustafasrepo

Thanks @alamb. I have left some minor comments. Other than these this PR is LGTM!.

mustafasrepo · 2023-06-12T07:17:51Z

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

+    fn memoize(&mut self, state: &mut WindowAggState) -> Result<()> {
+        let out = &state.out_col;
+        let size = out.len();
+        let (is_prunable, new_prunable) = match self.state.kind {


I know that is is not introduced here (and introduced by me) but while reading the code maybe instead of new_prunable we can use is_last flag. I think, that way intent will be more clear.

mustafasrepo · 2023-06-12T07:18:09Z

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

+            NthValueKind::First => {
+                let n_range =
+                    state.window_frame_range.end - state.window_frame_range.start;
+                (n_range > 0 && size > 0, true)


in case we use is_last flag, second flag should be false

mustafasrepo · 2023-06-12T07:18:32Z

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

+                    state.window_frame_range.end - state.window_frame_range.start;
+                (n_range > 0 && size > 0, true)
+            }
+            NthValueKind::Last => (true, false),


in case we use is_last flag, second flag should be true

mustafasrepo · 2023-06-12T07:19:13Z

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

+            NthValueKind::Nth(n) => {
+                let n_range =
+                    state.window_frame_range.end - state.window_frame_range.start;
+                (n_range >= (n as usize) && size >= (n as usize), true)


In case we use is_last flag, second flag should be false

mustafasrepo · 2023-06-12T07:19:52Z

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

+            }
+        };
+        if is_prunable {
+            if self.state.finalized_result.is_none() && new_prunable {


in case we use is_last flag, condition should be if self.state.finalized_result.is_none() && !is_last {

I agree the intent is clearer using the term is_last. Thank you for that suggestion . Done in 4542799

alamb added 2 commits June 9, 2023 16:51

Simplify and encapsulate window function state management

0e8ca7d

Fix docs

6decd54

github-actions bot added the physical-expr Changes to the physical-expr crates label Jun 9, 2023

Improve documentation

fb0c17e

alamb commented Jun 9, 2023

View reviewed changes

alamb requested a review from mustafasrepo June 9, 2023 21:50

mustafasrepo reviewed Jun 12, 2023

View reviewed changes

mustafasrepo approved these changes Jun 12, 2023

View reviewed changes

mustafasrepo reviewed Jun 12, 2023

View reviewed changes

alamb added 2 commits June 12, 2023 20:14

Merge remote-tracking branch 'apache/main' into alamb/simplify_state

7d3e361

Improve comment and readability

4542799

alamb merged commit a56ae74 into apache:main Jun 13, 2023

alamb deleted the alamb/simplify_state branch June 13, 2023 01:33

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Simplify and encapsulate window function state management #6621

Simplify and encapsulate window function state management #6621

alamb commented Jun 9, 2023

alamb Jun 9, 2023

alamb Jun 9, 2023

alamb Jun 9, 2023

alamb Jun 9, 2023

mustafasrepo Jun 12, 2023

alamb Jun 13, 2023

mustafasrepo Jun 12, 2023

alamb Jun 13, 2023

mustafasrepo Jun 13, 2023

mustafasrepo left a comment

mustafasrepo Jun 12, 2023

mustafasrepo Jun 12, 2023 •

edited

Loading

mustafasrepo Jun 12, 2023

mustafasrepo Jun 12, 2023

mustafasrepo Jun 12, 2023

alamb Jun 13, 2023

	/// PRECEDING), some functions such as FIRST_VALUE, LAST_VALUE and
	/// PRECEDING), for some functions such as FIRST_VALUE, LAST_VALUE and

Simplify and encapsulate window function state management #6621

Simplify and encapsulate window function state management #6621

Conversation

alamb commented Jun 9, 2023

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

mustafasrepo left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

mustafasrepo Jun 12, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

mustafasrepo Jun 12, 2023 •

edited

Loading