feat: support qualify clause

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

@@ -2777,6 +2777,16 @@ impl<'ast> Visitor<'ast> for AstFormatVisitor {
             children.push(window_list_node);
         }
 
+        if let Some(qualify) = &stmt.qualify {
+            self.visit_expr(qualify);
+            let qualify_child = self.children.pop().unwrap();
+            let qualify_name = "Qualify".to_string();
+            let qualify_format_ctx = AstFormatContext::with_children(qualify_name, 1);
+            let qualify_node =
+                FormatTreeNode::with_children(qualify_format_ctx, vec![qualify_child]);
+            children.push(qualify_node);
+        }
+
         let name = "SelectQuery".to_string();
         let format_ctx = AstFormatContext::with_children(name, children.len());
         let node = FormatTreeNode::with_children(format_ctx, children);

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

@@ -92,6 +92,8 @@ pub struct SelectStmt {
     pub having: Option<Expr>,
     // `WINDOW` clause
     pub window_list: Option<Vec<WindowDefinition>>,
+    // `QUALIFY` clause
+    pub qualify: Option<Expr>,
 }
 
 /// Group by Clause.

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

@@ -58,6 +58,7 @@ pub enum SetOperationElement {
         group_by: Option<GroupBy>,
         having: Box<Option<Expr>>,
         window_list: Option<Vec<WindowDefinition>>,
+        qualify: Box<Option<Expr>>,
     },
     SetOperation {
         op: SetOperator,
@@ -104,6 +105,7 @@ pub fn set_operation_element(i: Input) -> IResult<WithSpan<SetOperationElement>>
                 ~ ( GROUP ~ ^BY ~ ^#group_by_items )?
                 ~ ( HAVING ~ ^#expr )?
                 ~ ( WINDOW ~ ^#comma_separated_list1(window_clause) )?
+                ~ ( QUALIFY ~ ^#expr )?
         },
         |(
             _select,
@@ -115,6 +117,7 @@ pub fn set_operation_element(i: Input) -> IResult<WithSpan<SetOperationElement>>
             opt_group_by_block,
             opt_having_block,
             opt_window_block,
+            opt_qualify_block,
         )| {
             SetOperationElement::SelectStmt {
                 hints: opt_hints,
@@ -129,6 +132,7 @@ pub fn set_operation_element(i: Input) -> IResult<WithSpan<SetOperationElement>>
                 group_by: opt_group_by_block.map(|(_, _, group_by)| group_by),
                 having: Box::new(opt_having_block.map(|(_, having)| having)),
                 window_list: opt_window_block.map(|(_, windows)| windows),
+                qualify: Box::new(opt_qualify_block.map(|(_, qualify)| qualify)),
             }
         },
     );
@@ -142,6 +146,7 @@ pub fn set_operation_element(i: Input) -> IResult<WithSpan<SetOperationElement>>
                 ~ ( GROUP ~ ^BY ~ ^#group_by_items )?
                 ~ ( HAVING ~ ^#expr )?
                 ~ ( WINDOW ~ ^#comma_separated_list1(window_clause) )?
+                ~ ( QUALIFY ~ ^#expr )?
         },
         |(
             opt_from_block,
@@ -153,6 +158,7 @@ pub fn set_operation_element(i: Input) -> IResult<WithSpan<SetOperationElement>>
             opt_group_by_block,
             opt_having_block,
             opt_window_block,
+            opt_qualify_block,
         )| {
             SetOperationElement::SelectStmt {
                 hints: opt_hints,
@@ -167,6 +173,7 @@ pub fn set_operation_element(i: Input) -> IResult<WithSpan<SetOperationElement>>
                 group_by: opt_group_by_block.map(|(_, _, group_by)| group_by),
                 having: Box::new(opt_having_block.map(|(_, having)| having)),
                 window_list: opt_window_block.map(|(_, windows)| windows),
+                qualify: Box::new(opt_qualify_block.map(|(_, qualify)| qualify)),
             }
         },
     );
@@ -262,6 +269,7 @@ impl<'a, I: Iterator<Item = WithSpan<'a, SetOperationElement>>> PrattParser<I>
                 group_by,
                 having,
                 window_list,
+                qualify,
             } => SetExpr::Select(Box::new(SelectStmt {
                 span: transform_span(input.span.0),
                 hints,
@@ -272,6 +280,7 @@ impl<'a, I: Iterator<Item = WithSpan<'a, SetOperationElement>>> PrattParser<I>
                 group_by,
                 having: *having,
                 window_list,
+                qualify: *qualify,
             })),
             SetOperationElement::Values(values) => SetExpr::Values {
                 span: transform_span(input.span.0),

src/query/ast/src/parser/token.rs

@@ -807,6 +807,8 @@ pub enum TokenKind {
     PRESIGN,
     #[token("PRIVILEGES", ignore(ascii_case))]
     PRIVILEGES,
+    #[token("QUALIFY", ignore(ascii_case))]
+    QUALIFY,
     #[token("REMOVE", ignore(ascii_case))]
     REMOVE,
     #[token("RETAIN", ignore(ascii_case))]
@@ -1301,6 +1303,7 @@ impl TokenKind {
             | TokenKind::OF
             | TokenKind::ORDER
             | TokenKind::OVER
+            | TokenKind::QUALIFY
             | TokenKind::ROWS
             // | TokenKind::PRECISION
             // | TokenKind::RETURNING
@@ -1421,6 +1424,7 @@ impl TokenKind {
             | TokenKind::ORDER
             | TokenKind::OVER
             | TokenKind::PARTITION
+            | TokenKind::QUALIFY
             | TokenKind::ROWS
             | TokenKind::RANGE
             // | TokenKind::OVERLAPS

src/query/ast/src/visitors/visitor.rs

@@ -646,6 +646,8 @@ pub trait Visitor<'ast>: Sized {
             selection,
             group_by,
             having,
+            window_list,
+            qualify,
             ..
         } = stmt;
 
@@ -680,6 +682,16 @@ pub trait Visitor<'ast>: Sized {
         if let Some(having) = having {
             walk_expr(self, having);
         }
+
+        if let Some(window_list) = window_list {
+            for window_def in window_list {
+                walk_window_definition(self, window_def);
+            }
+        }
+
+        if let Some(qualify) = qualify {
+            walk_expr(self, qualify);
+        }
     }
 
     fn visit_select_target(&mut self, target: &'ast SelectTarget) {

src/query/ast/src/visitors/visitor_mut.rs

@@ -28,6 +28,7 @@ use super::walk_mut::walk_statement_mut;
 use super::walk_mut::walk_table_reference_mut;
 use super::walk_stream_point_mut;
 use super::walk_time_travel_point_mut;
+use super::walk_window_definition_mut;
 use crate::ast::*;
 use crate::visitors::walk_column_id_mut;
 
@@ -659,6 +660,8 @@ pub trait VisitorMut: Sized {
             selection,
             group_by,
             having,
+            window_list,
+            qualify,
             ..
         } = stmt;
 
@@ -693,6 +696,16 @@ pub trait VisitorMut: Sized {
         if let Some(having) = having {
             Self::visit_expr(self, having);
         }
+
+        if let Some(window_list) = window_list {
+            for window_def in window_list {
+                walk_window_definition_mut(self, window_def);
+            }
+        }
+
+        if let Some(qualify) = qualify {
+            Self::visit_expr(self, qualify);
+        }
     }
 
     fn visit_select_target(&mut self, target: &mut SelectTarget) {

src/query/ast/src/visitors/walk_mut.rs

@@ -198,6 +198,35 @@ pub fn walk_set_expr_mut<V: VisitorMut>(visitor: &mut V, set_expr: &mut SetExpr)
     }
 }
 
+pub fn walk_window_definition_mut<V: VisitorMut>(
+    visitor: &mut V,
+    window_definition: &mut WindowDefinition,
+) {
+    let WindowDefinition { name, spec: window } = window_definition;
+
+    visitor.visit_identifier(name);
+
+    let WindowSpec {
+        partition_by,
+        order_by,
+        window_frame,
+        ..
+    } = window;
+
+    for expr in partition_by {
+        visitor.visit_expr(expr);
+    }
+
+    for order_by in order_by {
+        visitor.visit_order_by(order_by);
+    }
+
+    if let Some(frame) = window_frame {
+        visitor.visit_frame_bound(&mut frame.start_bound);
+        visitor.visit_frame_bound(&mut frame.end_bound);
+    }
+}
+
 pub fn walk_select_target_mut<V: VisitorMut>(visitor: &mut V, target: &mut SelectTarget) {
     match target {
         SelectTarget::AliasedExpr { expr, alias } => {

src/query/ast/tests/it/testdata/query-error.txt

@@ -79,6 +79,6 @@ error:
   --> SQL:1:10
   |
 1 | select 1 1
-  |          ^ unexpected `1`, expecting <Ident>, <QuotedString>, `AS`, `,`, `FROM`, `WHERE`, `GROUP`, `HAVING`, `WINDOW`, `(`, `WITH`, `UNION`, `EXCEPT`, `INTERSECT`, `SELECT`, `VALUES`, `ORDER`, `LIMIT`, `OFFSET`, or `IGNORE_RESULT`
+  |          ^ unexpected `1`, expecting <Ident>, <QuotedString>, `AS`, `,`, `FROM`, `WHERE`, `GROUP`, `HAVING`, `WINDOW`, `QUALIFY`, `(`, `WITH`, `UNION`, `EXCEPT`, `INTERSECT`, `SELECT`, `VALUES`, `ORDER`, `LIMIT`, `OFFSET`, or `IGNORE_RESULT`