support some of pipe operators

Part of #1758

Still missing
- join
- union|intersect|except
- call
- tablesample
- pivot
- unpivot

Author: simonvandel

src/ast/mod.rs

@@ -68,15 +68,16 @@ pub use self::query::{
     JsonTableColumn, JsonTableColumnErrorHandling, JsonTableNamedColumn, JsonTableNestedColumn,
     LateralView, 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,
+    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,
 };
 
 pub use self::trigger::{

src/ast/query.rs

@@ -68,6 +68,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 {
@@ -103,6 +106,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(())
     }
 }
@@ -2407,6 +2413,98 @@ impl fmt::Display for OffsetRows {
     }
 }
 
+#[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 {
+    Limit {
+        expr: Expr,
+        offset: Option<Expr>,
+    },
+    Where {
+        expr: Expr,
+    },
+    OrderBy {
+        exprs: Vec<OrderByExpr>,
+    },
+    Select {
+        exprs: Vec<SelectItem>,
+    },
+    Extend {
+        exprs: Vec<SelectItem>,
+    },
+    Set {
+        assignments: Vec<Assignment>,
+    },
+    Drop {
+        columns: Vec<Ident>,
+    },
+    Alias {
+        alias: Ident,
+    },
+    Aggregate {
+        full_table_exprs: Vec<ExprWithAlias>,
+        group_by_exprs: Vec<ExprWithAlias>,
+    },
+}
+
+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::Alias { 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_exprs,
+            } => {
+                write!(f, "AGGREGATE")?;
+                if !full_table_exprs.is_empty() {
+                    write!(
+                        f,
+                        " {}",
+                        display_comma_separated(full_table_exprs.as_slice())
+                    )?;
+                }
+                if !group_by_exprs.is_empty() {
+                    write!(
+                        f,
+                        " GROUP BY {}",
+                        display_comma_separated(group_by_exprs.as_slice())
+                    )?;
+                }
+                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 {
             limit_by,
             offset,
             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

@@ -508,6 +508,20 @@ pub trait Dialect: Debug + Any {
         false
     }
 
+    /// Return true if the dialect supports pipe operator.
+    ///
+    /// Example:
+    /// ```sql
+    /// SELECT *
+    /// FROM table
+    /// |> limit 1
+    /// ```
+    ///
+    /// See "SQL Has Problems. We Can Fix Them: Pipe Syntax In SQL" https://research.google/pubs/sql-has-problems-we-can-fix-them-pipe-syntax-in-sql/
+    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,
@@ -335,6 +336,7 @@ define_keywords!(
     EXPLAIN,
     EXPLICIT,
     EXPORT,
+    EXTEND,
     EXTENDED,
     EXTENSION,
     EXTERNAL,

src/parser/mod.rs

@@ -10231,6 +10231,7 @@ impl<'a> Parser<'a> {
                 for_clause: None,
                 settings: None,
                 format_clause: None,
+                pipe_operators: vec![],
             }
             .into())
         } else if self.parse_keyword(Keyword::UPDATE) {
@@ -10246,6 +10247,7 @@ impl<'a> Parser<'a> {
                 for_clause: None,
                 settings: None,
                 format_clause: None,
+                pipe_operators: vec![],
             }
             .into())
         } else {
@@ -10319,6 +10321,104 @@ impl<'a> Parser<'a> {
                 None
             };
 
+            let mut pipe_operators = Vec::new();
+
+            // Syntax from "SQL Has Problems. We Can Fix Them: Pipe Syntax In SQL"
+            // https://storage.googleapis.com/gweb-research2023-media/pubtools/1004848.pdf
+            while self.consume_token(&Token::VerticalBarRightAngleBracket) {
+                let kw = self.expect_one_of_keywords(&[
+                    Keyword::SELECT,
+                    Keyword::EXTEND,
+                    Keyword::SET,
+                    Keyword::DROP,
+                    Keyword::AS,
+                    Keyword::WHERE,
+                    Keyword::LIMIT,
+                    Keyword::AGGREGATE,
+                    Keyword::ORDER,
+                ])?;
+                match kw {
+                    // SELECT <expr> [[AS] alias], ...
+                    Keyword::SELECT => {
+                        let exprs = self.parse_comma_separated(Parser::parse_select_item)?;
+                        pipe_operators.push(PipeOperator::Select { exprs })
+                    }
+                    // EXTEND <expr> [[AS] alias], ...
+                    Keyword::EXTEND => {
+                        let exprs = self.parse_comma_separated(Parser::parse_select_item)?;
+                        pipe_operators.push(PipeOperator::Extend { exprs })
+                    }
+                    // SET <column> = <expression>, ...
+                    Keyword::SET => {
+                        let assignments = self.parse_comma_separated(Parser::parse_assignment)?;
+                        pipe_operators.push(PipeOperator::Set { assignments })
+                    }
+                    // DROP <column>, ...
+                    Keyword::DROP => {
+                        let columns = self.parse_identifiers()?;
+                        pipe_operators.push(PipeOperator::Drop { columns })
+                    }
+                    // AS <alias>
+                    Keyword::AS => {
+                        let alias = self.parse_identifier()?;
+                        pipe_operators.push(PipeOperator::Alias { alias })
+                    }
+                    // WHERE <condition>
+                    Keyword::WHERE => {
+                        let expr = self.parse_expr()?;
+                        pipe_operators.push(PipeOperator::Where { expr })
+                    }
+                    // LIMIT <n> [OFFSET <m>]
+                    Keyword::LIMIT => {
+                        let expr = self.parse_expr()?;
+                        let offset = if self.parse_keyword(Keyword::OFFSET) {
+                            Some(self.parse_expr()?)
+                        } else {
+                            None
+                        };
+                        pipe_operators.push(PipeOperator::Limit { expr, offset })
+                    }
+                    // AGGREGATE <agg_expr> [[AS] alias], ...
+                    //
+                    // and
+                    //
+                    // AGGREGATE [<agg_expr> [[AS] alias], ...]
+                    // GROUP BY <grouping_expr> [AS alias], ...
+                    Keyword::AGGREGATE => {
+                        let full_table_exprs = self.parse_comma_separated0(
+                            |parser| {
+                                let expr = parser.parse_expr()?;
+                                let alias = parser.maybe_parse_select_item_alias()?;
+                                Ok(ExprWithAlias { expr, alias })
+                            },
+                            Token::make_keyword(keywords::GROUP),
+                        )?;
+
+                        let group_by_exprs = if self.parse_keywords(&[Keyword::GROUP, Keyword::BY])
+                        {
+                            self.parse_comma_separated(|parser| {
+                                let expr = parser.parse_expr()?;
+                                let alias = parser.maybe_parse_select_item_alias()?;
+                                Ok(ExprWithAlias { expr, alias })
+                            })?
+                        } else {
+                            vec![]
+                        };
+                        pipe_operators.push(PipeOperator::Aggregate {
+                            full_table_exprs,
+                            group_by_exprs,
+                        })
+                    }
+                    // ORDER BY <expr> [ASC|DESC], ...
+                    Keyword::ORDER => {
+                        self.expect_one_of_keywords(&[Keyword::BY])?;
+                        let exprs = self.parse_comma_separated(Parser::parse_order_by_expr)?;
+                        pipe_operators.push(PipeOperator::OrderBy { exprs })
+                    }
+                    _ => {}
+                }
+            }
+
             Ok(Query {
                 with,
                 body,
@@ -10331,6 +10431,7 @@ impl<'a> Parser<'a> {
                 for_clause,
                 settings,
                 format_clause,
+                pipe_operators,
             }
             .into())
         }
@@ -11688,6 +11789,7 @@ impl<'a> Parser<'a> {
                     for_clause: None,
                     settings: None,
                     format_clause: None,
+                    pipe_operators: vec![],
                 }),
                 alias,
             })

src/tokenizer.rs

@@ -246,6 +246,8 @@ pub enum Token {
     ShiftLeftVerticalBar,
     /// `|>> PostgreSQL/Redshift geometrical binary operator (Is strictly above?)
     VerticalBarShiftRight,
+    /// `|> BigQuery pipe operator
+    VerticalBarRightAngleBracket,
     /// `#>>`, extracts JSON sub-object at the specified path as text
     HashLongArrow,
     /// jsonb @> jsonb -> boolean: Test whether left json contains the right json
@@ -359,6 +361,7 @@ impl fmt::Display for Token {
             Token::AmpersandRightAngleBracket => f.write_str("&>"),
             Token::AmpersandLeftAngleBracketVerticalBar => f.write_str("&<|"),
             Token::VerticalBarAmpersandRightAngleBracket => f.write_str("|&>"),
+            Token::VerticalBarRightAngleBracket => f.write_str("|>"),
             Token::TwoWayArrow => f.write_str("<->"),
             Token::LeftAngleBracketCaret => f.write_str("<^"),
             Token::RightAngleBracketCaret => f.write_str(">^"),
@@ -1378,6 +1381,9 @@ impl<'a> Tokenizer<'a> {
                                 _ => self.start_binop_opt(chars, "|>", None),
                             }
                         }
+                        Some('>') if self.dialect.supports_pipe_operator() => {
+                            self.consume_for_binop(chars, "|>", Token::VerticalBarRightAngleBracket)
+                        }
                         // Bitshift '|' operator
                         _ => self.start_binop(chars, "|", Token::Pipe),
                     }