Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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",
]

Expand Down
2 changes: 1 addition & 1 deletion Makefile
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
2 changes: 2 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
12 changes: 9 additions & 3 deletions benchmarks/query_benchmark.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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();
})
});

Expand Down
62 changes: 62 additions & 0 deletions docs/features.md
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,68 @@ See [the ORM guide](../src/orm/README.md) for the full ORM guide, including:
- public ORM structs, enums, and traits
- related `ResultIter::orm::<M>()` integration

### Spill-backed Aggregation: `features = ["spill"]`

The native-only `spill` feature bounds memory usage for large grouped
aggregations. It is not enabled by default and cannot be combined with `wasm`.

For SQL, place the optimizer hint immediately after `SELECT`:

```sql
SELECT /*+ FORCE_AGG_SPILL */ user_id, COUNT(*)
FROM orders
GROUP BY user_id
ORDER BY user_id;
```

For ORM queries, enable both `orm` and `spill`, then call `force_spill()` before
building the projection and aggregate plan:

```rust,ignore
let grouped = database.bind(|ctx| {
ctx.from::<Order>()?
.force_spill()?
.project_tuple(|e| {
Ok(vec![e.column(Order::user_id())?, e.count_all()?])
})?
.group_by(|e| e.column(Order::user_id()))?
.order_by(Order::user_id())?
.finish()
})?;
```

KiteSQL still reuses an already ordered input when possible. Otherwise it adds
an external sort on the group keys and executes a streaming aggregate. The
hint is intended for grouped aggregates or `DISTINCT`; an aggregate without
group keys already uses constant-size accumulator state.

### Nested-loop Join Hint

Use `FORCE_NEST_LOOP_JOIN` to select the nested-loop implementation for joins
in the current `SELECT` query block:

```sql
SELECT /*+ FORCE_NEST_LOOP_JOIN */ orders.id, users.name
FROM orders
JOIN users ON orders.user_id = users.id;
```

ORM queries can select the same implementation before adding joins:

```rust,ignore
let joined = database.bind(|ctx| {
ctx.from::<Order>()?
.force_nested_loop()
.inner_join::<User, _>(|e| {
e.column(Order::user_id())?.eq(e.column(User::id())?)
})?
.finish()
})?;
```

It avoids Hash Join's build-side tuple materialization, but may perform
substantially more work on large inputs.

### User-Defined Function: `features = ["macros"]`
```rust
scala_function!(TestFunction::test(LogicalType::Integer, LogicalType::Integer) -> LogicalType::Integer => |v1: DataValue, v2: DataValue| {
Expand Down
1 change: 1 addition & 0 deletions src/binder/aggregate.rs
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,7 @@ impl<T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<'_, '_, T, A>
agg_calls,
groupby_exprs,
false,
self.force_spill,
))
}

Expand Down
1 change: 1 addition & 0 deletions src/binder/distinct.rs
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@ impl<T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<'_, '_, T, A>
vec![],
select_list,
true,
self.force_spill,
))
}

Expand Down
4 changes: 4 additions & 0 deletions src/binder/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -561,6 +561,8 @@ impl<'a, T: Transaction> BinderContext<'a, T> {
pub struct Binder<'a, 'parent, T: Transaction, A: AsRef<[(&'static str, DataValue)]>> {
pub(crate) context: BinderContext<'a, T>,
pub(crate) args: &'a A,
pub(crate) force_spill: bool,
pub(crate) force_nested_loop: bool,
with_pk: Option<TableName>,
pub(crate) parent: Option<&'parent BinderContext<'a, T>>,
}
Expand All @@ -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,
}
Expand Down
97 changes: 84 additions & 13 deletions src/binder/parser.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1364,6 +1364,7 @@ where
self.binder.bind_table_ref_sql(from, self.arena)?,
JoinCondition::None,
JoinType::Cross,
self.binder.force_nested_loop,
)
}
plan
Expand Down Expand Up @@ -2594,6 +2595,7 @@ impl<'a, 'parent, T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<
arena: &mut PlanArena,
) -> Result<LogicalPlan, DatabaseError> {
let Select {
optimizer_hint,
projection,
from,
selection,
Expand All @@ -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
Expand Down Expand Up @@ -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> {
Expand Down
17 changes: 11 additions & 6 deletions src/binder/select.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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,
))
}

Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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)),
))
}
Expand Down
1 change: 0 additions & 1 deletion src/db.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1376,7 +1376,6 @@ pub(crate) mod test {
while iter
.next_tuple(|_, tuple| {
println!("{tuple:#?}");
()
})?
.is_some()
{}
Expand Down
4 changes: 2 additions & 2 deletions src/errors.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 [
Expand All @@ -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:"));
Expand Down
Loading
Loading