PostgreSQL: Add support for LATERAL ... WITH ORDINALITY (#2304)

Author: yoavcloud

src/ast/query.rs

@@ -1522,6 +1522,8 @@ pub enum TableFactor {
         name: ObjectName,
         /// Arguments passed to the function.
         args: Vec<FunctionArg>,
+        /// Whether `WITH ORDINALITY` was specified to include ordinality.
+        with_ordinality: bool,
         /// Optional alias for the result of the function.
         alias: Option<TableAlias>,
     },
@@ -2277,13 +2279,17 @@ impl fmt::Display for TableFactor {
                 lateral,
                 name,
                 args,
+                with_ordinality,
                 alias,
             } => {
                 if *lateral {
                     write!(f, "LATERAL ")?;
                 }
                 write!(f, "{name}")?;
                 write!(f, "({})", display_comma_separated(args))?;
+                if *with_ordinality {
+                    write!(f, " WITH ORDINALITY")?;
+                }
                 if let Some(alias) = alias {
                     write!(f, " {alias}")?;
                 }

src/ast/spans.rs

@@ -1990,6 +1990,7 @@ impl Spanned for TableFactor {
                 lateral: _,
                 name,
                 args,
+                with_ordinality: _,
                 alias,
             } => union_spans(
                 name.0

src/parser/mod.rs

@@ -15898,11 +15898,13 @@ impl<'a> Parser<'a> {
                 let name = self.parse_object_name(false)?;
                 self.expect_token(&Token::LParen)?;
                 let args = self.parse_optional_args()?;
+                let with_ordinality = self.parse_keywords(&[Keyword::WITH, Keyword::ORDINALITY]);
                 let alias = self.maybe_parse_table_alias()?;
                 Ok(TableFactor::Function {
                     lateral: true,
                     name,
                     args,
+                    with_ordinality,
                     alias,
                 })
             }

tests/sqlparser_common.rs

@@ -9028,6 +9028,7 @@ fn lateral_function() {
                             vec![Ident::new("customer"), Ident::new("id")],
                         ))),
                     ],
+                    with_ordinality: false,
                     alias: None,
                 },
                 global: false,

tests/sqlparser_postgres.rs

@@ -6136,6 +6136,34 @@ fn test_table_function_with_ordinality() {
     }
 }
 
+#[test]
+fn test_lateral_function_with_ordinality_and_column_aliases() {
+    let from = pg()
+        .verified_only_select(
+            "SELECT * FROM tbl, \
+             LATERAL json_array_elements(c1::JSON) \
+             WITH ORDINALITY AS t (c1, index)",
+        )
+        .from;
+    assert_eq!(2, from.len());
+    match &from[1].relation {
+        TableFactor::Function {
+            lateral: true,
+            name,
+            with_ordinality: true,
+            alias: Some(alias),
+            ..
+        } => {
+            assert_eq!("json_array_elements", name.to_string().as_str());
+            assert_eq!("t", alias.name.value.as_str());
+            assert_eq!(2, alias.columns.len());
+            assert_eq!("c1", alias.columns[0].name.value.as_str());
+            assert_eq!("index", alias.columns[1].name.value.as_str());
+        }
+        _ => panic!("Expecting TableFactor::Function with ordinality and alias columns"),
+    }
+}
+
 #[test]
 fn test_table_unnest_with_ordinality() {
     let from = pg_and_generic()