Add repeat_row_non_negative as a new SQL function

Author: ggevay

src/expr/src/relation/func.rs

@@ -3289,6 +3289,22 @@ pub fn repeat_row(a: Datum) -> Option<(Row, Diff)> {
     }
 }
 
+pub fn repeat_row_non_negative<'a>(
+    a: Datum,
+) -> Result<Box<dyn Iterator<Item = (Row, Diff)> + 'a>, EvalError> {
+    let n = a.unwrap_int64();
+    if n < 0 {
+        Err(EvalError::InvalidParameterValue(
+            format!("repeat_row_non_negative got {}", n).into(),
+        ))
+    } else if n == 0 {
+        Ok(Box::new(iter::empty()))
+    } else {
+        // iterator with 1 element; n goes into the diff
+        Ok(Box::new(iter::once((Row::default(), n.into()))))
+    }
+}
+
 fn wrap<'a>(datums: &'a [Datum<'a>], width: usize) -> impl Iterator<Item = (Row, Diff)> + 'a {
     datums
         .chunks(width)
@@ -3396,7 +3412,16 @@ pub enum TableFunc {
     GuardSubquerySize {
         column_type: SqlScalarType,
     },
+    /// Repeats the input row the given number of times. Can even repeat a negative number of times,
+    /// which has some important consequences:
+    /// - can lead to negative accumulations downstream;
+    /// - can't be used in `WITH ORDINALITY` and other constructs that are implemented by
+    ///   `TableFunc::WithOrdinality`, e.g., `ROWS FROM`;
+    /// - output is non-monotonic.
     RepeatRow,
+    /// Same as `RepeatRow`, but errors on a negative count, and thereby avoids the above
+    /// peculiarities.
+    RepeatRowNonNegative,
     UnnestArray {
         el_typ: SqlScalarType,
     },
@@ -3471,7 +3496,8 @@ impl TableFunc {
             | TableFunc::GenerateSeriesTimestamp
             | TableFunc::GenerateSeriesTimestampTz
             | TableFunc::GuardSubquerySize { .. }
-            | TableFunc::RepeatRow
+            | TableFunc::RepeatRow // TODO: will move this to the not allowed category in the next commit
+            | TableFunc::RepeatRowNonNegative
             | TableFunc::UnnestArray { .. }
             | TableFunc::UnnestList { .. }
             | TableFunc::UnnestMap { .. }
@@ -3567,6 +3593,7 @@ impl TableFunc {
                 }
             }
             TableFunc::RepeatRow => Ok(Box::new(repeat_row(datums[0]).into_iter())),
+            TableFunc::RepeatRowNonNegative => repeat_row_non_negative(datums[0]),
             TableFunc::UnnestArray { .. } => Ok(Box::new(unnest_array(datums[0]))),
             TableFunc::UnnestList { .. } => Ok(Box::new(unnest_list(datums[0]))),
             TableFunc::UnnestMap { .. } => Ok(Box::new(unnest_map(datums[0]))),
@@ -3682,7 +3709,7 @@ impl TableFunc {
                 let keys = vec![];
                 (column_types, keys)
             }
-            TableFunc::RepeatRow => {
+            TableFunc::RepeatRow | TableFunc::RepeatRowNonNegative => {
                 let column_types = vec![];
                 let keys = vec![];
                 (column_types, keys)
@@ -3764,6 +3791,7 @@ impl TableFunc {
             TableFunc::GenerateSubscriptsArray => 1,
             TableFunc::GuardSubquerySize { .. } => 1,
             TableFunc::RepeatRow => 0,
+            TableFunc::RepeatRowNonNegative => 0,
             TableFunc::UnnestArray { .. } => 1,
             TableFunc::UnnestList { .. } => 1,
             TableFunc::UnnestMap { .. } => 2,
@@ -3791,6 +3819,7 @@ impl TableFunc {
             | TableFunc::RegexpExtract(_)
             | TableFunc::CsvExtract(_)
             | TableFunc::RepeatRow
+            | TableFunc::RepeatRowNonNegative
             | TableFunc::UnnestArray { .. }
             | TableFunc::UnnestList { .. }
             | TableFunc::UnnestMap { .. }
@@ -3822,6 +3851,7 @@ impl TableFunc {
             TableFunc::GenerateSeriesTimestampTz => true,
             TableFunc::GenerateSubscriptsArray => true,
             TableFunc::RepeatRow => false,
+            TableFunc::RepeatRowNonNegative => true,
             TableFunc::UnnestArray { .. } => true,
             TableFunc::UnnestList { .. } => true,
             TableFunc::UnnestMap { .. } => true,
@@ -3853,6 +3883,7 @@ impl fmt::Display for TableFunc {
             TableFunc::GenerateSubscriptsArray => f.write_str("generate_subscripts"),
             TableFunc::GuardSubquerySize { .. } => f.write_str("guard_subquery_size"),
             TableFunc::RepeatRow => f.write_str("repeat_row"),
+            TableFunc::RepeatRowNonNegative => f.write_str("repeat_row_non_negative"),
             TableFunc::UnnestArray { .. } => f.write_str("unnest_array"),
             TableFunc::UnnestList { .. } => f.write_str("unnest_list"),
             TableFunc::UnnestMap { .. } => f.write_str("unnest_map"),

src/pgrepr-consts/src/oid.rs

@@ -792,3 +792,4 @@ pub const VIEW_MZ_MCP_DATA_PRODUCT_DETAILS_OID: u32 = 17071;
 pub const VIEW_MZ_BUILTIN_MATERIALIZED_VIEWS_OID: u32 = 17072;
 pub const FUNC_PARSE_CATALOG_CREATE_SQL_OID: u32 = 17073;
 pub const FUNC_REDACT_SQL_OID: u32 = 17074;
+pub const FUNC_REPEAT_ROW_NON_NEGATIVE_OID: u32 = 17075;

src/sql/src/func.rs

@@ -4745,6 +4745,18 @@ pub static MZ_CATALOG_BUILTINS: LazyLock<BTreeMap<&'static str, Func>> = LazyLoc
                 })
             }) => ReturnType::none(true), oid::FUNC_REPEAT_ROW_OID;
         },
+        "repeat_row_non_negative" => Table {
+            params!(Int64) => Operation::unary(move |ecx, n| {
+                ecx.require_feature_flag(&vars::ENABLE_REPEAT_ROW_NON_NEGATIVE)?;
+                Ok(TableFuncPlan {
+                    imp: TableFuncImpl::CallTable {
+                        func: TableFunc::RepeatRowNonNegative,
+                        exprs: vec![n],
+                    },
+                    column_names: vec![]
+                })
+            }) => ReturnType::none(true), oid::FUNC_REPEAT_ROW_NON_NEGATIVE_OID;
+        },
         "seahash" => Scalar {
             params!(String) => UnaryFunc::SeahashString(func::SeahashString)
                 => UInt32, oid::FUNC_SEAHASH_STRING_OID;

src/sql/src/session/vars/definitions.rs

@@ -1880,6 +1880,12 @@ feature_flags!(
         default: false,
         enable_for_item_parsing: true,
     },
+    {
+        name: enable_repeat_row_non_negative,
+        desc: "the repeat_row_non_negative function",
+        default: false,
+        enable_for_item_parsing: true,
+    },
     {
         name: enable_replica_targeted_materialized_views,
         desc: "replica-targeted materialized views",

test/sqllogictest/table_func.slt

@@ -26,6 +26,11 @@ ALTER SYSTEM SET enable_repeat_row = true
 ----
 COMPLETE 0
 
+simple conn=mz_system,user=mz_system
+ALTER SYSTEM SET enable_repeat_row_non_negative = true
+----
+COMPLETE 0
+
 statement ok
 CREATE TABLE y (a JSONB);
 
@@ -462,6 +467,16 @@ SELECT * FROM generate_series(0,3), repeat_row(generate_series);
 3
 3
 
+query I
+SELECT * FROM generate_series(0,3), repeat_row_non_negative(generate_series);
+----
+1
+2
+2
+3
+3
+3
+
 query II
 SELECT * FROM generate_series(0,3), repeat_row(generate_series) WITH ORDINALITY;
 ----
@@ -472,6 +487,16 @@ SELECT * FROM generate_series(0,3), repeat_row(generate_series) WITH ORDINALITY;
 3  2
 3  3
 
+query II
+SELECT * FROM generate_series(0,3), repeat_row_non_negative(generate_series) WITH ORDINALITY;
+----
+1  1
+2  1
+2  2
+3  1
+3  2
+3  3
+
 query II
 SELECT * FROM generate_series(0,3), repeat_row(abs(generate_series - 1)) WITH ORDINALITY;
 ----
@@ -480,12 +505,34 @@ SELECT * FROM generate_series(0,3), repeat_row(abs(generate_series - 1)) WITH OR
 3  1
 3  2
 
+query II
+SELECT * FROM generate_series(0,3), repeat_row_non_negative(abs(generate_series - 1)) WITH ORDINALITY;
+----
+0  1
+2  1
+3  1
+3  2
+
 query I
 SELECT abs(generate_series) FROM generate_series(-1, 2), repeat_row(generate_series);
 ----
 2
 2
 
+query error repeat_row_non_negative got \-1
+SELECT abs(generate_series) FROM generate_series(-1, 2), repeat_row_non_negative(generate_series);
+
+query error repeat_row_non_negative got \-7
+SELECT 4 FROM repeat_row_non_negative(-7);
+
+query I
+SELECT 5 FROM repeat_row(null);
+----
+
+query I
+SELECT 5 FROM repeat_row_non_negative(null);
+----
+
 statement error  Negative multiplicity in constant result: -1
 SELECT * FROM (values ('a')), repeat_row(-1)
 
@@ -802,6 +849,31 @@ repeat_row
 ()
 (1)
 
+query TI
+SELECT repeat_row(5), generate_series(1,2);
+----
+()  NULL
+()  NULL
+()  NULL
+()  1
+()  2
+
+query TI
+SELECT repeat_row_non_negative(5), generate_series(1,2);
+----
+()  NULL
+()  NULL
+()  NULL
+()  1
+()  2
+
+query T colnames,rowsort
+SELECT repeat_row_non_negative FROM ROWS FROM (repeat_row_non_negative(2), generate_series(1, 1));
+----
+repeat_row_non_negative
+()
+(1)
+
 query T colnames
 SELECT g.a FROM generate_series(1, 1) AS g(a)
 ----