support select * except col_name

datafuselabs · Nov 29, 2022 · a876f0a · a876f0a
1 parent 3ac1e59
commit a876f0a
Show file tree

Hide file tree

Showing 13 changed files with 986 additions and 232 deletions.
diff --git a/src/query/ast/src/ast/format/ast_format.rs b/src/query/ast/src/ast/format/ast_format.rs
@@ -2093,7 +2093,7 @@ impl<'ast> Visitor<'ast> for AstFormatVisitor {
                 let node = FormatTreeNode::with_children(format_ctx, vec![child]);
                 self.children.push(node);
             }
-            SelectTarget::QualifiedName(_) => {
+            SelectTarget::QualifiedName { .. } => {
                 let name = format!("Target {}", target);
                 let format_ctx = AstFormatContext::new(name);
                 let node = FormatTreeNode::new(format_ctx);

diff --git a/src/query/ast/src/ast/format/syntax/query.rs b/src/query/ast/src/ast/format/syntax/query.rs
@@ -20,6 +20,7 @@ use crate::ast::format::syntax::inline_dot;
 use crate::ast::format::syntax::interweave_comma;
 use crate::ast::format::syntax::parenthenized;
 use crate::ast::format::syntax::NEST_FACTOR;
+use crate::ast::ExcludeCol;
 use crate::ast::Expr;
 use crate::ast::JoinCondition;
 use crate::ast::JoinOperator;
@@ -108,26 +109,70 @@ fn pretty_select_list(select_list: Vec<SelectTarget>) -> RcDoc {
     }
     .nest(NEST_FACTOR)
     .append(
-        interweave_comma(select_list.into_iter().map(|select_target| {
-            match select_target {
-                SelectTarget::AliasedExpr { expr, alias } => {
-                    pretty_expr(*expr).append(if let Some(alias) = alias {
-                        RcDoc::space()
-                            .append(RcDoc::text("AS"))
-                            .append(RcDoc::space())
-                            .append(RcDoc::text(alias.to_string()))
-                    } else {
-                        RcDoc::nil()
-                    })
-                }
-                SelectTarget::QualifiedName(object_name) => inline_dot(
-                    object_name
-                        .into_iter()
-                        .map(|indirection| RcDoc::text(indirection.to_string())),
-                )
-                .group(),
-            }
-        }))
+        interweave_comma(
+            select_list
+                .into_iter()
+                .map(|select_target| match select_target {
+                    SelectTarget::AliasedExpr { expr, alias } => {
+                        pretty_expr(*expr).append(if let Some(alias) = alias {
+                            RcDoc::space()
+                                .append(RcDoc::text("AS"))
+                                .append(RcDoc::space())
+                                .append(RcDoc::text(alias.to_string()))
+                        } else {
+                            RcDoc::nil()
+                        })
+                    }
+                    SelectTarget::QualifiedName {
+                        qualified: object_name,
+                        exclude,
+                    } => {
+                        let docs = inline_dot(
+                            object_name
+                                .into_iter()
+                                .map(|indirection| RcDoc::text(indirection.to_string())),
+                        )
+                        .group();
+                        docs.append(if let Some(exclude) = exclude {
+                            match exclude {
+                                ExcludeCol::Col(col) => RcDoc::line().append(
+                                    RcDoc::text("EXCEPT")
+                                        .append(RcDoc::space().nest(NEST_FACTOR))
+                                        .append(RcDoc::text(col.to_string())),
+                                ),
+                                ExcludeCol::Cols(cols) => {
+                                    if !cols.is_empty() {
+                                        RcDoc::line()
+                                            .append(
+                                                RcDoc::text("EXCEPT").append(
+                                                    if cols.len() > 1 {
+                                                        RcDoc::line()
+                                                    } else {
+                                                        RcDoc::space()
+                                                    }
+                                                    .nest(NEST_FACTOR),
+                                                ),
+                                            )
+                                            .append(
+                                                interweave_comma(cols.into_iter().map(|ident| {
+                                                    RcDoc::space()
+                                                        .append(RcDoc::space())
+                                                        .append(RcDoc::text(ident.to_string()))
+                                                }))
+                                                .nest(NEST_FACTOR)
+                                                .group(),
+                                            )
+                                    } else {
+                                        RcDoc::nil()
+                                    }
+                                }
+                            }
+                        } else {
+                            RcDoc::nil()
+                        })
+                    }
+                }),
+        )
         .nest(NEST_FACTOR)
         .group(),
     )

diff --git a/src/query/ast/src/ast/query.rs b/src/query/ast/src/ast/query.rs
@@ -119,9 +119,31 @@ pub enum SelectTarget<'a> {
         alias: Option<Identifier<'a>>,
     },
 
-    // Qualified name, e.g. `SELECT t.a, t.* FROM t`.
+    // Qualified name, e.g. `SELECT t.a, t.* exclude t.a FROM t`.
     // For simplicity, wildcard is involved.
-    QualifiedName(QualifiedName<'a>),
+    QualifiedName {
+        qualified: QualifiedName<'a>,
+        exclude: Option<ExcludeCol<'a>>,
+    },
+}
+
+#[derive(Debug, Clone, PartialEq)]
+pub enum ExcludeCol<'a> {
+    Col(Identifier<'a>),
+    Cols(Vec<Identifier<'a>>),
+}
+
+impl Display for ExcludeCol<'_> {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        match self {
+            ExcludeCol::Col(col) => write!(f, "{col}"),
+            ExcludeCol::Cols(cols) => {
+                write!(f, "(")?;
+                write_comma_separated_list(f, cols)?;
+                write!(f, ")")
+            }
+        }
+    }
 }
 
 pub type QualifiedName<'a> = Vec<Indirection<'a>>;
@@ -370,8 +392,24 @@ impl<'a> Display for SelectTarget<'a> {
                     write!(f, " AS {ident}")?;
                 }
             }
-            SelectTarget::QualifiedName(indirections) => {
-                write_period_separated_list(f, indirections)?;
+            SelectTarget::QualifiedName { qualified, exclude } => {
+                write_period_separated_list(f, qualified)?;
+                if let Some(exclude) = exclude {
+                    // ORDER BY clause
+                    match exclude {
+                        ExcludeCol::Col(col) => {
+                            write!(f, " EXCLUDE")?;
+                            write!(f, " {col}")?;
+                        }
+                        ExcludeCol::Cols(cols) => {
+                            if !cols.is_empty() {
+                                write!(f, " EXCLUDE (")?;
+                                write_comma_separated_list(f, cols)?;
+                                write!(f, ")")?;
+                            }
+                        }
+                    }
+                }
             }
         }
         Ok(())

diff --git a/src/query/ast/src/parser/query.rs b/src/query/ast/src/parser/query.rs
@@ -83,21 +83,48 @@ pub fn with(i: Input) -> IResult<With> {
     )(i)
 }
 
+pub fn exclude_col(i: Input) -> IResult<ExcludeCol> {
+    let var = map(
+        rule! {
+            #ident
+        },
+        ExcludeCol::Col,
+    );
+    let vars = map(
+        rule! {
+            "(" ~ ^#comma_separated_list1(ident) ~ ")"
+        },
+        |(_, cols, _)| ExcludeCol::Cols(cols),
+    );
+
+    rule!(
+        #var
+        | #vars
+    )(i)
+}
+
 pub fn select_target(i: Input) -> IResult<SelectTarget> {
     let qualified_wildcard = map(
         rule! {
-            ( #ident ~ "." ~ ( #ident ~ "." )? )? ~ "*"
+            ( #ident ~ "." ~ ( #ident ~ "." )? )? ~ "*" ~ ( ( EXCEPT | EXCLUDE ) ~ #exclude_col )?
         },
-        |(res, _)| match res {
-            Some((fst, _, Some((snd, _)))) => SelectTarget::QualifiedName(vec![
-                Indirection::Identifier(fst),
-                Indirection::Identifier(snd),
-                Indirection::Star,
-            ]),
-            Some((fst, _, None)) => {
-                SelectTarget::QualifiedName(vec![Indirection::Identifier(fst), Indirection::Star])
-            }
-            None => SelectTarget::QualifiedName(vec![Indirection::Star]),
+        |(res, _, opt_exclude)| match res {
+            Some((fst, _, Some((snd, _)))) => SelectTarget::QualifiedName {
+                qualified: vec![
+                    Indirection::Identifier(fst),
+                    Indirection::Identifier(snd),
+                    Indirection::Star,
+                ],
+                exclude: opt_exclude.map(|(_, exclude)| exclude),
+            },
+            Some((fst, _, None)) => SelectTarget::QualifiedName {
+                qualified: vec![Indirection::Identifier(fst), Indirection::Star],
+                exclude: opt_exclude.map(|(_, exclude)| exclude),
+            },
+            None => SelectTarget::QualifiedName {
+                qualified: vec![Indirection::Star],
+                exclude: opt_exclude.map(|(_, exclude)| exclude),
+            },
         },
     );
     let projection = map(

diff --git a/src/query/ast/src/parser/token.rs b/src/query/ast/src/parser/token.rs
@@ -367,6 +367,8 @@ pub enum TokenKind {
     DROP,
     #[token("EXCEPT", ignore(ascii_case))]
     EXCEPT,
+    #[token("EXCLUDE", ignore(ascii_case))]
+    EXCLUDE,
     #[token("ELSE", ignore(ascii_case))]
     ELSE,
     #[token("END", ignore(ascii_case))]

diff --git a/src/query/ast/src/visitors/walk.rs b/src/query/ast/src/visitors/walk.rs
@@ -185,7 +185,10 @@ pub fn walk_select_target<'a, V: Visitor<'a>>(visitor: &mut V, target: &'a Selec
                 visitor.visit_identifier(alias);
             }
         }
-        SelectTarget::QualifiedName(names) => {
+        SelectTarget::QualifiedName {
+            qualified: names,
+            exclude,
+        } => {
             for indirection in names {
                 match indirection {
                     Indirection::Identifier(ident) => {
@@ -194,6 +197,16 @@ pub fn walk_select_target<'a, V: Visitor<'a>>(visitor: &mut V, target: &'a Selec
                     Indirection::Star => {}
                 }
             }
+            if let Some(exclude) = exclude {
+                match exclude {
+                    ExcludeCol::Col(col) => visitor.visit_identifier(col),
+                    ExcludeCol::Cols(cols) => {
+                        for ident in cols.iter() {
+                            visitor.visit_identifier(ident);
+                        }
+                    }
+                }
+            }
         }
     }
 }

diff --git a/src/query/ast/src/visitors/walk_mut.rs b/src/query/ast/src/visitors/walk_mut.rs
@@ -185,7 +185,10 @@ pub fn walk_select_target_mut<'a, V: VisitorMut>(visitor: &mut V, target: &mut S
                 visitor.visit_identifier(alias);
             }
         }
-        SelectTarget::QualifiedName(names) => {
+        SelectTarget::QualifiedName {
+            qualified: names,
+            exclude,
+        } => {
             for indirection in names {
                 match indirection {
                     Indirection::Identifier(ident) => {
@@ -194,6 +197,16 @@ pub fn walk_select_target_mut<'a, V: VisitorMut>(visitor: &mut V, target: &mut S
                     Indirection::Star => {}
                 }
             }
+            if let Some(exclude) = exclude {
+                match exclude {
+                    ExcludeCol::Col(col) => visitor.visit_identifier(col),
+                    ExcludeCol::Cols(cols) => {
+                        for ident in cols {
+                            visitor.visit_identifier(ident);
+                        }
+                    }
+                }
+            }
         }
     }
 }

diff --git a/src/query/ast/tests/it/parser.rs b/src/query/ast/tests/it/parser.rs
@@ -418,6 +418,7 @@ fn test_query() {
     let mut mint = Mint::new("tests/it/testdata");
     let mut file = mint.new_goldenfile("query.txt").unwrap();
     let cases = &[
+        r#"select * except c1, b.* except (c2, c3, c4) from customer inner join orders on a = b limit 1"#,
         r#"select * from customer inner join orders"#,
         r#"select * from customer cross join orders"#,
         r#"select * from customer inner join orders on a = b limit 1"#,