feat: add spill-backed sorting and aggregation (#377)

* feat(execution): add spillable vector

* feat(execution): add spill-backed external sort

* feat(execution): add stream aggregation

* refactor(execution): simplify values and stream aggregation

* feat(optimizer): support forced spill aggregation

* chore: clean up spill docs and clippy warnings

* feat(optimizer): support forced nested-loop joins

* test: cover spill error paths

* refactor: keep external sort heads in cursors

* refactor: simplify spill state handling

Author: KKould

Cargo.toml

@@ -37,9 +37,12 @@ lmdb    = ["dep:lmdb", "dep:lmdb-sys"]
 pprof   = ["pprof/criterion", "pprof/flamegraph"]
 python  = ["parser", "dep:pyo3"]
 shell   = ["parser", "dep:comfy-table", "dep:rustyline"]
+spill   = []
 wasm    = [
     "dep:js-sys",
     "dep:wasm-bindgen",
+    "time",
+    "macros",
     "parser",
 ]
 

Makefile

@@ -21,7 +21,7 @@ COVERAGE_REPORT_ARGS ?= --llvm --ignore-not-existing --keep-only 'src/**' --igno
 
 ## Run default Rust tests in the current environment (non-WASM).
 test:
-	$(CARGO) test --all
+	$(CARGO) test --all --features spill
 
 ## Run Python binding API tests implemented with pyo3.
 test-python:

README.md

@@ -143,6 +143,8 @@ fn main() -> Result<(), DatabaseError> {
 - Cargo features:
   - `rocksdb` is enabled by default
   - `parser` is enabled by default and provides the SQL parser frontend
+  - `spill` optionally enables spill-backed external sorting and aggregation
+  - `spill` and `wasm` are mutually exclusive
   - `lmdb` is optional
   - `unsafe_txdb_checkpoint` enables experimental checkpoint support for RocksDB `TransactionDB`
   - `cargo check --no-default-features --features lmdb` builds an LMDB-only native configuration

benchmarks/query_benchmark.rs

@@ -112,9 +112,15 @@ fn query_on_execute(c: &mut Criterion) {
         let kite_label = format!("KiteSQL: {name} by '{case}'");
         c.bench_function(&kite_label, |b| {
             b.iter(|| {
-                for tuple in database.run(case).unwrap() {
-                    let _ = tuple.unwrap();
-                }
+                let mut tuples = database.run(case).unwrap();
+                while tuples
+                    .next_tuple(|_, tuple| {
+                        std::hint::black_box(tuple);
+                    })
+                    .unwrap()
+                    .is_some()
+                {}
+                tuples.done().unwrap();
             })
         });
 

docs/features.md

@@ -42,6 +42,68 @@ See [the ORM guide](../src/orm/README.md) for the full ORM guide, including:
 - public ORM structs, enums, and traits
 - related `ResultIter::orm::<M>()` integration
 
+### Spill-backed Aggregation: `features = ["spill"]`
+
+The native-only `spill` feature bounds memory usage for large grouped
+aggregations. It is not enabled by default and cannot be combined with `wasm`.
+
+For SQL, place the optimizer hint immediately after `SELECT`:
+
+```sql
+SELECT /*+ FORCE_AGG_SPILL */ user_id, COUNT(*)
+FROM orders
+GROUP BY user_id
+ORDER BY user_id;
+```
+
+For ORM queries, enable both `orm` and `spill`, then call `force_spill()` before
+building the projection and aggregate plan:
+
+```rust,ignore
+let grouped = database.bind(|ctx| {
+    ctx.from::<Order>()?
+        .force_spill()?
+        .project_tuple(|e| {
+            Ok(vec![e.column(Order::user_id())?, e.count_all()?])
+        })?
+        .group_by(|e| e.column(Order::user_id()))?
+        .order_by(Order::user_id())?
+        .finish()
+})?;
+```
+
+KiteSQL still reuses an already ordered input when possible. Otherwise it adds
+an external sort on the group keys and executes a streaming aggregate. The
+hint is intended for grouped aggregates or `DISTINCT`; an aggregate without
+group keys already uses constant-size accumulator state.
+
+### Nested-loop Join Hint
+
+Use `FORCE_NEST_LOOP_JOIN` to select the nested-loop implementation for joins
+in the current `SELECT` query block:
+
+```sql
+SELECT /*+ FORCE_NEST_LOOP_JOIN */ orders.id, users.name
+FROM orders
+JOIN users ON orders.user_id = users.id;
+```
+
+ORM queries can select the same implementation before adding joins:
+
+```rust,ignore
+let joined = database.bind(|ctx| {
+    ctx.from::<Order>()?
+        .force_nested_loop()
+        .inner_join::<User, _>(|e| {
+            e.column(Order::user_id())?.eq(e.column(User::id())?)
+        })?
+        .finish()
+})?;
+```
+
+It avoids Hash Join's build-side tuple materialization, but may perform
+substantially more work on large inputs.
+
 ### User-Defined Function: `features = ["macros"]`
 ```rust
 scala_function!(TestFunction::test(LogicalType::Integer, LogicalType::Integer) -> LogicalType::Integer => |v1: DataValue, v2: DataValue| {

src/binder/aggregate.rs

@@ -55,6 +55,7 @@ impl<T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<'_, '_, T, A>
             agg_calls,
             groupby_exprs,
             false,
+            self.force_spill,
         ))
     }
 

src/binder/distinct.rs

@@ -35,6 +35,7 @@ impl<T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<'_, '_, T, A>
             vec![],
             select_list,
             true,
+            self.force_spill,
         ))
     }
 

src/binder/mod.rs

@@ -561,6 +561,8 @@ impl<'a, T: Transaction> BinderContext<'a, T> {
 pub struct Binder<'a, 'parent, T: Transaction, A: AsRef<[(&'static str, DataValue)]>> {
     pub(crate) context: BinderContext<'a, T>,
     pub(crate) args: &'a A,
+    pub(crate) force_spill: bool,
+    pub(crate) force_nested_loop: bool,
     with_pk: Option<TableName>,
     pub(crate) parent: Option<&'parent BinderContext<'a, T>>,
 }
@@ -574,6 +576,8 @@ impl<'a, 'parent, T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<
         Binder {
             context,
             args,
+            force_spill: false,
+            force_nested_loop: false,
             with_pk: None,
             parent,
         }

src/binder/parser.rs

@@ -1364,6 +1364,7 @@ where
                     self.binder.bind_table_ref_sql(from, self.arena)?,
                     JoinCondition::None,
                     JoinType::Cross,
+                    self.binder.force_nested_loop,
                 )
             }
             plan
@@ -2594,6 +2595,7 @@ impl<'a, 'parent, T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<
         arena: &mut PlanArena,
     ) -> Result<LogicalPlan, DatabaseError> {
         let Select {
+            optimizer_hint,
             projection,
             from,
             selection,
@@ -2615,19 +2617,40 @@ impl<'a, 'parent, T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<
                 "QUALIFY is not supported".to_string(),
             ));
         }
-        Ok(self
-            .build_plan(arena)
-            .from_sql(from)?
-            .select_list_from_sql(projection)?
-            .where_sql(selection.as_ref())?
-            .aggregate_sql(group_by, having.as_ref(), orderby)?
-            .having()?
-            .window()?
-            .distinct_sql(distinct.as_ref())?
-            .order_by()?
-            .project()?
-            .select_into_sql(into.as_ref())?
-            .finish())
+        let has_hint = |expected: &str| {
+            optimizer_hint.as_ref().is_some_and(|hint| {
+                hint.text
+                    .split(|char: char| char.is_ascii_whitespace() || char == ',')
+                    .any(|hint| hint.eq_ignore_ascii_case(expected))
+            })
+        };
+        let force_spill = has_hint("FORCE_AGG_SPILL");
+        let force_nested_loop = has_hint("FORCE_NEST_LOOP_JOIN");
+        if force_spill && !cfg!(feature = "spill") {
+            return Err(DatabaseError::UnsupportedStmt(
+                "FORCE_AGG_SPILL requires the `spill` feature".to_string(),
+            ));
+        }
+        let previous_options = (self.force_spill, self.force_nested_loop);
+        self.force_spill = force_spill;
+        self.force_nested_loop = force_nested_loop;
+        let result = (|| {
+            Ok(self
+                .build_plan(arena)
+                .from_sql(from)?
+                .select_list_from_sql(projection)?
+                .where_sql(selection.as_ref())?
+                .aggregate_sql(group_by, having.as_ref(), orderby)?
+                .having()?
+                .window()?
+                .distinct_sql(distinct.as_ref())?
+                .order_by()?
+                .project()?
+                .select_into_sql(into.as_ref())?
+                .finish())
+        })();
+        (self.force_spill, self.force_nested_loop) = previous_options;
+        result
     }
 
     /// FIXME: temp values need to register BindContext.bind_table
@@ -3137,6 +3160,54 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn force_nest_loop_join_marks_join_operator() -> Result<(), DatabaseError> {
+        let tables = build_t1_table()?;
+        let plan =
+            tables.plan("select /*+ FORCE_NEST_LOOP_JOIN */ c1, c3 from t1 join t2 on c1 = c3")?;
+        let join = plan
+            .childrens
+            .iter()
+            .find_map(|plan| match &plan.operator {
+                Operator::Join(operator) => Some(operator),
+                _ => None,
+            })
+            .expect("query should contain a join");
+
+        assert!(join.force_nested_loop);
+        Ok(())
+    }
+
+    #[cfg(feature = "spill")]
+    #[test]
+    fn optimizer_hints_can_be_combined() -> Result<(), DatabaseError> {
+        let tables = build_t1_table()?;
+        let plan = tables.plan(
+            "select /*+ FORCE_AGG_SPILL, FORCE_NEST_LOOP_JOIN */ c1, count(c3) \
+             from t1 join t2 on c1 = c3 group by c1",
+        )?;
+        let aggregate_plan = plan
+            .childrens
+            .iter()
+            .find(|plan| matches!(plan.operator, Operator::Aggregate(_)))
+            .expect("query should contain an aggregate");
+        let Operator::Aggregate(aggregate) = &aggregate_plan.operator else {
+            unreachable!()
+        };
+        let join = aggregate_plan
+            .childrens
+            .iter()
+            .find_map(|plan| match &plan.operator {
+                Operator::Join(operator) => Some(operator),
+                _ => None,
+            })
+            .expect("aggregate input should contain a join");
+
+        assert!(aggregate.force_spill);
+        assert!(join.force_nested_loop);
+        Ok(())
+    }
+
     #[cfg(feature = "copy")]
     #[test]
     fn test_copy_file_format_options() -> Result<(), DatabaseError> {

src/binder/select.rs

@@ -884,6 +884,7 @@ impl<'a: 'b, 'b, T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<'
     }
 
     fn build_join_from_split_scope_predicates(
+        &self,
         mut children: LogicalPlan,
         mut plan: LogicalPlan,
         join_ty: JoinType,
@@ -931,6 +932,7 @@ impl<'a: 'b, 'b, T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<'
             plan,
             join_condition,
             join_ty,
+            self.force_nested_loop,
         ))
     }
 
@@ -1403,7 +1405,13 @@ impl<'a: 'b, 'b, T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<'
         )?;
         Self::localize_join_condition_from_join_scope(&mut on, left_len)?;
 
-        Ok(LJoinOperator::build(left, right, on, join_type))
+        Ok(LJoinOperator::build(
+            left,
+            right,
+            on,
+            join_type,
+            self.force_nested_loop,
+        ))
     }
 
     pub(crate) fn bind_where_expr(
@@ -1500,7 +1508,7 @@ impl<'a: 'b, 'b, T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<'
                                     .to_string(),
                             ));
                         }
-                        children = Self::build_join_from_split_scope_predicates(
+                        children = self.build_join_from_split_scope_predicates(
                             children,
                             plan,
                             JoinType::Inner,
@@ -1909,10 +1917,7 @@ impl<'a: 'b, 'b, T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<'
         self.context.step(QueryBindStep::Sort);
 
         Ok(LogicalPlan::new(
-            Operator::Sort(SortOperator {
-                sort_fields,
-                limit: None,
-            }),
+            Operator::Sort(SortOperator { sort_fields }),
             Childrens::Only(Box::new(children)),
         ))
     }