feat(frontend): implement WITH ORDINALITY clause

e2e_test/batch/basic/unnest.slt.part

@@ -83,3 +83,38 @@ select distinct unnest(array[1,1,2,3,1]) as x;
 1
 2
 3
+
+query I
+select * from unnest(array[0,1,2]) with ordinality;
+----
+0 1
+1 2
+2 3
+
+query I
+select * from unnest(array[0,1,2]) with ordinality, unnest(array[3,4]) with ordinality as unnest_2;
+----
+0 1 3 1
+0 1 4 2
+1 2 3 1
+1 2 4 2
+2 3 3 1
+2 3 4 2
+
+statement ok
+create table t(arr varchar[]);
+
+statement ok
+insert into t values (Array['a','b', 'c']), (Array['d','e']);
+
+query I rowsort
+select * from t cross join unnest(t.arr) WITH ORDINALITY AS x(elts, num);
+----
+{a,b,c} a 1
+{a,b,c} b 2
+{a,b,c} c 3
+{d,e}   d 1
+{d,e}   e 2
+
+statement ok
+drop table t;

e2e_test/streaming/project_set.slt

@@ -107,3 +107,38 @@ with cte as (SELECT 1 as v1, unnest(array[1,2,3,4,5]) AS v2) select v1 from cte;
 1
 1
 1
+
+statement ok
+create table t(arr varchar[]);
+
+statement ok
+create materialized view mv as select * from t cross join unnest(t.arr) WITH ORDINALITY AS x(elts, num);
+
+statement ok
+insert into t values (Array['a','b', 'c']), (Array['d','e']);
+
+query I rowsort
+select * from mv;
+----
+{a,b,c} a 1
+{a,b,c} b 2
+{a,b,c} c 3
+{d,e}   d 1
+{d,e}   e 2
+
+statement ok
+update t set arr = Array['a', 'c'] where arr = Array['a','b', 'c'];
+
+query I rowsort
+select * from mv;
+----
+{a,c} a 1
+{a,c} c 2
+{d,e} d 1
+{d,e} e 2
+
+statement ok
+drop materialized view mv;
+
+statement ok
+drop table t;

src/frontend/planner_test/tests/testdata/input/with_ordinality.yaml

@@ -0,0 +1,64 @@
+# constant FROM
+- sql: |
+    select * from unnest(array[1,2,3]) WITH ORDINALITY;
+  expected_outputs:
+    - batch_plan
+    - stream_plan
+# lateral join
+- sql: |
+    create table t(x int , arr int[]);
+    select * from t cross join unnest(arr) WITH ORDINALITY;
+  expected_outputs:
+    - batch_plan
+    - stream_plan
+- sql: |
+    create table t(x int , arr int[]);
+    select * from t cross join unnest(arr) WITH ORDINALITY as foo;
+  expected_outputs:
+    - batch_plan
+    - stream_plan
+- sql: |
+    create table t(x int , arr int[]);
+    select * from t cross join unnest(arr) WITH ORDINALITY as foo(a);
+  expected_outputs:
+    - batch_plan
+    - stream_plan
+- sql: |
+    create table t(x int , arr int[]);
+    select * from t cross join unnest(arr) WITH ORDINALITY as foo(a,ord);
+  expected_outputs:
+    - batch_plan
+    - stream_plan
+- name: use alias columns explicitlity
+  sql: |
+    create table t(x int , arr int[]);
+    select x, arr, a, ord from t cross join unnest(arr) WITH ORDINALITY as foo(a,ord);
+  expected_outputs:
+    - batch_plan
+    - stream_plan
+- sql: |
+    create table t(x int , arr int[]);
+    select * from t cross join unnest(arr) WITH ORDINALITY as foo(a,ord,bar);
+  expected_outputs:
+    - binder_error
+# multiple with ordinality
+- sql: |
+    create table t(x int , arr int[]);
+    select * from t cross join unnest(arr) WITH ORDINALITY, unnest(arr) WITH ORDINALITY AS unnest_2(arr_2,ordinality_2);
+  expected_outputs:
+    - batch_plan
+    - stream_plan
+# constant FROM (scalar function)
+- sql: |
+    select * from abs(1) WITH ORDINALITY;
+  expected_outputs:
+    - batch_plan
+    - stream_plan
+# lateral join (scalar function)
+# FIXME: currently this panics due to CorrelatedInputRef in Values https://github.com/risingwavelabs/risingwave/issues/12231
+- sql: |
+    create table t(x int , arr int[]);
+    select * from t, abs(x) WITH ORDINALITY;
+  expected_outputs:
+    - batch_plan
+    - stream_error

src/frontend/src/binder/relation/join.rs

@@ -59,7 +59,7 @@ impl Binder {
 
             let is_lateral = match &right {
                 Relation::Subquery(subquery) if subquery.lateral => true,
-                Relation::TableFunction(_) => true,
+                Relation::TableFunction { .. } => true,
                 _ => false,
             };
 
@@ -110,7 +110,7 @@ impl Binder {
 
             let is_lateral = match &right {
                 Relation::Subquery(subquery) if subquery.lateral => true,
-                Relation::TableFunction(_) => true,
+                Relation::TableFunction { .. } => true,
                 _ => false,
             };
 

src/frontend/src/binder/relation/mod.rs

@@ -55,7 +55,11 @@ pub enum Relation {
     Join(Box<BoundJoin>),
     Apply(Box<BoundJoin>),
     WindowTableFunction(Box<BoundWindowTableFunction>),
-    TableFunction(ExprImpl),
+    /// Table function or scalar function.
+    TableFunction {
+        expr: ExprImpl,
+        with_ordinality: bool,
+    },
     Watermark(Box<BoundWatermark>),
     Share(Box<BoundShare>),
 }
@@ -69,7 +73,9 @@ impl RewriteExprsRecursive for Relation {
             Relation::WindowTableFunction(inner) => inner.rewrite_exprs_recursive(rewriter),
             Relation::Watermark(inner) => inner.rewrite_exprs_recursive(rewriter),
             Relation::Share(inner) => inner.rewrite_exprs_recursive(rewriter),
-            Relation::TableFunction(inner) => *inner = rewriter.rewrite_expr(inner.take()),
+            Relation::TableFunction { expr: inner, .. } => {
+                *inner = rewriter.rewrite_expr(inner.take())
+            }
             _ => {}
         }
     }
@@ -113,7 +119,10 @@ impl Relation {
                 );
                 correlated_indices
             }
-            Relation::TableFunction(table_function) => table_function
+            Relation::TableFunction {
+                expr: table_function,
+                with_ordinality: _,
+            } => table_function
                 .collect_correlated_indices_by_depth_and_assign_id(depth + 1, correlated_id),
             _ => vec![],
         }
@@ -437,9 +446,14 @@ impl Binder {
                 alias,
                 for_system_time_as_of_proctime,
             } => self.bind_relation_by_name(name, alias, for_system_time_as_of_proctime),
-            TableFactor::TableFunction { name, alias, args } => {
+            TableFactor::TableFunction {
+                name,
+                alias,
+                args,
+                with_ordinality,
+            } => {
                 self.try_mark_lateral_as_visible();
-                let result = self.bind_table_function(name, alias, args);
+                let result = self.bind_table_function(name, alias, args, with_ordinality);
                 self.try_mark_lateral_as_invisible();
                 result
             }

src/frontend/src/binder/relation/table_function.rs

@@ -18,6 +18,7 @@ use itertools::Itertools;
 use risingwave_common::catalog::{
     Field, Schema, PG_CATALOG_SCHEMA_NAME, RW_INTERNAL_TABLE_FUNCTION_NAME,
 };
+use risingwave_common::error::ErrorCode;
 use risingwave_common::types::DataType;
 use risingwave_sqlparser::ast::{Function, FunctionArg, ObjectName, TableAlias};
 
@@ -34,23 +35,46 @@ impl Binder {
     ///
     /// Besides [`crate::expr::TableFunction`] expr, it can also be other things like window table
     /// functions, or scalar functions.
+    ///
+    /// `with_ordinality` is only supported for the `TableFunction` case now.
     pub(super) fn bind_table_function(
         &mut self,
         name: ObjectName,
         alias: Option<TableAlias>,
         args: Vec<FunctionArg>,
+        with_ordinality: bool,
     ) -> Result<Relation> {
         let func_name = &name.0[0].real_value();
         // internal/system table functions
         {
             if func_name.eq_ignore_ascii_case(RW_INTERNAL_TABLE_FUNCTION_NAME) {
+                if with_ordinality {
+                    return Err(ErrorCode::NotImplemented(
+                        format!(
+                            "WITH ORDINALITY for internal/system table function {}",
+                            func_name
+                        ),
+                        None.into(),
+                    )
+                    .into());
+                }
                 return self.bind_internal_table(args, alias);
             }
             if func_name.eq_ignore_ascii_case(PG_GET_KEYWORDS_FUNC_NAME)
                 || name.real_value().eq_ignore_ascii_case(
                     format!("{}.{}", PG_CATALOG_SCHEMA_NAME, PG_GET_KEYWORDS_FUNC_NAME).as_str(),
                 )
             {
+                if with_ordinality {
+                    return Err(ErrorCode::NotImplemented(
+                        format!(
+                            "WITH ORDINALITY for internal/system table function {}",
+                            func_name
+                        ),
+                        None.into(),
+                    )
+                    .into());
+                }
                 return self.bind_relation_by_name_inner(
                     Some(PG_CATALOG_SCHEMA_NAME),
                     PG_KEYWORDS_TABLE_NAME,
@@ -61,12 +85,25 @@ impl Binder {
         }
         // window table functions (tumble/hop)
         if let Ok(kind) = WindowTableFunctionKind::from_str(func_name) {
+            if with_ordinality {
+                return Err(ErrorCode::InvalidInputSyntax(format!(
+                    "WITH ORDINALITY for window table function {}",
+                    func_name
+                ))
+                .into());
+            }
             return Ok(Relation::WindowTableFunction(Box::new(
                 self.bind_window_table_function(alias, kind, args)?,
             )));
         }
         // watermark
         if is_watermark_func(func_name) {
+            if with_ordinality {
+                return Err(ErrorCode::InvalidInputSyntax(
+                    "WITH ORDINALITY for watermark".to_string(),
+                )
+                .into());
+            }
             return Ok(Relation::Watermark(Box::new(
                 self.bind_watermark(alias, args)?,
             )));
@@ -88,13 +125,14 @@ impl Binder {
         self.pop_context()?;
         let func = func?;
 
-        let columns = if let DataType::Struct(s) = func.return_type() {
-            // If the table function returns a struct, it's fields can be accessed just
-            // like a table's columns.
+        // bool indicates if the field is hidden
+        let mut columns = if let DataType::Struct(s) = func.return_type() {
+            // If the table function returns a struct, it will be flattened into multiple columns.
             let schema = Schema::from(&s);
             schema.fields.into_iter().map(|f| (false, f)).collect_vec()
         } else {
-            // If there is an table alias, we should use the alias as the table function's
+            // If there is an table alias (and it doesn't return a struct),
+            // we should use the alias as the table function's
             // column name. If column aliases are also provided, they
             // are handled in bind_table_to_context.
             //
@@ -112,9 +150,15 @@ impl Binder {
             };
             vec![(false, Field::with_name(func.return_type(), col_name))]
         };
+        if with_ordinality {
+            columns.push((false, Field::with_name(DataType::Int64, "ordinality")));
+        }
 
         self.bind_table_to_context(columns, func_name.clone(), alias)?;
 
-        Ok(Relation::TableFunction(func))
+        Ok(Relation::TableFunction {
+            expr: func,
+            with_ordinality,
+        })
     }
 }