Merge pull request #10676 from SkyFan2002/pivot

Author: sundy-li

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

@@ -2193,6 +2193,8 @@ impl<'ast> Visitor<'ast> for AstFormatVisitor {
                 table,
                 alias,
                 travel_point,
+                pivot,
+                unpivot,
             } => {
                 let mut name = String::new();
                 name.push_str("TableIdentifier ");
@@ -2206,6 +2208,16 @@ impl<'ast> Visitor<'ast> for AstFormatVisitor {
                 }
                 name.push_str(&table.to_string());
 
+                if let Some(pivot) = pivot {
+                    name.push(' ');
+                    name.push_str(&pivot.to_string());
+                }
+
+                if let Some(unpivot) = unpivot {
+                    name.push(' ');
+                    name.push_str(&unpivot.to_string());
+                }
+
                 let mut children = Vec::new();
                 if let Some(travel_point) = travel_point {
                     self.visit_time_travel_point(travel_point);

src/query/ast/src/ast/format/syntax/query.rs

@@ -275,6 +275,8 @@ pub(crate) fn pretty_table(table: TableReference) -> RcDoc<'static> {
             table,
             alias,
             travel_point,
+            pivot,
+            unpivot,
         } => if let Some(catalog) = catalog {
             RcDoc::text(catalog.to_string()).append(RcDoc::text("."))
         } else {
@@ -286,6 +288,16 @@ pub(crate) fn pretty_table(table: TableReference) -> RcDoc<'static> {
             RcDoc::nil()
         })
         .append(RcDoc::text(table.to_string()))
+        .append(if let Some(pivot) = pivot {
+            RcDoc::text(format!(" {pivot}"))
+        } else {
+            RcDoc::nil()
+        })
+        .append(if let Some(unpivot) = unpivot {
+            RcDoc::text(format!(" {unpivot}"))
+        } else {
+            RcDoc::nil()
+        })
         .append(if let Some(TimeTravelPoint::Snapshot(sid)) = travel_point {
             RcDoc::text(format!(" AT (SNAPSHOT => {sid})"))
         } else if let Some(TimeTravelPoint::Timestamp(ts)) = travel_point {

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

@@ -145,6 +145,26 @@ pub enum SelectTarget {
     },
 }
 
+impl SelectTarget {
+    pub fn is_star(&self) -> bool {
+        match self {
+            SelectTarget::AliasedExpr { .. } => false,
+            SelectTarget::QualifiedName { qualified, .. } => {
+                matches!(qualified.last(), Some(Indirection::Star(_)))
+            }
+        }
+    }
+
+    pub fn exclude(&mut self, exclude: Vec<Identifier>) {
+        match self {
+            SelectTarget::AliasedExpr { .. } => unreachable!(),
+            SelectTarget::QualifiedName { exclude: e, .. } => {
+                *e = Some(exclude);
+            }
+        }
+    }
+}
+
 pub type QualifiedName = Vec<Indirection>;
 
 /// Indirection of a select result, like a part of `db.table.column`.
@@ -164,6 +184,20 @@ pub enum TimeTravelPoint {
     Timestamp(Box<Expr>),
 }
 
+#[derive(Debug, Clone, PartialEq)]
+pub struct Pivot {
+    pub aggregate: Expr,
+    pub value_column: Identifier,
+    pub values: Vec<Expr>,
+}
+
+#[derive(Debug, Clone, PartialEq)]
+pub struct Unpivot {
+    pub value_column: Identifier,
+    pub column_name: Identifier,
+    pub names: Vec<Identifier>,
+}
+
 /// A table name or a parenthesized subquery with an optional alias
 #[derive(Debug, Clone, PartialEq)]
 pub enum TableReference {
@@ -175,6 +209,8 @@ pub enum TableReference {
         table: Identifier,
         alias: Option<TableAlias>,
         travel_point: Option<TimeTravelPoint>,
+        pivot: Option<Box<Pivot>>,
+        unpivot: Option<Box<Unpivot>>,
     },
     // `TABLE(expr)[ AS alias ]`
     TableFunction {
@@ -202,6 +238,22 @@ pub enum TableReference {
     },
 }
 
+impl TableReference {
+    pub fn pivot(&self) -> Option<&Pivot> {
+        match self {
+            TableReference::Table { pivot, .. } => pivot.as_ref().map(|b| b.as_ref()),
+            _ => None,
+        }
+    }
+
+    pub fn unpivot(&self) -> Option<&Unpivot> {
+        match self {
+            TableReference::Table { unpivot, .. } => unpivot.as_ref().map(|b| b.as_ref()),
+            _ => None,
+        }
+    }
+}
+
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct TableAlias {
     pub name: Identifier,
@@ -282,6 +334,28 @@ impl Display for TableAlias {
     }
 }
 
+impl Display for Pivot {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(f, "PIVOT({} FOR {} IN (", self.aggregate, self.value_column)?;
+        write_comma_separated_list(f, &self.values)?;
+        write!(f, "))")?;
+        Ok(())
+    }
+}
+
+impl Display for Unpivot {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(
+            f,
+            "UNPIVOT({} FOR {} IN (",
+            self.value_column, self.column_name
+        )?;
+        write_comma_separated_list(f, &self.names)?;
+        write!(f, "))")?;
+        Ok(())
+    }
+}
+
 impl Display for TableReference {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
         match self {
@@ -292,6 +366,8 @@ impl Display for TableReference {
                 table,
                 alias,
                 travel_point,
+                pivot,
+                unpivot,
             } => {
                 write_period_separated_list(
                     f,
@@ -309,6 +385,13 @@ impl Display for TableReference {
                 if let Some(alias) = alias {
                     write!(f, " AS {alias}")?;
                 }
+                if let Some(pivot) = pivot {
+                    write!(f, " {pivot}")?;
+                }
+
+                if let Some(unpivot) = unpivot {
+                    write!(f, " {unpivot}")?;
+                }
             }
             TableReference::TableFunction {
                 span: _,
@@ -545,6 +628,7 @@ impl Display for CTE {
         Ok(())
     }
 }
+
 impl Display for With {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
         if self.recursive {

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

@@ -264,6 +264,8 @@ pub enum TableReferenceElement {
         table: Identifier,
         alias: Option<TableAlias>,
         travel_point: Option<TimeTravelPoint>,
+        pivot: Option<Box<Pivot>>,
+        unpivot: Option<Box<Unpivot>>,
     },
     // `TABLE(expr)[ AS alias ]`
     TableFunction {
@@ -292,16 +294,42 @@ pub enum TableReferenceElement {
 }
 
 pub fn table_reference_element(i: Input) -> IResult<WithSpan<TableReferenceElement>> {
+    // PIVOT(expr FOR col IN (ident, ...))
+    let pivot = map(
+        rule! {
+           PIVOT ~ "(" ~ #expr ~ FOR ~ #ident ~ IN ~ "(" ~ #comma_separated_list1(expr) ~ ")" ~ ")"
+        },
+        |(_pivot, _, aggregate, _for, value_column, _in, _, values, _, _)| Pivot {
+            aggregate,
+            value_column,
+            values,
+        },
+    );
+    // UNPIVOT(ident for ident IN (ident, ...))
+    let unpivot = map(
+        rule! {
+            UNPIVOT ~ "(" ~ #ident ~ FOR ~ #ident ~ IN ~ "(" ~ #comma_separated_list1(ident) ~ ")" ~ ")"
+        },
+        |(_unpivot, _, value_column, _for, column_name, _in, _, names, _, _)| Unpivot {
+            value_column,
+            column_name,
+            names,
+        },
+    );
     let aliased_table = map(
         rule! {
-            #period_separated_idents_1_to_3 ~ (AT ~ #travel_point)? ~ #table_alias?
+            #period_separated_idents_1_to_3 ~ (AT ~ #travel_point)? ~ #table_alias? ~ #pivot? ~ #unpivot?
         },
-        |((catalog, database, table), travel_point_opt, alias)| TableReferenceElement::Table {
-            catalog,
-            database,
-            table,
-            alias,
-            travel_point: travel_point_opt.map(|p| p.1),
+        |((catalog, database, table), travel_point_opt, alias, pivot, unpivot)| {
+            TableReferenceElement::Table {
+                catalog,
+                database,
+                table,
+                alias,
+                travel_point: travel_point_opt.map(|p| p.1),
+                pivot: pivot.map(Box::new),
+                unpivot: unpivot.map(Box::new),
+            }
         },
     );
     let table_function = map(
@@ -426,13 +454,17 @@ impl<'a, I: Iterator<Item = WithSpan<'a, TableReferenceElement>>> PrattParser<I>
                 table,
                 alias,
                 travel_point,
+                pivot,
+                unpivot,
             } => TableReference::Table {
                 span: transform_span(input.span.0),
                 catalog,
                 database,
                 table,
                 alias,
                 travel_point,
+                pivot,
+                unpivot,
             },
             TableReferenceElement::TableFunction {
                 name,

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

@@ -1140,7 +1140,7 @@ pub fn statement(i: Input) -> IResult<StatementMsg> {
             | #show_file_formats: "`SHOW FILE FORMATS`"
             | #drop_file_format: "`DROP FILE FORMAT  [ IF EXISTS ] <format_name>`"
         ),
-        rule! (
+        rule!(
             #copy_into: "`COPY
                 INTO { internalStage | externalStage | externalLocation | [<database_name>.]<table_name> }
                 FROM { internalStage | externalStage | externalLocation | [<database_name>.]<table_name> | ( <query> ) }
@@ -1150,7 +1150,7 @@ pub fn statement(i: Input) -> IResult<StatementMsg> {
                 [ VALIDATION_MODE = RETURN_ROWS ]
                 [ copyOptions ]`"
         ),
-        rule! (
+        rule!(
             #call: "`CALL <procedure_name>(<parameter>, ...)`"
         ),
         rule!(
@@ -1889,6 +1889,8 @@ pub fn table_reference_only(i: Input) -> IResult<TableReference> {
             table,
             alias: None,
             travel_point: None,
+            pivot: None,
+            unpivot: None,
         },
     )(i)
 }

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

@@ -681,6 +681,10 @@ pub enum TokenKind {
     SECOND,
     #[token("SELECT", ignore(ascii_case))]
     SELECT,
+    #[token("PIVOT", ignore(ascii_case))]
+    PIVOT,
+    #[token("UNPIVOT", ignore(ascii_case))]
+    UNPIVOT,
     #[token("SEGMENT", ignore(ascii_case))]
     SEGMENT,
     #[token("SET", ignore(ascii_case))]
@@ -1136,6 +1140,8 @@ impl TokenKind {
             // | TokenKind::REFERENCES
             | TokenKind::RIGHT
             | TokenKind::SELECT
+            | TokenKind::PIVOT
+            | TokenKind::UNPIVOT
             // | TokenKind::SESSION_USER
             // | TokenKind::SIMILAR
             | TokenKind::SOME

src/query/ast/tests/it/parser.rs

@@ -485,6 +485,8 @@ fn test_query() {
         r#"select * from t1 union select * from t2 intersect select * from t3"#,
         r#"(select * from t1 union select * from t2) union select * from t3"#,
         r#"select * from t1 union (select * from t2 union select * from t3)"#,
+        r#"select * from monthly_sales pivot(sum(amount) for month in ('JAN', 'FEB', 'MAR', 'APR')) order by empid"#,
+        r#"select * from monthly_sales_1 unpivot(sales for month in (jan, feb, mar, april)) order by empid"#,
     ];
 
     for case in cases {