Support some of pipe operators (apache#1759)

Author: simonvandel

src/ast/mod.rs

@@ -66,23 +66,23 @@ pub use self::dml::{CreateIndex, CreateTable, Delete, IndexColumn, Insert};
 pub use self::operator::{BinaryOperator, UnaryOperator};
 pub use self::query::{
     AfterMatchSkip, ConnectBy, Cte, CteAsMaterialized, Distinct, EmptyMatchesMode,
-    ExceptSelectItem, ExcludeSelectItem, ExprWithAlias, Fetch, ForClause, ForJson, ForXml,
-    FormatClause, GroupByExpr, GroupByWithModifier, IdentWithAlias, IlikeSelectItem,
-    InputFormatClause, Interpolate, InterpolateExpr, Join, JoinConstraint, JoinOperator,
-    JsonTableColumn, JsonTableColumnErrorHandling, JsonTableNamedColumn, JsonTableNestedColumn,
-    LateralView, LimitClause, LockClause, LockType, MatchRecognizePattern, MatchRecognizeSymbol,
-    Measure, NamedWindowDefinition, NamedWindowExpr, NonBlock, Offset, OffsetRows,
-    OpenJsonTableColumn, OrderBy, OrderByExpr, OrderByKind, OrderByOptions, PivotValueSource,
-    ProjectionSelect, Query, RenameSelectItem, RepetitionQuantifier, ReplaceSelectElement,
-    ReplaceSelectItem, RowsPerMatch, Select, SelectFlavor, SelectInto, SelectItem,
-    SelectItemQualifiedWildcardKind, SetExpr, SetOperator, SetQuantifier, Setting,
-    SymbolDefinition, Table, TableAlias, TableAliasColumnDef, TableFactor, TableFunctionArgs,
-    TableIndexHintForClause, TableIndexHintType, TableIndexHints, TableIndexType, TableSample,
-    TableSampleBucket, TableSampleKind, TableSampleMethod, TableSampleModifier,
-    TableSampleQuantity, TableSampleSeed, TableSampleSeedModifier, TableSampleUnit, TableVersion,
-    TableWithJoins, Top, TopQuantity, UpdateTableFromKind, ValueTableMode, Values,
-    WildcardAdditionalOptions, With, WithFill, XmlNamespaceDefinition, XmlPassingArgument,
-    XmlPassingClause, XmlTableColumn, XmlTableColumnOption,
+    ExceptSelectItem, ExcludeSelectItem, ExprWithAlias, ExprWithAliasAndOrderBy, Fetch, ForClause,
+    ForJson, ForXml, FormatClause, GroupByExpr, GroupByWithModifier, IdentWithAlias,
+    IlikeSelectItem, InputFormatClause, Interpolate, InterpolateExpr, Join, JoinConstraint,
+    JoinOperator, JsonTableColumn, JsonTableColumnErrorHandling, JsonTableNamedColumn,
+    JsonTableNestedColumn, LateralView, LimitClause, LockClause, LockType, MatchRecognizePattern,
+    MatchRecognizeSymbol, Measure, NamedWindowDefinition, NamedWindowExpr, NonBlock, Offset,
+    OffsetRows, OpenJsonTableColumn, OrderBy, OrderByExpr, OrderByKind, OrderByOptions,
+    PipeOperator, PivotValueSource, ProjectionSelect, Query, RenameSelectItem,
+    RepetitionQuantifier, ReplaceSelectElement, ReplaceSelectItem, RowsPerMatch, Select,
+    SelectFlavor, SelectInto, SelectItem, SelectItemQualifiedWildcardKind, SetExpr, SetOperator,
+    SetQuantifier, Setting, SymbolDefinition, Table, TableAlias, TableAliasColumnDef, TableFactor,
+    TableFunctionArgs, TableIndexHintForClause, TableIndexHintType, TableIndexHints,
+    TableIndexType, TableSample, TableSampleBucket, TableSampleKind, TableSampleMethod,
+    TableSampleModifier, TableSampleQuantity, TableSampleSeed, TableSampleSeedModifier,
+    TableSampleUnit, TableVersion, TableWithJoins, Top, TopQuantity, UpdateTableFromKind,
+    ValueTableMode, Values, WildcardAdditionalOptions, With, WithFill, XmlNamespaceDefinition,
+    XmlPassingArgument, XmlPassingClause, XmlTableColumn, XmlTableColumnOption,
 };
 
 pub use self::trigger::{

src/ast/query.rs

@@ -62,6 +62,9 @@ pub struct Query {
     /// [ClickHouse](https://clickhouse.com/docs/en/sql-reference/statements/select/format)
     /// (ClickHouse-specific)
     pub format_clause: Option<FormatClause>,
+
+    /// Pipe operator
+    pub pipe_operators: Vec<PipeOperator>,
 }
 
 impl fmt::Display for Query {
@@ -92,6 +95,9 @@ impl fmt::Display for Query {
         if let Some(ref format) = self.format_clause {
             write!(f, " {}", format)?;
         }
+        for pipe_operator in &self.pipe_operators {
+            write!(f, " |> {}", pipe_operator)?;
+        }
         Ok(())
     }
 }
@@ -1004,6 +1010,26 @@ impl fmt::Display for ExprWithAlias {
     }
 }
 
+/// An expression optionally followed by an alias and order by options.
+///
+/// Example:
+/// ```sql
+/// 42 AS myint ASC
+/// ```
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub struct ExprWithAliasAndOrderBy {
+    pub expr: ExprWithAlias,
+    pub order_by: OrderByOptions,
+}
+
+impl fmt::Display for ExprWithAliasAndOrderBy {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(f, "{}{}", self.expr, self.order_by)
+    }
+}
+
 /// Arguments to a table-valued function
 #[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
@@ -2513,6 +2539,135 @@ impl fmt::Display for OffsetRows {
     }
 }
 
+/// Pipe syntax, first introduced in Google BigQuery.
+/// Example:
+///
+/// ```sql
+/// FROM Produce
+/// |> WHERE sales > 0
+/// |> AGGREGATE SUM(sales) AS total_sales, COUNT(*) AS num_sales
+///    GROUP BY item;
+/// ```
+///
+/// See <https://cloud.google.com/bigquery/docs/reference/standard-sql/pipe-syntax#pipe_syntax>
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub enum PipeOperator {
+    /// Limits the number of rows to return in a query, with an optional OFFSET clause to skip over rows.
+    ///
+    /// Syntax: `|> LIMIT <n> [OFFSET <m>]`
+    ///
+    /// See more at <https://cloud.google.com/bigquery/docs/reference/standard-sql/pipe-syntax#limit_pipe_operator>
+    Limit { expr: Expr, offset: Option<Expr> },
+    /// Filters the results of the input table.
+    ///
+    /// Syntax: `|> WHERE <condition>`
+    ///
+    /// See more at <https://cloud.google.com/bigquery/docs/reference/standard-sql/pipe-syntax#where_pipe_operator>
+    Where { expr: Expr },
+    /// `ORDER BY <expr> [ASC|DESC], ...`
+    OrderBy { exprs: Vec<OrderByExpr> },
+    /// Produces a new table with the listed columns, similar to the outermost SELECT clause in a table subquery in standard syntax.
+    ///
+    /// Syntax `|> SELECT <expr> [[AS] alias], ...`
+    ///
+    /// See more at <https://cloud.google.com/bigquery/docs/reference/standard-sql/pipe-syntax#select_pipe_operator>
+    Select { exprs: Vec<SelectItem> },
+    /// Propagates the existing table and adds computed columns, similar to SELECT *, new_column in standard syntax.
+    ///
+    /// Syntax: `|> EXTEND <expr> [[AS] alias], ...`
+    ///
+    /// See more at <https://cloud.google.com/bigquery/docs/reference/standard-sql/pipe-syntax#extend_pipe_operator>
+    Extend { exprs: Vec<SelectItem> },
+    /// Replaces the value of a column in the current table, similar to SELECT * REPLACE (expression AS column) in standard syntax.
+    ///
+    /// Syntax: `|> SET <column> = <expression>, ...`
+    ///
+    /// See more at <https://cloud.google.com/bigquery/docs/reference/standard-sql/pipe-syntax#set_pipe_operator>
+    Set { assignments: Vec<Assignment> },
+    /// Removes listed columns from the current table, similar to SELECT * EXCEPT (column) in standard syntax.
+    ///
+    /// Syntax: `|> DROP <column>, ...`
+    ///
+    /// See more at <https://cloud.google.com/bigquery/docs/reference/standard-sql/pipe-syntax#drop_pipe_operator>
+    Drop { columns: Vec<Ident> },
+    /// Introduces a table alias for the input table, similar to applying the AS alias clause on a table subquery in standard syntax.
+    ///
+    /// Syntax: `|> AS <alias>`
+    ///
+    /// See more at <https://cloud.google.com/bigquery/docs/reference/standard-sql/pipe-syntax#as_pipe_operator>
+    As { alias: Ident },
+    /// Performs aggregation on data across grouped rows or an entire table.
+    ///
+    /// Syntax: `|> AGGREGATE <agg_expr> [[AS] alias], ...`
+    ///
+    /// Syntax:
+    /// ```norust
+    /// |> AGGREGATE [<agg_expr> [[AS] alias], ...]
+    /// GROUP BY <grouping_expr> [AS alias], ...
+    /// ```
+    ///
+    /// See more at <https://cloud.google.com/bigquery/docs/reference/standard-sql/pipe-syntax#aggregate_pipe_operator>
+    Aggregate {
+        full_table_exprs: Vec<ExprWithAliasAndOrderBy>,
+        group_by_expr: Vec<ExprWithAliasAndOrderBy>,
+    },
+}
+
+impl fmt::Display for PipeOperator {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self {
+            PipeOperator::Select { exprs } => {
+                write!(f, "SELECT {}", display_comma_separated(exprs.as_slice()))
+            }
+            PipeOperator::Extend { exprs } => {
+                write!(f, "EXTEND {}", display_comma_separated(exprs.as_slice()))
+            }
+            PipeOperator::Set { assignments } => {
+                write!(f, "SET {}", display_comma_separated(assignments.as_slice()))
+            }
+            PipeOperator::Drop { columns } => {
+                write!(f, "DROP {}", display_comma_separated(columns.as_slice()))
+            }
+            PipeOperator::As { alias } => {
+                write!(f, "AS {}", alias)
+            }
+            PipeOperator::Limit { expr, offset } => {
+                write!(f, "LIMIT {}", expr)?;
+                if let Some(offset) = offset {
+                    write!(f, " OFFSET {}", offset)?;
+                }
+                Ok(())
+            }
+            PipeOperator::Aggregate {
+                full_table_exprs,
+                group_by_expr,
+            } => {
+                write!(f, "AGGREGATE")?;
+                if !full_table_exprs.is_empty() {
+                    write!(
+                        f,
+                        " {}",
+                        display_comma_separated(full_table_exprs.as_slice())
+                    )?;
+                }
+                if !group_by_expr.is_empty() {
+                    write!(f, " GROUP BY {}", display_comma_separated(group_by_expr))?;
+                }
+                Ok(())
+            }
+
+            PipeOperator::Where { expr } => {
+                write!(f, "WHERE {}", expr)
+            }
+            PipeOperator::OrderBy { exprs } => {
+                write!(f, "ORDER BY {}", display_comma_separated(exprs.as_slice()))
+            }
+        }
+    }
+}
+
 #[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
 #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]

src/ast/spans.rs

@@ -98,10 +98,11 @@ impl Spanned for Query {
             order_by,
             limit_clause,
             fetch,
-            locks: _,         // todo
-            for_clause: _,    // todo, mssql specific
-            settings: _,      // todo, clickhouse specific
-            format_clause: _, // todo, clickhouse specific
+            locks: _,          // todo
+            for_clause: _,     // todo, mssql specific
+            settings: _,       // todo, clickhouse specific
+            format_clause: _,  // todo, clickhouse specific
+            pipe_operators: _, // todo bigquery specific
         } = self;
 
         union_spans(

src/dialect/bigquery.rs

@@ -136,6 +136,10 @@ impl Dialect for BigQueryDialect {
     fn is_column_alias(&self, kw: &Keyword, _parser: &mut Parser) -> bool {
         !RESERVED_FOR_COLUMN_ALIAS.contains(kw)
     }
+
+    fn supports_pipe_operator(&self) -> bool {
+        true
+    }
 }
 
 impl BigQueryDialect {

src/dialect/mod.rs

@@ -518,6 +518,20 @@ pub trait Dialect: Debug + Any {
         false
     }
 
+    /// Return true if the dialect supports pipe operator.
+    ///
+    /// Example:
+    /// ```sql
+    /// SELECT *
+    /// FROM table
+    /// |> limit 1
+    /// ```
+    ///
+    /// See <https://cloud.google.com/bigquery/docs/pipe-syntax-guide#basic_syntax>
+    fn supports_pipe_operator(&self) -> bool {
+        false
+    }
+
     /// Does the dialect support MySQL-style `'user'@'host'` grantee syntax?
     fn supports_user_host_grantee(&self) -> bool {
         false

src/keywords.rs

@@ -83,6 +83,7 @@ define_keywords!(
     ADMIN,
     AFTER,
     AGAINST,
+    AGGREGATE,
     AGGREGATION,
     ALERT,
     ALGORITHM,
@@ -338,6 +339,7 @@ define_keywords!(
     EXPLAIN,
     EXPLICIT,
     EXPORT,
+    EXTEND,
     EXTENDED,
     EXTENSION,
     EXTERNAL,