Merge 'Window functions plumbing 1: add stubs for builtins, refactor row_number() to use VDBE aggregate machinery' from Jussi Saurio

1. Introduce stubs for all builtin window functions
2. rename `WindowFunctionKind` to `AccumulatorFunc`, because both window
functions and aggregate functions will use the same `op_agg_step`
mechanism
3. remove specialcased `row_number()` handling and route it through
`op_agg_step`, which is what other window functions will use as well

Reviewed-by: Mikaël Francoeur (@LeMikaelF)

Closes #7358

Author: jussisaurio

core/function.rs

@@ -460,29 +460,146 @@ pub enum AggFunc {
     External(Arc<ExtFunc>),
 }
 
-#[derive(Debug, Clone, Copy, PartialEq, Eq, strum::EnumIter)]
+#[derive(Debug, Clone, strum::EnumIter)]
 pub enum WindowFunc {
     RowNumber,
+    Rank,
+    DenseRank,
+    PercentRank,
+    CumeDist,
+    Ntile,
+    Lag,
+    Lead,
+    FirstValue,
+    LastValue,
+    NthValue,
+    #[strum(disabled)]
+    External(Arc<ExtFunc>),
 }
 
 impl WindowFunc {
+    /// SQL name of this window function. Matches the strings used by
+    /// `Display` so EXPLAIN output and error messages agree.
+    pub fn as_str(&self) -> &'static str {
+        match self {
+            Self::RowNumber => "row_number",
+            Self::Rank => "rank",
+            Self::DenseRank => "dense_rank",
+            Self::PercentRank => "percent_rank",
+            Self::CumeDist => "cume_dist",
+            Self::Ntile => "ntile",
+            Self::Lag => "lag",
+            Self::Lead => "lead",
+            Self::FirstValue => "first_value",
+            Self::LastValue => "last_value",
+            Self::NthValue => "nth_value",
+            Self::External(_) => unreachable!(
+                "WindowFunc::External is not constructible: ExtFunc has no Window variant"
+            ),
+        }
+    }
+
     pub fn arities(&self) -> &'static [i32] {
         match self {
-            Self::RowNumber => &[0],
+            Self::RowNumber | Self::Rank | Self::DenseRank | Self::PercentRank | Self::CumeDist => {
+                &[0]
+            }
+            Self::Ntile | Self::FirstValue | Self::LastValue => &[1],
+            Self::NthValue => &[2],
+            Self::Lag | Self::Lead => &[1, 2, 3],
+            Self::External(_) => unreachable!(
+                "WindowFunc::External is not constructible: ExtFunc has no Window variant"
+            ),
+        }
+    }
+
+    /// Whether name resolution + runtime dispatch are wired up. Stub variants
+    /// must not be advertised via `pragma_function_list`, or introspection
+    /// drifts ahead of the resolver and users get "no such function" when
+    /// they try to call them.
+    pub fn is_implemented(&self) -> bool {
+        matches!(self, Self::RowNumber)
+    }
+}
+
+impl PartialEq for WindowFunc {
+    fn eq(&self, other: &Self) -> bool {
+        match (self, other) {
+            (Self::RowNumber, Self::RowNumber)
+            | (Self::Rank, Self::Rank)
+            | (Self::DenseRank, Self::DenseRank)
+            | (Self::PercentRank, Self::PercentRank)
+            | (Self::CumeDist, Self::CumeDist)
+            | (Self::Ntile, Self::Ntile)
+            | (Self::Lag, Self::Lag)
+            | (Self::Lead, Self::Lead)
+            | (Self::FirstValue, Self::FirstValue)
+            | (Self::LastValue, Self::LastValue)
+            | (Self::NthValue, Self::NthValue) => true,
+            (Self::External(a), Self::External(b)) => Arc::ptr_eq(a, b),
+            _ => false,
         }
     }
 }
 
+impl Eq for WindowFunc {}
+
 impl Deterministic for WindowFunc {
     fn is_deterministic(&self) -> bool {
-        true
+        match self {
+            Self::RowNumber
+            | Self::Rank
+            | Self::DenseRank
+            | Self::PercentRank
+            | Self::CumeDist
+            | Self::Ntile
+            | Self::Lag
+            | Self::Lead
+            | Self::FirstValue
+            | Self::LastValue
+            | Self::NthValue => true,
+            Self::External(_) => unreachable!(
+                "WindowFunc::External is not constructible: ExtFunc has no Window variant"
+            ),
+        }
     }
 }
 
 impl std::fmt::Display for WindowFunc {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.write_str(self.as_str())
+    }
+}
+
+/// Function reference used by AggStep / AggValue / AggFinal opcodes.
+/// Aggregates used in window context and pure window functions share the same
+/// step/value dispatch path; this enum carries which side of that split a
+/// particular call belongs to.
+#[derive(Debug, Clone)]
+pub enum AccumulatorFunc {
+    Agg(AggFunc),
+    Window(WindowFunc),
+}
+
+impl AccumulatorFunc {
+    /// Extract the inner `AggFunc` when this kind is known to be an
+    /// aggregate. `unreachable!`s on `Window(...)` — the only opcodes
+    /// that carry an `AccumulatorFunc` are the AggStep / AggValue /
+    /// AggFinal trio, and the call sites that emit those wrap aggregates
+    /// only. A `Window` value reaching here is a planner bug.
+    pub fn expect_agg(&self) -> &AggFunc {
         match self {
-            Self::RowNumber => write!(f, "row_number"),
+            Self::Agg(f) => f,
+            Self::Window(f) => {
+                unreachable!("window function {f} reached an aggregate-only dispatch path")
+            }
+        }
+    }
+
+    pub fn as_str(&self) -> &'static str {
+        match self {
+            Self::Agg(f) => f.as_str(),
+            Self::Window(f) => f.as_str(),
         }
     }
 }
@@ -1721,8 +1838,11 @@ impl Func {
             push(f.to_string(), "w", f.arities(), f.is_deterministic());
         }
 
-        // Window functions.
+        // Window functions (skip stub variants until they're wired up).
         for f in WindowFunc::iter() {
+            if !f.is_implemented() {
+                continue;
+            }
             push(f.to_string(), "w", f.arities(), f.is_deterministic());
         }
 

core/translate/aggregation.rs

@@ -1,7 +1,7 @@
 use turso_parser::ast;
 
 use crate::{
-    function::AggFunc,
+    function::{AccumulatorFunc, AggFunc},
     schema::Table,
     sync::Arc,
     translate::collate::CollationSeq,
@@ -36,7 +36,7 @@ pub fn emit_ungrouped_aggregation<'a>(
         let agg_result_reg = agg_start_reg + i;
         program.emit_insn(Insn::AggFinal {
             register: agg_result_reg,
-            func: agg.func.clone(),
+            func: AccumulatorFunc::Agg(agg.func.clone()),
         });
     }
     // we now have the agg results in (agg_start_reg..agg_start_reg + aggregates.len() - 1)
@@ -364,7 +364,7 @@ pub fn translate_aggregation_step(
                 acc_reg: target_register,
                 col: expr_reg,
                 delimiter: 0,
-                func: AggFunc::Avg,
+                func: AccumulatorFunc::Agg(AggFunc::Avg),
                 comparator: None,
             });
             target_register
@@ -377,7 +377,7 @@ pub fn translate_aggregation_step(
                 acc_reg: target_register,
                 col: expr_reg,
                 delimiter: 0,
-                func: AggFunc::Count0,
+                func: AccumulatorFunc::Agg(AggFunc::Count0),
                 comparator: None,
             });
             target_register
@@ -392,7 +392,7 @@ pub fn translate_aggregation_step(
                 acc_reg: target_register,
                 col: expr_reg,
                 delimiter: 0,
-                func: AggFunc::Count,
+                func: AccumulatorFunc::Agg(AggFunc::Count),
                 comparator: None,
             });
             target_register
@@ -417,7 +417,7 @@ pub fn translate_aggregation_step(
                 acc_reg: target_register,
                 col: expr_reg,
                 delimiter: delimiter_reg,
-                func: AggFunc::GroupConcat,
+                func: AccumulatorFunc::Agg(AggFunc::GroupConcat),
                 comparator: None,
             });
 
@@ -437,7 +437,7 @@ pub fn translate_aggregation_step(
                 acc_reg: target_register,
                 col: expr_reg,
                 delimiter: 0,
-                func: AggFunc::Max,
+                func: AccumulatorFunc::Agg(AggFunc::Max),
                 comparator,
             });
             target_register
@@ -456,7 +456,7 @@ pub fn translate_aggregation_step(
                 acc_reg: target_register,
                 col: expr_reg,
                 delimiter: 0,
-                func: AggFunc::Min,
+                func: AccumulatorFunc::Agg(AggFunc::Min),
                 comparator,
             });
             target_register
@@ -474,7 +474,7 @@ pub fn translate_aggregation_step(
                 acc_reg: target_register,
                 col: expr_reg,
                 delimiter: value_reg,
-                func: AggFunc::JsonGroupObject,
+                func: AccumulatorFunc::Agg(AggFunc::JsonGroupObject),
                 comparator: None,
             });
             target_register
@@ -490,7 +490,7 @@ pub fn translate_aggregation_step(
                 acc_reg: target_register,
                 col: expr_reg,
                 delimiter: 0,
-                func: AggFunc::JsonGroupArray,
+                func: AccumulatorFunc::Agg(AggFunc::JsonGroupArray),
                 comparator: None,
             });
             target_register
@@ -508,7 +508,7 @@ pub fn translate_aggregation_step(
                 acc_reg: target_register,
                 col: expr_reg,
                 delimiter: delimiter_reg,
-                func: AggFunc::StringAgg,
+                func: AccumulatorFunc::Agg(AggFunc::StringAgg),
                 comparator: None,
             });
 
@@ -524,7 +524,7 @@ pub fn translate_aggregation_step(
                 acc_reg: target_register,
                 col: expr_reg,
                 delimiter: 0,
-                func: AggFunc::Sum,
+                func: AccumulatorFunc::Agg(AggFunc::Sum),
                 comparator: None,
             });
             target_register
@@ -539,7 +539,7 @@ pub fn translate_aggregation_step(
                 acc_reg: target_register,
                 col: expr_reg,
                 delimiter: 0,
-                func: AggFunc::Total,
+                func: AccumulatorFunc::Agg(AggFunc::Total),
                 comparator: None,
             });
             target_register
@@ -555,7 +555,7 @@ pub fn translate_aggregation_step(
                 acc_reg: target_register,
                 col: expr_reg,
                 delimiter: 0,
-                func: AggFunc::ArrayAgg,
+                func: AccumulatorFunc::Agg(AggFunc::ArrayAgg),
                 comparator: None,
             });
             target_register
@@ -573,7 +573,7 @@ pub fn translate_aggregation_step(
                 acc_reg: target_register,
                 col: value_reg,
                 delimiter: 0,
-                func: AggFunc::Mode,
+                func: AccumulatorFunc::Agg(AggFunc::Mode),
                 comparator: None,
             });
             target_register
@@ -595,7 +595,7 @@ pub fn translate_aggregation_step(
                 acc_reg: target_register,
                 col: value_reg,
                 delimiter: fraction_reg,
-                func: func.clone(),
+                func: AccumulatorFunc::Agg(func.clone()),
                 comparator: None,
             });
             target_register
@@ -630,11 +630,11 @@ pub fn translate_aggregation_step(
                 acc_reg: target_register,
                 col: expr_reg,
                 delimiter: 0,
-                func: AggFunc::External(if registered_argc < 0 {
+                func: AccumulatorFunc::Agg(AggFunc::External(if registered_argc < 0 {
                     Arc::new(func.with_aggregate_arg_count(num_args))
                 } else {
                     func.clone()
-                }),
+                })),
                 comparator: None,
             });
             target_register

core/translate/group_by.rs

@@ -6,6 +6,7 @@ use super::{
     plan::{Distinctness, GroupBy, SelectPlan, SubqueryEvalPhase, SubqueryOrigin},
     result_row::emit_select_result,
 };
+use crate::function::AccumulatorFunc;
 use crate::translate::{
     aggregation::{translate_aggregation_step, AggArgumentSource},
     order_by::{custom_type_comparator, EmitOrderBy},
@@ -1014,7 +1015,7 @@ pub fn group_by_emit_row_phase<'a>(
         let agg_result_reg = agg_start_reg + i;
         program.emit_insn(Insn::AggFinal {
             register: agg_result_reg,
-            func: agg.func.clone(),
+            func: AccumulatorFunc::Agg(agg.func.clone()),
         });
         t_ctx.resolver.cache_expr_reg(
             std::borrow::Cow::Owned(agg.original_expr.clone()),

core/translate/main_loop/body.rs

@@ -266,7 +266,7 @@ fn emit_loop_source<'a>(
                     acc_reg: start_reg,
                     col: expr_reg,
                     delimiter: 0,
-                    func: min_max.func.clone(),
+                    func: crate::function::AccumulatorFunc::Agg(min_max.func.clone()),
                     comparator,
                 });
                 program.emit_insn(Insn::Goto {

core/translate/plan.rs

@@ -1,6 +1,6 @@
 use crate::{
     alloc::{self, TursoIteratorExt},
-    function::{AggFunc, WindowFunc},
+    function::{AccumulatorFunc, AggFunc},
     schema::{
         BTreeTable, ColDef, Column, FromClauseSubquery, Index, Schema, Table, Type, ROWID_SENTINEL,
     },
@@ -3090,12 +3090,6 @@ impl Window {
     }
 }
 
-#[derive(Debug, Clone)]
-pub enum WindowFunctionKind {
-    Agg(AggFunc),
-    Window(WindowFunc),
-}
-
 /// One window function call belonging to a `Window`.
 ///
 /// Window queries are planned by wrapping the original FROM/WHERE in a
@@ -3108,7 +3102,7 @@ pub enum WindowFunctionKind {
 pub struct WindowFunction {
     /// The resolved function. Aggregate window functions and specialized window
     /// functions such as ROW_NUMBER() are supported.
-    pub func: WindowFunctionKind,
+    pub func: AccumulatorFunc,
     /// The expression from which the function was resolved. Used as the lookup
     /// key when matching SQL occurrences back to this entry during rewriting.
     pub original_expr: Expr,