feat(window): support named window (#10897)

* support window clause

* window clause ast part

* named window

* add tests

* regenerate files

* add comments

* remove unused trait impl

* fix review comments

Author: ariesdevil

Cargo.lock

@@ -16,6 +16,7 @@ common-io = { path = "../../common/io" }
 common-meta-app = { path = "../../meta/app" }
 
 # Crates.io dependencies
+enum-as-inner = "0.5.1"
 ethnum = "1.3.2"
 fast-float = "0.2.0"
 itertools = "0.10.5"

src/query/ast/Cargo.toml

@@ -21,6 +21,7 @@ use common_exception::Result;
 use common_exception::Span;
 use common_io::display_decimal_128;
 use common_io::display_decimal_256;
+use enum_as_inner::EnumAsInner;
 use ethnum::i256;
 
 use super::OrderByExpr;
@@ -155,7 +156,7 @@ pub enum Expr {
         name: Identifier,
         args: Vec<Expr>,
         params: Vec<Literal>,
-        window: Option<WindowSpec>,
+        window: Option<Window>,
     },
     /// `CASE ... WHEN ... ELSE ...` expression
     Case {
@@ -458,6 +459,23 @@ pub enum TrimWhere {
     Trailing,
 }
 
+#[derive(Debug, Clone, PartialEq, EnumAsInner)]
+pub enum Window {
+    WindowReference(WindowRef),
+    WindowSpec(WindowSpec),
+}
+
+#[derive(Debug, Clone, PartialEq)]
+pub struct WindowDefinition {
+    pub name: Identifier,
+    pub spec: WindowSpec,
+}
+
+#[derive(Debug, Clone, PartialEq)]
+pub struct WindowRef {
+    pub window_name: Identifier,
+}
+
 #[derive(Debug, Clone, PartialEq)]
 pub struct WindowSpec {
     pub partition_by: Vec<Expr>,
@@ -912,6 +930,28 @@ impl Display for Literal {
     }
 }
 
+impl Display for WindowDefinition {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(f, "WINDOW {} {}", self.name, self.spec)
+    }
+}
+
+impl Display for Window {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        let window_fmt = match *self {
+            Window::WindowSpec(ref window_spec) => format!("{}", window_spec),
+            Window::WindowReference(ref window_ref) => format!("{}", window_ref),
+        };
+        write!(f, "{}", window_fmt)
+    }
+}
+
+impl Display for WindowRef {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(f, "WINDOW {}", self.window_name)
+    }
+}
+
 impl Display for WindowSpec {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
         let mut first = true;
@@ -952,28 +992,22 @@ impl Display for WindowSpec {
                     write!(f, "RANGE")?;
                 }
             }
-            match (&frame.start_bound, &frame.end_bound) {
-                (WindowFrameBound::CurrentRow, WindowFrameBound::CurrentRow) => {
-                    write!(f, " CURRENT ROW")?
-                }
-                _ => {
-                    let format_frame = |frame: &WindowFrameBound| -> String {
-                        match frame {
-                            WindowFrameBound::CurrentRow => "CURRENT ROW".to_string(),
-                            WindowFrameBound::Preceding(None) => "UNBOUNDED PRECEDING".to_string(),
-                            WindowFrameBound::Following(None) => "UNBOUNDED FOLLOWING".to_string(),
-                            WindowFrameBound::Preceding(Some(n)) => format!("{} PRECEDING", n),
-                            WindowFrameBound::Following(Some(n)) => format!("{} FOLLOWING", n),
-                        }
-                    };
-                    write!(
-                        f,
-                        " BETWEEN {} AND {}",
-                        format_frame(&frame.start_bound),
-                        format_frame(&frame.end_bound)
-                    )?
+
+            let format_frame = |frame: &WindowFrameBound| -> String {
+                match frame {
+                    WindowFrameBound::CurrentRow => "CURRENT ROW".to_string(),
+                    WindowFrameBound::Preceding(None) => "UNBOUNDED PRECEDING".to_string(),
+                    WindowFrameBound::Following(None) => "UNBOUNDED FOLLOWING".to_string(),
+                    WindowFrameBound::Preceding(Some(n)) => format!("{} PRECEDING", n),
+                    WindowFrameBound::Following(Some(n)) => format!("{} FOLLOWING", n),
                 }
-            }
+            };
+            write!(
+                f,
+                " BETWEEN {} AND {}",
+                format_frame(&frame.start_bound),
+                format_frame(&frame.end_bound)
+            )?
         }
         Ok(())
     }

src/query/ast/src/ast/expr.rs

@@ -430,7 +430,7 @@ impl<'ast> Visitor<'ast> for AstFormatVisitor {
         name: &'ast Identifier,
         args: &'ast [Expr],
         _params: &'ast [Literal],
-        _over: &'ast Option<WindowSpec>,
+        _over: &'ast Option<Window>,
     ) {
         let mut children = Vec::with_capacity(args.len());
         for arg in args.iter() {
@@ -2215,6 +2215,20 @@ impl<'ast> Visitor<'ast> for AstFormatVisitor {
             children.push(having_node);
         }
 
+        if let Some(window_list) = &stmt.window_list {
+            let mut window_list_children = Vec::with_capacity(window_list.len());
+            for window in window_list {
+                self.visit_window_definition(window);
+                window_list_children.push(self.children.pop().unwrap());
+            }
+            let window_list_name = "WindowList".to_string();
+            let window_list_format_ctx =
+                AstFormatContext::with_children(window_list_name, window_list_children.len());
+            let window_list_node =
+                FormatTreeNode::with_children(window_list_format_ctx, window_list_children);
+            children.push(window_list_node);
+        }
+
         let name = "SelectQuery".to_string();
         let format_ctx = AstFormatContext::with_children(name, children.len());
         let node = FormatTreeNode::with_children(format_ctx, children);
@@ -2244,6 +2258,17 @@ impl<'ast> Visitor<'ast> for AstFormatVisitor {
         }
     }
 
+    fn visit_window_definition(&mut self, window: &'ast WindowDefinition) {
+        self.visit_identifier(&window.name);
+        let window_name = self.children.pop().unwrap();
+        self.visit_window(&Window::WindowSpec(window.spec.clone()));
+        let window = self.children.pop().unwrap();
+        let name = "Window".to_string();
+        let format_ctx = AstFormatContext::with_children(name, 2);
+        let node = FormatTreeNode::with_children(format_ctx, vec![window_name, window]);
+        self.children.push(node);
+    }
+
     fn visit_table_reference(&mut self, table: &'ast TableReference) {
         match table {
             TableReference::Table {

src/query/ast/src/ast/format/ast_format.rs

@@ -23,6 +23,7 @@ use crate::ast::Expr;
 use crate::ast::FileLocation;
 use crate::ast::Identifier;
 use crate::ast::SelectStageOptions;
+use crate::ast::WindowDefinition;
 
 /// Root node of a query tree
 #[derive(Debug, Clone, PartialEq)]
@@ -85,6 +86,8 @@ pub struct SelectStmt {
     pub group_by: Option<GroupBy>,
     // `HAVING` clause
     pub having: Option<Expr>,
+    // `WINDOW` clause
+    pub window_list: Option<Vec<WindowDefinition>>,
 }
 
 /// Group by Clause.

src/query/ast/src/ast/query.rs

@@ -284,7 +284,7 @@ pub enum ExprElement {
         distinct: bool,
         name: Identifier,
         args: Vec<Expr>,
-        window: Option<WindowSpec>,
+        window: Option<Window>,
         params: Vec<Literal>,
     },
     /// `CASE ... WHEN ... ELSE ...` expression
@@ -848,41 +848,6 @@ pub fn expr_element(i: Input) -> IResult<WithSpan<ExprElement>> {
         },
     );
 
-    let window_frame_between = alt((
-        map(
-            rule! { BETWEEN ~ #window_frame_bound ~ AND ~ #window_frame_bound },
-            |(_, s, _, e)| (s, e),
-        ),
-        map(rule! {#window_frame_bound}, |s| {
-            (s, WindowFrameBound::Following(None))
-        }),
-    ));
-
-    let window_spec = map(
-        rule! {
-            (PARTITION ~ ^BY ~ #comma_separated_list1(subexpr(0)))?
-            ~ ( ORDER ~ ^BY ~ ^#comma_separated_list1(order_by_expr) )?
-            ~ ((ROWS | RANGE) ~ #window_frame_between)?
-        },
-        |(opt_partition, opt_order, between)| WindowSpec {
-            partition_by: opt_partition.map(|x| x.2).unwrap_or_default(),
-            order_by: opt_order.map(|x| x.2).unwrap_or_default(),
-            window_frame: between.map(|x| {
-                let unit = match x.0.kind {
-                    ROWS => WindowFrameUnits::Rows,
-                    RANGE => WindowFrameUnits::Range,
-                    _ => unreachable!(),
-                };
-                let bw = x.1;
-                WindowFrame {
-                    units: unit,
-                    start_bound: bw.0,
-                    end_bound: bw.1,
-                }
-            }),
-        },
-    );
-
     let function_call = map(
         rule! {
             #function_name
@@ -901,14 +866,14 @@ pub fn expr_element(i: Input) -> IResult<WithSpan<ExprElement>> {
         rule! {
             #function_name
             ~ "(" ~ DISTINCT? ~ #comma_separated_list0(subexpr(0))? ~ ")"
-            ~ (OVER ~ "(" ~ #window_spec ~ ")")
+            ~ (OVER ~ #window_spec_ident)
         },
         |(name, _, opt_distinct, opt_args, _, window)| ExprElement::FunctionCall {
             distinct: opt_distinct.is_some(),
             name,
             args: opt_args.unwrap_or_default(),
             params: vec![],
-            window: Some(window.2),
+            window: Some(window.1),
         },
     );
 
@@ -1080,6 +1045,7 @@ pub fn expr_element(i: Input) -> IResult<WithSpan<ExprElement>> {
         },
         |(_, not, _, _)| ExprElement::IsDistinctFrom { not: not.is_some() },
     );
+
     let (rest, (span, elem)) = consumed(alt((
         // Note: each `alt` call supports maximum of 21 parsers
         rule!(
@@ -1125,26 +1091,6 @@ pub fn expr_element(i: Input) -> IResult<WithSpan<ExprElement>> {
     Ok((rest, WithSpan { span, elem }))
 }
 
-pub fn window_frame_bound(i: Input) -> IResult<WindowFrameBound> {
-    alt((
-        value(WindowFrameBound::CurrentRow, rule! { CURRENT ~ ROW }),
-        map(rule! { #subexpr(0) ~ PRECEDING }, |(expr, _)| {
-            WindowFrameBound::Preceding(Some(Box::new(expr)))
-        }),
-        value(
-            WindowFrameBound::Preceding(None),
-            rule! { UNBOUNDED ~ PRECEDING },
-        ),
-        map(rule! { #subexpr(0) ~ FOLLOWING }, |(expr, _)| {
-            WindowFrameBound::Following(Some(Box::new(expr)))
-        }),
-        value(
-            WindowFrameBound::Following(None),
-            rule! { UNBOUNDED ~ FOLLOWING },
-        ),
-    ))(i)
-}
-
 pub fn unary_op(i: Input) -> IResult<UnaryOperator> {
     // Plus and Minus are parsed as binary op at first.
     alt((