diff --git a/Cargo.toml b/Cargo.toml index eaaa021c..76eba6fa 100644 --- a/Cargo.toml +++ b/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", ] diff --git a/Makefile b/Makefile index 22ef8af8..f8600bb3 100644 --- a/Makefile +++ b/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: diff --git a/README.md b/README.md index 38e78357..a612aa6b 100755 --- a/README.md +++ b/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 diff --git a/benchmarks/query_benchmark.rs b/benchmarks/query_benchmark.rs index a3118747..8d61f727 100644 --- a/benchmarks/query_benchmark.rs +++ b/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(); }) }); diff --git a/docs/features.md b/docs/features.md index b66d5ef1..40e64121 100644 --- a/docs/features.md +++ b/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::()` 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::()? + .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::()? + .force_nested_loop() + .inner_join::(|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| { diff --git a/src/binder/aggregate.rs b/src/binder/aggregate.rs index 8216716e..b723fb44 100644 --- a/src/binder/aggregate.rs +++ b/src/binder/aggregate.rs @@ -55,6 +55,7 @@ impl> Binder<'_, '_, T, A> agg_calls, groupby_exprs, false, + self.force_spill, )) } diff --git a/src/binder/distinct.rs b/src/binder/distinct.rs index 02aa2441..d54286b9 100644 --- a/src/binder/distinct.rs +++ b/src/binder/distinct.rs @@ -35,6 +35,7 @@ impl> Binder<'_, '_, T, A> vec![], select_list, true, + self.force_spill, )) } diff --git a/src/binder/mod.rs b/src/binder/mod.rs index 15efb0fc..3901ba69 100644 --- a/src/binder/mod.rs +++ b/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, 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, } diff --git a/src/binder/parser.rs b/src/binder/parser.rs index 390a56fe..91129eb2 100644 --- a/src/binder/parser.rs +++ b/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 { 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> { diff --git a/src/binder/select.rs b/src/binder/select.rs index eace5f7b..055625ba 100644 --- a/src/binder/select.rs +++ b/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)), )) } diff --git a/src/db.rs b/src/db.rs index ebec405f..ec9e2bad 100644 --- a/src/db.rs +++ b/src/db.rs @@ -1376,7 +1376,6 @@ pub(crate) mod test { while iter .next_tuple(|_, tuple| { println!("{tuple:#?}"); - () })? .is_some() {} diff --git a/src/errors.rs b/src/errors.rs index 2ea780a0..e6ff28ca 100644 --- a/src/errors.rs +++ b/src/errors.rs @@ -774,7 +774,7 @@ mod tests { let bytes = vec![0xff]; let utf8: DatabaseError = std::str::from_utf8(&bytes).unwrap_err().into(); let from_utf8: DatabaseError = String::from_utf8(vec![0xff]).unwrap_err().into(); - let io: DatabaseError = std::io::Error::new(std::io::ErrorKind::Other, "disk").into(); + let io: DatabaseError = std::io::Error::other("disk").into(); let try_from_int: DatabaseError = u8::try_from(300_u16).unwrap_err().into(); for err in [ @@ -799,7 +799,7 @@ mod tests { let bytes = vec![0xff]; let utf8: DatabaseError = std::str::from_utf8(&bytes).unwrap_err().into(); let from_utf8: DatabaseError = String::from_utf8(vec![0xff]).unwrap_err().into(); - let io: DatabaseError = std::io::Error::new(std::io::ErrorKind::Other, "disk").into(); + let io: DatabaseError = std::io::Error::other("disk").into(); let try_from_int: DatabaseError = u8::try_from(300_u16).unwrap_err().into(); assert!(parse_int.to_string().starts_with("parser int:")); diff --git a/src/execution/dql/aggregate/hash_agg.rs b/src/execution/dql/aggregate/hash_agg.rs index fd5b38cf..32622451 100644 --- a/src/execution/dql/aggregate/hash_agg.rs +++ b/src/execution/dql/aggregate/hash_agg.rs @@ -13,7 +13,9 @@ // limitations under the License. use crate::errors::DatabaseError; -use crate::execution::dql::aggregate::{create_accumulators, Accumulator}; +use crate::execution::dql::aggregate::{ + create_accumulators, update_accumulators, write_aggregate_output, Accumulator, +}; use crate::execution::{ build_read, ExecArena, ExecId, ExecNode, ExecutionContext, ExecutorNode, ReadExecutor, }; @@ -21,34 +23,12 @@ use crate::expression::ScalarExpression; use crate::planner::operator::aggregate::AggregateOperator; use crate::planner::LogicalPlan; use crate::storage::Transaction; -use crate::types::tuple::Tuple; use crate::types::value::DataValue; use std::collections::hash_map::IntoIter as HashMapIntoIter; use std::collections::HashMap; type HashAggOutput = HashMapIntoIter, Vec>>; -fn update_accumulators( - accs: &mut [Box], - agg_calls: &[ScalarExpression], - tuple: &Tuple, -) -> Result<(), DatabaseError> { - for (acc, expr) in accs.iter_mut().zip(agg_calls.iter()) { - let ScalarExpression::AggCall { args, .. } = expr else { - unreachable!() - }; - if args.len() > 1 { - return Err(DatabaseError::UnsupportedStmt( - "currently aggregate functions only support a single Column as a parameter" - .to_string(), - )); - } - let value = args[0].eval(Some(tuple))?; - acc.update_value(&value)?; - } - Ok(()) -} - pub struct HashAggExecutor { agg_calls: Vec, groupby_exprs: Vec, @@ -118,17 +98,7 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for HashAggExecutor { return Ok(()); }; - let output = arena.result_tuple_mut(); - - output.pk = None; - output.values.clear(); - output.values.reserve(accs.len() + group_keys.len()); - - for mut acc in accs { - acc.evaluate()?; - output.values.push(acc.result_owned()); - } - output.values.extend(group_keys); + write_aggregate_output(arena.result_tuple_mut(), accs, group_keys)?; arena.resume(); Ok(()) } @@ -214,6 +184,7 @@ mod test { ty: LogicalType::Integer, }], is_distinct: false, + force_spill: false, }), Childrens::Only(Box::new(input)), ); diff --git a/src/execution/dql/aggregate/mod.rs b/src/execution/dql/aggregate/mod.rs index 3e125c1e..803cf65a 100644 --- a/src/execution/dql/aggregate/mod.rs +++ b/src/execution/dql/aggregate/mod.rs @@ -17,6 +17,7 @@ mod count; pub mod hash_agg; mod min_max; pub mod simple_agg; +pub mod stream_agg; pub mod stream_distinct; mod sum; @@ -28,6 +29,7 @@ use crate::execution::dql::aggregate::sum::{DistinctSumAccumulator, SumAccumulat use crate::expression::agg::AggKind; use crate::expression::ScalarExpression; use crate::iter_ext::Itertools; +use crate::types::tuple::Tuple; use crate::types::value::DataValue; use std::borrow::Cow; @@ -81,3 +83,40 @@ pub(crate) fn create_accumulators( }) .try_collect() } + +pub(crate) fn update_accumulators( + accs: &mut [Box], + agg_calls: &[ScalarExpression], + tuple: &Tuple, +) -> Result<(), DatabaseError> { + for (acc, expr) in accs.iter_mut().zip(agg_calls.iter()) { + let ScalarExpression::AggCall { args, .. } = expr else { + unreachable!() + }; + if args.len() > 1 { + return Err(DatabaseError::UnsupportedStmt( + "currently aggregate functions only support a single Column as a parameter" + .to_string(), + )); + } + let value = args[0].eval(Some(tuple))?; + acc.update_value(&value)?; + } + Ok(()) +} + +pub(crate) fn write_aggregate_output( + output: &mut Tuple, + accs: Vec>, + group_keys: Vec, +) -> Result<(), DatabaseError> { + output.pk = None; + output.values.clear(); + output.values.reserve(accs.len() + group_keys.len()); + for mut acc in accs { + acc.evaluate()?; + output.values.push(acc.result_owned()); + } + output.values.extend(group_keys); + Ok(()) +} diff --git a/src/execution/dql/aggregate/stream_agg.rs b/src/execution/dql/aggregate/stream_agg.rs new file mode 100644 index 00000000..5f281de6 --- /dev/null +++ b/src/execution/dql/aggregate/stream_agg.rs @@ -0,0 +1,237 @@ +// Copyright 2024 KipData/KiteSQL +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use crate::errors::DatabaseError; +use crate::execution::dql::aggregate::{ + create_accumulators, update_accumulators, write_aggregate_output, Accumulator, +}; +use crate::execution::{ + build_read, ExecArena, ExecId, ExecNode, ExecutionContext, ExecutorNode, ReadExecutor, +}; +use crate::expression::ScalarExpression; +use crate::planner::operator::aggregate::AggregateOperator; +use crate::planner::LogicalPlan; +use crate::storage::Transaction; +use crate::types::value::DataValue; +use std::mem; + +// The optimizer selects this executor only when equal group keys are contiguous in the input. +pub struct StreamAggExecutor { + agg_calls: Vec, + groupby_exprs: Vec, + group_keys: Option>, + accs: Vec>, + input: ExecId, +} + +impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for StreamAggExecutor { + type Input = (AggregateOperator, LogicalPlan); + + fn into_executor( + ( + AggregateOperator { + agg_calls, + groupby_exprs, + .. + }, + input, + ): Self::Input, + arena: &mut ExecArena<'a, T>, + plan_arena: &mut crate::planner::PlanArena<'a>, + cache: ExecutionContext<'_>, + transaction: &T, + ) -> ExecId { + let input = build_read(arena, plan_arena, input, cache, transaction); + arena.push(ExecNode::StreamAgg(StreamAggExecutor { + agg_calls, + groupby_exprs, + group_keys: None, + accs: Vec::new(), + input, + })) + } +} + +impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for StreamAggExecutor { + fn next_tuple( + &mut self, + arena: &mut ExecArena<'a, T>, + plan_arena: &mut crate::planner::PlanArena<'a>, + ) -> Result<(), DatabaseError> { + loop { + if !arena.next_tuple(self.input, plan_arena)? { + let Some(group_keys) = self.group_keys.take() else { + arena.finish(); + return Ok(()); + }; + write_aggregate_output( + arena.result_tuple_mut(), + mem::take(&mut self.accs), + group_keys, + )?; + arena.resume(); + return Ok(()); + } + + let tuple = arena.result_tuple(); + let mut group_keys = Vec::with_capacity(self.groupby_exprs.len()); + for expr in &self.groupby_exprs { + group_keys.push(expr.eval(Some(tuple))?); + } + + match &mut self.group_keys { + None => { + self.accs = create_accumulators(&self.agg_calls)?; + update_accumulators(&mut self.accs, &self.agg_calls, tuple)?; + self.group_keys = Some(group_keys); + } + Some(current_keys) if current_keys == &group_keys => { + update_accumulators(&mut self.accs, &self.agg_calls, tuple)?; + } + Some(current_keys) => { + let mut next_accs = create_accumulators(&self.agg_calls)?; + update_accumulators(&mut next_accs, &self.agg_calls, tuple)?; + mem::swap(current_keys, &mut group_keys); + let current_accs = mem::replace(&mut self.accs, next_accs); + write_aggregate_output(arena.result_tuple_mut(), current_accs, group_keys)?; + arena.resume(); + return Ok(()); + } + } + } + } +} + +#[cfg(all(test, not(target_arch = "wasm32")))] +mod tests { + use super::StreamAggExecutor; + use crate::catalog::{ColumnCatalog, ColumnDesc}; + use crate::errors::DatabaseError; + use crate::execution::{empty_context, execute_input, try_collect}; + use crate::expression::agg::AggKind; + use crate::expression::ScalarExpression; + use crate::planner::operator::aggregate::AggregateOperator; + use crate::planner::operator::values::ValuesOperator; + use crate::planner::operator::Operator; + use crate::planner::{Childrens, LogicalPlan}; + use crate::storage::memory::MemoryStorage; + use crate::storage::Storage; + use crate::types::value::DataValue; + use crate::types::LogicalType; + + #[test] + fn aggregates_sorted_groups() -> Result<(), DatabaseError> { + let table_arena = crate::planner::TableArenaCell::default(); + let mut plan_arena = crate::planner::PlanArena::new(&table_arena); + let desc = ColumnDesc::new(LogicalType::Integer, None, false, None)?; + let columns = ["group", "value"] + .map(|name| { + plan_arena.alloc_column(ColumnCatalog::new(name.to_string(), true, desc.clone())) + }) + .to_vec(); + let input = LogicalPlan::new( + Operator::Values(ValuesOperator { + rows: vec![ + vec![1.into(), 10.into()], + vec![1.into(), 20.into()], + vec![2.into(), 5.into()], + vec![2.into(), DataValue::Null], + vec![2.into(), 7.into()], + ], + schema_ref: columns.clone(), + }), + Childrens::None, + ); + let value = ScalarExpression::column_expr(columns[1], 1); + let operator = AggregateOperator { + groupby_exprs: vec![ScalarExpression::column_expr(columns[0], 0)], + agg_calls: vec![ + ScalarExpression::AggCall { + distinct: false, + kind: AggKind::Sum, + args: vec![value.clone()], + ty: LogicalType::Integer, + }, + ScalarExpression::AggCall { + distinct: false, + kind: AggKind::Count, + args: vec![value], + ty: LogicalType::Integer, + }, + ], + is_distinct: false, + force_spill: false, + }; + let table_cache = crate::storage::TableCache::default(); + let view_cache = crate::storage::ViewCache::default(); + let meta_cache = crate::storage::StatisticsMetaCache::default(); + let storage = MemoryStorage::new(); + let transaction = storage.transaction()?; + + let rows = try_collect(execute_input::<_, StreamAggExecutor>( + (operator, input), + empty_context(&table_cache, &view_cache, &meta_cache), + plan_arena, + &transaction, + ))?; + + assert_eq!( + rows.into_iter().map(|row| row.values).collect::>(), + vec![ + vec![30.into(), 2.into(), 1.into()], + vec![12.into(), 2.into(), 2.into()], + ] + ); + Ok(()) + } + + #[test] + fn empty_input_returns_no_groups() -> Result<(), DatabaseError> { + let table_arena = crate::planner::TableArenaCell::default(); + let mut plan_arena = crate::planner::PlanArena::new(&table_arena); + let column = plan_arena.alloc_column(ColumnCatalog::new( + "group".to_string(), + true, + ColumnDesc::new(LogicalType::Integer, None, false, None)?, + )); + let input = LogicalPlan::new( + Operator::Values(ValuesOperator { + rows: Vec::new(), + schema_ref: vec![column], + }), + Childrens::None, + ); + let operator = AggregateOperator { + groupby_exprs: vec![ScalarExpression::column_expr(column, 0)], + agg_calls: Vec::new(), + is_distinct: false, + force_spill: false, + }; + let table_cache = crate::storage::TableCache::default(); + let view_cache = crate::storage::ViewCache::default(); + let meta_cache = crate::storage::StatisticsMetaCache::default(); + let storage = MemoryStorage::new(); + let transaction = storage.transaction()?; + + let rows = try_collect(execute_input::<_, StreamAggExecutor>( + (operator, input), + empty_context(&table_cache, &view_cache, &meta_cache), + plan_arena, + &transaction, + ))?; + + assert!(rows.is_empty()); + Ok(()) + } +} diff --git a/src/execution/dql/aggregate/stream_distinct.rs b/src/execution/dql/aggregate/stream_distinct.rs index 8433c356..968f5988 100644 --- a/src/execution/dql/aggregate/stream_distinct.rs +++ b/src/execution/dql/aggregate/stream_distinct.rs @@ -164,6 +164,7 @@ mod tests { groupby_exprs: vec![ScalarExpression::column_expr(schema_ref[0], 0)], agg_calls: vec![], is_distinct: true, + force_spill: false, }; let plan = LogicalPlan::new(Operator::Aggregate(agg), Childrens::Only(Box::new(input))); let plan = optimize_exprs(plan, &mut plan_arena)?; @@ -220,6 +221,7 @@ mod tests { ], agg_calls: vec![], is_distinct: true, + force_spill: false, }; let plan = LogicalPlan::new(Operator::Aggregate(agg), Childrens::Only(Box::new(input))); let plan = optimize_exprs(plan, &mut plan_arena)?; diff --git a/src/execution/dql/external_sort.rs b/src/execution/dql/external_sort.rs new file mode 100644 index 00000000..c79501d0 --- /dev/null +++ b/src/execution/dql/external_sort.rs @@ -0,0 +1,477 @@ +// Copyright 2024 KipData/KiteSQL +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use crate::errors::DatabaseError; +use crate::execution::dql::sort::compare_sort_keys; +use crate::execution::spill::{SegmentOffset, SegmentReader, SortRow, SpillReader, SpillVec}; +use crate::execution::{ + build_read, ExecArena, ExecId, ExecNode, ExecutionContext, ExecutorNode, ReadExecutor, +}; +use crate::planner::operator::sort::{SortField, SortOperator}; +use crate::planner::LogicalPlan; +use crate::storage::Transaction; +use std::fs::File; +use std::io::BufReader; +use std::mem; + +const MERGE_FAN_IN: usize = 16; + +struct Run { + first_segment: SegmentOffset, + segment_count: usize, +} + +impl Run { + fn new(first_segment: SegmentOffset, segment_count: usize) -> Self { + assert!(segment_count > 0, "spill run must contain a segment"); + Self { + first_segment, + segment_count, + } + } +} + +pub struct ExternalSort { + rows: Option>, + sort_fields: Vec, + input: ExecId, +} + +impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for ExternalSort { + type Input = (SortOperator, LogicalPlan); + + fn into_executor( + (SortOperator { sort_fields }, input): Self::Input, + arena: &mut ExecArena<'a, T>, + plan_arena: &mut crate::planner::PlanArena<'a>, + cache: ExecutionContext<'_>, + transaction: &T, + ) -> ExecId { + let input = build_read(arena, plan_arena, input, cache, transaction); + arena.push(ExecNode::ExternalSort(ExternalSort { + rows: None, + sort_fields, + input, + })) + } +} + +impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for ExternalSort { + fn next_tuple( + &mut self, + arena: &mut ExecArena<'a, T>, + plan_arena: &mut crate::planner::PlanArena<'a>, + ) -> Result<(), DatabaseError> { + loop { + if let Some(rows) = &mut self.rows { + let Some(row) = rows.next().transpose()? else { + arena.finish(); + return Ok(()); + }; + arena.produce_tuple(row.tuple); + return Ok(()); + } + + // For R rows and B bytes of estimated SortRow data, sorting one run peaks at roughly + // B + O(R * SortRow) for the stable-sort scratch space. B includes each Tuple and its + // cached sort values, which are evaluated once here and spilled through every merge + // pass. By default R <= 1,024 and B is about 1 MiB; one oversized row makes this a + // soft bound. + let sort_fields = &self.sort_fields; + let mut rows = SpillVec::new().on_flush(move |rows| sort_segment(sort_fields, rows)); + let mut runs = Vec::new(); + while arena.next_tuple(self.input, plan_arena)? { + let tuple = mem::take(arena.result_tuple_mut()); + if let Some(segment) = rows.push(SortRow::new(sort_fields, tuple)?)? { + runs.push(Run::new(segment, 1)); + } + } + self.rows = Some(finish_sort(rows, runs, &self.sort_fields, MERGE_FAN_IN)?); + } + } +} + +#[inline] +fn sort_segment(sort_fields: &[SortField], rows: &mut [SortRow]) -> Result<(), DatabaseError> { + rows.sort_by(|left, right| { + compare_sort_keys( + sort_fields, + left.sort_values.iter(), + right.sort_values.iter(), + ) + }); + Ok(()) +} + +fn finish_sort<'on_flush>( + mut spill: SpillVec<'on_flush, SortRow>, + mut source_runs: Vec, + sort_fields: &[SortField], + fan_in: usize, +) -> Result, DatabaseError> { + assert!(fan_in > 1, "sort merge fan-in must be greater than one"); + if !spill.is_spilled() { + return Ok(spill.into_iter()); + } + + if let Some(segment) = spill.flush()? { + source_runs.push(Run::new(segment, 1)); + } + let mut target_runs = Vec::new(); + while source_runs.len() > 1 { + let source = spill.into_iter(); + spill = merge_pass( + source, + &mut source_runs, + &mut target_runs, + sort_fields, + fan_in, + )?; + mem::swap(&mut source_runs, &mut target_runs); + } + Ok(spill.into_iter()) +} + +struct RunCursor<'source> { + remaining_segments: usize, + reader: SegmentReader<'source, BufReader, SortRow>, + head: Option, +} + +impl<'source> RunCursor<'source> { + fn new(source: &'source SpillReader, run: &Run) -> Result { + let mut reader = source.open_segment_reader()?; + reader.reset(run.first_segment)?; + let mut cursor = Self { + remaining_segments: run.segment_count - 1, + reader, + head: None, + }; + let _ = cursor.next()?; + Ok(cursor) + } + + fn peek(&self) -> Option<&SortRow> { + self.head.as_ref() + } + + fn next(&mut self) -> Result, DatabaseError> { + let head = self.head.take(); + loop { + if let Some(row) = self.reader.next() { + self.head = Some(row?); + return Ok(head); + } + if self.remaining_segments == 0 { + return Ok(head); + } + if !self.reader.start_next_segment()? { + return Err(DatabaseError::InvalidValue( + "spill run ended before its segment count".to_string(), + )); + } + self.remaining_segments -= 1; + } + } +} + +// One merge pass groups source runs by fan-in and produces fewer target runs: +// +// source: [R0][R1] ... [Rk-1] | [Rk][Rk+1] ... +// \ k-way merge / \ k-way merge / +// target: [M0] [M1] +// +// Each cursor contributes one head tuple; the smallest head is appended to the target and then +// replaced from the same cursor. R/M is one logical run stored as contiguous spill segments. The +// target becomes the source of the next pass, until only one globally sorted run remains. +// +// For fan-in K, a merge keeps one output segment, K decoded SortRows, and K BufReader buffers: +// roughly 1 MiB + K * (SortRow + reader buffer). K = 16 therefore uses about 128 KiB for the +// current standard-library reader buffers. Run metadata adds one small entry per initial segment; +// no source segment is materialized in memory. +fn merge_pass<'on_flush>( + source: SpillReader, + source_runs: &mut Vec, + target_runs: &mut Vec, + sort_fields: &[SortField], + fan_in: usize, +) -> Result, DatabaseError> { + let mut target = SpillVec::new(); + target_runs.clear(); + target_runs.reserve(source_runs.len().div_ceil(fan_in)); + // Perf: a loser tree or binary heap would reduce head selection from O(K) to O(log K), but + // K is deliberately small, so a linear scan keeps the merge state and update path simpler. + let mut cursors = Vec::with_capacity(fan_in); + + for run_group in source_runs.chunks(fan_in) { + cursors.clear(); + for run in run_group { + cursors.push(RunCursor::new(&source, run)?); + } + let mut first_segment = None; + let mut segment_count = 0; + + loop { + let mut selected = None; + for (index, cursor) in cursors.iter().enumerate() { + let Some(row) = cursor.peek() else { + continue; + }; + let Some(current) = selected else { + selected = Some(index); + continue; + }; + let Some(current_row) = cursors[current].peek() else { + selected = Some(index); + continue; + }; + if compare_sort_keys( + sort_fields, + row.sort_values.iter(), + current_row.sort_values.iter(), + ) + .is_lt() + { + selected = Some(index); + } + } + let Some(selected) = selected else { + break; + }; + let Some(row) = cursors[selected].next()? else { + return Err(DatabaseError::InvalidValue( + "sort merge selected an empty run".to_string(), + )); + }; + if let Some(segment) = target.push(row)? { + first_segment.get_or_insert(segment); + segment_count += 1; + } + } + + if let Some(segment) = target.flush()? { + first_segment.get_or_insert(segment); + segment_count += 1; + } + let Some(first_segment) = first_segment else { + return Err(DatabaseError::InvalidValue( + "sort merge produced an empty run".to_string(), + )); + }; + target_runs.push(Run::new(first_segment, segment_count)); + } + source_runs.clear(); + + Ok(target) +} + +#[cfg(test)] +mod test { + use super::{finish_sort, sort_segment, Run, MERGE_FAN_IN}; + use crate::catalog::{ColumnCatalog, ColumnDesc}; + use crate::errors::DatabaseError; + use crate::execution::spill::{SortRow, SpillVec}; + use crate::expression::ScalarExpression; + use crate::planner::operator::sort::SortField; + use crate::types::tuple::Tuple; + use crate::types::value::DataValue; + use crate::types::LogicalType; + use std::cmp::Ordering; + + #[test] + fn finish_sort_orders_in_memory_rows() -> Result<(), DatabaseError> { + let table_arena = crate::planner::TableArenaCell::default(); + let mut plan_arena = crate::planner::PlanArena::new(&table_arena); + let sort_column = plan_arena.alloc_column(ColumnCatalog::new( + String::new(), + true, + ColumnDesc::new(LogicalType::Integer, None, false, None).unwrap(), + )); + let sort_fields = vec![SortField { + expr: ScalarExpression::ColumnRef { + column: sort_column, + position: 0, + }, + asc: true, + nulls_first: false, + }]; + + let mut rows = SpillVec::new() + .limit(4, usize::MAX) + .on_flush(|rows| sort_segment(&sort_fields, rows)); + for value in [DataValue::Int32(2), DataValue::Null, DataValue::Int32(1)] { + let _ = rows.push(SortRow::new(&sort_fields, Tuple::new(None, vec![value]))?)?; + } + + let values = finish_sort(rows, Vec::new(), &sort_fields, 2)? + .map(|row| row.map(|row| row.tuple.values[0].clone())) + .collect::, _>>()?; + assert_eq!( + values, + vec![DataValue::Int32(1), DataValue::Int32(2), DataValue::Null,] + ); + Ok(()) + } + + #[test] + fn merges_spilled_runs_and_preserves_ties() -> Result<(), DatabaseError> { + let table_arena = crate::planner::TableArenaCell::default(); + let mut plan_arena = crate::planner::PlanArena::new(&table_arena); + let sort_column = plan_arena.alloc_column(ColumnCatalog::new( + String::new(), + true, + ColumnDesc::new(LogicalType::Integer, None, false, None).unwrap(), + )); + let sort_fields = vec![SortField { + expr: ScalarExpression::ColumnRef { + column: sort_column, + position: 0, + }, + asc: false, + nulls_first: true, + }]; + let values = [ + DataValue::Int32(0), + DataValue::Null, + DataValue::Int32(2), + DataValue::Int32(1), + DataValue::Int32(2), + DataValue::Null, + DataValue::Int32(-1), + DataValue::Int32(3), + DataValue::Int32(0), + ]; + + let mut rows = SpillVec::new() + .limit(2, usize::MAX) + .on_flush(|rows| sort_segment(&sort_fields, rows)); + let mut runs = Vec::new(); + for (sequence, value) in values.into_iter().enumerate() { + let tuple = Tuple::new( + Some(DataValue::Int32(sequence as i32)), + vec![value, DataValue::Int32(sequence as i32)], + ); + if let Some(segment) = rows.push(SortRow::new(&sort_fields, tuple)?)? { + runs.push(Run::new(segment, 1)); + } + } + + let tuples = finish_sort(rows, runs, &sort_fields, 2)? + .map(|row| row.map(|row| row.tuple)) + .collect::, _>>()?; + let positions = tuples + .iter() + .map(|tuple| tuple.values[1].clone()) + .collect::>(); + assert_eq!( + positions, + vec![ + DataValue::Int32(1), + DataValue::Int32(5), + DataValue::Int32(7), + DataValue::Int32(2), + DataValue::Int32(4), + DataValue::Int32(3), + DataValue::Int32(0), + DataValue::Int32(8), + DataValue::Int32(6), + ] + ); + assert_eq!(tuples[0].pk, Some(DataValue::Int32(1))); + assert_eq!(tuples[1].pk, Some(DataValue::Int32(5))); + Ok(()) + } + + #[test] + fn merges_large_dataset_across_multiple_passes() -> Result<(), DatabaseError> { + const ROW_COUNT: usize = 20_000; + + let table_arena = crate::planner::TableArenaCell::default(); + let mut plan_arena = crate::planner::PlanArena::new(&table_arena); + let key_column = plan_arena.alloc_column(ColumnCatalog::new( + String::new(), + true, + ColumnDesc::new(LogicalType::Integer, None, false, None).unwrap(), + )); + let sequence_column = plan_arena.alloc_column(ColumnCatalog::new( + String::new(), + false, + ColumnDesc::new(LogicalType::Integer, None, false, None).unwrap(), + )); + let sort_fields = vec![ + SortField { + expr: ScalarExpression::ColumnRef { + column: key_column, + position: 0, + }, + asc: false, + nulls_first: true, + }, + SortField { + expr: ScalarExpression::ColumnRef { + column: sequence_column, + position: 1, + }, + asc: true, + nulls_first: false, + }, + ]; + + let mut rows = SpillVec::new().on_flush(|rows| sort_segment(&sort_fields, rows)); + let mut runs = Vec::new(); + for position in 0..ROW_COUNT { + let sequence = (position * 7919) % ROW_COUNT; + let key = if sequence % 113 == 0 { + DataValue::Null + } else { + DataValue::Int32((sequence % 257) as i32 - 128) + }; + let sequence = DataValue::Int32(sequence as i32); + let tuple = Tuple::new(Some(sequence.clone()), vec![key, sequence]); + if let Some(segment) = rows.push(SortRow::new(&sort_fields, tuple)?)? { + runs.push(Run::new(segment, 1)); + } + } + assert!(runs.len() > MERGE_FAN_IN); + + let tuples = finish_sort(rows, runs, &sort_fields, MERGE_FAN_IN)? + .map(|row| row.map(|row| row.tuple)) + .collect::, _>>()?; + let mut expected = (0..ROW_COUNT).collect::>(); + expected.sort_by(|left, right| { + match (left % 113 == 0, right % 113 == 0) { + (true, false) => Ordering::Less, + (false, true) => Ordering::Greater, + (false, false) => (right % 257).cmp(&(left % 257)), + (true, true) => Ordering::Equal, + } + .then_with(|| left.cmp(right)) + }); + let actual = tuples + .iter() + .map(|tuple| { + assert_eq!(tuple.pk.as_ref(), tuple.values.get(1)); + tuple.values[1].clone() + }) + .collect::>(); + assert_eq!( + actual, + expected + .into_iter() + .map(|value| DataValue::Int32(value as i32)) + .collect::>() + ); + Ok(()) + } +} diff --git a/src/execution/dql/join/hash_join.rs b/src/execution/dql/join/hash_join.rs index 8ba9a7cd..99bb9d66 100644 --- a/src/execution/dql/join/hash_join.rs +++ b/src/execution/dql/join/hash_join.rs @@ -68,7 +68,7 @@ enum HashJoinState { impl From<(JoinOperator, LogicalPlan, LogicalPlan)> for HashJoin { fn from( - (JoinOperator { on, join_type }, left_input, right_input): ( + (JoinOperator { on, join_type, .. }, left_input, right_input): ( JoinOperator, LogicalPlan, LogicalPlan, @@ -498,6 +498,7 @@ mod test { filter: None, }, join_type: JoinType::Inner, + force_nested_loop: false, }), Childrens::Twins { left: Box::new(left), @@ -555,6 +556,7 @@ mod test { filter: None, }, join_type: JoinType::LeftOuter, + force_nested_loop: false, }), Childrens::Twins { left: Box::new(left), @@ -618,6 +620,7 @@ mod test { filter: None, }, join_type: JoinType::RightOuter, + force_nested_loop: false, }), Childrens::Twins { left: Box::new(left), @@ -716,6 +719,7 @@ mod test { filter: Some(filter_expr), }, join_type: JoinType::RightOuter, + force_nested_loop: false, }), Childrens::Twins { left: Box::new(left), @@ -769,6 +773,7 @@ mod test { filter: None, }, join_type: JoinType::Full, + force_nested_loop: false, }), Childrens::Twins { left: Box::new(left), diff --git a/src/execution/dql/join/nested_loop_join.rs b/src/execution/dql/join/nested_loop_join.rs index 51ad5e0f..541395e8 100644 --- a/src/execution/dql/join/nested_loop_join.rs +++ b/src/execution/dql/join/nested_loop_join.rs @@ -102,7 +102,7 @@ struct ActiveLeftState { impl From<(JoinOperator, LogicalPlan, LogicalPlan)> for NestedLoopJoin { fn from( - (JoinOperator { on, join_type }, left_input, right_input): ( + (JoinOperator { on, join_type, .. }, left_input, right_input): ( JoinOperator, LogicalPlan, LogicalPlan, @@ -633,6 +633,7 @@ mod test { filter: Some(filter), }, join_type: JoinType::Inner, + force_nested_loop: false, }), Childrens::Twins { left: Box::new(left), @@ -685,6 +686,7 @@ mod test { filter: Some(filter), }, join_type: JoinType::LeftOuter, + force_nested_loop: false, }), Childrens::Twins { left: Box::new(left), @@ -766,6 +768,7 @@ mod test { filter: Some(filter), }, join_type: JoinType::Cross, + force_nested_loop: false, }), Childrens::Twins { left: Box::new(left), @@ -818,6 +821,7 @@ mod test { filter: None, }, join_type: JoinType::Cross, + force_nested_loop: false, }), Childrens::Twins { left: Box::new(left), @@ -885,6 +889,7 @@ mod test { filter: None, }, join_type: JoinType::Cross, + force_nested_loop: false, }), Childrens::Twins { left: Box::new(left), @@ -927,6 +932,7 @@ mod test { filter: Some(filter), }, join_type: JoinType::RightOuter, + force_nested_loop: false, }), Childrens::Twins { left: Box::new(left), @@ -1003,6 +1009,7 @@ mod test { filter: Some(filter), }, join_type: JoinType::Full, + force_nested_loop: false, }), Childrens::Twins { left: Box::new(left), @@ -1146,6 +1153,7 @@ mod test { filter: Some(filter_expr), }, join_type: JoinType::RightOuter, + force_nested_loop: false, }), Childrens::Twins { left: Box::new(left), diff --git a/src/execution/dql/mod.rs b/src/execution/dql/mod.rs index fbe27e87..9291ec57 100644 --- a/src/execution/dql/mod.rs +++ b/src/execution/dql/mod.rs @@ -16,6 +16,8 @@ pub(crate) mod aggregate; pub(crate) mod describe; pub(crate) mod dummy; pub(crate) mod explain; +#[cfg(feature = "spill")] +pub(crate) mod external_sort; pub(crate) mod filter; pub(crate) mod function_scan; pub(crate) mod index_scan; diff --git a/src/execution/dql/sort.rs b/src/execution/dql/sort.rs index c40bbe4f..69d24886 100644 --- a/src/execution/dql/sort.rs +++ b/src/execution/dql/sort.rs @@ -20,6 +20,7 @@ use crate::planner::operator::sort::{SortField, SortOperator}; use crate::planner::LogicalPlan; use crate::storage::Transaction; use crate::types::tuple::Tuple; +use crate::types::value::DataValue; use bumpalo::Bump; use std::cmp::Ordering; use std::mem::{self, transmute, MaybeUninit}; @@ -55,11 +56,11 @@ impl<'a, T> NullableVec<'a, T> { pub(crate) fn pop(&mut self) -> Option { self.0.pop().map(|item| unsafe { item.assume_init() }) } +} - pub(crate) fn truncate(&mut self, len: usize) { - while self.len() > len { - self.pop(); - } +impl Drop for NullableVec<'_, T> { + fn drop(&mut self) { + while self.pop().is_some() {} } } @@ -81,13 +82,6 @@ pub(crate) fn sort_tuples( sort_fields: &[SortField], tuples: &mut NullableVec<'_, (usize, Tuple)>, ) -> Result<(), DatabaseError> { - let fn_nulls_first = |nulls_first: bool| { - if nulls_first { - Ordering::Greater - } else { - Ordering::Less - } - }; // Extract the results of calculating SortFields to avoid double calculation // of data during comparison. let mut eval_values = vec![Vec::with_capacity(tuples.len()); sort_fields.len()]; @@ -101,46 +95,56 @@ pub(crate) fn sort_tuples( tuples.0.sort_by(|tuple_1, tuple_2| { let (i_1, _) = unsafe { tuple_1.assume_init_ref() }; let (i_2, _) = unsafe { tuple_2.assume_init_ref() }; - let mut ordering = Ordering::Equal; - - for ( - x, - SortField { - asc, nulls_first, .. - }, - ) in sort_fields.iter().enumerate() - { - let value_1 = &eval_values[x][*i_1]; - let value_2 = &eval_values[x][*i_2]; - - ordering = match (value_1.is_null(), value_2.is_null()) { - (false, true) => fn_nulls_first(*nulls_first), - (true, false) => fn_nulls_first(*nulls_first).reverse(), - _ => { - let mut ordering = value_1.partial_cmp(value_2).unwrap_or(Ordering::Equal); - if !*asc { - ordering = ordering.reverse(); - } - ordering - } - }; - if ordering != Ordering::Equal { - break; - } - } - - ordering + compare_sort_keys( + sort_fields, + eval_values.iter().map(|values| &values[*i_1]), + eval_values.iter().map(|values| &values[*i_2]), + ) }); drop(eval_values); Ok(()) } +pub(crate) fn compare_sort_keys<'a>( + sort_fields: &[SortField], + left: impl Iterator, + right: impl Iterator, +) -> Ordering { + for ( + (value_1, value_2), + SortField { + asc, nulls_first, .. + }, + ) in left.zip(right).zip(sort_fields.iter()) + { + let null_ordering = if *nulls_first { + Ordering::Greater + } else { + Ordering::Less + }; + let ordering = match (value_1.is_null(), value_2.is_null()) { + (false, true) => null_ordering, + (true, false) => null_ordering.reverse(), + _ => { + let mut ordering = value_1.partial_cmp(value_2).unwrap_or(Ordering::Equal); + if !*asc { + ordering = ordering.reverse(); + } + ordering + } + }; + if ordering != Ordering::Equal { + return ordering; + } + } + Ordering::Equal +} + pub struct Sort { rows: NullableVec<'static, (usize, Tuple)>, _arena: Box, sort_fields: Vec, - limit: Option, input: ExecId, } @@ -148,7 +152,7 @@ impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for Sort { type Input = (SortOperator, LogicalPlan); fn into_executor( - (SortOperator { sort_fields, limit }, input): Self::Input, + (SortOperator { sort_fields }, input): Self::Input, arena: &mut ExecArena<'a, T>, plan_arena: &mut crate::planner::PlanArena<'a>, cache: ExecutionContext<'_>, @@ -165,7 +169,6 @@ impl<'a, T: Transaction + 'a> ReadExecutor<'a, T> for Sort { rows, _arena: sort_arena, sort_fields, - limit, input, })) } @@ -191,7 +194,6 @@ impl<'a, T: Transaction + 'a> ExecutorNode<'a, T> for Sort { return Ok(()); } sort_tuples(&self.sort_fields, &mut self.rows)?; - self.rows.truncate(self.limit.unwrap_or(self.rows.len())); self.rows.reverse(); } } @@ -208,16 +210,39 @@ mod test { use crate::types::value::DataValue; use crate::types::LogicalType; use bumpalo::Bump; + use std::cell::Cell; + + #[test] + fn nullable_vec_drops_values() { + struct DropValue<'a>(&'a Cell); + + impl Drop for DropValue<'_> { + fn drop(&mut self) { + self.0.set(self.0.get() + 1); + } + } + + let dropped = Cell::new(0); + let arena = Bump::new(); + { + let mut values = NullableVec::new(&arena); + values.put(DropValue(&dropped)); + values.put(DropValue(&dropped)); + } + assert_eq!(dropped.get(), 2); + } fn sorted_rows<'a>( sort_fields: &[SortField], mut tuples: NullableVec<'a, (usize, Tuple)>, ) -> Result + 'a, DatabaseError> { sort_tuples(sort_fields, &mut tuples)?; - Ok(tuples.0.into_iter().map(|item| { - let (_, tuple) = unsafe { item.assume_init() }; - tuple - })) + let mut rows = Vec::with_capacity(tuples.len()); + while let Some((_, tuple)) = tuples.pop() { + rows.push(tuple); + } + rows.reverse(); + Ok(rows.into_iter()) } #[test] diff --git a/src/execution/mod.rs b/src/execution/mod.rs index 97cadb6c..1e046616 100644 --- a/src/execution/mod.rs +++ b/src/execution/mod.rs @@ -16,6 +16,8 @@ pub(crate) mod ddl; mod ddl_apply; pub(crate) mod dml; pub(crate) mod dql; +#[cfg(feature = "spill")] +pub(crate) mod spill; pub(crate) use ddl_apply::DDLApply; @@ -44,10 +46,13 @@ use crate::execution::dml::insert::Insert; use crate::execution::dml::update::Update; use crate::execution::dql::aggregate::hash_agg::HashAggExecutor; use crate::execution::dql::aggregate::simple_agg::SimpleAggExecutor; +use crate::execution::dql::aggregate::stream_agg::StreamAggExecutor; use crate::execution::dql::aggregate::stream_distinct::StreamDistinctExecutor; use crate::execution::dql::describe::Describe; use crate::execution::dql::dummy::Dummy; use crate::execution::dql::explain::Explain; +#[cfg(feature = "spill")] +use crate::execution::dql::external_sort::ExternalSort; use crate::execution::dql::filter::Filter; use crate::execution::dql::function_scan::FunctionScan; use crate::execution::dql::index_scan::IndexScan; @@ -180,6 +185,8 @@ pub(crate) enum ExecNode<'a, T: Transaction + 'a> { DropView(DropView), Dummy(Dummy), Explain(Explain), + #[cfg(feature = "spill")] + ExternalSort(ExternalSort), Filter(Filter), FunctionScan(FunctionScan), HashAgg(HashAggExecutor), @@ -198,6 +205,7 @@ pub(crate) enum ExecNode<'a, T: Transaction + 'a> { ShowViews(ShowViews<'a, T>), SimpleAgg(SimpleAggExecutor), Sort(Sort), + StreamAgg(StreamAggExecutor), StreamDistinct(StreamDistinctExecutor), TopK(TopK), Truncate(Truncate), @@ -273,6 +281,10 @@ impl<'a, T: Transaction + 'a> ExecNode<'a, T> { ExecNode::Explain(exec) => { >::next_tuple(exec, arena, plan_arena) } + #[cfg(feature = "spill")] + ExecNode::ExternalSort(exec) => { + >::next_tuple(exec, arena, plan_arena) + } ExecNode::Filter(exec) => { >::next_tuple(exec, arena, plan_arena) } @@ -327,6 +339,9 @@ impl<'a, T: Transaction + 'a> ExecNode<'a, T> { ExecNode::Sort(exec) => { >::next_tuple(exec, arena, plan_arena) } + ExecNode::StreamAgg(exec) => { + >::next_tuple(exec, arena, plan_arena) + } ExecNode::StreamDistinct(exec) => { >::next_tuple(exec, arena, plan_arena) } @@ -657,6 +672,20 @@ where cache, transaction, ) + } else if matches!( + physical_option, + Some(PhysicalOption { + plan: PlanImpl::StreamAggregate, + .. + }) + ) { + >::into_executor( + (op, input), + arena, + plan_arena, + cache, + transaction, + ) } else { >::into_executor( (op, input), @@ -777,13 +806,28 @@ where cache, transaction, ), - Operator::Sort(op) => >::into_executor( - (op, childrens.pop_only()), - arena, - plan_arena, - cache, - transaction, - ), + Operator::Sort(op) => { + #[cfg(feature = "spill")] + { + >::into_executor( + (op, childrens.pop_only()), + arena, + plan_arena, + cache, + transaction, + ) + } + #[cfg(not(feature = "spill"))] + { + >::into_executor( + (op, childrens.pop_only()), + arena, + plan_arena, + cache, + transaction, + ) + } + } Operator::Limit(op) => >::into_executor( (op, childrens.pop_only()), arena, diff --git a/src/execution/spill/codec.rs b/src/execution/spill/codec.rs new file mode 100644 index 00000000..233029d6 --- /dev/null +++ b/src/execution/spill/codec.rs @@ -0,0 +1,196 @@ +// Copyright 2024 KipData/KiteSQL +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use super::SpillCodec; +use crate::errors::DatabaseError; +use crate::planner::operator::sort::SortField; +use crate::types::tuple::Tuple; +use crate::types::value::DataValue; +use std::io::{Read, Write}; +use std::mem::size_of; + +pub(crate) struct SortRow { + pub(crate) sort_values: Vec, + pub(crate) tuple: Tuple, +} + +impl SortRow { + pub(crate) fn new(sort_fields: &[SortField], tuple: Tuple) -> Result { + let sort_values = sort_fields + .iter() + .map(|field| field.expr.eval(Some(&tuple))) + .collect::>()?; + Ok(Self { sort_values, tuple }) + } +} + +impl SpillCodec for SortRow { + fn encode(&self, writer: &mut W) -> Result<(), DatabaseError> { + self.sort_values.encode(writer)?; + self.tuple.encode(writer) + } + + fn decode(reader: &mut R) -> Result { + Ok(Self { + sort_values: Vec::::decode(reader)?, + tuple: Tuple::decode(reader)?, + }) + } + + fn estimated_size(&self) -> usize { + size_of::() + .saturating_add( + self.sort_values + .estimated_size() + .saturating_sub(size_of::>()), + ) + .saturating_add( + self.tuple + .estimated_size() + .saturating_sub(size_of::()), + ) + } +} + +impl SpillCodec for DataValue { + fn encode(&self, writer: &mut W) -> Result<(), DatabaseError> { + self.encode_reference_value(writer) + } + + fn decode(reader: &mut R) -> Result { + Self::decode_reference_value(reader) + } + + fn estimated_size(&self) -> usize { + size_of::().saturating_add(estimated_dynamic_value_size(self)) + } +} + +impl SpillCodec for Vec { + fn encode(&self, writer: &mut W) -> Result<(), DatabaseError> { + let len: u32 = self.len().try_into()?; + writer.write_all(&len.to_le_bytes())?; + for value in self { + value.encode(writer)?; + } + Ok(()) + } + + fn decode(reader: &mut R) -> Result { + let mut len = [0; size_of::()]; + reader.read_exact(&mut len)?; + let len = u32::from_le_bytes(len) as usize; + let mut values = Vec::with_capacity(len); + for _ in 0..len { + values.push(T::decode(reader)?); + } + Ok(values) + } + + fn estimated_size(&self) -> usize { + size_of::() + .saturating_add(self.capacity().saturating_mul(size_of::())) + .saturating_add( + self.iter() + .map(|value| value.estimated_size().saturating_sub(size_of::())) + .fold(0usize, usize::saturating_add), + ) + } +} + +impl SpillCodec for Option { + fn encode(&self, writer: &mut W) -> Result<(), DatabaseError> { + match self { + Some(value) => { + writer.write_all(&[1])?; + value.encode(writer) + } + None => { + writer.write_all(&[0])?; + Ok(()) + } + } + } + + fn decode(reader: &mut R) -> Result { + let mut tag = [0]; + reader.read_exact(&mut tag)?; + match tag[0] { + 0 => Ok(None), + 1 => Ok(Some(T::decode(reader)?)), + tag => Err(DatabaseError::InvalidValue(format!( + "invalid spill option tag: {tag}" + ))), + } + } + + fn estimated_size(&self) -> usize { + size_of::().saturating_add( + self.as_ref() + .map(|value| value.estimated_size().saturating_sub(size_of::())) + .unwrap_or_default(), + ) + } +} + +impl SpillCodec for Tuple { + fn encode(&self, writer: &mut W) -> Result<(), DatabaseError> { + self.pk.encode(writer)?; + self.values.encode(writer) + } + + fn decode(reader: &mut R) -> Result { + Ok(Self::new( + Option::::decode(reader)?, + Vec::::decode(reader)?, + )) + } + + fn estimated_size(&self) -> usize { + size_of::() + .saturating_add( + self.values + .capacity() + .saturating_mul(size_of::()), + ) + .saturating_add( + self.values + .iter() + .map(estimated_dynamic_value_size) + .fold(0usize, usize::saturating_add), + ) + .saturating_add( + self.pk + .as_ref() + .map(estimated_dynamic_value_size) + .unwrap_or_default(), + ) + } +} + +fn estimated_dynamic_value_size(value: &DataValue) -> usize { + match value { + DataValue::Utf8 { value, .. } => value.capacity(), + DataValue::Tuple(values, _) => values + .capacity() + .saturating_mul(size_of::()) + .saturating_add( + values + .iter() + .map(estimated_dynamic_value_size) + .fold(0usize, usize::saturating_add), + ), + _ => 0, + } +} diff --git a/src/execution/spill/mod.rs b/src/execution/spill/mod.rs new file mode 100644 index 00000000..e5eb82ad --- /dev/null +++ b/src/execution/spill/mod.rs @@ -0,0 +1,741 @@ +// Copyright 2024 KipData/KiteSQL +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use crate::errors::DatabaseError; +use std::fs::{File, OpenOptions}; +use std::io::{BufReader, Read, Seek, SeekFrom, Write}; +use std::marker::PhantomData; +use std::mem::size_of; +use std::path::PathBuf; +use std::sync::atomic::{AtomicU64, Ordering}; + +mod codec; + +pub(crate) use codec::SortRow; + +const DEFAULT_MAX_ROWS: usize = 1024; +const DEFAULT_MAX_BYTES: usize = 1024 * 1024; +static NEXT_SPILL_FILE_ID: AtomicU64 = AtomicU64::new(0); + +pub(crate) type SegmentOffset = u64; + +pub(crate) trait SpillCodec: Sized { + fn encode(&self, writer: &mut W) -> Result<(), DatabaseError>; + + fn decode(reader: &mut R) -> Result; + + fn estimated_size(&self) -> usize; +} + +pub(crate) struct SpillVec<'on_flush, T: SpillCodec> { + writer: Result, DatabaseError>, +} + +pub(crate) struct SpillReader { + state: ReadState, +} + +struct WriteState<'on_flush, T: SpillCodec> { + buffer: Vec, + buffer_bytes: usize, + file: Option, + on_flush: Option>, + max_rows: usize, + max_bytes: usize, +} + +type OnFlush<'on_flush, T> = Box) -> Result<(), DatabaseError> + 'on_flush>; + +enum ReadState { + Memory(std::vec::IntoIter), + Spilled { + reader: SegmentReader<'static, File, T>, + tail: std::vec::IntoIter, + _file_guard: SpillFileGuard, + }, + Exhausted(Option), +} + +impl<'on_flush, T: SpillCodec> SpillVec<'on_flush, T> { + pub(crate) fn new() -> Self { + Self { + writer: Ok(WriteState { + buffer: Vec::new(), + buffer_bytes: 0, + file: None, + on_flush: None, + max_rows: DEFAULT_MAX_ROWS, + max_bytes: DEFAULT_MAX_BYTES, + }), + } + } + + #[cfg(test)] + pub(crate) fn limit(mut self, max_rows: usize, max_bytes: usize) -> Self { + assert!(max_rows > 0, "spill row limit must be positive"); + assert!(max_bytes > 0, "spill byte limit must be positive"); + if let Ok(state) = &mut self.writer { + state.max_rows = max_rows; + state.max_bytes = max_bytes; + } + self + } + + pub(crate) fn on_flush(mut self, on_flush: F) -> Self + where + F: FnMut(&mut Vec) -> Result<(), DatabaseError> + 'on_flush, + { + if let Ok(state) = &mut self.writer { + state.on_flush = Some(Box::new(on_flush)); + } + self + } + + pub(crate) fn push(&mut self, value: T) -> Result, DatabaseError> { + let state = self.writer.as_mut().map_err(|_| { + DatabaseError::InvalidValue("cannot append to a failed SpillVec".to_string()) + })?; + state.push(value) + } + + pub(crate) fn is_spilled(&self) -> bool { + matches!(&self.writer, Ok(state) if state.file.is_some()) + } + + pub(crate) fn flush(&mut self) -> Result, DatabaseError> { + let state = self.writer.as_mut().map_err(|_| { + DatabaseError::InvalidValue("cannot flush a failed SpillVec".to_string()) + })?; + if state.buffer.is_empty() { + return Ok(None); + } + state.start_spilling()?; + state.flush() + } +} + +impl<'on_flush, T: SpillCodec> From> for SpillVec<'on_flush, T> { + fn from(values: Vec) -> Self { + let mut result = Self::new(); + for value in values { + if let Err(error) = result.push(value) { + result.writer = Err(error); + break; + } + } + result + } +} + +impl IntoIterator for SpillVec<'_, T> { + type Item = Result; + type IntoIter = SpillReader; + + fn into_iter(self) -> Self::IntoIter { + let state = match self.writer { + Ok(writer) => writer + .into_read() + .unwrap_or_else(|error| ReadState::Exhausted(Some(error))), + Err(error) => ReadState::Exhausted(Some(error)), + }; + SpillReader { state } + } +} + +impl Iterator for SpillReader { + type Item = Result; + + fn next(&mut self) -> Option { + let result = match &mut self.state { + ReadState::Memory(rows) => Ok(rows.next()), + ReadState::Spilled { reader, tail, .. } => loop { + match reader.next() { + Some(Ok(value)) => break Ok(Some(value)), + Some(Err(error)) => break Err(error), + None => match reader.start_next_segment() { + Ok(true) => continue, + Ok(false) => break Ok(tail.next()), + Err(error) => break Err(error), + }, + } + }, + ReadState::Exhausted(error) => return error.take().map(Err), + }; + match result { + Ok(Some(value)) => Some(Ok(value)), + Ok(None) => { + self.state = ReadState::Exhausted(None); + None + } + Err(error) => { + self.state = ReadState::Exhausted(None); + Some(Err(error)) + } + } + } +} + +impl SpillReader { + pub(crate) fn open_segment_reader<'source>( + &'source self, + ) -> Result, T>, DatabaseError> { + let ReadState::Spilled { _file_guard, .. } = &self.state else { + return Err(DatabaseError::InvalidValue( + "cannot open a segment reader for an in-memory SpillVec".to_string(), + )); + }; + Ok(SegmentReader::new(BufReader::new(File::open( + &_file_guard.path, + )?))) + } +} + +pub(crate) struct SegmentReader<'source, R, T: SpillCodec> { + reader: R, + remaining_rows: usize, + marker: PhantomData, + source: PhantomData<&'source SpillFileGuard>, +} + +impl SegmentReader<'_, R, T> { + fn new(reader: R) -> Self { + Self { + reader, + remaining_rows: 0, + marker: PhantomData, + source: PhantomData, + } + } + + fn is_exhausted(&self) -> bool { + self.remaining_rows == 0 + } +} + +impl SegmentReader<'_, R, T> { + pub(crate) fn start_next_segment(&mut self) -> Result { + debug_assert!(self.is_exhausted()); + let mut row_count = [0; size_of::()]; + if self.reader.read(&mut row_count[..1])? == 0 { + return Ok(false); + } + self.reader.read_exact(&mut row_count[1..])?; + self.remaining_rows = u64::from_le_bytes(row_count).try_into()?; + if self.remaining_rows == 0 { + return Err(DatabaseError::InvalidValue( + "spill segment cannot be empty".to_string(), + )); + } + Ok(true) + } +} + +impl SegmentReader<'_, R, T> { + pub(crate) fn reset(&mut self, offset: SegmentOffset) -> Result<(), DatabaseError> { + self.reader.seek(SeekFrom::Start(offset))?; + self.remaining_rows = 0; + if !self.start_next_segment()? { + return Err(DatabaseError::InvalidValue( + "spill segment offset points to end of file".to_string(), + )); + } + Ok(()) + } +} + +impl Iterator for SegmentReader<'_, R, T> { + type Item = Result; + + fn next(&mut self) -> Option { + if self.remaining_rows == 0 { + return None; + } + match T::decode(&mut self.reader) { + Ok(value) => { + self.remaining_rows -= 1; + Some(Ok(value)) + } + Err(error) => { + self.remaining_rows = 0; + Some(Err(error)) + } + } + } +} + +impl WriteState<'_, T> { + fn push(&mut self, value: T) -> Result, DatabaseError> { + let value_size = value.estimated_size(); + self.buffer.push(value); + self.buffer_bytes = self.buffer_bytes.saturating_add(value_size); + + if self.buffer.len() >= self.max_rows || self.buffer_bytes >= self.max_bytes { + self.start_spilling()?; + return self.flush(); + } + Ok(None) + } + + fn start_spilling(&mut self) -> Result<(), DatabaseError> { + if self.file.is_none() { + self.file = Some(SpillFileWriter::new()?); + } + Ok(()) + } + + fn flush(&mut self) -> Result, DatabaseError> { + if self.buffer.is_empty() { + return Ok(None); + } + if let Some(on_flush) = &mut self.on_flush { + on_flush(&mut self.buffer)?; + } + let Some(file) = self.file.as_mut() else { + return Err(DatabaseError::InvalidValue( + "cannot flush without a spill file".to_string(), + )); + }; + let segment = file.append_segment(&self.buffer)?; + self.buffer.clear(); + self.buffer_bytes = 0; + Ok(Some(segment)) + } + + fn into_read(mut self) -> Result, DatabaseError> { + if let Some(on_flush) = &mut self.on_flush { + on_flush(&mut self.buffer)?; + } + let Some(file) = self.file.take() else { + return Ok(ReadState::Memory(self.buffer.into_iter())); + }; + file.into_reader(self.buffer.into_iter()) + } +} + +struct SpillFileGuard { + path: PathBuf, +} + +impl Drop for SpillFileGuard { + fn drop(&mut self) { + let _ = std::fs::remove_file(&self.path); + } +} + +struct SpillFileWriter { + file: File, + file_guard: SpillFileGuard, +} + +impl SpillFileWriter { + fn new() -> Result { + let (file, path) = create_spill_file()?; + Ok(Self { + file, + file_guard: SpillFileGuard { path }, + }) + } + + fn append_segment( + &mut self, + rows: &[T], + ) -> Result { + let offset = self.file.stream_position()?; + // The row-count header makes segment boundaries discoverable without an in-memory index. + let row_count: u64 = rows.len().try_into()?; + self.file.write_all(&row_count.to_le_bytes())?; + for row in rows { + row.encode(&mut self.file)?; + } + let end = self.file.stream_position()?; + debug_assert!(end > offset); + Ok(offset) + } + + fn into_reader( + mut self, + tail: std::vec::IntoIter, + ) -> Result, DatabaseError> { + self.file.flush()?; + self.file.seek(SeekFrom::Start(0))?; + // Flushed segments are always a prefix; the in-memory buffer is its ordered tail. + Ok(ReadState::Spilled { + reader: SegmentReader::new(self.file), + tail, + _file_guard: self.file_guard, + }) + } +} + +fn create_spill_file() -> Result<(File, PathBuf), DatabaseError> { + loop { + let id = NEXT_SPILL_FILE_ID.fetch_add(1, Ordering::Relaxed); + let path = + std::env::temp_dir().join(format!("kitesql-spill-{}-{id}.tmp", std::process::id())); + match OpenOptions::new() + .read(true) + .write(true) + .create_new(true) + .open(&path) + { + Ok(file) => return Ok((file, path)), + Err(error) if error.kind() == std::io::ErrorKind::AlreadyExists => continue, + Err(error) => return Err(error.into()), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::types::tuple::Tuple; + use crate::types::value::DataValue; + + fn row(value: i32) -> Vec { + vec![DataValue::Int32(value)] + } + + #[test] + fn small_spill_vec_transitions_to_memory_reading() -> Result<(), DatabaseError> { + let mut values = SpillVec::new().limit(2, usize::MAX); + let _ = values.push(row(1))?; + + let mut reader = values.into_iter(); + assert!(matches!(reader.state, ReadState::Memory(_))); + assert_eq!(reader.next().transpose()?, Some(row(1))); + assert_eq!(reader.next().transpose()?, None); + Ok(()) + } + + #[test] + fn push_automatically_spills_and_next_preserves_order() -> Result<(), DatabaseError> { + let mut values = SpillVec::new().limit(2, usize::MAX); + for value in 0..5 { + let _ = values.push(row(value))?; + } + assert!(matches!( + &values.writer, + Ok(WriteState { file: Some(_), .. }) + )); + + let reader = values.into_iter(); + assert!(matches!( + &reader.state, + ReadState::Spilled { tail, .. } if tail.len() == 1 + )); + let restored = reader.collect::, _>>()?; + assert_eq!(restored, (0..5).map(row).collect::>()); + Ok(()) + } + + #[test] + fn spill_reader_stays_exhausted() -> Result<(), DatabaseError> { + let mut reader = SpillVec::from(vec![row(1)]).into_iter(); + assert_eq!(reader.next().transpose()?, Some(row(1))); + assert_eq!(reader.next().transpose()?, None); + assert_eq!(reader.next().transpose()?, None); + Ok(()) + } + + #[test] + fn spill_reader_opens_independent_segment_readers() -> Result<(), DatabaseError> { + let mut values = SpillVec::new().limit(usize::MAX, usize::MAX); + let _ = values.push(row(1))?; + let _ = values.push(row(2))?; + let first = values.flush()?.ok_or_else(|| { + DatabaseError::InvalidValue("expected first spill segment".to_string()) + })?; + let _ = values.push(row(3))?; + let second = values.flush()?.ok_or_else(|| { + DatabaseError::InvalidValue("expected second spill segment".to_string()) + })?; + + let source = values.into_iter(); + let mut first_reader = source.open_segment_reader()?; + first_reader.reset(first)?; + let mut second_reader = source.open_segment_reader()?; + second_reader.reset(second)?; + + assert_eq!(first_reader.next().transpose()?, Some(row(1))); + assert_eq!(second_reader.next().transpose()?, Some(row(3))); + assert_eq!(first_reader.next().transpose()?, Some(row(2))); + assert_eq!(second_reader.next().transpose()?, None); + assert_eq!(first_reader.next().transpose()?, None); + Ok(()) + } + + #[test] + fn spill_vec_rejects_operations_after_failure() { + let mut failed = SpillVec { + writer: Err(DatabaseError::InvalidValue("failed spill".to_string())), + }; + + assert!(matches!( + failed.push(row(1)), + Err(DatabaseError::InvalidValue(message)) + if message == "cannot append to a failed SpillVec" + )); + assert!(matches!( + failed.flush(), + Err(DatabaseError::InvalidValue(message)) + if message == "cannot flush a failed SpillVec" + )); + + let mut reader = failed.into_iter(); + assert!(matches!( + reader.next(), + Some(Err(DatabaseError::InvalidValue(message))) if message == "failed spill" + )); + assert!(reader.next().is_none()); + } + + #[test] + fn spill_reader_rejects_segment_reader_for_memory_values() { + let values = SpillVec::from(vec![row(1)]).into_iter(); + + assert!(matches!( + values.open_segment_reader(), + Err(DatabaseError::InvalidValue(message)) + if message == "cannot open a segment reader for an in-memory SpillVec" + )); + } + + #[test] + fn segment_reader_rejects_invalid_segments() { + let empty_segment_bytes = 0_u64.to_le_bytes().to_vec(); + let mut empty_segment = + SegmentReader::<_, Vec>::new(empty_segment_bytes.as_slice()); + assert!(matches!( + empty_segment.start_next_segment(), + Err(DatabaseError::InvalidValue(message)) + if message == "spill segment cannot be empty" + )); + + let truncated_segment_bytes = 1_u64.to_le_bytes().to_vec(); + let mut truncated_segment = + SegmentReader::<_, Vec>::new(truncated_segment_bytes.as_slice()); + assert!(truncated_segment.start_next_segment().unwrap()); + assert!(truncated_segment.next().transpose().is_err()); + + let mut end_offset = + SegmentReader::<_, Vec>::new(std::io::Cursor::new(Vec::new())); + assert!(matches!( + end_offset.reset(0), + Err(DatabaseError::InvalidValue(message)) + if message == "spill segment offset points to end of file" + )); + } + + #[test] + fn write_state_rejects_flush_without_spill_file() -> Result<(), DatabaseError> { + let mut empty = WriteState { + buffer: Vec::>::new(), + buffer_bytes: 0, + file: None, + on_flush: None, + max_rows: DEFAULT_MAX_ROWS, + max_bytes: DEFAULT_MAX_BYTES, + }; + assert_eq!(empty.flush()?, None); + + let mut missing_file = WriteState { + buffer: vec![row(1)], + buffer_bytes: 0, + file: None, + on_flush: None, + max_rows: DEFAULT_MAX_ROWS, + max_bytes: DEFAULT_MAX_BYTES, + }; + assert!(matches!( + missing_file.flush(), + Err(DatabaseError::InvalidValue(message)) + if message == "cannot flush without a spill file" + )); + Ok(()) + } + + #[test] + fn from_records_push_failure_for_later_read() { + struct FailingEncode; + + impl SpillCodec for FailingEncode { + fn encode(&self, _: &mut W) -> Result<(), DatabaseError> { + Err(DatabaseError::InvalidValue("encode failed".to_string())) + } + + fn decode(_: &mut R) -> Result { + Ok(Self) + } + + fn estimated_size(&self) -> usize { + 0 + } + } + + let values = SpillVec::from( + std::iter::repeat_with(|| FailingEncode) + .take(DEFAULT_MAX_ROWS) + .collect::>(), + ); + assert!(matches!( + values.writer, + Err(DatabaseError::InvalidValue(message)) if message == "encode failed" + )); + let _ = FailingEncode::decode(&mut [].as_slice()).unwrap(); + } + + #[test] + fn spill_reader_reports_decode_errors() -> Result<(), DatabaseError> { + struct FailingDecode; + + impl SpillCodec for FailingDecode { + fn encode(&self, _: &mut W) -> Result<(), DatabaseError> { + Ok(()) + } + + fn decode(_: &mut R) -> Result { + Err(DatabaseError::InvalidValue("decode failed".to_string())) + } + + fn estimated_size(&self) -> usize { + 0 + } + } + + let mut values = SpillVec::new().limit(1, usize::MAX); + let _ = values.push(FailingDecode)?; + let mut reader = values.into_iter(); + assert!(matches!( + reader.next(), + Some(Err(DatabaseError::InvalidValue(message))) if message == "decode failed" + )); + assert!(reader.next().is_none()); + Ok(()) + } + + #[test] + fn spill_reader_reports_truncated_segment_header() -> Result<(), DatabaseError> { + let path = std::env::temp_dir().join(format!( + "kitesql-spill-test-{}-{}.tmp", + std::process::id(), + NEXT_SPILL_FILE_ID.fetch_add(1, Ordering::Relaxed) + )); + { + let mut file = File::create(&path)?; + file.write_all(&1_u64.to_le_bytes())?; + row(1).encode(&mut file)?; + file.write_all(&[0])?; + } + + let mut reader = SpillReader { + state: ReadState::Spilled { + reader: SegmentReader::new(File::open(&path)?), + tail: Vec::>::new().into_iter(), + _file_guard: SpillFileGuard { path }, + }, + }; + + assert_eq!(reader.next().transpose()?, Some(row(1))); + assert!(reader.next().transpose().is_err()); + assert!(reader.next().is_none()); + Ok(()) + } + + #[test] + fn open_segment_reader_propagates_file_open_errors() -> Result<(), DatabaseError> { + let mut values = SpillVec::new().limit(1, usize::MAX); + let _ = values.push(row(1))?; + let source = values.into_iter(); + let mut path = None; + if let ReadState::Spilled { _file_guard, .. } = &source.state { + path = Some(_file_guard.path.clone()); + } + std::fs::remove_file(path.expect("expected spilled reader"))?; + + assert!(source.open_segment_reader().is_err()); + Ok(()) + } + + #[test] + fn on_flush_runs_before_memory_read_and_segment_flush() -> Result<(), DatabaseError> { + fn reverse(rows: &mut [Vec]) -> Result<(), DatabaseError> { + rows.reverse(); + Ok(()) + } + + let mut memory = SpillVec::new() + .limit(3, usize::MAX) + .on_flush(|rows| reverse(rows)); + let _ = memory.push(row(1))?; + let _ = memory.push(row(2))?; + assert_eq!( + memory.into_iter().collect::, _>>()?, + vec![row(2), row(1)] + ); + + let mut spilled = SpillVec::new() + .limit(2, usize::MAX) + .on_flush(|rows| reverse(rows)); + for value in 1..=4 { + let _ = spilled.push(row(value))?; + } + assert_eq!( + spilled.into_iter().collect::, _>>()?, + vec![row(2), row(1), row(4), row(3)] + ); + Ok(()) + } + + #[test] + fn tuple_codec_preserves_primary_key_and_values() -> Result<(), DatabaseError> { + let tuple = Tuple::new( + Some(DataValue::Int32(7)), + vec![DataValue::Null, DataValue::Int32(11)], + ); + let mut bytes = Vec::new(); + tuple.encode(&mut bytes)?; + assert_eq!(Tuple::decode(&mut bytes.as_slice())?, tuple); + Ok(()) + } + + #[test] + fn codec_handles_option_and_nested_tuple_edges() -> Result<(), DatabaseError> { + let mut encoded_none = Vec::new(); + Option::::None.encode(&mut encoded_none)?; + assert_eq!( + Option::::decode(&mut encoded_none.as_slice())?, + None + ); + + assert!(matches!( + Option::::decode(&mut [2].as_slice()), + Err(DatabaseError::InvalidValue(message)) + if message == "invalid spill option tag: 2" + )); + + let some = Some(DataValue::new_utf8("spill".to_string())); + let none = Option::::None; + assert!(some.estimated_size() > none.estimated_size()); + + let nested = DataValue::Tuple( + vec![ + DataValue::new_utf8("outer".to_string()), + DataValue::Tuple(vec![DataValue::new_utf8("inner".to_string())], false), + ], + false, + ); + assert!(nested.estimated_size() > std::mem::size_of::()); + Ok(()) + } +} diff --git a/src/lib.rs b/src/lib.rs index 715f3044..5027a76a 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -102,6 +102,11 @@ #![allow(unused_doc_comments)] extern crate core; +#[cfg(all(feature = "spill", feature = "wasm", not(clippy)))] +compile_error!("features `spill` and `wasm` are mutually exclusive"); +#[cfg(all(feature = "spill", target_arch = "wasm32"))] +compile_error!("feature `spill` is not supported on wasm32 targets"); + pub mod binder; pub mod catalog; pub mod db; diff --git a/src/optimizer/heuristic/optimizer.rs b/src/optimizer/heuristic/optimizer.rs index a6c5ee8b..06a2cb1a 100644 --- a/src/optimizer/heuristic/optimizer.rs +++ b/src/optimizer/heuristic/optimizer.rs @@ -25,7 +25,6 @@ use crate::optimizer::rule::normalization::{ apply_annotated_post_rules, apply_scan_order_hint, constant_calculation_current, evaluator_bind_current, NormalizationRuleImpl, OrderHintKind, ScanOrderHint, WholeTreePassKind, }; -use crate::planner::operator::join::JoinCondition; use crate::planner::operator::table_scan::TableScanOperator; use crate::planner::operator::{Operator, PhysicalOption, PlanImpl, SortOption}; use crate::planner::{Childrens, LogicalPlan, PlanArena}; @@ -77,14 +76,14 @@ impl<'a> HepOptimizer<'a> { if self.implementation_index.is_empty().not() { let apply_no_sort_hints = |_scan_op: &mut TableScanOperator, _arena: &PlanArena| Ok(()); - let apply_no_stream_distinct_hints = + let apply_no_stream_aggregate_hints = |_scan_op: &mut TableScanOperator, _arena: &PlanArena| Ok(()); Self::annotate_hints_and_physical_options( &mut self.plan, loader, self.implementation_index, &apply_no_sort_hints, - &apply_no_stream_distinct_hints, + &apply_no_stream_aggregate_hints, arena, )?; } @@ -247,12 +246,12 @@ impl<'a> HepOptimizer<'a> { loader: &StatisticMetaLoader<'_>, implementation_index: &ImplementationRuleIndex, inherited_sort_hints: &'plan ScanHintApplier<'plan>, - inherited_stream_distinct_hints: &'plan ScanHintApplier<'plan>, + inherited_stream_aggregate_hints: &'plan ScanHintApplier<'plan>, arena: &mut PlanArena, ) -> Result<(), DatabaseError> { if let Operator::TableScan(scan_op) = &mut plan.operator { inherited_sort_hints(scan_op, arena)?; - inherited_stream_distinct_hints(scan_op, arena)?; + inherited_stream_aggregate_hints(scan_op, arena)?; } { @@ -281,16 +280,16 @@ impl<'a> HepOptimizer<'a> { .. } = plan; Self::with_child_sort_hints(operator, inherited_sort_hints, |child_sort_hints| { - Self::with_child_stream_distinct_hints( + Self::with_child_stream_aggregate_hints( operator, - inherited_stream_distinct_hints, - |child_stream_distinct_hints| match &mut **childrens { + inherited_stream_aggregate_hints, + |child_stream_aggregate_hints| match &mut **childrens { Childrens::Only(child) => Self::annotate_hints_and_physical_options( child, loader, implementation_index, child_sort_hints, - child_stream_distinct_hints, + child_stream_aggregate_hints, arena, ), Childrens::Twins { left, right } => { @@ -299,7 +298,7 @@ impl<'a> HepOptimizer<'a> { loader, implementation_index, child_sort_hints, - child_stream_distinct_hints, + child_stream_aggregate_hints, arena, )?; Self::annotate_hints_and_physical_options( @@ -307,7 +306,7 @@ impl<'a> HepOptimizer<'a> { loader, implementation_index, child_sort_hints, - child_stream_distinct_hints, + child_stream_aggregate_hints, arena, ) } @@ -357,9 +356,9 @@ impl<'a> HepOptimizer<'a> { } } - fn with_child_stream_distinct_hints<'plan, R>( + fn with_child_stream_aggregate_hints<'plan, R>( operator: &'plan Operator, - inherited_stream_distinct_hints: &'plan ScanHintApplier<'plan>, + inherited_stream_aggregate_hints: &'plan ScanHintApplier<'plan>, f: impl for<'b> FnOnce(&'b ScanHintApplier<'plan>) -> R, ) -> R { let propagate_hints = matches!( @@ -372,25 +371,23 @@ impl<'a> HepOptimizer<'a> { ); match operator { - Operator::Aggregate(op) - if op.is_distinct && op.agg_calls.is_empty() && !op.groupby_exprs.is_empty() => - { - let child_stream_distinct_hints = + Operator::Aggregate(op) if !op.groupby_exprs.is_empty() => { + let child_stream_aggregate_hints = |scan_op: &mut TableScanOperator, arena: &PlanArena| { apply_scan_order_hint( scan_op, - ScanOrderHint::distinct_groupby(&op.groupby_exprs), - OrderHintKind::StreamDistinct, + ScanOrderHint::groupby(&op.groupby_exprs), + OrderHintKind::StreamAggregate, arena, ) }; - f(&child_stream_distinct_hints) + f(&child_stream_aggregate_hints) } - _ if propagate_hints => f(inherited_stream_distinct_hints), + _ if propagate_hints => f(inherited_stream_aggregate_hints), _ => { - let no_stream_distinct_hints = + let no_stream_aggregate_hints = |_scan_op: &mut TableScanOperator, _arena: &PlanArena| Ok(()); - f(&no_stream_distinct_hints) + f(&no_stream_aggregate_hints) } } } @@ -504,11 +501,7 @@ impl ImplementationRuleIndex { Some(PhysicalOption::new(PlanImpl::Filter, SortOption::Follow)) } Operator::Join(join_op) if self.contains(ImplementationRuleImpl::HashJoin) => { - let plan = match &join_op.on { - JoinCondition::On { on, .. } if !on.is_empty() => PlanImpl::HashJoin, - _ => PlanImpl::NestLoopJoin, - }; - Some(PhysicalOption::new(plan, SortOption::None)) + Some(PhysicalOption::new(join_op.plan_impl(), SortOption::None)) } Operator::Limit(_) if self.contains(ImplementationRuleImpl::Limit) => { Some(PhysicalOption::new(PlanImpl::Limit, SortOption::Follow)) diff --git a/src/optimizer/rule/implementation/dql/join.rs b/src/optimizer/rule/implementation/dql/join.rs index e81be2fd..61310f5d 100644 --- a/src/optimizer/rule/implementation/dql/join.rs +++ b/src/optimizer/rule/implementation/dql/join.rs @@ -16,8 +16,7 @@ use crate::errors::DatabaseError; use crate::optimizer::core::pattern::{Pattern, PatternChildrenPredicate}; use crate::optimizer::core::rule::{BestPhysicalOption, ImplementationRule, MatchPattern}; use crate::optimizer::core::statistics_meta::StatisticMetaLoader; -use crate::planner::operator::join::{JoinCondition, JoinOperator}; -use crate::planner::operator::{Operator, PhysicalOption, PlanImpl, SortOption}; +use crate::planner::operator::{Operator, PhysicalOption, SortOption}; use std::sync::LazyLock; static JOIN_PATTERN: LazyLock = LazyLock::new(|| Pattern { @@ -42,17 +41,10 @@ impl ImplementationRule for JoinImplementation { _: &StatisticMetaLoader<'_>, best_physical_option: &mut BestPhysicalOption, ) -> Result<(), DatabaseError> { - let mut physical_option = PhysicalOption::new(PlanImpl::NestLoopJoin, SortOption::None); - - if let Operator::Join(JoinOperator { - on: JoinCondition::On { on, .. }, - .. - }) = op - { - if !on.is_empty() { - physical_option.plan = PlanImpl::HashJoin; - } - } + let Operator::Join(operator) = op else { + unreachable!() + }; + let physical_option = PhysicalOption::new(operator.plan_impl(), SortOption::None); crate::optimizer::core::rule::keep_best_physical_option( best_physical_option, physical_option, diff --git a/src/optimizer/rule/implementation/dql/table_scan.rs b/src/optimizer/rule/implementation/dql/table_scan.rs index b32bdc9f..bf1cdc3a 100644 --- a/src/optimizer/rule/implementation/dql/table_scan.rs +++ b/src/optimizer/rule/implementation/dql/table_scan.rs @@ -125,7 +125,7 @@ impl ImplementationRule for IndexScanImplementation { .map(|hint| hint.cover_num()) .unwrap_or(0) + index_info - .stream_distinct_hint + .stream_aggregate_hint .map(|hint| hint.cover_num()) .unwrap_or(0); if hint_sum > 0 { diff --git a/src/optimizer/rule/implementation/mod.rs b/src/optimizer/rule/implementation/mod.rs index a93e6dca..f2fcb9ff 100644 --- a/src/optimizer/rule/implementation/mod.rs +++ b/src/optimizer/rule/implementation/mod.rs @@ -326,7 +326,7 @@ mod tests { use crate::storage::StatisticsMetaCache; use crate::types::value::DataValue; - fn find_operator<'a, F>(plan: &'a LogicalPlan, predicate: F) -> Option<&'a Operator> + fn find_operator(plan: &LogicalPlan, predicate: F) -> Option<&Operator> where F: Fn(&Operator) -> bool + Copy, { diff --git a/src/optimizer/rule/normalization/combine_operators.rs b/src/optimizer/rule/normalization/combine_operators.rs index edcc853d..7c7e2712 100644 --- a/src/optimizer/rule/normalization/combine_operators.rs +++ b/src/optimizer/rule/normalization/combine_operators.rs @@ -438,9 +438,10 @@ mod tests { vec![], vec![column_expr(&mut arena, "c2", 1)], false, + false, ); let expr = column_expr(&mut arena, "c2", 0); - let mut plan = AggregateOperator::build(child, vec![], vec![expr], true); + let mut plan = AggregateOperator::build(child, vec![], vec![expr], true, false); assert!(CollapseGroupByAgg.apply(&mut plan, &mut arena)?); let Operator::Aggregate(op) = &plan.operator else { diff --git a/src/optimizer/rule/normalization/compilation_in_advance.rs b/src/optimizer/rule/normalization/compilation_in_advance.rs index a62f2e1f..6b092fbc 100644 --- a/src/optimizer/rule/normalization/compilation_in_advance.rs +++ b/src/optimizer/rule/normalization/compilation_in_advance.rs @@ -135,7 +135,6 @@ mod tests { }), Operator::Sort(SortOperator { sort_fields: vec![SortField::from(expr())], - limit: None, }), Operator::TopK(TopKOperator { sort_fields: vec![SortField::from(expr())], @@ -205,6 +204,7 @@ mod tests { let mut join = LogicalPlan::new( Operator::Join(JoinOperator { join_type: JoinType::Inner, + force_nested_loop: false, on: JoinCondition::On { on: vec![(expr(), expr())], filter: Some(expr()), diff --git a/src/optimizer/rule/normalization/elimination.rs b/src/optimizer/rule/normalization/elimination.rs index 527f5053..f1bde9a1 100644 --- a/src/optimizer/rule/normalization/elimination.rs +++ b/src/optimizer/rule/normalization/elimination.rs @@ -15,9 +15,9 @@ use crate::errors::DatabaseError; use crate::expression::ScalarExpression; use crate::optimizer::core::rule::NormalizationRule; -use crate::optimizer::plan_utils::{only_child_mut, replace_with_only_child}; +use crate::optimizer::plan_utils::{only_child_mut, replace_with_only_child, wrap_child_with}; use crate::planner::operator::limit::LimitOperator; -use crate::planner::operator::sort::SortField; +use crate::planner::operator::sort::{SortField, SortOperator}; use crate::planner::operator::table_scan::TableScanOperator; use crate::planner::operator::{Operator, PhysicalOption, PlanImpl, SortOption}; use crate::planner::{Childrens, LogicalPlan}; @@ -45,7 +45,7 @@ impl NormalizationRule for EliminateRedundantSort { None => return Ok(false), }; mark_sort_preserving_indexes(child, &sort_fields, arena)?; - let can_remove = ensure_index_order(child, &sort_fields, arena); + let can_remove = ensure_order(child, &sort_fields, arena); if !can_remove { return Ok(false); @@ -116,13 +116,13 @@ fn mark_sort_preserving_indexes( #[derive(Copy, Clone)] pub(crate) enum OrderHintKind { SortElimination, - StreamDistinct, + StreamAggregate, } #[derive(Copy, Clone)] pub(crate) enum ScanOrderHint<'a> { SortFields(&'a [SortField]), - DistinctGroupBy(&'a [ScalarExpression]), + GroupBy(&'a [ScalarExpression]), } impl<'a> ScanOrderHint<'a> { @@ -130,8 +130,8 @@ impl<'a> ScanOrderHint<'a> { Self::SortFields(fields) } - pub(crate) fn distinct_groupby(groupby_exprs: &'a [ScalarExpression]) -> Self { - Self::DistinctGroupBy(groupby_exprs) + pub(crate) fn groupby(groupby_exprs: &'a [ScalarExpression]) -> Self { + Self::GroupBy(groupby_exprs) } } @@ -173,7 +173,7 @@ pub(crate) fn apply_scan_order_hint( for index in 0..hint_len(required) { let expr = match required { ScanOrderHint::SortFields(fields) => &fields[index].expr, - ScanOrderHint::DistinctGroupBy(groupby_exprs) => &groupby_exprs[index], + ScanOrderHint::GroupBy(groupby_exprs) => &groupby_exprs[index], }; if !expr.all_referenced_columns(arena, |arena, column| { scan_op @@ -199,11 +199,11 @@ pub(crate) fn apply_scan_order_hint( index_info.sort_elimination_hint = Some(IndexOrderHint::new(covered)); } } - OrderHintKind::StreamDistinct => { - if let Some(hint) = &mut index_info.stream_distinct_hint { + OrderHintKind::StreamAggregate => { + if let Some(hint) = &mut index_info.stream_aggregate_hint { hint.merge_cover_num(covered); } else { - index_info.stream_distinct_hint = Some(IndexOrderHint::new(covered)); + index_info.stream_aggregate_hint = Some(IndexOrderHint::new(covered)); } } } @@ -215,7 +215,7 @@ pub(crate) fn apply_scan_order_hint( fn hint_len(required: ScanOrderHint<'_>) -> usize { match required { ScanOrderHint::SortFields(fields) => fields.len(), - ScanOrderHint::DistinctGroupBy(groupby_exprs) => groupby_exprs.len(), + ScanOrderHint::GroupBy(groupby_exprs) => groupby_exprs.len(), } } @@ -228,15 +228,13 @@ fn hint_covers( ScanOrderHint::SortFields(fields) => covers(fields, provided, |required, provided| { sort_field_matches(required, provided, arena) }), - ScanOrderHint::DistinctGroupBy(groupby_exprs) => { - covers(groupby_exprs, provided, |expr, field| { - field.asc && !field.nulls_first && expr.eq_ignore_colref_pos(&field.expr, arena) - }) - } + ScanOrderHint::GroupBy(groupby_exprs) => covers(groupby_exprs, provided, |expr, field| { + field.asc && !field.nulls_first && expr.eq_ignore_colref_pos(&field.expr, arena) + }), } } -pub(crate) fn distinct_sort_fields(groupby_exprs: &[ScalarExpression]) -> Vec { +pub(crate) fn groupby_sort_fields(groupby_exprs: &[ScalarExpression]) -> Vec { groupby_exprs .iter() .cloned() @@ -244,9 +242,9 @@ pub(crate) fn distinct_sort_fields(groupby_exprs: &[ScalarExpression]) -> Vec child, None => return Ok(false), }; - if !ensure_stream_distinct_order(child, &required, arena) { + if !ensure_order(child, &required, arena) { return Ok(false); } - plan.physical_option = Some(PhysicalOption::new( - PlanImpl::StreamDistinct, - SortOption::Follow, - )); + plan.physical_option = Some(PhysicalOption::new(implementation, SortOption::Follow)); + Ok(true) + } +} + +pub struct ForceSpillAggregate; + +impl NormalizationRule for ForceSpillAggregate { + fn apply( + &self, + plan: &mut LogicalPlan, + _: &mut crate::planner::PlanArena, + ) -> Result { + let (implementation, sort_fields) = match (&plan.operator, &plan.physical_option) { + ( + Operator::Aggregate(op), + Some(PhysicalOption { + plan: PlanImpl::HashAggregate, + .. + }), + ) if op.force_spill && !op.groupby_exprs.is_empty() => { + let implementation = if op.is_distinct && op.agg_calls.is_empty() { + PlanImpl::StreamDistinct + } else { + PlanImpl::StreamAggregate + }; + (implementation, groupby_sort_fields(&op.groupby_exprs)) + } + _ => return Ok(false), + }; + + if !cfg!(feature = "spill") { + return Err(DatabaseError::UnsupportedStmt( + "FORCE_AGG_SPILL requires the `spill` feature".to_string(), + )); + } + + let sort_option = SortOption::OrderBy { + fields: sort_fields.clone(), + ignore_prefix_len: 0, + }; + if !wrap_child_with(plan, 0, Operator::Sort(SortOperator { sort_fields })) { + return Ok(false); + } + let sort = only_child_mut(plan).expect("aggregate child was wrapped with sort"); + sort.physical_option = Some(PhysicalOption::new(PlanImpl::Sort, sort_option)); + plan.physical_option = Some(PhysicalOption::new(implementation, SortOption::Follow)); Ok(true) } } @@ -297,14 +343,18 @@ pub(crate) fn apply_annotated_post_rules( if EliminateIndexFilter.apply(plan, arena)? { changed = true; } - if UseStreamDistinct.apply(plan, arena)? { + if UseStreamAggregate.apply(plan, arena)? { + changed = true; + } + // Run last so an existing ordered child is reused before a forced spill adds a Sort. + if ForceSpillAggregate.apply(plan, arena)? { changed = true; } Ok(changed) } -fn ensure_stream_distinct_order( +fn ensure_order( plan: &mut LogicalPlan, required: &[SortField], arena: &crate::planner::PlanArena, @@ -323,19 +373,16 @@ fn ensure_stream_distinct_order( if let Some(physical_option) = plan.physical_option.as_ref() { match physical_option.sort_option() { - SortOption::OrderBy { .. } - if covers( + SortOption::OrderBy { .. } => { + return covers( required, physical_option.sort_option(), |required, provided| sort_field_matches(required, provided, arena), - ) => - { - return true + ); } - SortOption::OrderBy { .. } => {} SortOption::Follow => { if let Childrens::Only(child) = plan.childrens.as_mut() { - if ensure_stream_distinct_order(child, required, arena) { + if ensure_order(child, required, arena) { return true; } } @@ -347,36 +394,6 @@ fn ensure_stream_distinct_order( false } -fn ensure_index_order( - plan: &mut LogicalPlan, - required: &[SortField], - arena: &crate::planner::PlanArena, -) -> bool { - if let Some(PhysicalOption { - plan: PlanImpl::IndexScan(index_info), - .. - }) = plan.physical_option.as_ref() - { - if covers(required, &index_info.sort_option, |required, provided| { - sort_field_matches(required, provided, arena) - }) { - return true; - } - } - - if let Some(physical_option) = plan.physical_option.as_ref() { - if matches!(physical_option.sort_option(), SortOption::Follow) { - if let Childrens::Only(child) = plan.childrens.as_mut() { - if ensure_index_order(child, required, arena) { - return true; - } - } - } - } - - false -} - fn sort_field_matches( required: &SortField, provided: &SortField, @@ -426,7 +443,9 @@ pub(crate) fn covers( #[cfg(all(test, not(target_arch = "wasm32")))] mod tests { - use super::{EliminateIndexFilter, EliminateRedundantSort, UseStreamDistinct}; + use super::{ + EliminateIndexFilter, EliminateRedundantSort, ForceSpillAggregate, UseStreamAggregate, + }; use crate::catalog::{ColumnCatalog, TableName}; use crate::errors::DatabaseError; use crate::expression::range_detacher::Range; @@ -485,7 +504,6 @@ mod tests { LogicalPlan::new( Operator::Sort(SortOperator { sort_fields: required_fields, - limit: None, }), Childrens::Only(Box::new(filter)), ) @@ -521,7 +539,7 @@ mod tests { covered_deserializers: None, cover_mapping: None, sort_elimination_hint: None, - stream_distinct_hint: None, + stream_aggregate_hint: None, }; (index_info, sort_option) } @@ -600,7 +618,7 @@ mod tests { covered_deserializers: None, cover_mapping: None, sort_elimination_hint: None, - stream_distinct_hint: None, + stream_aggregate_hint: None, }; let scan = LogicalPlan::new( @@ -619,6 +637,7 @@ mod tests { groupby_exprs: vec![ScalarExpression::column_expr(c1, 0)], agg_calls: vec![], is_distinct: true, + force_spill: false, }), Childrens::Only(Box::new(scan)), ); @@ -792,7 +811,6 @@ mod tests { let mut plan = LogicalPlan::new( Operator::Sort(SortOperator { sort_fields: vec![sort_field], - limit: None, }), Childrens::Only(Box::new(table_scan)), ); @@ -818,19 +836,19 @@ mod tests { } #[test] - fn annotate_sets_stream_distinct_hint_on_table_scan() -> Result<(), DatabaseError> { + fn annotate_sets_stream_aggregate_hint_on_table_scan() -> Result<(), DatabaseError> { let table_arena = crate::planner::TableArenaCell::default(); let mut arena = crate::planner::PlanArena::new(&table_arena); let (mut plan, _) = build_distinct_scan_plan(&mut arena); let required = match &plan.operator { - Operator::Aggregate(op) => super::distinct_sort_fields(&op.groupby_exprs), + Operator::Aggregate(op) => super::groupby_sort_fields(&op.groupby_exprs), _ => unreachable!("expected aggregate operator"), }; if let Childrens::Only(child) = plan.childrens.as_mut() { super::mark_order_hint( child, &required, - super::OrderHintKind::StreamDistinct, + super::OrderHintKind::StreamAggregate, &arena, )?; } @@ -843,7 +861,7 @@ mod tests { assert_eq!(scan_op.index_infos.len(), 1); assert_eq!( scan_op.index_infos[0] - .stream_distinct_hint + .stream_aggregate_hint .map(|hint| hint.cover_num()), Some(1) ); @@ -869,7 +887,96 @@ mod tests { SortOption::None, )); - let rule = UseStreamDistinct; + let rule = UseStreamAggregate; + assert!(rule.apply(&mut plan, &mut arena)?); + assert!(matches!( + plan.physical_option, + Some(PhysicalOption { + plan: PlanImpl::StreamDistinct, + .. + }) + )); + Ok(()) + } + + #[test] + fn use_stream_aggregate_only_when_order_satisfied() -> Result<(), DatabaseError> { + let table_arena = crate::planner::TableArenaCell::default(); + let mut arena = crate::planner::PlanArena::new(&table_arena); + let (mut plan, sort_option) = build_distinct_scan_plan(&mut arena); + let Operator::Aggregate(op) = &mut plan.operator else { + unreachable!() + }; + op.is_distinct = false; + op.force_spill = true; + let Childrens::Only(child) = plan.childrens.as_mut() else { + unreachable!() + }; + let Operator::TableScan(scan_op) = &child.operator else { + unreachable!() + }; + child.physical_option = Some(PhysicalOption::new( + PlanImpl::IndexScan(Box::new(scan_op.index_infos[0].clone())), + sort_option, + )); + plan.physical_option = Some(PhysicalOption::new( + PlanImpl::HashAggregate, + SortOption::None, + )); + + let rule = UseStreamAggregate; + assert!(rule.apply(&mut plan, &mut arena)?); + assert!(matches!( + plan.physical_option, + Some(PhysicalOption { + plan: PlanImpl::StreamAggregate, + .. + }) + )); + let force_spill = ForceSpillAggregate; + assert!(!force_spill.apply(&mut plan, &mut arena)?); + assert!(!matches!( + plan.childrens.as_ref(), + Childrens::Only(child) if matches!(child.operator, Operator::Sort(_)) + )); + + let (mut unordered, _) = build_distinct_scan_plan(&mut arena); + let Operator::Aggregate(op) = &mut unordered.operator else { + unreachable!() + }; + op.is_distinct = false; + unordered.physical_option = Some(PhysicalOption::new( + PlanImpl::HashAggregate, + SortOption::None, + )); + assert!(!rule.apply(&mut unordered, &mut arena)?); + assert!(matches!( + unordered.physical_option, + Some(PhysicalOption { + plan: PlanImpl::HashAggregate, + .. + }) + )); + Ok(()) + } + + #[cfg(feature = "spill")] + #[test] + fn force_spill_distinct_sorts_unordered_input() -> Result<(), DatabaseError> { + let table_arena = crate::planner::TableArenaCell::default(); + let mut arena = crate::planner::PlanArena::new(&table_arena); + let (mut plan, _) = build_distinct_scan_plan(&mut arena); + let Operator::Aggregate(op) = &mut plan.operator else { + unreachable!() + }; + op.force_spill = true; + let expected_fields = super::groupby_sort_fields(&op.groupby_exprs); + plan.physical_option = Some(PhysicalOption::new( + PlanImpl::HashAggregate, + SortOption::None, + )); + + let rule = ForceSpillAggregate; assert!(rule.apply(&mut plan, &mut arena)?); assert!(matches!( plan.physical_option, @@ -878,6 +985,24 @@ mod tests { .. }) )); + let Childrens::Only(sort) = plan.childrens.as_ref() else { + unreachable!() + }; + assert!(matches!( + &sort.operator, + Operator::Sort(SortOperator { sort_fields }) if sort_fields == &expected_fields + )); + assert!(matches!( + sort.physical_option, + Some(PhysicalOption { + plan: PlanImpl::Sort, + .. + }) + )); + assert!(matches!( + sort.childrens.as_ref(), + Childrens::Only(child) if matches!(child.operator, Operator::TableScan(_)) + )); Ok(()) } @@ -946,7 +1071,6 @@ mod tests { let mut plan = LogicalPlan::new( Operator::Sort(SortOperator { sort_fields: vec![sort_field], - limit: None, }), Childrens::Only(Box::new(filter)), ); diff --git a/src/optimizer/rule/normalization/pushdown_predicates.rs b/src/optimizer/rule/normalization/pushdown_predicates.rs index e9e93ed6..5494cc52 100644 --- a/src/optimizer/rule/normalization/pushdown_predicates.rs +++ b/src/optimizer/rule/normalization/pushdown_predicates.rs @@ -275,7 +275,7 @@ impl NormalizationRule for PushPredicateIntoScan { cover_mapping, sort_option, sort_elimination_hint: _, - stream_distinct_hint: _, + stream_aggregate_hint: _, } in &mut scan_op.index_infos { if lookup.is_some() { @@ -847,7 +847,7 @@ mod tests { covered_deserializers: None, cover_mapping: None, sort_elimination_hint: None, - stream_distinct_hint: None, + stream_aggregate_hint: None, }, IndexInfo { meta: index_meta_aligned, @@ -860,7 +860,7 @@ mod tests { covered_deserializers: None, cover_mapping: None, sort_elimination_hint: None, - stream_distinct_hint: None, + stream_aggregate_hint: None, }, ], with_pk: false, diff --git a/src/orm/mod.rs b/src/orm/mod.rs index c2f21582..d0df33f8 100644 --- a/src/orm/mod.rs +++ b/src/orm/mod.rs @@ -1845,6 +1845,23 @@ where } } + /// Forces subsequent grouped aggregate or distinct operations to use spill-backed execution. + pub fn force_spill(self) -> Result { + if !cfg!(feature = "spill") { + return Err(DatabaseError::UnsupportedStmt( + "force_spill requires the `spill` feature".to_string(), + )); + } + self.binder.force_spill = true; + Ok(self) + } + + /// Forces subsequent joins in this query to use nested-loop execution. + pub fn force_nested_loop(self) -> Self { + self.binder.force_nested_loop = true; + self + } + pub fn filter( mut self, build: impl for<'scope> FnOnce( @@ -3997,5 +4014,83 @@ mod tests { Ok(()) } + + #[test] + fn query_builder_force_nested_loop_join() -> Result<(), DatabaseError> { + let database = build_orm_unit_database()?; + + let plan = database.explain(|ctx| { + ctx.from::()? + .force_nested_loop() + .inner_join::(|e| { + e.column(OrmUnitUser::id())? + .eq(e.column(OrmUnitOrder::user_id())?) + })? + .project_scalar(OrmUnitUser::id())? + .finish() + })?; + assert_eq!( + plan, + concat!( + "Projection [#1] [Project => (Sort Option: Follow)] ", + "Inner Join On #1 = #5 [NestLoopJoin => (Sort Option: None)] ", + "TableScan orm_unit_users -> [#1] [SeqScan => (Sort Option: None)] ", + "TableScan orm_unit_orders -> [#5] [SeqScan => (Sort Option: None)]" + ), + "{plan}" + ); + + Ok(()) + } + + #[cfg(feature = "spill")] + #[test] + fn query_builder_force_spill_aggregate_and_distinct() -> Result<(), DatabaseError> { + let database = build_orm_unit_database()?; + + let plan = database.explain(|ctx| { + ctx.from::()? + .force_spill()? + .project_tuple(|e| { + Ok(vec![ + e.column(OrmUnitOrder::user_id())?, + e.aggregate(AggKind::Sum, [e.column(OrmUnitOrder::amount())?])?, + ]) + })? + .group_by(|e| e.column(OrmUnitOrder::user_id()))? + .order_by_expr(|e| Ok(e.column(OrmUnitOrder::user_id())?.asc()))? + .finish() + })?; + assert_eq!( + plan, + concat!( + "Projection [#5, #7] [Project => (Sort Option: Follow)] ", + "Aggregate [Sum(#6)] -> Group By [#5] [StreamAggregate => (Sort Option: Follow)] ", + "Sort By #5 Asc Nulls Last [Sort => (Sort Option: OrderBy: (#5 Asc Nulls Last) ignore_prefix_len: 0)] ", + "TableScan orm_unit_orders -> [#5, #6] [SeqScan => (Sort Option: None)]" + ), + "{plan}" + ); + + let distinct_plan = database.explain(|ctx| { + ctx.from::()? + .force_spill()? + .project_scalar(OrmUnitOrder::user_id())? + .distinct()? + .finish() + })?; + assert_eq!( + distinct_plan, + concat!( + "Projection [#5] [Project => (Sort Option: Follow)] ", + "Aggregate [] -> Group By [#5] [StreamDistinct => (Sort Option: Follow)] ", + "Sort By #5 Asc Nulls Last [Sort => (Sort Option: OrderBy: (#5 Asc Nulls Last) ignore_prefix_len: 0)] ", + "TableScan orm_unit_orders -> [#5] [SeqScan => (Sort Option: None)]" + ), + "{distinct_plan}" + ); + + Ok(()) + } } // GRCOV_EXCL_STOP diff --git a/src/parser/mod.rs b/src/parser/mod.rs index 7bb2659b..1ed6d108 100644 --- a/src/parser/mod.rs +++ b/src/parser/mod.rs @@ -13,9 +13,9 @@ // limitations under the License. use sqlparser::parser::ParserError; -use sqlparser::{ast::Statement, dialect::PostgreSqlDialect, parser::Parser}; +use sqlparser::{ast::Statement, dialect::GenericDialect, parser::Parser}; -const DIALECT: PostgreSqlDialect = PostgreSqlDialect {}; +const DIALECT: GenericDialect = GenericDialect {}; /// Parse a string to a collection of statements. /// @@ -32,3 +32,25 @@ const DIALECT: PostgreSqlDialect = PostgreSqlDialect {}; pub fn parse_sql>(sql: S) -> Result, ParserError> { Parser::parse_sql(&DIALECT, sql.as_ref()) } + +#[cfg(test)] +mod tests { + use super::parse_sql; + use sqlparser::ast::{SetExpr, Statement}; + + #[test] + fn parses_optimizer_hint() { + let statements = parse_sql("SELECT /*+ FORCE_AGG_SPILL */ a FROM t").unwrap(); + let Statement::Query(query) = &statements[0] else { + panic!("expected query statement"); + }; + let SetExpr::Select(select) = query.body.as_ref() else { + panic!("expected select query"); + }; + + assert_eq!( + select.optimizer_hint.as_ref().map(|hint| hint.text.trim()), + Some("FORCE_AGG_SPILL") + ); + } +} diff --git a/src/planner/arena.rs b/src/planner/arena.rs index 439c4154..5a1749db 100644 --- a/src/planner/arena.rs +++ b/src/planner/arena.rs @@ -626,7 +626,7 @@ mod tests { let dummy = plan_arena.alloc_dummy("TABLE"); assert_eq!(plan_arena.column(dummy).name(), "TABLE"); - assert!(format!("{:?}", table_arena).contains("columns_len")); + assert!(format!("{table_arena:?}").contains("columns_len")); assert_eq!(plan_arena.temp_table().to_string(), "_temp_table_0_"); let first_index = plan_arena.alloc_index(index_meta("idx_a")); diff --git a/src/planner/operator/aggregate.rs b/src/planner/operator/aggregate.rs index a8b836b6..899d9c02 100644 --- a/src/planner/operator/aggregate.rs +++ b/src/planner/operator/aggregate.rs @@ -24,6 +24,7 @@ pub struct AggregateOperator { pub groupby_exprs: Vec, pub agg_calls: Vec, pub is_distinct: bool, + pub force_spill: bool, } impl AggregateOperator { @@ -32,12 +33,14 @@ impl AggregateOperator { agg_calls: Vec, groupby_exprs: Vec, is_distinct: bool, + force_spill: bool, ) -> LogicalPlan { LogicalPlan::new( Operator::Aggregate(Self { groupby_exprs, agg_calls, is_distinct, + force_spill, }), Childrens::Only(Box::new(children)), ) diff --git a/src/planner/operator/join.rs b/src/planner/operator/join.rs index 528b88dd..0a2e671c 100644 --- a/src/planner/operator/join.rs +++ b/src/planner/operator/join.rs @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -use super::Operator; +use super::{Operator, PlanImpl}; use crate::expression::ScalarExpression; use crate::iter_ext::Itertools; use crate::planner::{Childrens, LogicalPlan}; @@ -50,6 +50,7 @@ pub enum JoinCondition { pub struct JoinOperator { pub on: JoinCondition, pub join_type: JoinType, + pub force_nested_loop: bool, } impl JoinOperator { @@ -58,15 +59,28 @@ impl JoinOperator { right: LogicalPlan, on: JoinCondition, join_type: JoinType, + force_nested_loop: bool, ) -> LogicalPlan { LogicalPlan::new( - Operator::Join(JoinOperator { on, join_type }), + Operator::Join(JoinOperator { + on, + join_type, + force_nested_loop, + }), Childrens::Twins { left: Box::new(left), right: Box::new(right), }, ) } + + pub(crate) fn plan_impl(&self) -> PlanImpl { + match (&self.on, self.force_nested_loop) { + (_, true) => PlanImpl::NestLoopJoin, + (JoinCondition::On { on, .. }, false) if !on.is_empty() => PlanImpl::HashJoin, + _ => PlanImpl::NestLoopJoin, + } + } } impl fmt::Display for JoinType { @@ -115,3 +129,24 @@ impl fmt::Display for JoinCondition { Ok(()) } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn forced_nested_loop_overrides_equi_join() { + let mut operator = JoinOperator { + on: JoinCondition::On { + on: vec![(1_i32.into(), 2_i32.into())], + filter: None, + }, + join_type: JoinType::Inner, + force_nested_loop: false, + }; + assert_eq!(operator.plan_impl(), PlanImpl::HashJoin); + + operator.force_nested_loop = true; + assert_eq!(operator.plan_impl(), PlanImpl::NestLoopJoin); + } +} diff --git a/src/planner/operator/mod.rs b/src/planner/operator/mod.rs index 0296a200..13e9473e 100644 --- a/src/planner/operator/mod.rs +++ b/src/planner/operator/mod.rs @@ -172,6 +172,7 @@ pub enum PlanImpl { Dummy, SimpleAggregate, HashAggregate, + StreamAggregate, StreamDistinct, ScalarApply, MarkApply, @@ -527,6 +528,7 @@ impl fmt::Display for PlanImpl { PlanImpl::Dummy => write!(f, "Dummy"), PlanImpl::SimpleAggregate => write!(f, "SimpleAggregate"), PlanImpl::HashAggregate => write!(f, "HashAggregate"), + PlanImpl::StreamAggregate => write!(f, "StreamAggregate"), PlanImpl::StreamDistinct => write!(f, "StreamDistinct"), PlanImpl::ScalarApply => write!(f, "ScalarApply"), PlanImpl::MarkApply => write!(f, "MarkApply"), @@ -605,7 +607,7 @@ mod tests { covered_deserializers: None, cover_mapping: None, sort_elimination_hint: None, - stream_distinct_hint: None, + stream_aggregate_hint: None, } } @@ -653,6 +655,7 @@ mod tests { (PlanImpl::Dummy, "Dummy"), (PlanImpl::SimpleAggregate, "SimpleAggregate"), (PlanImpl::HashAggregate, "HashAggregate"), + (PlanImpl::StreamAggregate, "StreamAggregate"), (PlanImpl::StreamDistinct, "StreamDistinct"), (PlanImpl::ScalarApply, "ScalarApply"), (PlanImpl::MarkApply, "MarkApply"), @@ -736,6 +739,7 @@ mod tests { agg_calls: vec![column_expr(a, 0)], groupby_exprs: vec![column_expr(b, 1)], is_distinct: false, + force_spill: false, }); assert_eq!(referenced_columns(&aggregate, &mut arena)?, vec![a, b]); @@ -755,6 +759,7 @@ mod tests { let join = Operator::Join(JoinOperator { join_type: join::JoinType::Inner, + force_nested_loop: false, on: JoinCondition::On { on: vec![(column_expr(a, 0), column_expr(b, 1))], filter: Some(column_expr(c, 2)), @@ -838,7 +843,6 @@ mod tests { let sort = Operator::Sort(SortOperator { sort_fields: vec![SortField::from(column_expr(a, 0))], - limit: None, }); assert_eq!(referenced_columns(&sort, &mut arena)?, vec![a]); @@ -1093,11 +1097,10 @@ mod tests { let sort = Operator::Sort(SortOperator { sort_fields: vec![descending_nulls_first.clone(), ascending_nulls_last.clone()], - limit: Some(10), }); assert_eq!( sort.to_string(), - "Sort By 9 Desc Nulls First, 1 Asc Nulls Last, Limit 10" + "Sort By 9 Desc Nulls First, 1 Asc Nulls Last" ); let child = LogicalPlan::new(Operator::ShowTable, Childrens::None); diff --git a/src/planner/operator/sort.rs b/src/planner/operator/sort.rs index 2e51e5ff..a680591a 100644 --- a/src/planner/operator/sort.rs +++ b/src/planner/operator/sort.rs @@ -64,8 +64,6 @@ impl From for SortField { #[derive(Debug, PartialEq, Eq, Clone, Hash, ReferenceSerialization)] pub struct SortOperator { pub sort_fields: Vec, - /// Support push down limit to sort plan. - pub limit: Option, } impl fmt::Display for SortOperator { @@ -75,13 +73,7 @@ impl fmt::Display for SortOperator { .iter() .map(|sort_field| format!("{sort_field}")) .join(", "); - write!(f, "Sort By {sort_fields}")?; - - if let Some(limit) = self.limit { - write!(f, ", Limit {limit}")?; - } - - Ok(()) + write!(f, "Sort By {sort_fields}") } } diff --git a/src/planner/operator/table_scan.rs b/src/planner/operator/table_scan.rs index a337fea5..5ad7e552 100644 --- a/src/planner/operator/table_scan.rs +++ b/src/planner/operator/table_scan.rs @@ -74,7 +74,7 @@ impl TableScanOperator { covered_deserializers: None, cover_mapping: None, sort_elimination_hint: None, - stream_distinct_hint: None, + stream_aggregate_hint: None, }); } diff --git a/src/planner/operator/visitor.rs b/src/planner/operator/visitor.rs index 47d5e6cf..09b1cabe 100644 --- a/src/planner/operator/visitor.rs +++ b/src/planner/operator/visitor.rs @@ -391,7 +391,7 @@ pub(crate) mod tests { covered_deserializers: None, cover_mapping: None, sort_elimination_hint: None, - stream_distinct_hint: None, + stream_aggregate_hint: None, } } @@ -417,6 +417,7 @@ pub(crate) mod tests { groupby_exprs: vec![1_i32.into()], agg_calls: vec![2_i32.into()], is_distinct: false, + force_spill: false, }), Operator::ScalarApply(ScalarApplyOperator), Operator::MarkApply(mark_apply), @@ -431,6 +432,7 @@ pub(crate) mod tests { filter: Some(8_i32.into()), }, join_type: JoinType::Inner, + force_nested_loop: false, }), Operator::Project(ProjectOperator { exprs: vec![9_i32.into()], @@ -446,7 +448,6 @@ pub(crate) mod tests { Operator::FunctionScan(FunctionScanOperator { table_function }), Operator::Sort(SortOperator { sort_fields: vec![SortField::from(ScalarExpression::from(13_i32))], - limit: None, }), Operator::Limit(LimitOperator { offset: None, diff --git a/src/types/index.rs b/src/types/index.rs index 3c474c4c..27eb1432 100644 --- a/src/types/index.rs +++ b/src/types/index.rs @@ -77,7 +77,7 @@ pub struct IndexInfo { pub(crate) covered_deserializers: Option>, pub(crate) cover_mapping: Option>, pub(crate) sort_elimination_hint: Option, - pub(crate) stream_distinct_hint: Option, + pub(crate) stream_aggregate_hint: Option, } #[derive(Debug, Clone, Copy, Eq, PartialEq, Hash, ReferenceSerialization)] @@ -200,7 +200,7 @@ mod tests { covered_deserializers: None, cover_mapping: None, sort_elimination_hint: None, - stream_distinct_hint: None, + stream_aggregate_hint: None, } } @@ -363,7 +363,7 @@ mod tests { covered_deserializers: Some(vec![LogicalType::Integer.serializable()]), cover_mapping: Some(vec![0]), sort_elimination_hint: Some(IndexOrderHint::new(1)), - stream_distinct_hint: Some(IndexOrderHint::new(1)), + stream_aggregate_hint: Some(IndexOrderHint::new(1)), }; assert_eq!( diff --git a/tests/slt/group_by.slt b/tests/slt/group_by.slt index 6caec916..d7fdc691 100644 --- a/tests/slt/group_by.slt +++ b/tests/slt/group_by.slt @@ -58,4 +58,4 @@ select v2, v2 + 1, sum(v1) from t group by v2 + 1, v2 order by v2 # 7 statement ok -drop table t \ No newline at end of file +drop table t diff --git a/tests/slt/join.slt b/tests/slt/join.slt index 8a765660..439ab503 100644 --- a/tests/slt/join.slt +++ b/tests/slt/join.slt @@ -28,6 +28,11 @@ select a, b, c, d from x join y on a = c; 1 2 1 6 1 3 1 6 +query T +explain select /*+ FORCE_NEST_LOOP_JOIN */ a, b, c, d from x join y on a = c; +---- +Projection [#10, #11, #12, #13] [Project => (Sort Option: Follow)] Inner Join On #2 = #5 [NestLoopJoin => (Sort Option: None)] TableScan x -> [#2, #3] [SeqScan => (Sort Option: None)] TableScan y -> [#5, #6] [SeqScan => (Sort Option: None)] + query IIII select a, b, c, d, e, f from x join y on a = c and c < 5 join z on e = a and f = 5; ---- diff --git a/tests/slt/order_by.slt b/tests/slt/order_by.slt index fe7e75be..5ac306f5 100644 --- a/tests/slt/order_by.slt +++ b/tests/slt/order_by.slt @@ -25,6 +25,21 @@ select v1 from t order by v1 desc statement ok drop table t +statement ok +create table t(id int primary key, v1 int, v2 int) + +statement ok +copy t from 'tests/data/row_20000.csv' ( DELIMITER '|' ) + +query III +select * from t order by v1 % 10 desc, v2 desc offset 19997 +---- +69 70 71 +39 40 41 +9 10 11 + +statement ok +drop table t statement ok create table t(id int primary key, v1 int, v2 int) diff --git a/tests/slt/stream_distinct.slt b/tests/slt/stream_distinct.slt index 54def2c8..ff355d7f 100644 --- a/tests/slt/stream_distinct.slt +++ b/tests/slt/stream_distinct.slt @@ -24,6 +24,20 @@ select distinct c1 from distinct_t where c1 < 10 and c1 > 0; 8 9 +# stream aggregate +query II rowsort +select c1, count(c2) from distinct_t where c1 < 10 and c1 > 0 group by c1; +---- +1 10000 +2 10000 +3 10000 +4 10000 +5 10000 +6 10000 +7 10000 +8 10000 +9 10000 + statement ok drop index distinct_t.distinct_t_c1_index; @@ -41,5 +55,19 @@ select distinct c1 from distinct_t where c1 < 10 and c1 > 0; 8 9 +# forced spill aggregate +query II rowsort +select /*+ FORCE_AGG_SPILL */ c1, count(c2) from distinct_t where c1 < 10 and c1 > 0 group by c1; +---- +1 10000 +2 10000 +3 10000 +4 10000 +5 10000 +6 10000 +7 10000 +8 10000 +9 10000 + statement ok drop table distinct_t; diff --git a/tests/slt/stream_distinct_explain.slt b/tests/slt/stream_distinct_explain.slt index e7a960fb..5bbec0a9 100644 --- a/tests/slt/stream_distinct_explain.slt +++ b/tests/slt/stream_distinct_explain.slt @@ -16,6 +16,18 @@ explain select distinct c1 from distinct_t where c1 < 10 and c1 > 0; ---- Projection [#2] [Project => (Sort Option: Follow)] Aggregate [] -> Group By [#2] [StreamDistinct => (Sort Option: Follow)] TableScan distinct_t -> [#2] [IndexScan By #1 => (0, 10) Covered => (Sort Option: OrderBy: (#2 Asc Nulls Last) ignore_prefix_len: 0)] +# stream aggregate +query T +explain select c1, count(c2) from distinct_t where c1 < 10 and c1 > 0 group by c1; +---- +Projection [#2, #4] [Project => (Sort Option: Follow)] Aggregate [Count(#3)] -> Group By [#2] [StreamAggregate => (Sort Option: Follow)] TableScan distinct_t -> [#2, #3] [IndexScan By #1 => (0, 10) => (Sort Option: OrderBy: (#2 Asc Nulls Last) ignore_prefix_len: 0)] + +# forced spill reuses ordered input +query T +explain select /*+ FORCE_AGG_SPILL */ c1, count(c2) from distinct_t where c1 < 10 and c1 > 0 group by c1; +---- +Projection [#2, #4] [Project => (Sort Option: Follow)] Aggregate [Count(#3)] -> Group By [#2] [StreamAggregate => (Sort Option: Follow)] TableScan distinct_t -> [#2, #3] [IndexScan By #1 => (0, 10) => (Sort Option: OrderBy: (#2 Asc Nulls Last) ignore_prefix_len: 0)] + statement ok drop index distinct_t.distinct_t_c1_index; @@ -25,5 +37,17 @@ explain select distinct c1 from distinct_t where c1 < 10 and c1 > 0; ---- Projection [#2] [Project => (Sort Option: Follow)] Aggregate [] -> Group By [#2] [HashAggregate => (Sort Option: None)] Filter ((#2 < 10) && (#2 > 0)), Is Having: false [Filter => (Sort Option: Follow)] TableScan distinct_t -> [#2] [SeqScan => (Sort Option: None)] +# forced spill aggregate +query T +explain select /*+ FORCE_AGG_SPILL */ c1, count(c2) from distinct_t where c1 < 10 and c1 > 0 group by c1; +---- +Projection [#2, #4] [Project => (Sort Option: Follow)] Aggregate [Count(#3)] -> Group By [#2] [StreamAggregate => (Sort Option: Follow)] Sort By #2 Asc Nulls Last [Sort => (Sort Option: OrderBy: (#2 Asc Nulls Last) ignore_prefix_len: 0)] Filter ((#2 < 10) && (#2 > 0)), Is Having: false [Filter => (Sort Option: Follow)] TableScan distinct_t -> [#2, #3] [SeqScan => (Sort Option: None)] + +# forced spill sort also satisfies order by +query T +explain select /*+ FORCE_AGG_SPILL */ c1, count(c2) from distinct_t where c1 < 10 and c1 > 0 group by c1 order by c1; +---- +Projection [#2, #4] [Project => (Sort Option: Follow)] Aggregate [Count(#3)] -> Group By [#2] [StreamAggregate => (Sort Option: Follow)] Sort By #2 Asc Nulls Last [Sort => (Sort Option: OrderBy: (#2 Asc Nulls Last) ignore_prefix_len: 0)] Filter ((#2 < 10) && (#2 > 0)), Is Having: false [Filter => (Sort Option: Follow)] TableScan distinct_t -> [#2, #3] [SeqScan => (Sort Option: None)] + statement ok drop table distinct_t; diff --git a/tests/sqllogictest/Cargo.toml b/tests/sqllogictest/Cargo.toml index 675dbb06..a090bda6 100644 --- a/tests/sqllogictest/Cargo.toml +++ b/tests/sqllogictest/Cargo.toml @@ -5,7 +5,7 @@ edition = "2021" [dependencies] clap = { version = "4" } -"kite_sql" = { path = "../..", features = ["copy", "decimal", "orm"] } +"kite_sql" = { path = "../..", features = ["copy", "decimal", "orm", "spill"] } glob = { version = "0.3" } sqllogictest = { version = "0.14" } tempfile = { version = "3.10" }