feat(query): Support values Clause (#12380)

* feat(query): Support values Clause * fix * fix * add tests * fix --------- Co-authored-by: xudong.w <wxd963996380@gmail.com>
databendlabs · Aug 9, 2023 · 120669f · 120669f · vercel · Aug 9, 2023
1 parent 33f9b77
commit 120669f
Show file tree

Hide file tree

Showing 34 changed files with 1,729 additions and 572 deletions.
diff --git a/src/query/ast/src/ast/format/ast_format.rs b/src/query/ast/src/ast/format/ast_format.rs
@@ -2353,14 +2353,37 @@ impl<'ast> Visitor<'ast> for AstFormatVisitor {
     fn visit_with(&mut self, with: &'ast With) {
         let mut children = Vec::with_capacity(with.ctes.len());
         for cte in with.ctes.iter() {
-            self.visit_query(&cte.query);
-            let query_child = self.children.pop().unwrap();
+            let source_child = match &cte.source {
+                CTESource::Query { query, .. } => {
+                    self.visit_query(query);
+                    self.children.pop().unwrap()
+                }
+                CTESource::Values(values) => {
+                    let mut value_children = Vec::with_capacity(values.len());
+                    for (i, row_values) in values.iter().enumerate() {
+                        let mut row_children = Vec::with_capacity(row_values.len());
+                        for value in row_values {
+                            self.visit_expr(value);
+                            row_children.push(self.children.pop().unwrap());
+                        }
+                        let row_name = format!("Row {}", i);
+                        let row_format_ctx =
+                            AstFormatContext::with_children(row_name, row_children.len());
+                        let row_node = FormatTreeNode::with_children(row_format_ctx, row_children);
+                        value_children.push(row_node);
+                    }
+                    let value_format_ctx =
+                        AstFormatContext::with_children("Values".to_string(), value_children.len());
+                    FormatTreeNode::with_children(value_format_ctx, value_children)
+                }
+            };
+
             let cte_format_ctx = AstFormatContext::with_children_alias(
                 "CTE".to_string(),
                 1,
                 Some(format!("{}", cte.alias)),
             );
-            let cte_node = FormatTreeNode::with_children(cte_format_ctx, vec![query_child]);
+            let cte_node = FormatTreeNode::with_children(cte_format_ctx, vec![source_child]);
             children.push(cte_node);
         }
 
@@ -2719,6 +2742,36 @@ impl<'ast> Visitor<'ast> for AstFormatVisitor {
                 let node = FormatTreeNode::with_children(format_ctx, children);
                 self.children.push(node)
             }
+            TableReference::Values {
+                span: _,
+                values,
+                alias,
+            } => {
+                let mut children = Vec::with_capacity(values.len());
+                for (i, row_values) in values.iter().enumerate() {
+                    let mut row_children = Vec::with_capacity(row_values.len());
+                    for value in row_values {
+                        self.visit_expr(value);
+                        row_children.push(self.children.pop().unwrap());
+                    }
+                    let row_name = format!("Row {}", i);
+                    let row_format_ctx =
+                        AstFormatContext::with_children(row_name, row_children.len());
+                    let row_node = FormatTreeNode::with_children(row_format_ctx, row_children);
+                    children.push(row_node);
+                }
+                let format_ctx = if let Some(alias) = alias {
+                    AstFormatContext::with_children_alias(
+                        "Values".to_string(),
+                        children.len(),
+                        Some(format!("{}", alias)),
+                    )
+                } else {
+                    AstFormatContext::with_children("Values".to_string(), children.len())
+                };
+                let node = FormatTreeNode::with_children(format_ctx, children);
+                self.children.push(node);
+            }
         }
     }
 

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::parenthesized;
 use crate::ast::format::syntax::NEST_FACTOR;
+use crate::ast::CTESource;
 use crate::ast::Expr;
 use crate::ast::GroupBy;
 use crate::ast::JoinCondition;
@@ -65,9 +66,26 @@ fn pretty_with(with: Option<With>) -> RcDoc<'static> {
 }
 
 fn pretty_cte(cte: CTE) -> RcDoc<'static> {
-    RcDoc::text(format!("{} AS", cte.alias))
+    RcDoc::text(format!("{} AS ", cte.alias))
         .append(RcDoc::softline())
-        .append(parenthesized(pretty_query(cte.query)))
+        .append(match cte.source {
+            CTESource::Query {
+                materialized,
+                query,
+            } => if materialized {
+                RcDoc::text("MATERIALIZED ".to_string())
+            } else {
+                RcDoc::nil()
+            }
+            .append(parenthesized(pretty_query(*query))),
+            CTESource::Values(values) => RcDoc::text("(VALUES")
+                .append(inline_comma(values.into_iter().map(|row_values| {
+                    RcDoc::text("(")
+                        .append(inline_comma(row_values.into_iter().map(pretty_expr)))
+                        .append(RcDoc::text(")"))
+                })))
+                .append(RcDoc::text(")")),
+        })
 }
 
 fn pretty_body(body: SetExpr) -> RcDoc<'static> {
@@ -416,6 +434,22 @@ pub(crate) fn pretty_table(table: TableReference) -> RcDoc<'static> {
             } else {
                 RcDoc::nil()
             }),
+        TableReference::Values {
+            span: _,
+            values,
+            alias,
+        } => RcDoc::text("(VALUES")
+            .append(inline_comma(values.into_iter().map(|row_values| {
+                RcDoc::text("(")
+                    .append(inline_comma(row_values.into_iter().map(pretty_expr)))
+                    .append(RcDoc::text(")"))
+            })))
+            .append(RcDoc::text(")"))
+            .append(if let Some(alias) = alias {
+                RcDoc::text(format!(" {alias}"))
+            } else {
+                RcDoc::nil()
+            }),
     }
 }
 

diff --git a/src/query/ast/src/ast/query.rs b/src/query/ast/src/ast/query.rs
@@ -59,8 +59,16 @@ pub struct With {
 pub struct CTE {
     pub span: Span,
     pub alias: TableAlias,
-    pub materialized: bool,
-    pub query: Query,
+    pub source: CTESource,
+}
+
+#[derive(Debug, Clone, PartialEq)]
+pub enum CTESource {
+    Query {
+        materialized: bool,
+        query: Box<Query>,
+    },
+    Values(Vec<Vec<Expr>>),
 }
 
 #[derive(Debug, Clone, PartialEq)]
@@ -267,6 +275,12 @@ pub enum TableReference {
         options: SelectStageOptions,
         alias: Option<TableAlias>,
     },
+    // A set of values generates a temporary constant table that can be used for queries
+    Values {
+        span: Span,
+        values: Vec<Vec<Expr>>,
+        alias: Option<TableAlias>,
+    },
 }
 
 impl TableReference {
@@ -531,6 +545,25 @@ impl Display for TableReference {
                     write!(f, " AS {alias}")?;
                 }
             }
+            TableReference::Values {
+                span: _,
+                values,
+                alias,
+            } => {
+                write!(f, "(VALUES")?;
+                for (i, value) in values.iter().enumerate() {
+                    if i > 0 {
+                        write!(f, ", ")?;
+                    }
+                    write!(f, "(")?;
+                    write_comma_separated_list(f, value)?;
+                    write!(f, ")")?;
+                }
+                write!(f, ")")?;
+                if let Some(alias) = alias {
+                    write!(f, " {alias}")?;
+                }
+            }
         }
         Ok(())
     }
@@ -676,7 +709,36 @@ impl Display for SetExpr {
 
 impl Display for CTE {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
-        write!(f, "{} AS ({})", self.alias, self.query)?;
+        write!(f, "{} AS {}", self.alias, self.source)?;
+        Ok(())
+    }
+}
+
+impl Display for CTESource {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        match self {
+            CTESource::Query {
+                materialized,
+                query,
+            } => {
+                if *materialized {
+                    write!(f, "MATERIALIZED ")?;
+                }
+                write!(f, "({query})")?;
+            }
+            CTESource::Values(values) => {
+                write!(f, "(VALUES")?;
+                for (i, value) in values.iter().enumerate() {
+                    if i > 0 {
+                        write!(f, ", ")?;
+                    }
+                    write!(f, "(")?;
+                    write_comma_separated_list(f, value)?;
+                    write!(f, ")")?;
+                }
+                write!(f, ")")?;
+            }
+        }
         Ok(())
     }
 }

diff --git a/src/query/ast/src/parser/query.rs b/src/query/ast/src/parser/query.rs
@@ -304,16 +304,42 @@ impl<'a, I: Iterator<Item = WithSpan<'a, SetOperationElement>>> PrattParser<I>
     }
 }
 
+pub fn row_values(i: Input) -> IResult<Vec<Expr>> {
+    map(
+        rule! {"(" ~ #comma_separated_list1(expr) ~ ")"},
+        |(_, row_values, _)| row_values,
+    )(i)
+}
+
+pub fn cte_source(i: Input) -> IResult<CTESource> {
+    alt((
+        map(
+            rule! {
+                MATERIALIZED? ~ "(" ~ #query ~ ")"
+            },
+            |(materialized, _, query, _)| CTESource::Query {
+                materialized: materialized.is_some(),
+                query: Box::new(query),
+            },
+        ),
+        map(
+            rule! {
+                "(" ~ VALUES ~ ^#comma_separated_list1(row_values) ~ ")"
+            },
+            |(_, _, values, _)| CTESource::Values(values),
+        ),
+    ))(i)
+}
+
 pub fn with(i: Input) -> IResult<With> {
     let cte = map(
         consumed(rule! {
-            #table_alias ~ AS ~ MATERIALIZED? ~ "(" ~ #query ~ ")"
+            #table_alias ~ AS ~ #cte_source
         }),
-        |(span, (table_alias, _, materialized, _, query, _))| CTE {
+        |(span, (table_alias, _, source))| CTE {
             span: transform_span(span.0),
             alias: table_alias,
-            materialized: materialized.is_some(),
-            query,
+            source,
         },
     );
 
@@ -526,6 +552,10 @@ pub enum TableReferenceElement {
         options: Vec<SelectStageOption>,
         alias: Option<TableAlias>,
     },
+    Values {
+        values: Vec<Vec<Expr>>,
+        alias: Option<TableAlias>,
+    },
 }
 
 pub fn table_reference_element(i: Input) -> IResult<WithSpan<TableReferenceElement>> {
@@ -633,11 +663,19 @@ pub fn table_reference_element(i: Input) -> IResult<WithSpan<TableReferenceEleme
         },
     );
 
+    let values = map(
+        rule! {
+            "(" ~ VALUES ~ ^#comma_separated_list1(row_values) ~ ")" ~ #table_alias?
+        },
+        |(_, _, values, _, alias)| TableReferenceElement::Values { values, alias },
+    );
+
     let (rest, (span, elem)) = consumed(rule! {
         #aliased_stage
         | #table_function
         | #aliased_table
         | #subquery
+        | #values
         | #group
         | #join
         | #join_condition_on
@@ -730,6 +768,11 @@ impl<'a, I: Iterator<Item = WithSpan<'a, TableReferenceElement>>> PrattParser<I>
                     alias,
                 }
             }
+            TableReferenceElement::Values { values, alias } => TableReference::Values {
+                span: transform_span(input.span.0),
+                values,
+                alias,
+            },
             _ => unreachable!(),
         };
         Ok(table_ref)

diff --git a/src/query/ast/src/visitors/walk.rs b/src/query/ast/src/visitors/walk.rs
@@ -263,6 +263,16 @@ pub fn walk_table_reference<'a, V: Visitor<'a>>(visitor: &mut V, table_ref: &'a
             visitor.visit_join(join);
         }
         TableReference::Stage { .. } => {}
+        TableReference::Values { values, alias, .. } => {
+            for row_values in values {
+                for value in row_values {
+                    visitor.visit_expr(value);
+                }
+            }
+            if let Some(alias) = alias {
+                visitor.visit_identifier(&alias.name);
+            }
+        }
     }
 }
 
@@ -287,10 +297,21 @@ pub fn walk_join_condition<'a, V: Visitor<'a>>(visitor: &mut V, join_cond: &'a J
 }
 
 pub fn walk_cte<'a, V: Visitor<'a>>(visitor: &mut V, cte: &'a CTE) {
-    let CTE { alias, query, .. } = cte;
+    let CTE { alias, source, .. } = cte;
 
     visitor.visit_identifier(&alias.name);
-    visitor.visit_query(query);
+    match source {
+        CTESource::Query { query, .. } => {
+            visitor.visit_query(query);
+        }
+        CTESource::Values(values) => {
+            for row_values in values {
+                for value in row_values {
+                    visitor.visit_expr(value);
+                }
+            }
+        }
+    }
 }
 
 pub fn walk_window_definition<'a, V: Visitor<'a>>(