refactor(over window): simplify find_affected_ranges and correct so…

…me corner cases (#14580)

Signed-off-by: Richard Chien <[email protected]>
Co-authored-by: Noel Kwan <[email protected]>

src/expr/core/src/window_function/call.rs

@@ -111,7 +111,7 @@ impl Frame {
     }
 }
 
-#[derive(Display, Debug, Clone, Eq, PartialEq, Hash)]
+#[derive(Display, Debug, Clone, Eq, PartialEq, Hash, EnumAsInner)]
 #[display("{0}")]
 pub enum FrameBounds {
     Rows(RowsFrameBounds),
@@ -154,6 +154,40 @@ impl RowsFrameBounds {
     fn validate(&self) -> Result<()> {
         FrameBound::validate_bounds(&self.start, &self.end)
     }
+
+    /// Check if the `ROWS` frame is canonical.
+    ///
+    /// A canonical `ROWS` frame is defined as:
+    ///
+    /// - Its bounds are valid (see [`Self::validate`]).
+    /// - It contains the current row.
+    pub fn is_canonical(&self) -> bool {
+        self.validate().is_ok() && {
+            let start = self.start.to_offset();
+            let end = self.end.to_offset();
+            start.unwrap_or(0) <= 0 && end.unwrap_or(0) >= 0
+        }
+    }
+
+    /// Get the number of preceding rows.
+    pub fn n_preceding_rows(&self) -> Option<usize> {
+        match (&self.start, &self.end) {
+            (UnboundedPreceding, _) => None,
+            (Preceding(n1), Preceding(n2)) => Some(*n1.max(n2)),
+            (Preceding(n), _) => Some(*n),
+            (CurrentRow | Following(_) | UnboundedFollowing, _) => Some(0),
+        }
+    }
+
+    /// Get the number of following rows.
+    pub fn n_following_rows(&self) -> Option<usize> {
+        match (&self.start, &self.end) {
+            (_, UnboundedFollowing) => None,
+            (Following(n1), Following(n2)) => Some(*n1.max(n2)),
+            (_, Following(n)) => Some(*n),
+            (_, CurrentRow | Preceding(_) | UnboundedPreceding) => Some(0),
+        }
+    }
 }
 
 #[derive(Display, Debug, Clone, Eq, PartialEq, Hash, EnumAsInner)]
@@ -233,16 +267,6 @@ impl FrameBound<usize> {
             Following(n) => Some(*n as isize),
         }
     }
-
-    /// View the bound as frame start, and get the number of preceding rows.
-    pub fn n_preceding_rows(&self) -> Option<usize> {
-        self.to_offset().map(|x| x.min(0).unsigned_abs())
-    }
-
-    /// View the bound as frame end, and get the number of following rows.
-    pub fn n_following_rows(&self) -> Option<usize> {
-        self.to_offset().map(|x| x.max(0) as usize)
-    }
 }
 
 #[derive(Display, Debug, Copy, Clone, Eq, PartialEq, Hash, Default, EnumAsInner)]
@@ -272,3 +296,92 @@ impl FrameExclusion {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+
+    use super::*;
+
+    #[test]
+    fn test_rows_frame_bounds() {
+        let bounds = RowsFrameBounds {
+            start: Preceding(1),
+            end: CurrentRow,
+        };
+        assert!(bounds.validate().is_ok());
+        assert!(bounds.is_canonical());
+        assert_eq!(bounds.start.to_offset(), Some(-1));
+        assert_eq!(bounds.end.to_offset(), Some(0));
+        assert_eq!(bounds.n_preceding_rows(), Some(1));
+        assert_eq!(bounds.n_following_rows(), Some(0));
+
+        let bounds = RowsFrameBounds {
+            start: CurrentRow,
+            end: Following(1),
+        };
+        assert!(bounds.validate().is_ok());
+        assert!(bounds.is_canonical());
+        assert_eq!(bounds.start.to_offset(), Some(0));
+        assert_eq!(bounds.end.to_offset(), Some(1));
+        assert_eq!(bounds.n_preceding_rows(), Some(0));
+        assert_eq!(bounds.n_following_rows(), Some(1));
+
+        let bounds = RowsFrameBounds {
+            start: UnboundedPreceding,
+            end: Following(10),
+        };
+        assert!(bounds.validate().is_ok());
+        assert!(bounds.is_canonical());
+        assert_eq!(bounds.start.to_offset(), None);
+        assert_eq!(bounds.end.to_offset(), Some(10));
+        assert_eq!(bounds.n_preceding_rows(), None);
+        assert_eq!(bounds.n_following_rows(), Some(10));
+
+        let bounds = RowsFrameBounds {
+            start: Preceding(10),
+            end: UnboundedFollowing,
+        };
+        assert!(bounds.validate().is_ok());
+        assert!(bounds.is_canonical());
+        assert_eq!(bounds.start.to_offset(), Some(-10));
+        assert_eq!(bounds.end.to_offset(), None);
+        assert_eq!(bounds.n_preceding_rows(), Some(10));
+        assert_eq!(bounds.n_following_rows(), None);
+
+        let bounds = RowsFrameBounds {
+            start: Preceding(1),
+            end: Preceding(10),
+        };
+        assert!(bounds.validate().is_ok());
+        assert!(!bounds.is_canonical());
+        assert_eq!(bounds.start.to_offset(), Some(-1));
+        assert_eq!(bounds.end.to_offset(), Some(-10));
+        assert_eq!(bounds.n_preceding_rows(), Some(10));
+        assert_eq!(bounds.n_following_rows(), Some(0));
+
+        let bounds = RowsFrameBounds {
+            start: Following(10),
+            end: Following(1),
+        };
+        assert!(bounds.validate().is_ok());
+        assert!(!bounds.is_canonical());
+        assert_eq!(bounds.start.to_offset(), Some(10));
+        assert_eq!(bounds.end.to_offset(), Some(1));
+        assert_eq!(bounds.n_preceding_rows(), Some(0));
+        assert_eq!(bounds.n_following_rows(), Some(10));
+
+        let bounds = RowsFrameBounds {
+            start: UnboundedFollowing,
+            end: Following(10),
+        };
+        assert!(bounds.validate().is_err());
+        assert!(!bounds.is_canonical());
+
+        let bounds = RowsFrameBounds {
+            start: Preceding(10),
+            end: UnboundedPreceding,
+        };
+        assert!(bounds.validate().is_err());
+        assert!(!bounds.is_canonical());
+    }
+}