Add support of table function WITH ORDINALITY modifier for Postgre Parser (apache#1337)

Author: git-hulk

src/ast/query.rs

@@ -913,6 +913,10 @@ pub enum TableFactor {
         /// Optional version qualifier to facilitate table time-travel, as
         /// supported by BigQuery and MSSQL.
         version: Option<TableVersion>,
+        //  Optional table function modifier to generate the ordinality for column.
+        /// For example, `SELECT * FROM generate_series(1, 10) WITH ORDINALITY AS t(a, b);`
+        /// [WITH ORDINALITY](https://www.postgresql.org/docs/current/functions-srf.html), supported by Postgres.
+        with_ordinality: bool,
         /// [Partition selection](https://dev.mysql.com/doc/refman/8.0/en/partitioning-selection.html), supported by MySQL.
         partitions: Vec<Ident>,
     },
@@ -953,6 +957,7 @@ pub enum TableFactor {
         array_exprs: Vec<Expr>,
         with_offset: bool,
         with_offset_alias: Option<Ident>,
+        with_ordinality: bool,
     },
     /// The `JSON_TABLE` table-valued function.
     /// Part of the SQL standard, but implemented only by MySQL, Oracle, and DB2.
@@ -1298,6 +1303,7 @@ impl fmt::Display for TableFactor {
                 with_hints,
                 version,
                 partitions,
+                with_ordinality,
             } => {
                 write!(f, "{name}")?;
                 if !partitions.is_empty() {
@@ -1306,6 +1312,9 @@ impl fmt::Display for TableFactor {
                 if let Some(args) = args {
                     write!(f, "({})", display_comma_separated(args))?;
                 }
+                if *with_ordinality {
+                    write!(f, " WITH ORDINALITY")?;
+                }
                 if let Some(alias) = alias {
                     write!(f, " AS {alias}")?;
                 }
@@ -1366,9 +1375,14 @@ impl fmt::Display for TableFactor {
                 array_exprs,
                 with_offset,
                 with_offset_alias,
+                with_ordinality,
             } => {
                 write!(f, "UNNEST({})", display_comma_separated(array_exprs))?;
 
+                if *with_ordinality {
+                    write!(f, " WITH ORDINALITY")?;
+                }
+
                 if let Some(alias) = alias {
                     write!(f, " AS {alias}")?;
                 }

src/keywords.rs

@@ -518,6 +518,7 @@ define_keywords!(
     OR,
     ORC,
     ORDER,
+    ORDINALITY,
     OUT,
     OUTER,
     OUTPUTFORMAT,

src/parser/mod.rs

@@ -9224,6 +9224,7 @@ impl<'a> Parser<'a> {
             let array_exprs = self.parse_comma_separated(Parser::parse_expr)?;
             self.expect_token(&Token::RParen)?;
 
+            let with_ordinality = self.parse_keywords(&[Keyword::WITH, Keyword::ORDINALITY]);
             let alias = match self.parse_optional_table_alias(keywords::RESERVED_FOR_TABLE_ALIAS) {
                 Ok(Some(alias)) => Some(alias),
                 Ok(None) => None,
@@ -9250,6 +9251,7 @@ impl<'a> Parser<'a> {
                 array_exprs,
                 with_offset,
                 with_offset_alias,
+                with_ordinality,
             })
         } else if self.parse_keyword_with_tokens(Keyword::JSON_TABLE, &[Token::LParen]) {
             let json_expr = self.parse_expr()?;
@@ -9293,6 +9295,8 @@ impl<'a> Parser<'a> {
                 None
             };
 
+            let with_ordinality = self.parse_keywords(&[Keyword::WITH, Keyword::ORDINALITY]);
+
             let alias = self.parse_optional_table_alias(keywords::RESERVED_FOR_TABLE_ALIAS)?;
 
             // MSSQL-specific table hints:
@@ -9314,6 +9318,7 @@ impl<'a> Parser<'a> {
                 with_hints,
                 version,
                 partitions,
+                with_ordinality,
             };
 
             while let Some(kw) = self.parse_one_of_keywords(&[Keyword::PIVOT, Keyword::UNPIVOT]) {

src/test_utils.rs

@@ -309,6 +309,7 @@ pub fn table(name: impl Into<String>) -> TableFactor {
         with_hints: vec![],
         version: None,
         partitions: vec![],
+        with_ordinality: false,
     }
 }
 
@@ -323,6 +324,7 @@ pub fn table_with_alias(name: impl Into<String>, alias: impl Into<String>) -> Ta
         with_hints: vec![],
         version: None,
         partitions: vec![],
+        with_ordinality: false,
     }
 }
 

tests/sqlparser_bigquery.rs

@@ -224,6 +224,7 @@ fn parse_delete_statement() {
                     with_hints: vec![],
                     version: None,
                     partitions: vec![],
+                    with_ordinality: false,
                 },
                 from[0].relation
             );
@@ -1353,6 +1354,7 @@ fn parse_table_identifiers() {
                     with_hints: vec![],
                     version: None,
                     partitions: vec![],
+                    with_ordinality: false,
                 },
                 joins: vec![]
             },]
@@ -1525,6 +1527,7 @@ fn parse_table_time_travel() {
                     Value::SingleQuotedString(version)
                 ))),
                 partitions: vec![],
+                with_ordinality: false,
             },
             joins: vec![]
         },]
@@ -1551,7 +1554,8 @@ fn parse_join_constraint_unnest_alias() {
                     Ident::new("a")
                 ])],
                 with_offset: false,
-                with_offset_alias: None
+                with_offset_alias: None,
+                with_ordinality: false,
             },
             join_operator: JoinOperator::Inner(JoinConstraint::On(Expr::BinaryOp {
                 left: Box::new(Expr::Identifier("c1".into())),
@@ -1620,6 +1624,7 @@ fn parse_merge() {
                     with_hints: Default::default(),
                     version: Default::default(),
                     partitions: Default::default(),
+                    with_ordinality: false,
                 },
                 table
             );
@@ -1634,6 +1639,7 @@ fn parse_merge() {
                     with_hints: Default::default(),
                     version: Default::default(),
                     partitions: Default::default(),
+                    with_ordinality: false,
                 },
                 source
             );

tests/sqlparser_clickhouse.rs

@@ -59,6 +59,7 @@ fn parse_map_access_expr() {
                     with_hints: vec![],
                     version: None,
                     partitions: vec![],
+                    with_ordinality: false,
                 },
                 joins: vec![],
             }],
@@ -162,6 +163,7 @@ fn parse_delimited_identifiers() {
             args,
             with_hints,
             version,
+            with_ordinality: _,
             partitions: _,
         } => {
             assert_eq!(vec![Ident::with_quote('"', "a table")], name.0);