From 429f9bc5bce265bf0fcdf969cd2a5719b30a6f4b Mon Sep 17 00:00:00 2001
From: Mark Nguyen <marknguyen@retool.com>
Date: Sat, 6 Jun 2026 13:11:30 -0700
Subject: [PATCH] fix(postgres): support select with ordinality

Postgres supports `WITH ORDINALITY`: https://www.postgresql.org/docs/18/sql-select.html

Example:

```
SELECT * FROM unnest(ARRAY['a','b','c','d','e','f']) WITH ORDINALITY;
 unnest | ordinality
--------+----------
 a      |        1
 b      |        2
 c      |        3
 d      |        4
 e      |        5
 f      |        6

```
---
 ast/postgresql.ts      | 8 +++++++-
 pegjs/postgresql.pegjs | 6 +++---
 src/tables.js          | 3 ++-
 test/postgres.spec.js  | 7 +++++++
 4 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/ast/postgresql.ts b/ast/postgresql.ts
index e291a7626..b55568e0a 100644
--- a/ast/postgresql.ts
+++ b/ast/postgresql.ts
@@ -1024,7 +1024,7 @@ export type table_join = table_base & {join: join_op; using: ident_without_kw_ty
       on?: on_clause;
     };
 
-export type table_base = { type: 'dual' } | { expr: value_clause; as?: alias_clause; } | { prefix?: string; expr: union_stmt | value_clause; as?: alias_clause; } | { prefix?: string; expr: table_ref_list; as?: alias_clause; } | { prefix?: string; type: 'expr'; expr: expr; as?: alias_clause; } | table_name & { expr: expr, repeatable: literal_numeric; as?: alias_clause;} | table_name & { as?: alias_clause; };
+export type table_base = { type: 'dual' } | { expr: value_clause; as?: alias_clause; } | { prefix?: string; expr: union_stmt | value_clause; as?: alias_clause; } | { prefix?: string; expr: table_ref_list; as?: alias_clause; } | { prefix?: string; type: 'expr'; expr: expr; ordinality?: boolean; as?: alias_clause; } | table_name & { expr: expr, repeatable: literal_numeric; as?: alias_clause;} | table_name & { as?: alias_clause; };
 
 
 
@@ -1851,6 +1851,8 @@ type KW_TIMESTAMP_TZ = never;
 
 type KW_TIMESTAMP_NTZ = never;
 
+type KW_TIMESTAMP_LTZ = never;
+
 type KW_YEAR = never;
 
 type KW_GEOMETRY = never;
@@ -2013,6 +2015,10 @@ type KW_VAR_PRE_DOLLAR_DOUBLE = never;
 
 type KW_VAR_PRE = never;
 
+type KW_VARIANT = never;
+
+type KW_OBJECT = never;
+
 type KW_ASSIGN = never;
 
 type KW_ASSIGIN_EQUAL = never;
diff --git a/pegjs/postgresql.pegjs b/pegjs/postgresql.pegjs
index 7322b7020..c603f4e44 100644
--- a/pegjs/postgresql.pegjs
+++ b/pegjs/postgresql.pegjs
@@ -3705,9 +3705,9 @@ table_base
       as: alias
     };
   }
-  / l:('LATERAL'i)? __ e:func_call __ alias:alias_clause? {
-    // => { prefix?: string; type: 'expr'; expr: expr; as?: alias_clause; }
-      return { prefix: l, type: 'expr', expr: e, as: alias };
+  / l:('LATERAL'i)? __ e:func_call __ wo:(KW_WITH __ 'ORDINALITY'i)? __ alias:alias_clause? {
+    // => { prefix?: string; type: 'expr'; expr: expr; ordinality?: boolean; as?: alias_clause; }
+      return { prefix: l, type: 'expr', expr: e, ordinality: !!wo, as: alias };
     }
   / t:table_name __ 'TABLESAMPLE'i __ f:func_call __ re:('REPEATABLE'i __ LPAREN __ literal_numeric __ RPAREN)? __ alias:alias_clause? {
     // => table_name & { expr: expr, repeatable: literal_numeric; as?: alias_clause;}
diff --git a/src/tables.js b/src/tables.js
index 7c9e9a1c6..f1d9bc649 100644
--- a/src/tables.js
+++ b/src/tables.js
@@ -114,7 +114,7 @@ function generateVirtualTable(stmt) {
 
 function tableToSQL(tableInfo) {
   if (toUpper(tableInfo.type) === 'UNNEST') return unnestToSQL(tableInfo)
-  const { table, db, as, expr, operator, prefix: prefixStr, schema, server, suffix, tablesample, temporal_table, table_hint, surround = {} } = tableInfo
+  const { table, db, as, expr, operator, ordinality, prefix: prefixStr, schema, server, suffix, tablesample, temporal_table, table_hint, surround = {} } = tableInfo
   const serverName = identifierToSql(server, false, surround.server)
   const database = identifierToSql(db, false, surround.db)
   const schemaStr = identifierToSql(schema, false, surround.schema)
@@ -147,6 +147,7 @@ function tableToSQL(tableInfo) {
     const tableSampleSQL = ['TABLESAMPLE', exprToSQL(tablesample.expr), literalToSQL(tablesample.repeatable)].filter(hasVal).join(' ')
     result.push(tableSampleSQL)
   }
+  if (ordinality) result.push('WITH ORDINALITY')
   result.push(temporalTableToSQL(temporal_table), commonOptionConnector('AS', typeof as === 'string' ? identifierToSql : exprToSQL, as), operatorToSQL(operator))
   if (table_hint) result.push(toUpper(table_hint.keyword), `(${table_hint.expr.map(tableHintToSQL).filter(hasVal).join(', ')})`)
   const tableSQL = result.filter(hasVal).join(' ')
diff --git a/test/postgres.spec.js b/test/postgres.spec.js
index b85da66ae..7f978413b 100644
--- a/test/postgres.spec.js
+++ b/test/postgres.spec.js
@@ -2072,6 +2072,13 @@ describe('Postgres', () => {
         'CREATE TABLE "test_table" (id INT NOT NULL, vector TSVECTOR)'
       ]
     },
+    {
+      title: 'select from set-returning function WITH ORDINALITY',
+      sql: [
+        `SELECT t.v, t.ord FROM unnest(ARRAY[1, 2, 3]) WITH ORDINALITY AS t`,
+        'SELECT "t".v, "t".ord FROM unnest(ARRAY[1,2,3]) WITH ORDINALITY AS "t"'
+      ]
+    },
   ]
   function neatlyNestTestedSQL(sqlList){
     sqlList.forEach(sqlInfo => {