use std::collections::{BTreeMap, BTreeSet};
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex, MutexGuard, PoisonError, RwLock, RwLockReadGuard, RwLockWriteGuard};
use std::time::Instant;
use rudb_bind::{Bound, Parameters, Write};
use rudb_catalog::{Catalog, DEFAULT_CATALOG, Entry, QualifiedName, View};
use rudb_common::stat::Provenance;
use rudb_common::{
Cancel, Clustering, Error, Field, LogicalType, Memory, Result, Rule, Session, Value,
};
use rudb_io::{Filesystem, RealFilesystem};
use rudb_metrics::{Document, LoadProfile, Report, Span, Stage};
use rudb_native::graph::Edge;
use rudb_parse::ast::{self, Ast};
use rudb_pipeline::{Lease, Morsel, Pool, Progress, Sink, keep_pages};
use rudb_plan::{Expr, Node, Plan};
use rudb_vector::{Chunk, Data, Form, Selection, Vector};
use crate::config::Config;
use crate::connection::{Connection, single};
use crate::prepared::Prepared;
use crate::result::QueryResult;
use crate::settings::{COMPILED_ENGINE, Settings};
use crate::{foreign, upsert};
const MEMORY: &str = ":memory:";
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum NativeExtremaValues {
Integer { low: i128, high: i128 },
Date { low: i32, high: i32 },
}
fn native_simple_identifier(text: &str) -> bool {
let mut bytes = text.bytes();
matches!(bytes.next(), Some(b'a'..=b'z' | b'A'..=b'Z' | b'_'))
&& bytes.all(|byte| byte.is_ascii_alphanumeric() || byte == b'_')
}
fn native_simple_unquoted_identifier(text: &str) -> bool {
native_simple_identifier(text)
&& rudb_parse::classes(rudb_parse::lookup(text)) & rudb_parse::RESERVED == 0
}
fn native_simple_average_statement(sql: &str) -> Option<(&str, &str)> {
let statement = sql.trim();
let statement = statement.strip_suffix(';').unwrap_or(statement).trim_end();
let mut words = statement.split_ascii_whitespace();
let select = words.next()?;
let aggregate = words.next()?;
let from = words.next()?;
let table = words.next()?;
if words.next().is_some()
|| !select.eq_ignore_ascii_case("select")
|| !from.eq_ignore_ascii_case("from")
{
return None;
}
let prefix = aggregate.get(..4)?;
if !prefix.eq_ignore_ascii_case("avg(") || !aggregate.ends_with(')') {
return None;
}
let column = &aggregate[4..aggregate.len() - 1];
(native_simple_identifier(column) && native_simple_identifier(table)).then_some((table, column))
}
fn native_simple_nonzero_statement(sql: &str) -> Option<(&str, &str)> {
let statement = sql.trim();
let statement = statement.strip_suffix(';').unwrap_or(statement).trim_end();
let mut words = statement.split_ascii_whitespace();
let select = words.next()?;
let count = words.next()?;
let from = words.next()?;
let table = words.next()?;
let where_keyword = words.next()?;
let column = words.next()?;
let comparison = words.next()?;
let zero = words.next()?;
(words.next().is_none()
&& select.eq_ignore_ascii_case("select")
&& count.eq_ignore_ascii_case("count(*)")
&& from.eq_ignore_ascii_case("from")
&& where_keyword.eq_ignore_ascii_case("where")
&& comparison == "<>"
&& zero == "0"
&& native_simple_unquoted_identifier(table)
&& native_simple_unquoted_identifier(column))
.then_some((table, column))
}
fn native_simple_distinct_statement(sql: &str) -> Option<(&str, &str)> {
let statement = sql.trim();
let statement = statement.strip_suffix(';').unwrap_or(statement).trim_end();
let mut words = statement.split_ascii_whitespace();
let select = words.next()?;
let aggregate = words.next()?;
let column = words.next()?.strip_suffix(')')?;
let from = words.next()?;
let table = words.next()?;
if words.next().is_some()
|| !select.eq_ignore_ascii_case("select")
|| !aggregate.eq_ignore_ascii_case("count(distinct")
|| !from.eq_ignore_ascii_case("from")
|| !native_simple_identifier(column)
|| !native_simple_identifier(table)
{
return None;
}
Some((table, column))
}
fn native_simple_extrema_statement(sql: &str) -> Option<(&str, &str)> {
let statement = sql.trim();
let statement = statement.strip_suffix(';').unwrap_or(statement).trim_end();
let mut words = statement.split_ascii_whitespace();
let select = words.next()?;
let minimum = words.next()?;
let maximum = words.next()?;
let from = words.next()?;
let table = words.next()?;
if words.next().is_some()
|| !select.eq_ignore_ascii_case("select")
|| !from.eq_ignore_ascii_case("from")
|| !native_simple_identifier(table)
{
return None;
}
let min_prefix = minimum.get(..4)?;
let max_prefix = maximum.get(..4)?;
if !min_prefix.eq_ignore_ascii_case("min(")
|| !max_prefix.eq_ignore_ascii_case("max(")
|| !minimum.ends_with("),")
|| !maximum.ends_with(')')
{
return None;
}
let column = &minimum[4..minimum.len() - 2];
let other = &maximum[4..maximum.len() - 1];
(native_simple_identifier(column) && column.eq_ignore_ascii_case(other))
.then_some((table, column))
}
fn native_simple_three_statement(sql: &str) -> Option<(&str, &str, &str)> {
let statement = sql.trim();
let statement = statement.strip_suffix(';').unwrap_or(statement).trim_end();
let mut words = statement.split_ascii_whitespace();
let select = words.next()?;
let sum = words.next()?;
let count = words.next()?;
let average = words.next()?;
let from = words.next()?;
let table = words.next()?;
if words.next().is_some()
|| !select.eq_ignore_ascii_case("select")
|| !sum.get(..4)?.eq_ignore_ascii_case("sum(")
|| !count.eq_ignore_ascii_case("count(*),")
|| !average.get(..4)?.eq_ignore_ascii_case("avg(")
|| !from.eq_ignore_ascii_case("from")
|| !native_simple_unquoted_identifier(table)
{
return None;
}
let sum_column = sum.strip_suffix("),")?.get(4..)?;
let average_column = average.strip_suffix(')')?.get(4..)?;
(native_simple_unquoted_identifier(sum_column)
&& native_simple_unquoted_identifier(average_column))
.then_some((table, sum_column, average_column))
}
fn native_simple_group_count_statement(sql: &str) -> Option<(&str, &str)> {
let statement = sql.trim();
let statement = statement.strip_suffix(';').unwrap_or(statement).trim_end();
let mut words = statement.split_ascii_whitespace();
let select = words.next()?;
let key = words.next()?.strip_suffix(',')?;
let count = words.next()?;
let from = words.next()?;
let table = words.next()?;
let where_keyword = words.next()?;
let filtered = words.next()?;
let comparison = words.next()?;
let zero = words.next()?;
let group = words.next()?;
let group_by = words.next()?;
let grouped = words.next()?;
let order = words.next()?;
let order_by = words.next()?;
let order_count = words.next()?;
let descending = words.next()?;
(words.next().is_none()
&& select.eq_ignore_ascii_case("select")
&& count.eq_ignore_ascii_case("count(*)")
&& from.eq_ignore_ascii_case("from")
&& where_keyword.eq_ignore_ascii_case("where")
&& comparison == "<>"
&& zero == "0"
&& group.eq_ignore_ascii_case("group")
&& group_by.eq_ignore_ascii_case("by")
&& order.eq_ignore_ascii_case("order")
&& order_by.eq_ignore_ascii_case("by")
&& order_count.eq_ignore_ascii_case("count(*)")
&& descending.eq_ignore_ascii_case("desc")
&& native_simple_unquoted_identifier(table)
&& native_simple_unquoted_identifier(key)
&& filtered.eq_ignore_ascii_case(key)
&& grouped.eq_ignore_ascii_case(key))
.then_some((table, key))
}
fn native_nonzero_shape(ast: &Ast) -> Option<(&str, &str, &str)> {
use ast::{BinaryOp, Distinct, Expr, LiteralKind, QueryBody, Source, Statement};
use rudb_parse::NONE;
let [Statement::Query(query_ref)] = ast.statements.as_slice() else {
return None;
};
let query = ast.query(*query_ref);
if query.ctes.len != 0
|| query.order_by.len != 0
|| query.order_by_all
|| query.limit != NONE
|| query.offset != NONE
|| query.limit_percent
{
return None;
}
let QueryBody::Select(select_ref) = query.body else {
return None;
};
let select = ast.select(select_ref);
if select.distinct != Distinct::No
|| select.group_by.len != 0
|| select.group_by_all
|| select.having != NONE
{
return None;
}
let [target] = ast.target_list(select.targets) else {
return None;
};
let Expr::Function { name, args, distinct: false, filter: NONE } = ast.expr(target.expr) else {
return None;
};
if name.len != 1 {
return None;
}
let function = ast.name(name).next()?;
if !function.eq_ignore_ascii_case("count") {
return None;
}
let [arg] = ast.expr_list(args) else {
return None;
};
if !matches!(ast.expr(*arg), Expr::Star { qualifier, replacements } if qualifier.len == 0 && replacements.len == 0)
{
return None;
}
let [source] = ast.source_list(select.from) else {
return None;
};
let Source::Table { name, alias: NONE, columns } = ast.source(*source) else {
return None;
};
if columns.len != 0 {
return None;
}
if name.len != 1 {
return None;
}
let table = ast.name(name).next()?;
let Expr::Binary { op: BinaryOp::NotEq, left, right } = ast.expr(select.filter) else {
return None;
};
let (column, zero) = match (ast.expr(left), ast.expr(right)) {
(Expr::Column { name }, Expr::Literal { kind: LiteralKind::Number, text }) => (name, text),
(Expr::Literal { kind: LiteralKind::Number, text }, Expr::Column { name }) => (name, text),
_ => return None,
};
if ast.string(zero) != "0" {
return None;
}
if column.len != 1 {
return None;
}
let column = ast.name(column).next()?;
Some((
table,
column,
if target.alias == NONE { "count_star()" } else { ast.string(target.alias) },
))
}
fn native_column_aggregate<'a>(
ast: &'a Ast,
expr: ast::ExprRef,
function: &str,
) -> Option<&'a str> {
use rudb_parse::NONE;
let ast::Expr::Function { name, args, distinct: false, filter: NONE } = ast.expr(expr) else {
return None;
};
if name.len != 1 || !ast.name(name).next()?.eq_ignore_ascii_case(function) {
return None;
}
let [argument] = ast.expr_list(args) else { return None };
let ast::Expr::Column { name } = ast.expr(*argument) else { return None };
(name.len == 1).then(|| ast.name(name).next()).flatten()
}
fn native_three_aggregate_shape(ast: &Ast) -> Option<(&str, &str, &str, [String; 3])> {
use ast::{Distinct, Expr, QueryBody, Source, Statement};
use rudb_parse::NONE;
let [Statement::Query(query_ref)] = ast.statements.as_slice() else { return None };
let query = ast.query(*query_ref);
if query.ctes.len != 0
|| query.order_by.len != 0
|| query.order_by_all
|| query.limit != NONE
|| query.offset != NONE
|| query.limit_percent
{
return None;
}
let QueryBody::Select(select_ref) = query.body else { return None };
let select = ast.select(select_ref);
if select.distinct != Distinct::No
|| select.filter != NONE
|| select.group_by.len != 0
|| select.group_by_all
|| select.having != NONE
{
return None;
}
let [sum, count, average] = ast.target_list(select.targets) else { return None };
let sum_column = native_column_aggregate(ast, sum.expr, "sum")?;
let average_column = native_column_aggregate(ast, average.expr, "avg")?;
let Expr::Function { name, args, distinct: false, filter: NONE } = ast.expr(count.expr) else {
return None;
};
if name.len != 1 || !ast.name(name).next()?.eq_ignore_ascii_case("count") {
return None;
}
let [argument] = ast.expr_list(args) else { return None };
if !matches!(ast.expr(*argument), Expr::Star { qualifier, replacements } if qualifier.len == 0 && replacements.len == 0)
{
return None;
}
let [source] = ast.source_list(select.from) else { return None };
let Source::Table { name, alias: NONE, columns } = ast.source(*source) else { return None };
if name.len != 1 || columns.len != 0 {
return None;
}
let table = ast.name(name).next()?;
let names = [
if sum.alias == NONE { format!("sum({sum_column})") } else { ast.string(sum.alias).into() },
if count.alias == NONE { "count_star()".into() } else { ast.string(count.alias).into() },
if average.alias == NONE {
format!("avg({average_column})")
} else {
ast.string(average.alias).into()
},
];
Some((table, sum_column, average_column, names))
}
fn native_single_average_shape(ast: &Ast) -> Option<(&str, &str, String)> {
use ast::{Distinct, QueryBody, Source, Statement};
use rudb_parse::NONE;
let [Statement::Query(query_ref)] = ast.statements.as_slice() else { return None };
let query = ast.query(*query_ref);
if query.ctes.len != 0
|| query.order_by.len != 0
|| query.order_by_all
|| query.limit != NONE
|| query.offset != NONE
|| query.limit_percent
{
return None;
}
let QueryBody::Select(select_ref) = query.body else { return None };
let select = ast.select(select_ref);
if select.distinct != Distinct::No
|| select.filter != NONE
|| select.group_by.len != 0
|| select.group_by_all
|| select.having != NONE
{
return None;
}
let [target] = ast.target_list(select.targets) else { return None };
let column = native_column_aggregate(ast, target.expr, "avg")?;
let [source] = ast.source_list(select.from) else { return None };
let Source::Table { name, alias: NONE, columns } = ast.source(*source) else { return None };
if name.len != 1 || columns.len != 0 {
return None;
}
let table = ast.name(name).next()?;
let name = if target.alias == NONE {
format!("avg({column})")
} else {
ast.string(target.alias).into()
};
Some((table, column, name))
}
fn native_single_distinct_shape(ast: &Ast) -> Option<(&str, &str, String)> {
use ast::{Distinct, Expr, QueryBody, Source, Statement};
use rudb_parse::NONE;
let [Statement::Query(query_ref)] = ast.statements.as_slice() else { return None };
let query = ast.query(*query_ref);
if query.ctes.len != 0
|| query.order_by.len != 0
|| query.order_by_all
|| query.limit != NONE
|| query.offset != NONE
|| query.limit_percent
{
return None;
}
let QueryBody::Select(select_ref) = query.body else { return None };
let select = ast.select(select_ref);
if select.distinct != Distinct::No
|| select.filter != NONE
|| select.group_by.len != 0
|| select.group_by_all
|| select.having != NONE
{
return None;
}
let [target] = ast.target_list(select.targets) else { return None };
let Expr::Function { name, args, distinct: true, filter: NONE } = ast.expr(target.expr) else {
return None;
};
if name.len != 1 || !ast.name(name).next()?.eq_ignore_ascii_case("count") {
return None;
}
let [argument] = ast.expr_list(args) else { return None };
let Expr::Column { name } = ast.expr(*argument) else { return None };
if name.len != 1 {
return None;
}
let column = ast.name(name).next()?;
let [source] = ast.source_list(select.from) else { return None };
let Source::Table { name, alias: NONE, columns } = ast.source(*source) else { return None };
if name.len != 1 || columns.len != 0 {
return None;
}
let table = ast.name(name).next()?;
let name = if target.alias == NONE {
format!("count(DISTINCT {column})")
} else {
ast.string(target.alias).into()
};
Some((table, column, name))
}
fn native_extrema_shape(ast: &Ast) -> Option<(&str, &str, [String; 2])> {
use ast::{Distinct, QueryBody, Source, Statement};
use rudb_parse::NONE;
let [Statement::Query(query_ref)] = ast.statements.as_slice() else { return None };
let query = ast.query(*query_ref);
if query.ctes.len != 0
|| query.order_by.len != 0
|| query.order_by_all
|| query.limit != NONE
|| query.offset != NONE
|| query.limit_percent
{
return None;
}
let QueryBody::Select(select_ref) = query.body else { return None };
let select = ast.select(select_ref);
if select.distinct != Distinct::No
|| select.filter != NONE
|| select.group_by.len != 0
|| select.group_by_all
|| select.having != NONE
{
return None;
}
let [minimum, maximum] = ast.target_list(select.targets) else { return None };
let column = native_column_aggregate(ast, minimum.expr, "min")?;
let other = native_column_aggregate(ast, maximum.expr, "max")?;
if !column.eq_ignore_ascii_case(other) {
return None;
}
let [source] = ast.source_list(select.from) else { return None };
let Source::Table { name, alias: NONE, columns } = ast.source(*source) else { return None };
if name.len != 1 || columns.len != 0 {
return None;
}
let table = ast.name(name).next()?;
let names = [
if minimum.alias == NONE {
format!("min({column})")
} else {
ast.string(minimum.alias).into()
},
if maximum.alias == NONE {
format!("max({column})")
} else {
ast.string(maximum.alias).into()
},
];
Some((table, column, names))
}
fn native_integer_value(ty: &LogicalType, value: i128) -> Option<Value> {
Some(match ty {
LogicalType::TinyInt => Value::TinyInt(i8::try_from(value).ok()?),
LogicalType::SmallInt => Value::SmallInt(i16::try_from(value).ok()?),
LogicalType::Integer => Value::Integer(i32::try_from(value).ok()?),
LogicalType::BigInt => Value::BigInt(i64::try_from(value).ok()?),
LogicalType::UTinyInt => Value::UTinyInt(u8::try_from(value).ok()?),
LogicalType::USmallInt => Value::USmallInt(u16::try_from(value).ok()?),
LogicalType::UInteger => Value::UInteger(u32::try_from(value).ok()?),
LogicalType::UBigInt => Value::UBigInt(u64::try_from(value).ok()?),
LogicalType::Date => Value::Date(i32::try_from(value).ok()?),
_ => return None,
})
}
#[derive(Debug, Clone)]
pub struct Database {
shared: Shared,
}
#[derive(Debug, Clone)]
pub(crate) struct Shared {
inner: Arc<Inner>,
}
#[derive(Debug)]
struct Inner {
catalog: RwLock<Catalog>,
writer: Mutex<()>,
open: Mutex<Option<Open>>,
#[cfg(test)]
loading: Mutex<Option<std::sync::mpsc::Sender<()>>>,
path: Option<PathBuf>,
writable: bool,
settings: Settings,
memory: Memory,
pool: Pool,
pages: rudb_native::PagePool,
facts: Mutex<Arc<rudb_opt::estimate::Facts>>,
relationships: Mutex<(u64, String, Arc<Vec<rudb_opt::link::Linked>>)>,
declined: Mutex<BTreeSet<(String, bool)>>,
settings_revision: AtomicU64,
native_aggregate_plan: Mutex<Option<CachedNativeAggregate>>,
refusals: Mutex<Vec<String>>,
}
const REFUSALS_KEPT: usize = 1000;
#[derive(Debug)]
struct CachedNativeAggregate {
sql: String,
catalog_generation: u64,
settings_revision: u64,
plan: Arc<Plan>,
}
#[derive(Debug)]
struct Open {
before: Catalog,
aborted: bool,
read_only: bool,
}
impl Drop for Inner {
fn drop(&mut self) {
let catalog = self.catalog.get_mut().unwrap_or_else(PoisonError::into_inner);
if let Some(open) = self.open.get_mut().unwrap_or_else(PoisonError::into_inner).take() {
catalog.restore(open.before);
}
if let Some(path) = self.path.as_ref().filter(|_| self.writable) {
let _ = persist(path, catalog, &self.pages, DEFAULT_CATALOG);
}
for (name, path) in attached_files(catalog) {
let _ = persist(&path, catalog, &self.pages, &name);
}
}
}
impl Default for Database {
fn default() -> Self {
Self::new()
}
}
fn page_budget(limit: Option<u64>) -> usize {
limit.map_or(usize::MAX, |limit| usize::try_from(limit / 2).unwrap_or(usize::MAX))
}
fn dictionary_budget(limit: Option<u64>) -> u64 {
limit.map_or(u64::MAX, |limit| limit / 4)
}
fn runtime(config: &Config) -> Pool {
keep_pages();
Pool::new(config.threads())
}
impl Database {
pub fn query_native_nonzero_value_once(path: &str, sql: &str) -> Result<Option<i64>> {
let Some((table, column)) = native_simple_nonzero_statement(sql) else { return Ok(None) };
let native = rudb_native::Catalog::open(path)?;
let Some(table) = native.names().find(|name| name.eq_ignore_ascii_case(table)) else {
return Ok(None);
};
let Some(fields) = native.table_fields(table) else { return Ok(None) };
let Some(index) = fields.iter().position(|field| field.name.eq_ignore_ascii_case(column))
else {
return Ok(None);
};
if !matches!(
fields[index].ty,
LogicalType::TinyInt
| LogicalType::SmallInt
| LogicalType::Integer
| LogicalType::BigInt
| LogicalType::UTinyInt
| LogicalType::USmallInt
| LogicalType::UInteger
| LogicalType::UBigInt
) {
return Ok(None);
}
Ok(native.nonzero_count(table, index)?.and_then(|count| i64::try_from(count).ok()))
}
pub fn query_native_group_count_once(
path: &str,
sql: &str,
) -> Result<Option<Vec<(i128, i64)>>> {
let Some((table, column)) = native_simple_group_count_statement(sql) else {
return Ok(None);
};
let catalog = rudb_native::Catalog::open(path)?;
let Some(name) = catalog.names().find(|name| name.eq_ignore_ascii_case(table)) else {
return Ok(None);
};
let Some(fields) = catalog.table_fields(name) else { return Ok(None) };
let Some(index) = fields.iter().position(|field| field.name.eq_ignore_ascii_case(column))
else {
return Ok(None);
};
if !matches!(
fields[index].ty,
LogicalType::TinyInt
| LogicalType::SmallInt
| LogicalType::Integer
| LogicalType::BigInt
| LogicalType::UTinyInt
| LogicalType::USmallInt
| LogicalType::UInteger
| LogicalType::UBigInt
) {
return Ok(None);
}
let distinct = catalog.distinct_count(name, index)?;
if distinct.is_some_and(|count| count > 4096) {
return Ok(None);
}
let mut dense = match catalog.integer_extremes(name, index)? {
Some(rudb_native::IntegerExtremes::Values { low, high })
if (0..=4096).contains(&(high - low)) =>
{
Some((low, vec![0_u64; usize::try_from(high - low + 1).unwrap_or(0)]))
}
_ => None,
};
if distinct.is_none() && dense.is_none() {
return Ok(None);
}
let mut sparse = BTreeMap::<i128, u64>::new();
let folded = catalog.integer_fold(name, index, |value, count| {
if value == 0 {
return Ok(());
}
let value = i128::from(value);
if let Some((low, dense_counts)) = &mut dense {
let at = usize::try_from(value - *low).unwrap_or(usize::MAX);
if let Some(held) = dense_counts.get_mut(at) {
*held += count;
return Ok(());
}
}
*sparse.entry(value).or_default() += count;
Ok(())
})?;
if folded.is_none() {
let reader = catalog.table(name)?;
for part in 0..reader.parts() {
if let Some(counts) = reader.integer_tally(part, index)? {
for (value, count) in counts {
if value == 0 {
continue;
}
let value = i128::from(value);
if let Some((low, dense_counts)) = &mut dense {
let at = usize::try_from(value - *low).unwrap_or(usize::MAX);
if let Some(held) = dense_counts.get_mut(at) {
*held += count;
continue;
}
}
*sparse.entry(value).or_default() += count;
}
continue;
}
let chunk = reader.read(part, &[index])?;
let column = chunk
.into_columns()
.into_iter()
.next()
.ok_or_else(|| Error::internal("native column scan returned no column"))?
.into_flat()?;
let validity = column.validity();
macro_rules! count_values {
($values:expr) => {
if let Some((low, counts)) = &mut dense {
if column.none_null() {
for &value in $values.as_slice() {
let value = i128::from(value);
if value != 0 {
let at =
usize::try_from(value - *low).unwrap_or(usize::MAX);
if let Some(held) = counts.get_mut(at) {
*held += 1;
} else {
*sparse.entry(value).or_default() += 1;
}
}
}
} else {
for (row, &value) in $values.as_slice().iter().enumerate() {
let value = i128::from(value);
if value != 0 && validity.is_valid(row) {
let at =
usize::try_from(value - *low).unwrap_or(usize::MAX);
if let Some(held) = counts.get_mut(at) {
*held += 1;
} else {
*sparse.entry(value).or_default() += 1;
}
}
}
}
} else {
for (row, &value) in $values.as_slice().iter().enumerate() {
let value = i128::from(value);
if value != 0 && validity.is_valid(row) {
*sparse.entry(value).or_default() += 1;
}
}
}
};
}
match column.data() {
Some(Data::Int8(values)) => count_values!(values),
Some(Data::Int16(values)) => count_values!(values),
Some(Data::Int32(values)) => count_values!(values),
Some(Data::Int64(values)) => count_values!(values),
Some(Data::UInt8(values)) => count_values!(values),
Some(Data::UInt16(values)) => count_values!(values),
Some(Data::UInt32(values)) => count_values!(values),
Some(Data::UInt64(values)) => count_values!(values),
_ => return Ok(None),
}
}
}
if let Some((low, counts)) = dense {
for (at, count) in counts.into_iter().enumerate() {
if count != 0 {
sparse.insert(low + at as i128, count);
}
}
}
let mut groups = sparse
.into_iter()
.map(|(value, count)| i64::try_from(count).map(|count| (value, count)))
.collect::<std::result::Result<Vec<_>, _>>()
.map_err(|_| Error::internal("grouped count exceeds BIGINT"))?;
groups.sort_unstable_by(|(left_value, left_count), (right_value, right_count)| {
right_count.cmp(left_count).then_with(|| left_value.cmp(right_value))
});
Ok(Some(groups))
}
pub fn query_native_three_values_once(
path: &str,
sql: &str,
) -> Result<Option<(i128, i64, f64)>> {
let Some((table_name, sum_column, avg_column)) = native_simple_three_statement(sql) else {
return Ok(None);
};
let catalog = rudb_native::Catalog::open(path)?;
let Some(name) = catalog.names().find(|name| name.eq_ignore_ascii_case(table_name)) else {
return Ok(None);
};
let Some(fields) = catalog.table_fields(name) else { return Ok(None) };
let Some(sum_column) =
fields.iter().position(|field| field.name.eq_ignore_ascii_case(sum_column))
else {
return Ok(None);
};
let Some(avg_column) =
fields.iter().position(|field| field.name.eq_ignore_ascii_case(avg_column))
else {
return Ok(None);
};
let Some(sums) = catalog.aggregate_sums(name, &[sum_column, avg_column])? else {
return Ok(None);
};
let Ok(rows) = i64::try_from(sums.rows) else { return Ok(None) };
let (sum, sum_count) = sums.columns[0];
let (avg_sum, avg_count) = sums.columns[1];
if sum_count == 0 || avg_count == 0 {
return Ok(None);
}
let average = avg_sum as f64 / avg_count as f64;
Ok(Some((sum, rows, average)))
}
pub fn query_native_average_value_once(path: &str, sql: &str) -> Result<Option<f64>> {
let Some((table, column)) = native_simple_average_statement(sql) else {
return Ok(None);
};
let catalog = rudb_native::Catalog::open(path)?;
let Some(name) = catalog.names().find(|name| name.eq_ignore_ascii_case(table)) else {
return Ok(None);
};
let Some(fields) = catalog.table_fields(name) else { return Ok(None) };
let Some(index) = fields.iter().position(|field| field.name.eq_ignore_ascii_case(column))
else {
return Ok(None);
};
let Some(sums) = catalog.aggregate_sums(name, &[index])? else {
return Ok(None);
};
let (sum, count) = sums.columns[0];
if count == 0 {
return Ok(None);
}
Ok(Some(sum as f64 / count as f64))
}
pub fn query_native_distinct_value_once(path: &str, sql: &str) -> Result<Option<i64>> {
let Some((table, column)) = native_simple_distinct_statement(sql) else {
return Ok(None);
};
let catalog = rudb_native::Catalog::open(path)?;
let Some(name) = catalog.names().find(|name| name.eq_ignore_ascii_case(table)) else {
return Ok(None);
};
let Some(fields) = catalog.table_fields(name) else { return Ok(None) };
let Some(index) = fields.iter().position(|field| field.name.eq_ignore_ascii_case(column))
else {
return Ok(None);
};
let Some(count) = catalog.distinct_count(name, index)? else {
return Ok(None);
};
Ok(i64::try_from(count).ok())
}
pub fn query_native_extrema_values_once(
path: &str,
sql: &str,
) -> Result<Option<NativeExtremaValues>> {
let Some((table, column)) = native_simple_extrema_statement(sql) else {
return Ok(None);
};
let catalog = rudb_native::Catalog::open(path)?;
let Some(name) = catalog.names().find(|name| name.eq_ignore_ascii_case(table)) else {
return Ok(None);
};
let Some(fields) = catalog.table_fields(name) else { return Ok(None) };
let Some(index) = fields.iter().position(|field| field.name.eq_ignore_ascii_case(column))
else {
return Ok(None);
};
let Some(rudb_native::IntegerExtremes::Values { low, high }) =
catalog.integer_extremes(name, index)?
else {
return Ok(None);
};
let ty = &fields[index].ty;
if native_integer_value(ty, low).is_none() || native_integer_value(ty, high).is_none() {
return Ok(None);
}
Ok(Some(match ty {
LogicalType::Date => {
let (Ok(low), Ok(high)) = (i32::try_from(low), i32::try_from(high)) else {
return Ok(None);
};
NativeExtremaValues::Date { low, high }
}
_ => NativeExtremaValues::Integer { low, high },
}))
}
pub fn query_native_once(path: &str, sql: &str) -> Result<Option<QueryResult>> {
let ast = rudb_parse::parse_ast(sql)?;
let nonzero = native_nonzero_shape(&ast);
let three = native_three_aggregate_shape(&ast);
let average = native_single_average_shape(&ast);
let distinct = native_single_distinct_shape(&ast);
let extrema = native_extrema_shape(&ast);
let Some(table) = nonzero
.map(|(table, _, _)| table)
.or_else(|| three.as_ref().map(|(table, _, _, _)| *table))
.or_else(|| average.as_ref().map(|(table, _, _)| *table))
.or_else(|| distinct.as_ref().map(|(table, _, _)| *table))
.or_else(|| extrema.as_ref().map(|(table, _, _)| *table))
else {
return Ok(None);
};
let native = rudb_native::Catalog::open(path)?;
let Some((stored_name, fields)) = native
.names()
.find(|stored| stored.eq_ignore_ascii_case(table))
.and_then(|stored| native.table_fields(stored).map(|fields| (stored, fields)))
else {
return Ok(None);
};
if let Some((_, column, name)) = nonzero {
let Some(index) =
fields.iter().position(|field| field.name.eq_ignore_ascii_case(column))
else {
return Ok(None);
};
let Some(count) = native.nonzero_count(stored_name, index)? else {
return Ok(None);
};
let Ok(count) = i64::try_from(count) else {
return Ok(None);
};
let vector = Vector::from_values(LogicalType::BigInt, &[Value::BigInt(count)])?;
let chunk = Chunk::new(vec![vector])?;
return Ok(Some(QueryResult::new(
vec![name.to_string()],
vec![LogicalType::BigInt],
vec![chunk],
Memory::unlimited().reservation(),
)));
}
if let Some((_, column, name)) = average {
let Some(index) =
fields.iter().position(|field| field.name.eq_ignore_ascii_case(column))
else {
return Ok(None);
};
let Some(sums) = native.aggregate_sums(stored_name, &[index])? else {
return Ok(None);
};
let (sum, count) = sums.columns[0];
let value =
if count == 0 { Value::Null } else { Value::Double(sum as f64 / count as f64) };
let vector = Vector::from_values(LogicalType::Double, &[value])?;
let chunk = Chunk::new(vec![vector])?;
return Ok(Some(QueryResult::new(
vec![name],
vec![LogicalType::Double],
vec![chunk],
Memory::unlimited().reservation(),
)));
}
if let Some((_, column, name)) = distinct {
let Some(index) =
fields.iter().position(|field| field.name.eq_ignore_ascii_case(column))
else {
return Ok(None);
};
let Some(count) = native.distinct_count(stored_name, index)? else {
return Ok(None);
};
let Ok(count) = i64::try_from(count) else {
return Ok(None);
};
let vector = Vector::from_values(LogicalType::BigInt, &[Value::BigInt(count)])?;
let chunk = Chunk::new(vec![vector])?;
return Ok(Some(QueryResult::new(
vec![name],
vec![LogicalType::BigInt],
vec![chunk],
Memory::unlimited().reservation(),
)));
}
if let Some((_, column, names)) = extrema {
let Some(index) =
fields.iter().position(|field| field.name.eq_ignore_ascii_case(column))
else {
return Ok(None);
};
let Some(ends) = native.integer_extremes(stored_name, index)? else {
return Ok(None);
};
let ty = fields[index].ty.clone();
let (low, high) = match ends {
rudb_native::IntegerExtremes::Null => (Value::Null, Value::Null),
rudb_native::IntegerExtremes::Values { low, high } => {
let Some(low) = native_integer_value(&ty, low) else { return Ok(None) };
let Some(high) = native_integer_value(&ty, high) else { return Ok(None) };
(low, high)
}
};
let chunk = Chunk::new(vec![
Vector::from_values(ty.clone(), &[low])?,
Vector::from_values(ty.clone(), &[high])?,
])?;
return Ok(Some(QueryResult::new(
names.into(),
vec![ty.clone(), ty],
vec![chunk],
Memory::unlimited().reservation(),
)));
}
let Some((_, sum_column, avg_column, names)) = three else {
return Ok(None);
};
let Some(sum_index) =
fields.iter().position(|field| field.name.eq_ignore_ascii_case(sum_column))
else {
return Ok(None);
};
let Some(avg_index) =
fields.iter().position(|field| field.name.eq_ignore_ascii_case(avg_column))
else {
return Ok(None);
};
let Some(sums) = native.aggregate_sums(stored_name, &[sum_index, avg_index])? else {
return Ok(None);
};
let Ok(rows) = i64::try_from(sums.rows) else {
return Ok(None);
};
let (sum, sum_count) = sums.columns[0];
let (avg_sum, avg_count) = sums.columns[1];
let types = vec![LogicalType::HugeInt, LogicalType::BigInt, LogicalType::Double];
let values = [
if sum_count == 0 { Value::Null } else { Value::HugeInt(sum) },
Value::BigInt(rows),
if avg_count == 0 {
Value::Null
} else {
Value::Double(avg_sum as f64 / avg_count as f64)
},
];
let vectors = types
.iter()
.cloned()
.zip(values)
.map(|(ty, value)| Vector::from_values(ty, &[value]))
.collect::<Result<Vec<_>>>()?;
let chunk = Chunk::new(vectors)?;
Ok(Some(QueryResult::new(
names.into(),
types,
vec![chunk],
Memory::unlimited().reservation(),
)))
}
#[must_use]
pub fn new() -> Self {
Self::with_config(Config::default())
}
#[must_use]
pub fn with_config(config: Config) -> Self {
let memory = Memory::new(config.memory_limit());
let pool = runtime(&config);
let writable = !config.read_only();
let settings = Settings::new(config);
let inner = Inner {
catalog: RwLock::new(Catalog::new()),
writer: Mutex::default(),
open: Mutex::default(),
#[cfg(test)]
loading: Mutex::default(),
path: None,
writable,
settings,
memory,
pool,
pages: rudb_native::PagePool::default(),
facts: Mutex::default(),
relationships: Mutex::default(),
declined: Mutex::default(),
settings_revision: AtomicU64::new(0),
native_aggregate_plan: Mutex::default(),
refusals: Mutex::default(),
};
Self { shared: Shared { inner: Arc::new(inner) } }
}
#[must_use]
pub fn config(&self) -> Config {
self.shared.inner.settings.config()
}
#[must_use]
pub fn opened_with(&self) -> Config {
self.shared.inner.settings.defaults()
}
#[must_use]
pub fn refusals(&self) -> Vec<String> {
self.shared.inner.refusals.lock().unwrap_or_else(PoisonError::into_inner).clone()
}
pub fn setting(&self, name: &str) -> Result<String> {
if crate::settings::is_clustering(name) {
return Ok(self.shared.read().clustering());
}
match crate::settings::search_setting(name) {
Some(true) => return Ok(self.shared.read().default_schema().to_string()),
Some(false) => return Ok(self.shared.read().search_path()),
None => {}
}
self.shared.inner.settings.value(name)
}
#[must_use]
pub fn seams(&self) -> rudb_seam::Settings {
self.shared.inner.settings.seams()
}
pub fn seams_for(&self, sql: &str) -> Result<rudb_seam::Settings> {
self.shared.seams(sql)
}
#[must_use]
pub fn memory(&self) -> &Memory {
&self.shared.inner.memory
}
pub fn open(path: &str) -> Result<Self> {
Self::open_with(path, Config::default())
}
pub fn open_with(path: &str, config: Config) -> Result<Self> {
if path.is_empty() || path == MEMORY {
return Ok(Self::with_config(config));
}
let path = PathBuf::from(path);
let mut catalog = Catalog::new();
let pages = rudb_native::PagePool::new(page_budget(config.memory_limit()));
if path.exists() {
let native = rudb_native::Catalog::open_in(&path, &pages)?;
let names = native.names().map(str::to_string).collect::<Vec<_>>();
for name in names {
catalog.create_native_table(native.table(&name)?)?;
}
let views = native.views().cloned().collect::<Vec<_>>();
for view in &views {
catalog.create_native_view(view)?;
}
}
let memory = Memory::new(config.memory_limit());
let pool = runtime(&config);
let writable = !config.read_only();
let settings = Settings::new(config);
let inner = Inner {
catalog: RwLock::new(catalog),
writer: Mutex::default(),
open: Mutex::default(),
#[cfg(test)]
loading: Mutex::default(),
path: Some(path),
writable,
settings,
memory,
pool,
pages,
facts: Mutex::default(),
relationships: Mutex::default(),
declined: Mutex::default(),
settings_revision: AtomicU64::new(0),
native_aggregate_plan: Mutex::default(),
refusals: Mutex::default(),
};
Ok(Self { shared: Shared { inner: Arc::new(inner) } })
}
#[must_use]
pub fn connect(&self) -> Connection {
Connection::new(self.shared.clone())
}
pub fn close(self) -> Result<()> {
let Some(path) = self.shared.inner.path.as_ref().filter(|_| self.shared.inner.writable)
else {
return Ok(());
};
let path = path.clone();
let _writing = self.shared.writing();
let mut catalog = self.shared.write();
persist(&path, &mut catalog, &self.shared.inner.pages, DEFAULT_CATALOG)?;
for (name, path) in attached_files(&catalog) {
persist(&path, &mut catalog, &self.shared.inner.pages, &name)?;
}
Ok(())
}
pub fn prepare(&self, sql: &str) -> Result<Prepared> {
Prepared::new(self.shared.clone(), sql).map_err(|error| self.shared.process_error(error))
}
pub fn with_catalog<T>(&self, read: impl FnOnce(&Catalog) -> T) -> T {
read(&self.shared.read())
}
pub fn with_catalog_mut<T>(&self, write: impl FnOnce(&mut Catalog) -> T) -> T {
let _writing = self.shared.writing();
write(&mut self.shared.write())
}
pub fn create_table(&self, name: &str, columns: Vec<Field>) -> Result<()> {
let parts: Vec<&str> = name.split('.').collect();
let _writing = self.shared.writing();
let mut catalog = self.shared.write();
let resolved = catalog.resolve_for_create(&parts)?;
catalog.create_table(resolved, columns)
}
pub fn drop_table(&self, name: &str) -> Result<()> {
let parts: Vec<&str> = name.split('.').collect();
let _writing = self.shared.writing();
let mut catalog = self.shared.write();
let resolved = catalog.resolve(&parts)?;
catalog.drop_table(&resolved)
}
pub fn append(&self, name: &str, rows: &[Vec<Value>]) -> Result<()> {
let parts: Vec<&str> = name.split('.').collect();
let _writing = self.shared.writing();
let mut catalog = self.shared.write();
let resolved = catalog.resolve(&parts)?;
let table = catalog.table_mut(&resolved)?;
let types: Vec<LogicalType> =
table.columns().iter().map(|field| field.ty.clone()).collect();
let fits = |value: &Value, ty: &LogicalType| value.is_null() || &value.logical_type() == ty;
if rows
.iter()
.all(|row| row.len() == types.len() && row.iter().zip(&types).all(|(v, t)| fits(v, t)))
{
return table.append_rows(rows);
}
let mut converted = Vec::with_capacity(rows.len());
for (index, row) in rows.iter().enumerate() {
if row.len() != types.len() {
return Err(Error::invalid_input(format!(
"row {index} has {} values and the table has {} columns",
row.len(),
types.len()
)));
}
let row = row
.iter()
.zip(&types)
.map(|(value, ty)| rudb_kernels::cast::cast_value(value, ty, false))
.collect::<Result<Vec<Value>>>()?;
converted.push(row);
}
table.append_rows(&converted)
}
pub fn table_len(&self, name: &str) -> Result<usize> {
let parts: Vec<&str> = name.split('.').collect();
let catalog = self.shared.read();
let resolved = catalog.resolve(&parts)?;
Ok(catalog.table(&resolved)?.rows().len())
}
#[must_use]
pub fn table_names(&self) -> Vec<String> {
self.shared.read().tables().map(|table| table.name().table.clone()).collect()
}
pub fn table_sql(&self, name: &str) -> Result<String> {
let parts: Vec<&str> = name.split('.').collect();
let catalog = self.shared.read();
let resolved = catalog.resolve(&parts)?;
let table = catalog.table(&resolved)?;
let columns: Vec<String> = table
.columns()
.iter()
.map(|field| {
let null = if field.not_null { " NOT NULL" } else { "" };
format!("{} {}{null}", field.name, field.ty)
})
.collect();
Ok(format!("CREATE TABLE {}({});", resolved.table, columns.join(", ")))
}
pub fn query(&self, sql: &str) -> Result<QueryResult> {
self.shared
.query(sql, &self.shared.token())
.map_err(|error| self.shared.process_error(error))
}
pub fn execute(&self, sql: &str) -> Result<QueryResult> {
self.shared
.execute(sql, &self.shared.token())
.map_err(|error| self.shared.process_error(error))
}
pub fn plan(&self, sql: &str) -> Result<String> {
self.shared.plan(sql).map_err(|error| self.shared.process_error(error))
}
pub fn value(&self, sql: &str) -> Result<Value> {
single(&self.query(sql)?)
}
}
fn persist(
path: &Path,
catalog: &mut Catalog,
pages: &rudb_native::PagePool,
database: &str,
) -> Result<()> {
let names =
catalog.stored_tables_in(database).map(|table| table.name().clone()).collect::<Vec<_>>();
let views = views(catalog, database);
let clean = catalog
.stored_tables_in(database)
.all(|table| table.rows().is_native() && table.clustering_is_stored());
let held = committed(path)?;
if clean
&& held
.as_ref()
.is_some_and(|held| held.tables == wanted(&names) && same_views(&held.views, &views))
{
return Ok(());
}
if names.is_empty() {
let temporary = scratch(path)?;
rudb_native::Writer::empty(&temporary, &views)?;
return rename(&temporary, path);
}
if clean && held.is_some_and(|held| held.tables == wanted(&names)) {
rudb_native::Writer::restate(path, &views)?;
return rebind(path, catalog, &names, pages);
}
if appended(path, catalog, &names, &views)? {
return rebind(path, catalog, &names, pages);
}
let temporary = scratch(path)?;
let mut writer: Option<rudb_native::Writer> = None;
for name in &names {
let table = catalog.table(name)?;
let fields = table.columns().to_vec();
let columns = (0..fields.len()).collect::<Vec<_>>();
let mut open = match writer.take() {
None => rudb_native::Writer::create(&temporary, name.table.clone(), fields)?,
Some(writer) => writer.next(name.table.clone(), fields)?,
};
if let Some(clustering) = table.clustering() {
open = open.declare(clustering.clone())?;
}
let constraints = table.stored_constraints()?;
if !constraints.is_empty() {
open = open.constrain(constraints)?;
}
for at in 0..table.rows().chunk_count() {
open.append(&table.rows().read(at, &columns)?)?;
}
writer = Some(open);
}
let writer = writer.ok_or_else(|| Error::internal("a catalog with tables wrote none"))?;
writer.with_views(views).finish()?;
rename(&temporary, path)?;
rebind(path, catalog, &names, pages)
}
fn scratch(path: &Path) -> Result<PathBuf> {
let temporary = path.with_extension(format!("{}.tmp", std::process::id()));
if temporary.exists() {
std::fs::remove_file(&temporary).map_err(|error| Error::io(error.to_string()))?;
}
Ok(temporary)
}
fn rename(temporary: &Path, path: &Path) -> Result<()> {
publish(&RealFilesystem::new(), temporary, path)
}
pub(crate) fn publish(fs: &dyn Filesystem, temporary: &Path, path: &Path) -> Result<()> {
fs.rename(temporary, path)?;
fs.sync_dir(directory_of(path))
}
fn directory_of(path: &Path) -> &Path {
match path.parent() {
Some(parent) if !parent.as_os_str().is_empty() => parent,
_ => Path::new("."),
}
}
fn committed(path: &Path) -> Result<Option<Held>> {
if !path.exists() {
return Ok(None);
}
let held = rudb_native::Catalog::open(path)?;
Ok(Some(Held {
tables: held.names().map(str::to_string).collect(),
views: held.views().cloned().collect(),
}))
}
struct Held {
tables: BTreeSet<String>,
views: Vec<rudb_native::ViewEntry>,
}
fn held_rows(path: &Path) -> Result<Option<BTreeMap<String, usize>>> {
if !path.exists() {
return Ok(None);
}
let held = rudb_native::Catalog::open(path)?;
Ok(Some(held.rows().map(|(name, rows)| (name.to_string(), rows)).collect()))
}
fn wanted(names: &[QualifiedName]) -> BTreeSet<String> {
names.iter().map(|name| name.table.clone()).collect()
}
fn views(catalog: &Catalog, database: &str) -> Vec<rudb_native::ViewEntry> {
catalog
.stored_views_in(database)
.map(|view| rudb_native::ViewEntry {
name: view.name().table.clone(),
sql: view.sql().to_string(),
statement: view.statement().to_string(),
aliases: view.aliases().to_vec(),
columns: view.columns(),
})
.collect()
}
fn written_database(bound: &Bound) -> Option<(&'static str, &str)> {
match bound {
Bound::CreateTable(create) => Some(("CREATE", &create.name.catalog)),
Bound::CreateView(create) => Some(("CREATE", &create.name.catalog)),
Bound::DropTable(drop) => Some(("DROP", &drop.names.first()?.catalog)),
Bound::Insert(insert) => Some((
match insert.write {
Write::Append => "INSERT",
Write::Update => "UPDATE",
Write::Delete => "DELETE",
},
&insert.name.catalog,
)),
Bound::Alter(alter) => Some(("ALTER", &alter.name.as_ref()?.catalog)),
_ => None,
}
}
fn attach_database(
attach: rudb_bind::Attach,
catalog: &mut Catalog,
pages: &rudb_native::PagePool,
) -> Result<()> {
let memory = attach.path.is_empty() || attach.path == MEMORY;
let mut access = None;
for (name, value) in &attach.options {
let name = name.to_ascii_lowercase();
match name.as_str() {
"read_only" | "readonly" => {
access = Some(value.as_ref().map_or(Ok(true), option_flag)?)
}
"read_write" | "readwrite" => {
access = Some(!value.as_ref().map_or(Ok(true), option_flag)?);
}
"block_size" => {
let size = value.as_ref().map(Value::to_string).unwrap_or_default();
let size = size.parse::<u64>().map_err(|_| {
Error::invalid_input(format!("the block size must be a number, got {size}"))
})?;
check_block_size(size)?;
}
"type" => {
let kind = value.as_ref().map(Value::to_string).unwrap_or_default();
if !kind.eq_ignore_ascii_case("duckdb") {
return Err(Error::not_implemented(format!(
"ATTACH of a database of type {kind}, since only native files are read so far"
)));
}
}
"storage_version"
| "row_group_size"
| "compress"
| "io_mode"
| "mmap_reserve_size"
| "recovery_mode"
| "vacuum_rebuild_indexes"
| "hidden" => {}
"encryption_key" | "encryption_cipher" | "default_table" => {
return Err(Error::not_implemented(format!("the {name} option of ATTACH")));
}
_ => return Err(Error::binder(format!("Unrecognized option for attach \"{name}\""))),
}
}
let name = match attach.alias {
Some(alias) => alias,
None if memory => DEFAULT_CATALOG.to_string(),
None => Path::new(&attach.path)
.file_stem()
.map_or_else(|| attach.path.clone(), |stem| stem.to_string_lossy().into_owned()),
};
let read_only = access.unwrap_or(false);
if let Some(held) = catalog.attached(&name) {
if (attach.if_not_exists || attach.or_replace)
&& access.is_some_and(|asked| asked != held.read_only())
{
let mode = |read_only: bool| if read_only { "READ_ONLY" } else { "READ_WRITE" };
return Err(Error::binder(format!(
"Database \"{name}\" is already attached in {} mode, cannot re-attach in {} mode",
mode(held.read_only()),
mode(read_only)
)));
}
if attach.if_not_exists {
return Ok(());
}
let place = held.path().unwrap_or(MEMORY);
let asked = if memory { MEMORY } else { attach.path.as_str() };
if attach.or_replace && place == asked {
return Ok(());
}
if attach.or_replace && !held.internal() {
let file = held.path().filter(|_| !held.read_only()).map(PathBuf::from);
let held = held.name().to_string();
if let Some(path) = file {
persist(&path, catalog, pages, &held)?;
}
catalog.detach(&held)?;
}
}
if memory {
if read_only {
return Err(Error::catalog("Cannot launch in-memory database in read-only mode!"));
}
return catalog.attach_file(&name, None, false);
}
let path = PathBuf::from(&attach.path);
let same = |held: &str| {
let held = Path::new(held);
held == path
|| std::fs::canonicalize(held)
.ok()
.zip(std::fs::canonicalize(&path).ok())
.is_some_and(|(held, path)| held == path)
};
if let Some(holder) = catalog.databases().iter().find(|held| held.path().is_some_and(same)) {
return Err(Error::resource_in_use(format!(
"Unique file handle conflict: Cannot attach \"{name}\" - the database file \"{}\" is \
already attached by database \"{}\"",
attach.path,
holder.name()
)));
}
if !path.exists() {
if read_only {
return Err(Error::io(format!(
"Cannot open database \"{}\" in read-only mode: database does not exist",
attach.path
)));
}
rudb_native::Writer::empty(&path, &[])?;
}
let native = rudb_native::Catalog::open_in(&path, pages)?;
catalog.attach_file(&name, Some(attach.path.clone()), read_only)?;
let tables = native.names().map(str::to_string).collect::<Vec<_>>();
for table in tables {
catalog.create_native_table_in(&name, native.table(&table)?)?;
}
for view in native.views().cloned().collect::<Vec<_>>() {
catalog.create_native_view_in(&name, &view)?;
}
Ok(())
}
fn several(sql: &str) -> Option<Vec<&str>> {
if !sql.contains(';') {
return None;
}
let found = crate::statements(sql).ok()?;
(found.len() > 1).then(|| found.iter().map(crate::Statement::sql).collect())
}
fn check_block_size(size: u64) -> Result<()> {
const SMALLEST: u64 = 16384;
const LARGEST: u64 = 262_144;
if !size.is_power_of_two() {
return Err(Error::invalid_input(format!(
"the block size must be a power of two, got {size}"
)));
}
if size < SMALLEST {
return Err(Error::invalid_input(format!(
"the block size must be greater or equal than the minimum block size of {SMALLEST}, \
got {size}"
)));
}
if size > LARGEST {
return Err(Error::invalid_input(format!(
"the block size must be lesser or equal than the maximum block size of {LARGEST}, \
got {size}"
)));
}
Ok(())
}
fn option_flag(value: &Value) -> Result<bool> {
match value {
Value::Boolean(flag) => Ok(*flag),
Value::Null => Ok(false),
other => match other.to_string().to_ascii_lowercase().as_str() {
"true" | "1" | "on" => Ok(true),
"false" | "0" | "off" => Ok(false),
text => Err(Error::binder(format!("Could not read \"{text}\" as a boolean"))),
},
}
}
fn holds_a_file(inner: &Inner, catalog: &Catalog, database: &str) -> bool {
if database.eq_ignore_ascii_case(DEFAULT_CATALOG) {
return inner.path.is_some() && inner.writable;
}
catalog.attached(database).is_some_and(|held| held.path().is_some() && !held.read_only())
}
fn attached_files(catalog: &Catalog) -> Vec<(String, PathBuf)> {
catalog
.databases()
.iter()
.filter(|database| !database.read_only())
.filter_map(|database| Some((database.name().to_string(), database.path()?.into())))
.collect()
}
fn same_views(held: &[rudb_native::ViewEntry], wanted: &[rudb_native::ViewEntry]) -> bool {
held.len() == wanted.len()
&& held.iter().zip(wanted).all(|(held, wanted)| {
held.name == wanted.name
&& held.sql == wanted.sql
&& held.statement == wanted.statement
&& held.aliases == wanted.aliases
})
}
fn index(
path: &Path,
catalog: &mut Catalog,
links: &str,
pages: &rudb_native::PagePool,
) -> Result<()> {
let declared = rudb_exec::declared(catalog, links);
if declared.is_empty() {
return Ok(());
}
let mut wanted: Vec<(QualifiedName, Vec<usize>)> = Vec::new();
for link in &declared {
let [column] = &link.parent.columns[..] else { continue };
let Some(table) = catalog
.tables()
.find(|table| table.name().table.eq_ignore_ascii_case(&link.parent.table))
else {
continue;
};
let Some(at) = table.column_index(column) else { continue };
let name = table.name().clone();
match wanted.iter_mut().find(|(held, _)| held == &name) {
Some((_, columns)) if columns.contains(&at) => {}
Some((_, columns)) => columns.push(at),
None => wanted.push((name, vec![at])),
}
}
if wanted.is_empty() {
return Ok(());
}
for (name, columns) in &wanted {
rudb_native::graph::build_key_maps(path, &name.table, columns)?;
}
let edges = edges_of(catalog, &declared);
let mut names = wanted.into_iter().map(|(name, _)| name).collect::<Vec<_>>();
if !edges.is_empty() {
rudb_native::graph::build_links(path, &edges)?;
for edge in &edges {
let Some(name) = catalog
.tables()
.find(|table| table.name().table == edge.child)
.map(|table| table.name().clone())
else {
continue;
};
if !names.contains(&name) {
names.push(name);
}
}
}
rebind(path, catalog, &names, pages)
}
fn sketch(path: &Path, catalog: &mut Catalog, pages: &rudb_native::PagePool) -> Result<()> {
let native = rudb_native::Catalog::open_in(path, pages)?;
let mut stale = Vec::new();
for name in native.names() {
if !rudb_native::grams::current(&native.table(name)?) {
stale.push(name.to_string());
}
}
drop(native);
let mut names = Vec::new();
for table in &stale {
rudb_native::grams::build_text_grams(path, table)?;
if let Some(name) = catalog
.tables()
.find(|held| held.name().table == *table)
.map(|held| held.name().clone())
{
names.push(name);
}
}
if names.is_empty() {
return Ok(());
}
rebind(path, catalog, &names, pages)
}
fn edges_of(catalog: &Catalog, declared: &[rudb_graph::Relationship]) -> Vec<Edge> {
let mut edges = Vec::new();
for link in declared {
let find = |name: &str| {
catalog.tables().find(|table| table.name().table.eq_ignore_ascii_case(name))
};
let (Some(child_table), Some(parent_table)) =
(find(&link.child.table), find(&link.parent.table))
else {
continue;
};
let (Some(child_column), Some(parent_column)) =
(key_in(child_table, &link.child.columns), key_in(parent_table, &link.parent.columns))
else {
continue;
};
edges.push(Edge {
child: child_table.name().table.clone(),
child_column,
parent: parent_table.name().table.clone(),
parent_column,
});
}
edges
}
fn key_in(table: &rudb_catalog::Table, columns: &[String]) -> Option<usize> {
let at = columns.iter().map(|column| table.column_index(column)).collect::<Option<Vec<_>>>()?;
rudb_native::graph::key_of(&at)
}
fn rebind(
path: &Path,
catalog: &mut Catalog,
names: &[QualifiedName],
pages: &rudb_native::PagePool,
) -> Result<()> {
let native = rudb_native::Catalog::open_in(path, pages)?;
for name in names {
let reader = native.table(&name.table)?;
catalog.table_mut(name)?.rebind_native(reader)?;
}
Ok(())
}
fn appended(
path: &Path,
catalog: &mut Catalog,
names: &[QualifiedName],
views: &[rudb_native::ViewEntry],
) -> Result<bool> {
let Some(held) = held_rows(path)? else { return Ok(false) };
if held.is_empty() {
return Ok(false);
}
let native = names
.iter()
.filter(|name| catalog.table(name).is_ok_and(|table| table.rows().is_native()))
.map(|name| name.table.clone())
.collect::<BTreeSet<_>>();
let dirty =
names.iter().filter(|name| !native.contains(&name.table)).cloned().collect::<Vec<_>>();
let replaced = |name: &String| held.get(name).is_some_and(|rows| *rows == 0);
let carried = held.keys().all(|name| {
native.contains(name) || (replaced(name) && dirty.iter().any(|d| &d.table == name))
});
if !carried || !native.iter().all(|name| held.contains_key(name)) {
return Ok(false);
}
let mut writer: Option<rudb_native::Writer> = None;
for name in &dirty {
let table = catalog.table(name)?;
let fields = table.columns().to_vec();
let columns = (0..fields.len()).collect::<Vec<_>>();
let mut open = match writer.take() {
None => rudb_native::Writer::open(path, name.table.clone(), fields)?,
Some(writer) => writer.next(name.table.clone(), fields)?,
};
if let Some(clustering) = table.clustering() {
open = open.declare(clustering.clone())?;
}
let constraints = table.stored_constraints()?;
if !constraints.is_empty() {
open = open.constrain(constraints)?;
}
for at in 0..table.rows().chunk_count() {
open.append(&table.rows().read(at, &columns)?)?;
}
writer = Some(open);
}
let Some(writer) = writer else { return Ok(false) };
writer.with_views(views.to_vec()).finish()?;
Ok(true)
}
fn split(
chunks: Vec<Chunk>,
delete: bool,
wanted: bool,
) -> Result<(Vec<Chunk>, Vec<Chunk>, usize)> {
let mut kept = Vec::with_capacity(chunks.len());
let mut changed = Vec::new();
let mut count = 0;
for chunk in chunks {
let width = chunk.width().saturating_sub(1);
let hit = Selection::from_predicate(chunk.len(), |row| {
chunk.value_at(row, width) == Value::Boolean(true)
});
count += hit.len();
let columns: Vec<usize> = (0..width).collect();
let chunk = chunk.project(&columns)?;
if wanted && !hit.is_empty() {
changed.push(chunk.clone().select(&hit)?);
}
if delete {
let rest = hit.complement(chunk.len());
if !rest.is_empty() {
kept.push(chunk.select(&rest)?);
}
} else {
kept.push(chunk);
}
}
Ok((kept, changed, count))
}
fn appendable(path: &Path, catalog: &Catalog, target: &QualifiedName) -> Result<bool> {
let Some(held) = held_rows(path)? else { return Ok(false) };
if held.get(&target.table).is_some_and(|rows| *rows > 0) {
return Ok(false);
}
let carried =
held.keys().filter(|name| *name != &target.table).cloned().collect::<BTreeSet<_>>();
let others = catalog.stored_tables().filter(|table| table.name() != target).count();
let native = catalog
.stored_tables()
.filter(|table| table.name() != target && table.rows().is_native())
.map(|table| table.name().table.clone())
.collect::<BTreeSet<_>>();
Ok(carried == native && others == native.len())
}
#[derive(Debug, Default)]
struct NativePlace {
morsel: u64,
chunk: u64,
held: Vec<((u64, u64), Chunk)>,
holding: u64,
building: Option<rudb_native::Building>,
fed: usize,
started: Option<Span>,
inside_wall: u64,
inside_cpu: u64,
rows: u64,
bytes: u64,
}
impl NativePlace {
fn start(&mut self) {
if self.started.is_none() {
self.started = Some(Span::start());
}
}
}
fn declared(
writer: rudb_native::Writer,
clustering: Option<Clustering>,
) -> Result<rudb_native::Writer> {
match clustering {
None => Ok(writer),
Some(clustering) => writer.declare(clustering),
}
}
const GATHER_ROWS: usize = 131_072;
const FEED_PARTS: usize = 8;
#[derive(Debug)]
struct NativeSink {
writer: Mutex<Option<rudb_native::Writer>>,
preparer: rudb_native::Preparer,
merger: rudb_native::Merger,
temporary: Option<PathBuf>,
target: PathBuf,
table: String,
fields: Vec<Field>,
profile: Arc<LoadProfile>,
}
impl NativeSink {
fn keep(&self, chunk: Chunk, place: &mut NativePlace) -> Result<Progress> {
for (at, field) in self.fields.iter().enumerate().filter(|(_, field)| field.not_null) {
let vector = chunk.column(at)?;
let null = match vector.form() {
Form::Dictionary | Form::Rle => (0..vector.len()).any(|row| vector.is_null_at(row)),
_ => vector.validity().has_nulls(vector.len()),
};
if null {
return Err(Error::constraint(format!(
"NOT NULL constraint failed: {}.{}",
self.table, field.name
)));
}
}
place.start();
place.rows = place.rows.saturating_add(chunk.len() as u64);
let footprint = chunk.footprint() as u64;
place.bytes = place.bytes.saturating_add(footprint);
place.holding = place.holding.saturating_add(footprint);
self.profile.hold(footprint);
place.held.push(((place.morsel, place.chunk), chunk));
place.chunk = place.chunk.saturating_add(1);
if place.held.len() == FEED_PARTS {
self.feed(place)?;
}
if place.fed.saturating_add(place.held.len()) == rudb_native::STRIPE_PARTS {
self.hand_over(place)?;
}
Ok(Progress::More)
}
fn create(
target: &Path,
name: String,
fields: Vec<Field>,
clustering: Option<Clustering>,
limit: Option<u64>,
) -> Result<Self> {
let temporary = target.with_extension(format!("{}.tmp", std::process::id()));
if temporary.exists() {
std::fs::remove_file(&temporary).map_err(|error| Error::io(error.to_string()))?;
}
let profile = LoadProfile::begin(name.clone());
let writer = rudb_native::Writer::create(&temporary, name.clone(), fields.clone())?
.with_profile(Arc::clone(&profile))
.with_dictionary_cap(dictionary_budget(limit));
let mut writer = declared(writer, clustering)?;
Ok(Self {
preparer: writer.preparer(),
merger: writer.merger()?,
writer: Mutex::new(Some(writer)),
temporary: Some(temporary),
target: target.to_path_buf(),
table: name,
fields,
profile,
})
}
fn open(
target: &Path,
name: String,
fields: Vec<Field>,
clustering: Option<Clustering>,
limit: Option<u64>,
) -> Result<Self> {
let profile = LoadProfile::begin(name.clone());
let writer = rudb_native::Writer::open(target, name.clone(), fields.clone())?
.with_profile(Arc::clone(&profile))
.with_dictionary_cap(dictionary_budget(limit));
let mut writer = declared(writer, clustering)?;
Ok(Self {
preparer: writer.preparer(),
merger: writer.merger()?,
writer: Mutex::new(Some(writer)),
temporary: None,
target: target.to_path_buf(),
table: name,
fields,
profile,
})
}
fn locked<T>(&self, work: impl FnOnce(&mut rudb_native::Writer) -> Result<T>) -> Result<T> {
let waiting = Instant::now();
let mut writer =
self.writer.lock().map_err(|_| Error::internal("native writer panicked"))?;
self.profile.waited(Stage::Write, elapsed_ns(waiting));
work(
writer
.as_mut()
.ok_or_else(|| Error::internal("native writer was already committed"))?,
)
}
fn feed(&self, place: &mut NativePlace) -> Result<()> {
if place.held.is_empty() {
return Ok(());
}
let parts = std::mem::take(&mut place.held);
let holding = std::mem::take(&mut place.holding);
place.fed = place.fed.saturating_add(parts.len());
let inside = Span::start();
let building = place.building.get_or_insert_with(|| self.preparer.start());
let fed = self.preparer.feed(building, parts);
self.profile.release(holding);
let (wall, cpu) = inside.stop();
place.inside_wall = place.inside_wall.saturating_add(wall);
place.inside_cpu = place.inside_cpu.saturating_add(cpu);
fed
}
fn hand_over(&self, place: &mut NativePlace) -> Result<()> {
self.feed(place)?;
place.fed = 0;
let Some(building) = place.building.take() else { return Ok(()) };
let inside = Span::start();
let appended = self.preparer.finish(building).and_then(|prepared| {
let merged = self.merger.merge(prepared)?;
let mut paged = merged.pages()?;
self.merger.give_back(&mut paged)?;
self.locked(|writer| writer.write(paged))
});
let (wall, cpu) = inside.stop();
place.inside_wall = place.inside_wall.saturating_add(wall);
place.inside_cpu = place.inside_cpu.saturating_add(cpu);
rudb_common::heap::release();
appended
}
}
impl Sink for NativeSink {
type Local = NativePlace;
fn parallel(&self) -> bool {
true
}
fn local(&self) -> Self::Local {
NativePlace::default()
}
fn gather(&self) -> usize {
GATHER_ROWS
}
fn at(&self, morsel: &Morsel, place: &mut Self::Local) -> Result<()> {
place.start();
self.hand_over(place)?;
place.morsel = morsel.index();
place.chunk = 0;
Ok(())
}
fn sink(&self, chunk: &Chunk, place: &mut Self::Local) -> Result<Progress> {
self.keep(chunk.clone(), place)
}
fn sink_taking(&self, chunk: &mut Chunk, place: &mut Self::Local) -> Result<Progress> {
self.keep(std::mem::replace(chunk, Chunk::empty(&[])), place)
}
fn combine(&self, mut local: Self::Local) -> Result<()> {
let handed = self.hand_over(&mut local);
if let Some(started) = local.started.take() {
let (wall, cpu) = started.stop();
self.profile.charge(
Stage::Convert,
wall.saturating_sub(local.inside_wall),
cpu.saturating_sub(local.inside_cpu),
);
self.profile.moved(Stage::Convert, 0, local.bytes, local.rows);
}
handed
}
fn finalize(&self, _threads: &Lease<'_>) -> Result<()> {
let writer = self
.writer
.lock()
.map_err(|_| Error::internal("native writer panicked"))?
.take()
.ok_or_else(|| Error::internal("native writer was already committed"))?;
writer.finish()?;
let Some(temporary) = &self.temporary else {
self.profile.finish();
return Ok(());
};
let renamed = {
let _timing = self.profile.span(Stage::Publish);
publish(&RealFilesystem::new(), temporary, &self.target)
};
self.profile.finish();
renamed
}
}
impl Drop for NativeSink {
fn drop(&mut self) {
self.profile.finish();
}
}
fn elapsed_ns(since: Instant) -> u64 {
u64::try_from(since.elapsed().as_nanos()).unwrap_or(u64::MAX)
}
impl Shared {
pub(crate) fn process_error(&self, error: Error) -> Error {
if self.session().semantics().errors_as_json() { error.into_json() } else { error }
}
fn read(&self) -> RwLockReadGuard<'_, Catalog> {
self.inner.catalog.read().unwrap_or_else(PoisonError::into_inner)
}
fn write(&self) -> RwLockWriteGuard<'_, Catalog> {
self.inner.catalog.write().unwrap_or_else(PoisonError::into_inner)
}
fn writing(&self) -> MutexGuard<'_, ()> {
self.inner.writer.lock().unwrap_or_else(PoisonError::into_inner)
}
fn refresh_device_card(&self) -> Result<()> {
let Some(path) = self.inner.path.as_ref().filter(|_| self.inner.writable) else {
return Err(Error::invalid_input(
"PRAGMA device_card_refresh measures the device a database file is on, and this \
database has no file it can write",
));
};
let dir = path.parent().filter(|dir| !dir.as_os_str().is_empty());
let dir = dir.unwrap_or(Path::new("."));
rudb_io::device::card(dir, Some(rudb_io::device::Options::default().iterations))?;
if path.exists() {
rudb_native::Writer::keep_device_card(path)?;
}
Ok(())
}
fn budget(&self) -> Budget<'_> {
Budget { memory: &self.inner.memory, pool: &self.inner.pool }
}
pub(crate) fn session(&self) -> Session {
self.inner.settings.session()
}
pub(crate) fn query(&self, sql: &str, cancel: &Cancel) -> Result<QueryResult> {
if let Some(script) = several(sql) {
return self.run_script(&script, cancel, Self::query);
}
self.in_transaction(sql, || {
if let Some(answer) = self.cached_native_aggregate(sql, cancel)? {
return Ok(answer);
}
self.query_mirrored(sql, cancel, true)
})
}
fn in_transaction(
&self,
sql: &str,
run: impl FnOnce() -> Result<QueryResult>,
) -> Result<QueryResult> {
let aborted = self.open().as_ref().is_some_and(|open| open.aborted);
if aborted && crate::syntax::statement_kind(sql) != Some("TransactionStatement") {
return Err(Error::transaction("Current transaction is aborted (please ROLLBACK)"));
}
let result = run();
let aborts = result.as_ref().err().is_some_and(|error| {
!matches!(
error.code(),
rudb_common::ErrorCode::Parser | rudb_common::ErrorCode::NotImplemented
)
});
if aborts && let Some(open) = self.open().as_mut() {
open.aborted = true;
}
result
}
fn open(&self) -> MutexGuard<'_, Option<Open>> {
self.inner.open.lock().unwrap_or_else(PoisonError::into_inner)
}
fn transaction(&self, kind: ast::Transaction, catalog: &mut Catalog) -> Result<QueryResult> {
let mut open = self.open();
match kind {
ast::Transaction::Begin { read_only } => {
if open.is_some() {
return Err(Error::transaction(
"cannot start a transaction within a transaction",
));
}
*open = Some(Open { before: catalog.clone(), aborted: false, read_only });
}
ast::Transaction::Commit => {
let Some(closed) = open.take() else {
return Err(Error::transaction("cannot commit - no transaction is active"));
};
if closed.aborted {
catalog.restore(closed.before);
}
}
ast::Transaction::Rollback => {
let Some(closed) = open.take() else {
return Err(Error::transaction("cannot rollback - no transaction is active"));
};
catalog.restore(closed.before);
}
}
Ok(QueryResult::empty())
}
fn transacting(&self) -> bool {
self.open().is_some()
}
fn cached_native_aggregate(&self, sql: &str, cancel: &Cancel) -> Result<Option<QueryResult>> {
let catalog = self.read();
let revision = self.inner.settings_revision.load(Ordering::Relaxed);
let cached = self
.inner
.native_aggregate_plan
.lock()
.unwrap_or_else(PoisonError::into_inner)
.as_ref()
.filter(|cached| {
cached.sql == sql
&& cached.catalog_generation == catalog.generation()
&& cached.settings_revision == revision
})
.map(|cached| Arc::clone(&cached.plan));
let Some(plan) = cached else { return Ok(None) };
let seams = self.seams(sql)?;
let context = self.optimizer(&catalog)?;
let session = self.session();
let under = Under::new(self.budget(), context.facts(), &seams, &session, Rows::ForACaller);
run(sql, &plan, &catalog, cancel, under).map(Some)
}
fn refused(&self, sql: &str, refusal: &rudb_qc::Refusal) {
let mut log = self.inner.refusals.lock().unwrap_or_else(PoisonError::into_inner);
if log.len() == REFUSALS_KEPT {
log.remove(0);
}
let sql = sql.split_whitespace().collect::<Vec<_>>().join(" ");
log.push(format!("{refusal} | {sql}"));
}
fn remember_native_aggregate(&self, sql: &str, ast: &Ast, plan: &Plan, catalog: &Catalog) {
if !is_native_summary_aggregate(ast, plan, catalog) {
return;
}
*self.inner.native_aggregate_plan.lock().unwrap_or_else(PoisonError::into_inner) =
Some(CachedNativeAggregate {
sql: sql.to_string(),
catalog_generation: catalog.generation(),
settings_revision: self.inner.settings_revision.load(Ordering::Relaxed),
plan: Arc::new(plan.clone()),
});
}
fn query_mirrored(&self, sql: &str, cancel: &Cancel, mirror: bool) -> Result<QueryResult> {
let catalog = self.read();
let seams = self.seams(sql)?;
let context = self.optimizer(&catalog)?;
let session = self.session();
let (ast, parse_ns) =
timed(|| rudb_parse::parse_ast_with_case(sql, session.semantics().identifier_case()))?;
let outlined = mirror && self.inner.settings.config().parquet_mirror();
let (bound, bind_ns) = timed(|| {
if outlined {
rudb_bind::bind_statement_outlined(&ast, &catalog, &Parameters::new(), &session)
} else {
rudb_bind::bind_statement_with(&ast, &catalog, &Parameters::new(), &session)
}
})?;
if mirror {
let wanted = self.wanted_mirrors(&bound);
if !wanted.is_empty() || (outlined && asked_for_mirrors(&bound)) {
drop(bound);
drop(catalog);
self.mirror(&wanted);
return self.query_mirrored(sql, cancel, false);
}
}
match bound {
Bound::Query(mut plan) => {
let ((), optimize_ns) = timed(|| rudb_opt::optimize_with(&mut plan, &context))?;
self.remember_native_aggregate(sql, &ast, &plan, &catalog);
let budget = self.budget();
let under = Under::new(budget, context.facts(), &seams, &session, Rows::ForACaller)
.after(Planning { parse_ns, bind_ns, optimize_ns });
if self.inner.settings.engine() == COMPILED_ENGINE {
match rudb_qc::compile(&plan, cancel) {
Ok(compiled) => {
return run_compiled(sql, &plan, &catalog, cancel, compiled, under);
}
Err(refusal) => self.refused(sql, &refusal),
}
}
run(sql, &plan, &catalog, cancel, under)
}
Bound::Explain { mut plan, analyze, statistics, codegen } => {
let ((), optimize_ns) = timed(|| rudb_opt::optimize_with(&mut plan, &context))?;
if codegen {
return explained_codegen(&plan, cancel);
}
let seams = rudb_opt::explain::Seams::new(&seams, rudb_exec::registries());
explaining(
&plan,
&catalog,
cancel,
self.budget(),
&context,
seams,
&session,
Asked { analyze, statistics },
sql,
Planning { parse_ns, bind_ns, optimize_ns },
)
}
_ => Err(Error::not_implemented("a statement that is not a query, on the query path")),
}
}
fn wanted_mirrors(&self, bound: &Bound) -> Vec<(String, bool)> {
let Bound::Query(plan) = bound else { return Vec::new() };
if plan.wanted_mirrors().is_empty() {
return Vec::new();
}
let config = self.inner.settings.config();
if !config.parquet_mirror() {
return Vec::new();
}
let declined = self.inner.declined.lock().unwrap_or_else(PoisonError::into_inner);
plan.wanted_mirrors()
.iter()
.filter(|(_, _, rows)| *rows >= config.mirror_rows())
.map(|(path, binary_as_string, _)| (path.clone(), *binary_as_string))
.filter(|wanted| !declined.contains(wanted))
.collect()
}
fn mirror(&self, wanted: &[(String, bool)]) {
let config = self.inner.settings.config();
for (path, binary_as_string) in wanted {
let added = crate::mirror::ensure(path, *binary_as_string, config).and_then(|found| {
let Some((stamp, reader)) = found else { return Ok(false) };
self.write().add_mirror(path, *binary_as_string, stamp, reader)?;
Ok(true)
});
if !matches!(added, Ok(true)) {
self.inner
.declined
.lock()
.unwrap_or_else(PoisonError::into_inner)
.insert((path.clone(), *binary_as_string));
}
}
}
pub(crate) fn seams(&self, sql: &str) -> Result<rudb_seam::Settings> {
let mut seams = self.inner.settings.seams();
for hint in rudb_parse::hints(sql)? {
seams.hint(hint)?;
}
Ok(seams)
}
fn optimizer(&self, catalog: &Catalog) -> Result<rudb_opt::pass::Context> {
let mut context =
rudb_opt::pass::Context::without(&self.inner.settings.disabled_optimizers())?;
context.measure(self.estimates(catalog));
context.relate(self.relationships(catalog));
context.size(self.inner.settings.sizes());
context.govern(self.inner.settings.rules());
Ok(context)
}
fn relationships(&self, catalog: &Catalog) -> Arc<Vec<rudb_opt::link::Linked>> {
if !self.inner.settings.rules().enabled(Rule::GraphSections) {
return Arc::default();
}
let declared = self.inner.settings.links();
if declared.is_empty() && catalog.tables().all(|table| table.foreign().is_empty()) {
return Arc::default();
}
let generation = catalog.generation();
let mut held = match self.inner.relationships.lock() {
Ok(held) => held,
Err(poisoned) => poisoned.into_inner(),
};
if held.0 != generation || held.1 != declared {
let found = Arc::new(Self::related(catalog, &declared));
*held = (generation, declared, found);
}
Arc::clone(&held.2)
}
fn related(catalog: &Catalog, declared: &str) -> Vec<rudb_opt::link::Linked> {
let mut found = Vec::new();
for link in rudb_exec::declared(catalog, declared) {
let (child_keys, parent_keys) = (&link.child.columns, &link.parent.columns);
let (Some(child_key), Some(parent_key)) = (child_keys.first(), parent_keys.first())
else {
continue;
};
let second = match (&child_keys[1..], &parent_keys[1..]) {
([], []) => None,
([child], [parent]) => Some((child, parent)),
_ => continue,
};
let built = stored_link(
catalog,
(&link.child.table, child_keys),
(&link.parent.table, parent_keys),
);
let sides = (&link.child.table, child_key, &link.parent.table, parent_key);
let linked = match built {
Some(head) if head.linked == head.children => {
rudb_opt::link::Linked::verified(sides.0, sides.1, sides.2, sides.3)
}
Some(_) => rudb_opt::link::Linked::built(sides.0, sides.1, sides.2, sides.3),
None => rudb_opt::link::Linked::declared(sides.0, sides.1, sides.2, sides.3),
};
let linked = match built {
Some(head) if head.form == rudb_graph::link::Form::Monotone => linked.monotone(),
_ => linked,
};
found.push(match second {
Some((child, parent)) => linked.and(child, parent),
None if parent_keyed(catalog, &link.parent.table, parent_keys) => linked.keyed(),
None => linked,
});
}
found
}
fn facts(&self, catalog: &Catalog) -> Arc<rudb_opt::estimate::Facts> {
let generation = catalog.generation();
let mut held = match self.inner.facts.lock() {
Ok(held) => held,
Err(poisoned) => poisoned.into_inner(),
};
if held.generation() != generation {
*held = Arc::new(Self::measured(catalog, generation));
}
Arc::clone(&held)
}
fn estimates(&self, catalog: &Catalog) -> Arc<rudb_opt::estimate::Facts> {
let held = self.facts(catalog);
if self.inner.settings.rules().enabled(Rule::StatsAll) {
return held;
}
Arc::new(held.without_distincts())
}
fn measured(catalog: &Catalog, generation: u64) -> rudb_opt::estimate::Facts {
let mut facts = rudb_opt::estimate::Facts::at(generation);
for table in catalog.tables().chain(catalog.mirrored_tables()) {
let name = table.name();
let rows = u64::try_from(table.rows().len()).unwrap_or(u64::MAX);
facts.record(&name.catalog, &name.schema, &name.table, rows);
let provenance = match table.rows() {
rudb_catalog::table::Rows::Memory(_) => Provenance::Sketch,
rudb_catalog::table::Rows::Native(_) => Provenance::Dictionary,
rudb_catalog::table::Rows::Grown(_, _) => Provenance::Sketch,
};
for (at, column) in table.columns().iter().enumerate() {
let Ok(Some(distinct)) = table.rows().distinct_values(at) else {
continue;
};
facts.record_distinct(
&name.catalog,
&name.schema,
&name.table,
&column.name,
distinct,
provenance,
);
}
}
facts
}
pub(crate) fn timeout(&self) -> Option<std::time::Duration> {
self.inner.settings.config().query_timeout()
}
pub(crate) fn token(&self) -> Cancel {
match self.inner.settings.config().query_timeout() {
Some(timeout) => Cancel::after(timeout),
None => Cancel::new(),
}
}
pub(crate) fn plan(&self, sql: &str) -> Result<String> {
let catalog = self.read();
let context = self.optimizer(&catalog)?;
Ok(planned(sql, &catalog, &context, &self.session())?.to_string())
}
pub(crate) fn execute(&self, sql: &str, cancel: &Cancel) -> Result<QueryResult> {
if let Some(script) = several(sql) {
return self.run_script(&script, cancel, Self::execute);
}
self.in_transaction(sql, || {
if let Some(answer) = self.cached_native_aggregate(sql, cancel)? {
return Ok(answer);
}
let session = self.session();
let (ast, parse_ns) = timed(|| {
rudb_parse::parse_ast_with_case(sql, session.semantics().identifier_case())
})?;
self.execute_ast(&ast, sql, &Parameters::new(), cancel, parse_ns)
})
}
fn run_script(
&self,
script: &[&str],
cancel: &Cancel,
last: impl Fn(&Self, &str, &Cancel) -> Result<QueryResult>,
) -> Result<QueryResult> {
let Some((final_statement, before)) = script.split_last() else {
return Err(Error::binder("no statement to bind"));
};
for statement in before {
self.execute(statement, cancel)?;
}
last(self, final_statement, cancel)
}
pub(crate) fn execute_ast(
&self,
ast: &Ast,
sql: &str,
parameters: &Parameters,
cancel: &Cancel,
parse_ns: u64,
) -> Result<QueryResult> {
let _writing = self.writing();
self.execute_mirrored(ast, sql, parameters, cancel, parse_ns, true)
}
fn execute_mirrored(
&self,
ast: &Ast,
sql: &str,
parameters: &Parameters,
cancel: &Cancel,
parse_ns: u64,
mirror: bool,
) -> Result<QueryResult> {
let seams = self.seams(sql)?;
let mut catalog = self.write();
let context = self.optimizer(&catalog)?;
let session = self.session();
let outlined = mirror && self.inner.settings.config().parquet_mirror();
let (bound, bind_ns) = timed(|| {
if outlined {
rudb_bind::bind_statement_outlined(ast, &catalog, parameters, &session)
} else {
rudb_bind::bind_statement_with(ast, &catalog, parameters, &session)
}
})?;
if mirror {
let wanted = self.wanted_mirrors(&bound);
if !wanted.is_empty() || (outlined && asked_for_mirrors(&bound)) {
drop(bound);
drop(catalog);
self.mirror(&wanted);
return self.execute_mirrored(ast, sql, parameters, cancel, parse_ns, false);
}
}
let writes = matches!(
bound,
Bound::CreateTable(_)
| Bound::CreateView(_)
| Bound::DropTable(_)
| Bound::Schema(_)
| Bound::Sequence(_)
| Bound::Type(_)
| Bound::Alter(_)
| Bound::Index(_)
| Bound::Insert(_)
);
if let Some((kind, database)) = written_database(&bound)
&& let Some(held) = catalog.attached(database).filter(|held| held.read_only())
{
return Err(Error::invalid_input(format!(
"Cannot execute statement of type \"{kind}\" on database \"{}\" which is \
attached in read-only mode!",
held.name()
)));
}
if writes && self.open().as_ref().is_some_and(|open| open.read_only) {
return Err(Error::transaction(format!(
"Cannot write to database \"\"{}\"\" - transaction is launched in read-only mode",
catalog.default_catalog()
)));
}
match bound {
Bound::Query(mut plan) => {
let ((), optimize_ns) = timed(|| rudb_opt::optimize_with(&mut plan, &context))?;
if parameters.is_empty() {
self.remember_native_aggregate(sql, ast, &plan, &catalog);
}
let budget = self.budget();
let under = Under::new(budget, context.facts(), &seams, &session, Rows::ForACaller)
.after(Planning { parse_ns, bind_ns, optimize_ns });
run(sql, &plan, &catalog, cancel, under)
}
Bound::Explain { mut plan, analyze, statistics, codegen } => {
let ((), optimize_ns) = timed(|| rudb_opt::optimize_with(&mut plan, &context))?;
if codegen {
return explained_codegen(&plan, cancel);
}
let seams = rudb_opt::explain::Seams::new(&seams, rudb_exec::registries());
explaining(
&plan,
&catalog,
cancel,
self.budget(),
&context,
seams,
&session,
Asked { analyze, statistics },
sql,
Planning { parse_ns, bind_ns, optimize_ns },
)
}
Bound::Setting(setting)
if setting.pragma && setting.name.eq_ignore_ascii_case("device_card_refresh") =>
{
self.refresh_device_card()?;
Ok(QueryResult::empty())
}
Bound::Setting(setting) if setting.pragma => {
self.inner.settings.toggle(&setting.name)?;
self.inner.settings_revision.fetch_add(1, Ordering::Relaxed);
Ok(QueryResult::empty())
}
Bound::Setting(setting) => {
let value = setting.value.as_ref();
self.inner.settings.apply(
&self.inner.memory,
&self.inner.pool,
&mut catalog,
&setting.name,
setting.scope,
value,
)?;
self.inner.settings_revision.fetch_add(1, Ordering::Relaxed);
Ok(QueryResult::empty())
}
Bound::Transaction(kind) => self.transaction(kind, &mut catalog),
Bound::Checkpoint(name) => {
if let Some(name) = name.filter(|name| !name.eq_ignore_ascii_case(DEFAULT_CATALOG))
{
let Some(database) = catalog.attached(&name).filter(|held| !held.internal())
else {
return Err(Error::binder(format!("Database \"{name}\" not found")));
};
let file = database.path().filter(|_| !database.read_only()).map(PathBuf::from);
let name = database.name().to_string();
if let Some(path) = file {
persist(&path, &mut catalog, &self.inner.pages, &name)?;
}
return Ok(QueryResult::empty());
}
if let Some(path) = self.inner.path.as_ref().filter(|_| self.inner.writable) {
persist(path, &mut catalog, &self.inner.pages, DEFAULT_CATALOG)?;
index(path, &mut catalog, &self.inner.settings.links(), &self.inner.pages)?;
sketch(path, &mut catalog, &self.inner.pages)?;
}
Ok(QueryResult::empty())
}
Bound::Attach(attach) => {
attach_database(attach, &mut catalog, &self.inner.pages)?;
Ok(QueryResult::empty())
}
Bound::Detach { name, if_exists } => {
let Some(database) = catalog.attached(&name).filter(|held| !held.internal()) else {
if if_exists {
return Ok(QueryResult::empty());
}
catalog.detach(&name)?;
return Ok(QueryResult::empty());
};
let file = database.path().filter(|_| !database.read_only()).map(PathBuf::from);
let name = database.name().to_string();
if let Some(path) =
file.filter(|_| !name.eq_ignore_ascii_case(catalog.default_catalog()))
{
persist(&path, &mut catalog, &self.inner.pages, &name)?;
}
catalog.detach(&name)?;
Ok(QueryResult::empty())
}
Bound::CreateTable(mut create) => {
let writable = self.inner.writable
&& create.name.catalog.eq_ignore_ascii_case(DEFAULT_CATALOG)
&& !self.transacting();
if let Some(path) = self.inner.path.as_ref().filter(|_| writable) {
let fresh = create.source.is_some() && catalog.table(&create.name).is_err();
let alone = fresh && !path.exists() && catalog.stored_tables().count() == 0;
if fresh && (alone || appendable(path, &catalog, &create.name)?) {
let plan = create.source.as_mut().expect("a source, asked for above");
rudb_opt::optimize_with(plan, &context)?;
let table = create.name.table.clone();
let fields = create.columns.clone();
let limit = self.inner.memory.limit();
let sink = Arc::new(if alone {
NativeSink::create(path, table.clone(), fields, None, limit)?
} else {
NativeSink::open(path, table.clone(), fields, None, limit)?
});
drop(catalog);
let reading = self.read();
#[cfg(test)]
if let Some(told) = self.inner.loading.lock().unwrap().take() {
let _ = told.send(());
}
let query = rudb_exec::build_measured_into(
plan,
&reading,
cancel,
&self.inner.memory,
&seams,
&session,
sink,
)?;
query.run(cancel, &self.inner.pool)?;
drop(query);
drop(reading);
let reader = rudb_native::Catalog::open(path)?.table(&table)?;
let mut catalog = self.write();
catalog.create_table(create.name.clone(), create.columns)?;
catalog.table_mut(&create.name)?.commit_native(reader)?;
return Ok(QueryResult::empty());
}
}
create_table(
sql,
create,
&mut catalog,
cancel,
self.budget(),
&context,
&seams,
&session,
)?;
Ok(QueryResult::empty())
}
Bound::CreateView(create) => {
create_view(create, &mut catalog)?;
Ok(QueryResult::empty())
}
Bound::DropTable(drop) => {
for name in &drop.names {
match drop.kind {
Entry::Table => catalog.drop_table(name)?,
Entry::View => catalog.drop_view(name)?,
}
}
Ok(QueryResult::empty())
}
Bound::Schema(change) => {
let there = catalog.has_schema(&change.catalog, &change.name);
if change.drop {
if there || !change.quiet {
catalog.drop_schema(&change.catalog, &change.name, change.cascade)?;
}
return Ok(QueryResult::empty());
}
if holds_a_file(&self.inner, &catalog, &change.catalog) {
return Err(Error::not_implemented(
"CREATE SCHEMA in a database file, which holds only the main schema so far",
));
}
if there && change.quiet {
return Ok(QueryResult::empty());
}
if there && change.or_replace {
catalog.drop_schema(&change.catalog, &change.name, false)?;
}
catalog.create_schema(&change.catalog, &change.name)?;
Ok(QueryResult::empty())
}
Bound::Sequence(change) => {
let Some(name) = change.name else { return Ok(QueryResult::empty()) };
if let Some(owner) = change.owner {
catalog.own_sequence(&name, owner)?;
return Ok(QueryResult::empty());
}
if change.drop {
catalog.drop_sequence(&name, change.cascade)?;
return Ok(QueryResult::empty());
}
if holds_a_file(&self.inner, &catalog, &name.catalog) {
return Err(Error::not_implemented(
"CREATE SEQUENCE in a database file, which cannot hold one so far",
));
}
let counter = rudb_common::sequence::Counter::register(&name.table, change.options);
catalog.create_sequence(name, counter, change.or_replace, change.if_not_exists)?;
Ok(QueryResult::empty())
}
Bound::Type(change) => {
let Some(name) = change.name else { return Ok(QueryResult::empty()) };
let Some(ty) = change.ty else {
catalog.drop_type(&name, change.cascade)?;
return Ok(QueryResult::empty());
};
if holds_a_file(&self.inner, &catalog, &name.catalog) {
return Err(Error::not_implemented(
"CREATE TYPE in a database file, which cannot hold one so far",
));
}
catalog.create_type(
name,
ty,
change.uses,
change.or_replace,
change.if_not_exists,
)?;
Ok(QueryResult::empty())
}
Bound::Alter(alter) => {
let (Some(name), Some(alteration)) = (alter.name, alter.alteration) else {
return Ok(QueryResult::empty());
};
let rows = match alter.rewrite {
Some(mut plan) => {
let ((), optimize_ns) =
timed(|| rudb_opt::optimize_with(&mut plan, &context))?;
let facts = context.facts();
let under =
Under::new(self.budget(), facts, &seams, &session, Rows::ForATable)
.after(Planning { parse_ns, bind_ns, optimize_ns });
Some(run(sql, &plan, &catalog, cancel, under)?.into_chunks())
}
None => None,
};
catalog.alter(&name, alteration, rows, self.inner.pool.threads())?;
Ok(QueryResult::empty())
}
Bound::Index(change) => {
match (change.table, change.index) {
(Some(table), Some(index)) => {
catalog.create_index(&table, index, change.quiet)?;
}
_ => {
let parts: Vec<&str> = change.name.iter().map(String::as_str).collect();
catalog.drop_index(&parts, change.quiet)?;
}
}
Ok(QueryResult::empty())
}
Bound::Insert(mut insert) => {
let ((), optimize_ns) =
timed(|| rudb_opt::optimize_with(&mut insert.source, &context))?;
let writable = self.inner.writable
&& insert.name.catalog.eq_ignore_ascii_case(DEFAULT_CATALOG)
&& !self.transacting();
if let Some(path) = self.inner.path.as_ref().filter(|_| {
writable
&& insert.write == Write::Append
&& insert.returning.is_none()
&& insert.checks.is_none()
&& catalog.table(&insert.name).is_ok_and(|table| {
table.guards().is_empty() && table.foreign().is_empty()
})
}) {
let target = catalog.table(&insert.name)?;
let alone = !path.exists() && catalog.stored_tables().count() == 1;
let empty = target.rows().is_empty();
if empty && (alone || appendable(path, &catalog, &insert.name)?) {
let table = target.name().table.clone();
let fields = target.columns().to_vec();
let clustering = target.clustering().cloned();
let limit = self.inner.memory.limit();
let sink = Arc::new(if alone {
NativeSink::create(path, table.clone(), fields, clustering, limit)?
} else {
NativeSink::open(path, table.clone(), fields, clustering, limit)?
});
let query = rudb_exec::build_measured_into(
&insert.source,
&catalog,
cancel,
&self.inner.memory,
&seams,
&session,
sink,
)?;
query.run(cancel, &self.inner.pool)?;
drop(query);
let reader = rudb_native::Catalog::open(path)?.table(&table)?;
let added = reader.table().rows();
catalog.table_mut(&insert.name)?.commit_native(reader)?;
return QueryResult::changed(added);
}
}
let facts = context.facts();
let under = Under::new(self.budget(), facts, &seams, &session, Rows::ForATable)
.after(Planning { parse_ns, bind_ns, optimize_ns });
let result = run(sql, &insert.source, &catalog, cancel, under)?;
let workers = self.inner.pool.threads();
let chunks = result.into_chunks();
let wanted = insert.returning.is_some();
let place = (cancel, &seams, &session);
let mut checks = insert.checks.take();
let (count, written) = match insert.write {
Write::Append if insert.conflict.is_some() => {
let conflict = insert.conflict.take().expect("asked just above");
let name = insert.name.clone();
let upsert = (conflict, checks.as_mut());
self.upsert(sql, &mut catalog, place, &name, upsert, chunks)?
}
Write::Append => {
if let Some(checks) = checks.as_mut() {
self.check(sql, &mut catalog, place, &insert.name, checks, &chunks)?;
}
foreign::missing(&catalog, &insert.name, &chunks)?;
let added = chunks.iter().map(Chunk::len).sum();
let written = if wanted { chunks.clone() } else { Vec::new() };
catalog.table_mut(&insert.name)?.append_all(chunks, workers)?;
(added, written)
}
Write::Update | Write::Delete => {
let delete = insert.write == Write::Delete;
let plain = catalog.table(&insert.name)?.foreign().is_empty();
let (kept, changed, count) =
split(chunks, delete, wanted || checks.is_some() || !plain)?;
if let Some(checks) = checks.as_mut() {
self.check(sql, &mut catalog, place, &insert.name, checks, &changed)?;
}
if !delete {
foreign::missing(&catalog, &insert.name, &changed)?;
}
foreign::lost(&catalog, &insert.name, &kept)?;
catalog.table_mut(&insert.name)?.replace_all(kept, workers)?;
(count, changed)
}
};
let Some(mut returning) = insert.returning else {
return QueryResult::changed(count);
};
let held = catalog.table_mut(&insert.name)?.stand_in(written, workers)?;
let answer = (|| {
let context = self.optimizer(&catalog)?;
rudb_opt::optimize_with(&mut returning, &context)?;
let under = Under::new(
self.budget(),
context.facts(),
&seams,
&session,
Rows::ForACaller,
);
run(sql, &returning, &catalog, cancel, under)
})();
catalog.table_mut(&insert.name)?.put_back(held);
answer
}
}
}
}
impl Shared {
fn check(
&self,
sql: &str,
catalog: &mut Catalog,
(cancel, seams, session): (&Cancel, &rudb_seam::Settings, &Session),
name: &QualifiedName,
checks: &mut rudb_bind::Checks,
rows: &[Chunk],
) -> Result<()> {
if rows.iter().all(|chunk| chunk.is_empty()) {
return Ok(());
}
let workers = self.inner.pool.threads();
let before = catalog.table_mut(name)?.stand_in(rows.to_vec(), workers)?;
let answer = (|| {
let context = self.optimizer(catalog)?;
rudb_opt::optimize_with(&mut checks.plan, &context)?;
let under =
Under::new(self.budget(), context.facts(), seams, session, Rows::ForACaller);
run(sql, &checks.plan, catalog, cancel, under)
})();
catalog.table_mut(name)?.put_back(before);
let chunks = answer?.into_chunks();
for (at, message) in checks.messages.iter().enumerate() {
let failed = chunks.iter().any(|chunk| {
(0..chunk.len()).any(|row| chunk.value_at(row, at) == Value::Boolean(true))
});
if failed {
return Err(Error::constraint(message.clone()));
}
}
Ok(())
}
fn upsert(
&self,
sql: &str,
catalog: &mut Catalog,
(cancel, seams, session): (&Cancel, &rudb_seam::Settings, &Session),
name: &QualifiedName,
(conflict, checks): (rudb_bind::Conflict, Option<&mut rudb_bind::Checks>),
chunks: Vec<Chunk>,
) -> Result<(usize, Vec<Chunk>)> {
let workers = self.inner.pool.threads();
let table = catalog.table(name)?;
let types = table.types();
let fields = table.columns().to_vec();
let keys = table.keys().to_vec();
let all: Vec<usize> = (0..fields.len()).collect();
let mut stored = Vec::with_capacity(table.rows().chunk_count());
for at in 0..table.rows().chunk_count() {
stored.push(table.rows().read(at, &all)?);
}
let mut held = upsert::rows_of(&stored);
let new = upsert::rows_of(&chunks);
let arrivals = upsert::arrivals(&keys, conflict.key, &held, &new);
let mut added = Vec::new();
let mut clashes = Vec::new();
for (row, arrival) in new.into_iter().zip(&arrivals) {
match arrival {
upsert::Arrival::New => added.push(row),
upsert::Arrival::Held(at) => clashes.push((*at, row)),
upsert::Arrival::Dropped => {}
}
}
let mut updated = Vec::new();
match conflict.action {
rudb_bind::ConflictAction::Nothing => {}
rudb_bind::ConflictAction::Replace(columns) => {
for (at, row) in clashes {
for &column in &columns {
held[at][column] = row[column].clone();
}
updated.push(at);
}
}
rudb_bind::ConflictAction::Update { columns, mut plan } if !clashes.is_empty() => {
let matched: Vec<Vec<Value>> =
clashes.iter().map(|(at, _)| held[*at].clone()).collect();
let incoming: Vec<Vec<Value>> =
clashes.iter().map(|(_, row)| row.clone()).collect();
let excluded = QualifiedName::excluded();
let loose =
fields.iter().map(|field| Field { not_null: false, ..field.clone() }).collect();
catalog.create_table(excluded.clone(), loose)?;
let answer = (|| {
let rows = upsert::chunks_of(&types, &incoming)?;
catalog.table_mut(&excluded)?.append_all(rows, workers)?;
let rows = upsert::chunks_of(&types, &matched)?;
let before = catalog.table_mut(name)?.stand_in(rows, workers)?;
let answer = (|| {
let context = self.optimizer(catalog)?;
rudb_opt::optimize_with(&mut plan, &context)?;
let under = Under::new(
self.budget(),
context.facts(),
seams,
session,
Rows::ForACaller,
);
run(sql, &plan, catalog, cancel, under)
})();
catalog.table_mut(name)?.put_back(before);
answer
})();
catalog.drop_table(&excluded)?;
let values = upsert::rows_of(&answer?.into_chunks());
for ((at, _), row) in clashes.iter().zip(&values) {
if row.last() != Some(&Value::Boolean(true)) {
continue;
}
for (&column, value) in columns.iter().zip(row) {
held[*at][column] = value.clone();
}
updated.push(*at);
}
}
rudb_bind::ConflictAction::Update { .. } => {}
}
let count = updated.len() + added.len();
let mut written: Vec<Vec<Value>> = updated.iter().map(|&at| held[at].clone()).collect();
written.extend(added.iter().cloned());
let rows = upsert::chunks_of(&types, &written)?;
if let Some(checks) = checks {
self.check(sql, catalog, (cancel, seams, session), name, checks, &rows)?;
}
foreign::missing(catalog, name, &rows)?;
let table = catalog.table_mut(name)?;
if updated.is_empty() {
table.append_all(upsert::chunks_of(&types, &added)?, workers)?;
} else {
held.extend(added);
table.replace_all(upsert::chunks_of(&types, &held)?, workers)?;
}
Ok((count, upsert::chunks_of(&types, &written)?))
}
}
fn planned(
sql: &str,
catalog: &Catalog,
context: &rudb_opt::pass::Context,
session: &Session,
) -> Result<Plan> {
let mut plan = rudb_bind::bind_sql_with(sql, catalog, session)?;
rudb_opt::optimize_with(&mut plan, context)?;
Ok(plan)
}
fn stored_link(
catalog: &Catalog,
child: (&str, &[String]),
parent: (&str, &[String]),
) -> Option<rudb_graph::link::Counts> {
let child_table = table_named(catalog, child.0)?;
let parent_table = table_named(catalog, parent.0)?;
let (
rudb_catalog::table::Rows::Native(child_rows),
rudb_catalog::table::Rows::Native(parent_rows),
) = (child_table.rows(), parent_table.rows())
else {
return None;
};
let (Some(child_column), Some(parent_column)) =
(key_in(child_table, child.1), key_in(parent_table, parent.1))
else {
return None;
};
let edge = Edge {
child: child_table.name().table.clone(),
child_column,
parent: parent_table.name().table.clone(),
parent_column,
};
rudb_native::graph::stored_link_counts(child_rows, parent_rows, &edge)
}
fn parent_keyed(catalog: &Catalog, table: &str, columns: &[String]) -> bool {
let Some(table) = table_named(catalog, table) else { return false };
let rudb_catalog::table::Rows::Native(rows) = table.rows() else { return false };
key_in(table, columns).is_some_and(|column| rudb_native::graph::holds_key_map(rows, column))
}
fn table_named<'a>(catalog: &'a Catalog, name: &str) -> Option<&'a rudb_catalog::Table> {
catalog.tables().into_iter().find(|table| table.name().table.eq_ignore_ascii_case(name))
}
#[derive(Clone, Copy)]
struct Budget<'a> {
memory: &'a Memory,
pool: &'a Pool,
}
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
enum Rows {
ForACaller,
ForATable,
}
#[derive(Clone, Copy, Default, Debug)]
struct Planning {
parse_ns: u64,
bind_ns: u64,
optimize_ns: u64,
}
impl Planning {
fn total_ns(self) -> u64 {
self.parse_ns.saturating_add(self.bind_ns).saturating_add(self.optimize_ns)
}
}
fn asked_for_mirrors(bound: &Bound) -> bool {
matches!(bound, Bound::Query(plan) if !plan.wanted_mirrors().is_empty())
}
fn timed<T>(what: impl FnOnce() -> Result<T>) -> Result<(T, u64)> {
let span = Span::start();
let out = what()?;
Ok((out, span.stop().0))
}
#[derive(Clone, Copy)]
struct Under<'a> {
budget: Budget<'a>,
facts: &'a rudb_opt::estimate::Facts,
seams: &'a rudb_seam::Settings,
session: &'a Session,
going: Rows,
planning: Planning,
}
impl<'a> Under<'a> {
fn new(
budget: Budget<'a>,
facts: &'a rudb_opt::estimate::Facts,
seams: &'a rudb_seam::Settings,
session: &'a Session,
going: Rows,
) -> Self {
Self { budget, facts, seams, session, going, planning: Planning::default() }
}
fn after(mut self, planning: Planning) -> Self {
self.planning = planning;
self
}
}
fn is_native_summary_aggregate(ast: &Ast, plan: &Plan, catalog: &Catalog) -> bool {
let Node::Project { input, exprs, .. } = *plan.node(plan.root()) else { return false };
let Node::Aggregate { input, index, groups, aggregates } = *plan.node(input) else {
return false;
};
let projected = plan.expr_list(exprs);
let aggregates = plan.expr_list(aggregates);
if projected.is_empty()
|| projected.len() != aggregates.len()
|| !plan.expr_list(groups).is_empty()
|| !projected.iter().enumerate().all(|(position, expr)| {
matches!(plan.expr(*expr), Expr::Column(column)
if column.table == index && column.column as usize == position)
})
{
return false;
}
let filtered = matches!(plan.node(input), Node::Filter { .. });
let input = match *plan.node(input) {
Node::Filter { input, .. } if simple_literal_filter(ast) => input,
Node::Filter { .. } => return false,
_ => input,
};
let Node::Get { catalog: source_catalog, schema, table, index: source_index, columns, .. } =
*plan.node(input)
else {
return false;
};
let direct_aggregate = |expr, expected, arguments| {
let Expr::Aggregate { name, args, distinct: false, filter: None } = plan.expr(expr) else {
return false;
};
if plan.string(*name) != expected {
return false;
}
let args = plan.expr_list(*args);
args.len() == arguments
&& args.iter().all(|arg| {
matches!(plan.expr(*arg), Expr::Column(column) if column.table == source_index)
})
};
let supported = match aggregates {
[aggregate] => {
(direct_aggregate(*aggregate, "count_star", 0)
&& (filtered || plan.field_list(columns).is_empty()))
|| (!filtered && direct_aggregate(*aggregate, "avg", 1))
}
[sum, count, avg] if !filtered => {
direct_aggregate(*sum, "sum", 1)
&& direct_aggregate(*count, "count_star", 0)
&& direct_aggregate(*avg, "avg", 1)
}
_ => false,
};
if !supported {
return false;
}
let source =
QualifiedName::new(plan.string(source_catalog), plan.string(schema), plan.string(table));
catalog.table(&source).is_ok_and(|table| table.rows().is_native())
}
fn simple_literal_filter(ast: &Ast) -> bool {
let [ast::Statement::Query(reference)] = ast.statements.as_slice() else { return false };
let query = ast.query(*reference);
if !query.ctes.is_empty()
|| !query.order_by.is_empty()
|| query.order_by_all
|| query.limit != rudb_parse::NONE
|| query.offset != rudb_parse::NONE
{
return false;
}
let ast::QueryBody::Select(reference) = query.body else { return false };
let select = ast.select(reference);
if select.distinct != ast::Distinct::No
|| !select.group_by.is_empty()
|| select.group_by_all
|| select.filter == rudb_parse::NONE
|| select.having != rudb_parse::NONE
{
return false;
}
let [source] = ast.source_list(select.from) else { return false };
if !matches!(ast.source(*source), ast::Source::Table { .. }) {
return false;
}
let ast::Expr::Binary { op: ast::BinaryOp::NotEq, left, right } = ast.expr(select.filter)
else {
return false;
};
matches!(ast.expr(left), ast::Expr::Column { .. })
&& matches!(ast.expr(right), ast::Expr::Literal { kind: ast::LiteralKind::Number, .. })
}
fn run(
sql: &str,
plan: &Plan,
catalog: &Catalog,
cancel: &Cancel,
under: Under<'_>,
) -> Result<QueryResult> {
let Under { budget: Budget { memory, pool }, facts, seams, session, going, planning } = under;
memory.forget_peak();
let report = Report::new();
let building = Span::start();
let query = rudb_exec::build_measured(plan, catalog, cancel, memory, seams, session, &report)?;
let (built_wall, built_cpu) = building.stop();
if going == Rows::ForACaller {
query.for_a_caller();
}
let names = query.schema().names();
let types = query.schema().types();
let mut held = memory.reservation();
let mut chunks = Vec::new();
let driving = Span::start();
query.run(cancel, pool)?;
while let Some(chunk) = query.next_chunk()? {
if chunk.is_empty() {
continue;
}
let chunk = match going {
Rows::ForACaller => chunk.into_flat()?,
Rows::ForATable => chunk,
};
held.grow(u64::try_from(chunk.footprint()).unwrap_or(u64::MAX))?;
chunks.push(chunk);
}
let (ran_wall, ran_cpu) = driving.stop();
let mut metrics = Document::new(sql);
metrics.settings.memory_limit = memory.limit();
metrics.settings.threads = u32::try_from(pool.threads()).unwrap_or(u32::MAX);
metrics.timing.parse_ns = planning.parse_ns;
metrics.timing.bind_ns = planning.bind_ns;
metrics.timing.optimize_ns = planning.optimize_ns;
metrics.timing.physical_ns = built_wall;
metrics.timing.execute_ns = ran_wall;
metrics.timing.total_ns =
planning.total_ns().saturating_add(built_wall).saturating_add(ran_wall);
let ran_cpu = ran_cpu.saturating_add(query.worker_cpu_ns());
metrics.resource.cpu_ns = built_cpu.saturating_add(ran_cpu);
metrics.resource.build_cpu_ns = built_cpu;
metrics.resource.peak_bytes = memory.peak();
report.fill(&mut metrics);
rudb_opt::explain::record_estimates(plan, facts, &mut metrics);
Ok(QueryResult::new(names, types, chunks, held).in_session(session.clone()).measured(metrics))
}
fn run_compiled(
sql: &str,
plan: &Plan,
catalog: &Catalog,
cancel: &Cancel,
compiled: rudb_qc::Compiled,
under: Under<'_>,
) -> Result<QueryResult> {
let Under { budget: Budget { memory, pool }, seams, session, planning, .. } = under;
memory.forget_peak();
let driving = Span::start();
let qc = rudb_qc::Under { catalog, cancel, memory, seams, session, pool };
let answer = compiled.run(plan, qc)?;
let (ran_wall, ran_cpu) = driving.stop();
let mut held = memory.reservation();
let mut chunks = Vec::with_capacity(answer.chunks.len());
for chunk in answer.chunks {
let chunk = chunk.into_flat()?;
held.grow(u64::try_from(chunk.footprint()).unwrap_or(u64::MAX))?;
chunks.push(chunk);
}
let mut metrics = Document::new(sql);
metrics.settings.memory_limit = memory.limit();
metrics.settings.threads = u32::try_from(pool.threads()).unwrap_or(u32::MAX);
metrics.timing.parse_ns = planning.parse_ns;
metrics.timing.bind_ns = planning.bind_ns;
metrics.timing.optimize_ns = planning.optimize_ns;
metrics.timing.execute_ns = ran_wall;
metrics.timing.total_ns = planning.total_ns().saturating_add(ran_wall);
metrics.resource.cpu_ns = ran_cpu;
metrics.resource.peak_bytes = memory.peak();
Ok(QueryResult::new(answer.names, answer.types, chunks, held)
.in_session(session.clone())
.measured(metrics))
}
#[allow(clippy::too_many_arguments)]
fn explaining(
plan: &Plan,
catalog: &Catalog,
cancel: &Cancel,
budget: Budget<'_>,
context: &rudb_opt::pass::Context,
seams: rudb_opt::explain::Seams<'_>,
session: &Session,
asked: Asked,
sql: &str,
planning: Planning,
) -> Result<QueryResult> {
let facts = context.facts();
let statistics = asked.statistics();
if !asked.analyze {
let text = rudb_opt::explain::explain_with(plan, context, seams, statistics);
return explained("logical_plan", &text);
}
let mut profiled = session.clone();
profiled.set("enable_profiling", "query_tree");
let under =
Under::new(budget, facts, seams.settings(), &profiled, Rows::ForACaller).after(planning);
let result = run(sql, plan, catalog, cancel, under)?;
let measured = result.metrics().expect("a query that ran reports what it did");
let text = rudb_opt::explain::analyzed(plan, context, seams, measured, statistics);
explained("analyzed_plan", &text)
}
fn explained_codegen(plan: &Plan, cancel: &Cancel) -> Result<QueryResult> {
match rudb_qc::compile(plan, cancel) {
Ok(compiled) => explained("codegen", &compiled.explain()),
Err(refusal) => explained("codegen", &format!("refused: {refusal}")),
}
}
#[derive(Debug, Clone, Copy)]
struct Asked {
analyze: bool,
statistics: bool,
}
impl Asked {
fn statistics(self) -> rudb_opt::explain::Statistics {
if self.statistics {
rudb_opt::explain::Statistics::Asked
} else {
rudb_opt::explain::Statistics::NotAsked
}
}
}
fn explained(key: &str, text: &str) -> Result<QueryResult> {
let key = Vector::from_values(LogicalType::Varchar, &[Value::Varchar(key.to_owned())])?;
let value = Vector::from_values(LogicalType::Varchar, &[Value::Varchar(text.to_owned())])?;
Ok(QueryResult::new(
vec!["explain_key".to_owned(), "explain_value".to_owned()],
vec![LogicalType::Varchar, LogicalType::Varchar],
vec![Chunk::new(vec![key, value])?],
Memory::unlimited().reservation(),
))
}
fn create_view(create: rudb_bind::CreateView, catalog: &mut Catalog) -> Result<()> {
if create.if_not_exists && catalog.entry(&create.name).is_ok() {
return Ok(());
}
if create.or_replace && catalog.view(&create.name).is_ok() {
catalog.drop_view(&create.name)?;
}
catalog.create_view(View::new(
create.name,
create.sql,
create.statement,
create.aliases,
create.columns,
))
}
#[allow(clippy::too_many_arguments)]
fn create_table(
sql: &str,
mut create: rudb_bind::CreateTable,
catalog: &mut Catalog,
cancel: &Cancel,
budget: Budget<'_>,
context: &rudb_opt::pass::Context,
seams: &rudb_seam::Settings,
session: &Session,
) -> Result<()> {
if create.if_not_exists && catalog.table(&create.name).is_ok() {
return Ok(());
}
let rows = match &mut create.source {
Some(plan) => {
rudb_opt::optimize_with(plan, context)?;
let under = Under::new(budget, context.facts(), seams, session, Rows::ForATable);
Some(run(sql, plan, catalog, cancel, under)?)
}
None => None,
};
if create.or_replace && catalog.table(&create.name).is_ok() {
catalog.drop_table(&create.name)?;
}
catalog.create_table(create.name.clone(), create.columns)?;
if !create.keys.is_empty() {
catalog.table_mut(&create.name)?.set_keys(create.keys)?;
}
if create.defaults.iter().any(Option::is_some) {
catalog.table_mut(&create.name)?.set_defaults(create.defaults);
}
if !create.sequences.is_empty() {
catalog.table_mut(&create.name)?.set_sequences(create.sequences);
}
if !create.checks.is_empty() {
catalog.table_mut(&create.name)?.set_checks(create.checks);
}
if !create.order.is_empty() {
catalog.table_mut(&create.name)?.set_order(create.order);
}
if !create.foreign.is_empty() {
catalog.table_mut(&create.name)?.set_foreign(create.foreign);
}
if let Some(rows) = rows {
catalog.table_mut(&create.name)?.append_all(rows.into_chunks(), budget.pool.threads())?;
}
Ok(())
}
#[cfg(test)]
mod tests {
use std::path::Path;
use std::sync::mpsc;
use std::time::Duration;
use rudb_common::Value;
use rudb_io::{Filesystem, Op, OpenMode, SimFilesystem};
use super::{
Database, NativeExtremaValues, native_extrema_shape, native_nonzero_shape,
native_simple_average_statement, native_simple_distinct_statement,
native_simple_extrema_statement, native_simple_three_statement,
native_single_average_shape, native_single_distinct_shape, native_three_aggregate_shape,
publish,
};
#[test]
fn grouped_counts_read_rows_instead_of_returning_a_saved_frequency_list() {
let path =
std::env::temp_dir().join(format!("rudb-grouped-rows-{}.rdb", std::process::id()));
let name = path.to_str().unwrap();
let database = Database::open(name).unwrap();
database.execute("CREATE TABLE events (source_id SMALLINT)").unwrap();
database.execute("INSERT INTO events VALUES (2), (2), (3), (0), (NULL)").unwrap();
database.execute("CREATE TABLE wide_events (source_id BIGINT)").unwrap();
database
.execute("INSERT INTO wide_events VALUES (-10000), (-10000), (10000), (0), (NULL)")
.unwrap();
drop(database);
let sql = "SELECT source_id, COUNT(*) FROM events WHERE source_id <> 0 GROUP BY source_id ORDER BY COUNT(*) DESC";
assert!(Database::query_native_once(name, sql).unwrap().is_none());
assert_eq!(
Database::query_native_group_count_once(name, sql).unwrap(),
Some(vec![(2, 2), (3, 1)])
);
assert_eq!(
Database::query_native_group_count_once(
name,
"SELECT source_id, COUNT(*) FROM wide_events WHERE source_id <> 0 GROUP BY source_id ORDER BY COUNT(*) DESC"
)
.unwrap(),
Some(vec![(-10000, 2), (10000, 1)])
);
assert!(
Database::query_native_group_count_once(name, "SELECT COUNT(*) FROM events")
.unwrap()
.is_none()
);
for changed_sql in [
"SELECT source_id, COUNT(*) FROM events WHERE source_id <> 1 GROUP BY source_id ORDER BY COUNT(*) DESC",
"SELECT source_id, COUNT(*) FROM events WHERE source_id <> 0 GROUP BY source_id ORDER BY source_id DESC",
"SELECT source_id, COUNT(*) FROM events WHERE source_id <> 0 GROUP BY source_id ORDER BY COUNT(*) DESC LIMIT 1",
] {
assert!(Database::query_native_group_count_once(name, changed_sql).unwrap().is_none());
}
let database = Database::open(name).unwrap();
let rows = database.query(sql).unwrap().rows().collect::<Vec<_>>();
assert_eq!(
rows,
vec![
vec![Value::SmallInt(2), Value::BigInt(2)],
vec![Value::SmallInt(3), Value::BigInt(1)]
]
);
drop(database);
std::fs::remove_file(path).unwrap();
}
#[test]
fn covering_projection_runs_grouped_distinct_from_bound_columns() {
let path = std::env::temp_dir()
.join(format!("rudb-covering-distinct-plan-{}.rdb", std::process::id()));
let name = path.to_str().unwrap();
let database = Database::open(name).unwrap();
database
.execute("CREATE TABLE events (person BIGINT NOT NULL, zone INTEGER NOT NULL)")
.unwrap();
database
.execute(
"INSERT INTO events VALUES (9, 7), (2, 1), (9, 7), (2, 2), \
(2, 2), (5, 1), (9, 1), (7, 2)",
)
.unwrap();
drop(database);
rudb_native::build_run_projection(name, "events", "person", "zone").unwrap();
let database = Database::open(name).unwrap();
let plain = "SELECT zone, COUNT(DISTINCT person) FROM events GROUP BY zone ORDER BY zone";
let result = database.query(plain).unwrap();
assert_eq!(
result.rows().collect::<Vec<_>>(),
vec![
vec![Value::Integer(1), Value::BigInt(3)],
vec![Value::Integer(2), Value::BigInt(2)],
vec![Value::Integer(7), Value::BigInt(1)],
]
);
assert!(
result.metrics().unwrap().operators.iter().any(|operator| {
operator.detail.as_deref() == Some("covering grouped distinct")
})
);
let ranked = database
.query("SELECT zone AS z, COUNT(DISTINCT person) AS n FROM events GROUP BY zone ORDER BY n DESC LIMIT 2")
.unwrap();
assert_eq!(
ranked.rows().collect::<Vec<_>>(),
vec![
vec![Value::Integer(1), Value::BigInt(3)],
vec![Value::Integer(2), Value::BigInt(2)],
]
);
assert!(
ranked.metrics().unwrap().operators.iter().any(|operator| {
operator.detail.as_deref() == Some("covering grouped distinct")
})
);
let filtered = database
.query("SELECT zone, COUNT(DISTINCT person) FROM events WHERE zone > 0 GROUP BY zone ORDER BY zone")
.unwrap();
assert_eq!(filtered.rows().collect::<Vec<_>>(), result.rows().collect::<Vec<_>>());
assert!(
!filtered.metrics().unwrap().operators.iter().any(|operator| {
operator.detail.as_deref() == Some("covering grouped distinct")
})
);
database.execute("INSERT INTO events VALUES (3, 1)").unwrap();
let stale = database.query(plain).unwrap();
assert_eq!(stale.rows().next().unwrap(), vec![Value::Integer(1), Value::BigInt(4)]);
assert!(
!stale.metrics().unwrap().operators.iter().any(|operator| {
operator.detail.as_deref() == Some("covering grouped distinct")
})
);
drop(database);
std::fs::remove_file(path).unwrap();
}
#[test]
fn cold_extrema_shape_accepts_only_direct_bounds() {
let parsed =
rudb_parse::parse_ast("SELECT MIN(EventDate), MAX(EventDate) FROM hits").unwrap();
assert_eq!(
native_extrema_shape(&parsed),
Some(("hits", "EventDate", ["min(EventDate)".into(), "max(EventDate)".into()]))
);
for sql in [
"SELECT MIN(EventDate), MAX(EventDate) FROM hits WHERE EventDate > 0",
"SELECT MIN(EventDate), MAX(EventDate) FROM hits LIMIT 1",
"SELECT MIN(EventDate + 1), MAX(EventDate) FROM hits",
"SELECT MIN(EventDate), MAX(OtherDate) FROM hits",
"SELECT MIN(DISTINCT EventDate), MAX(EventDate) FROM hits",
] {
let parsed = rudb_parse::parse_ast(sql).unwrap();
assert_eq!(native_extrema_shape(&parsed), None, "{sql}");
}
}
#[test]
fn simple_extrema_statement_rejects_other_sql() {
assert_eq!(
native_simple_extrema_statement(" SELECT MIN(EventDate), MAX(EventDate) FROM hits; "),
Some(("hits", "EventDate"))
);
for sql in [
"SELECT MIN(EventDate), MAX(EventDate) FROM hits WHERE EventDate > 0",
"SELECT MIN(EventDate), MAX(OtherDate) FROM hits",
"SELECT MIN(EventDate + 1), MAX(EventDate) FROM hits",
"SELECT MIN(EventDate), MAX(EventDate) FROM hits; SELECT 1",
"SELECT MIN(EventDate), MAX(EventDate) FROM hits GROUP BY RegionID",
] {
assert_eq!(native_simple_extrema_statement(sql), None, "{sql}");
}
}
#[test]
fn cold_extrema_matches_regular_execution_for_dates_and_nulls() {
let path = std::env::temp_dir().join(format!("rudb-q7-{}.rdb", std::process::id()));
let name = path.to_str().unwrap();
let database = Database::open(name).unwrap();
database.execute("CREATE TABLE hits (EventDate DATE)").unwrap();
database
.execute("INSERT INTO hits VALUES (DATE '2013-07-31'), (NULL), (DATE '2013-07-02')")
.unwrap();
database.execute("CREATE TABLE small_hits (EventDate USMALLINT)").unwrap();
database.execute("INSERT INTO small_hits VALUES (15917), (15888), (NULL)").unwrap();
database.execute("CREATE TABLE empty_hits (EventDate DATE)").unwrap();
database.execute("CREATE TABLE null_hits (EventDate DATE)").unwrap();
database.execute("INSERT INTO null_hits VALUES (NULL)").unwrap();
let cases = [
"SELECT MIN(EventDate), MAX(EventDate) FROM hits",
"SELECT MIN(EventDate), MAX(EventDate) FROM small_hits",
"SELECT MIN(EventDate), MAX(EventDate) FROM empty_hits",
"SELECT MIN(EventDate), MAX(EventDate) FROM null_hits",
];
let expected = cases.map(|sql| database.query(sql).unwrap().rows().collect::<Vec<_>>());
drop(database);
for (sql, expected) in cases.into_iter().zip(expected) {
let actual = Database::query_native_once(name, sql).unwrap().unwrap();
assert_eq!(actual.rows().collect::<Vec<_>>(), expected, "{sql}");
}
assert_eq!(
Database::query_native_extrema_values_once(
name,
"SELECT MIN(EventDate), MAX(EventDate) FROM hits"
)
.unwrap(),
Some(NativeExtremaValues::Date { low: 15888, high: 15917 })
);
assert_eq!(
Database::query_native_extrema_values_once(
name,
"SELECT MIN(EventDate), MAX(EventDate) FROM small_hits"
)
.unwrap(),
Some(NativeExtremaValues::Integer { low: 15888, high: 15917 })
);
for table in ["empty_hits", "null_hits"] {
let sql = format!("SELECT MIN(EventDate), MAX(EventDate) FROM {table}");
assert_eq!(Database::query_native_extrema_values_once(name, &sql).unwrap(), None);
}
std::fs::remove_file(path).unwrap();
}
#[test]
fn cold_distinct_shape_accepts_only_a_direct_count() {
let parsed = rudb_parse::parse_ast("SELECT COUNT(DISTINCT UserID) FROM hits").unwrap();
assert_eq!(
native_single_distinct_shape(&parsed),
Some(("hits", "UserID", "count(DISTINCT UserID)".into()))
);
for sql in [
"SELECT COUNT(DISTINCT UserID) FROM hits WHERE UserID > 0",
"SELECT COUNT(DISTINCT UserID + 1) FROM hits",
"SELECT COUNT(UserID) FROM hits",
"SELECT COUNT(DISTINCT UserID) FROM hits LIMIT 1",
"SELECT COUNT(DISTINCT UserID) FROM hits GROUP BY RegionID",
] {
let parsed = rudb_parse::parse_ast(sql).unwrap();
assert_eq!(native_single_distinct_shape(&parsed), None, "{sql}");
}
}
#[test]
fn simple_distinct_statement_rejects_other_sql() {
assert_eq!(
native_simple_distinct_statement(" SELECT COUNT(DISTINCT UserID) FROM hits; "),
Some(("hits", "UserID"))
);
for sql in [
"SELECT COUNT(DISTINCT UserID) FROM hits WHERE UserID > 0",
"SELECT COUNT(DISTINCT UserID + 1) FROM hits",
"SELECT COUNT(UserID) FROM hits",
"SELECT COUNT(DISTINCT UserID) FROM hits; SELECT 1",
"SELECT COUNT(DISTINCT UserID) FROM hits GROUP BY RegionID",
"SELECT COUNT(DISTINCT UserID) FROM hits; ;",
] {
assert_eq!(native_simple_distinct_statement(sql), None, "{sql}");
}
}
#[test]
fn cold_distinct_matches_regular_execution_with_nulls() {
let path = std::env::temp_dir().join(format!("rudb-q5-{}.rdb", std::process::id()));
let name = path.to_str().unwrap();
let database = Database::open(name).unwrap();
database.execute("CREATE TABLE hits (UserID BIGINT)").unwrap();
database.execute("INSERT INTO hits VALUES (1), (2), (1), (NULL)").unwrap();
database.execute("CREATE TABLE empty_hits (UserID BIGINT)").unwrap();
database.execute("CREATE TABLE null_hits (UserID BIGINT)").unwrap();
database.execute("INSERT INTO null_hits VALUES (NULL)").unwrap();
let cases = [
"SELECT COUNT(DISTINCT UserID) FROM hits",
"SELECT COUNT(DISTINCT UserID) FROM empty_hits",
"SELECT COUNT(DISTINCT UserID) FROM null_hits",
];
let expected = cases.map(|sql| database.query(sql).unwrap().rows().collect::<Vec<_>>());
drop(database);
for (sql, expected) in cases.into_iter().zip(expected) {
let actual = Database::query_native_once(name, sql).unwrap().unwrap();
assert_eq!(actual.rows().collect::<Vec<_>>(), expected, "{sql}");
}
assert_eq!(
Database::query_native_distinct_value_once(
name,
"SELECT COUNT(DISTINCT UserID) FROM hits"
)
.unwrap(),
Some(2)
);
for sql in [
"SELECT COUNT(DISTINCT UserID) FROM empty_hits",
"SELECT COUNT(DISTINCT UserID) FROM null_hits",
] {
assert_eq!(Database::query_native_distinct_value_once(name, sql).unwrap(), Some(0));
}
std::fs::remove_file(path).unwrap();
}
#[test]
fn cold_average_shape_accepts_only_a_direct_aggregate() {
let parsed = rudb_parse::parse_ast("SELECT AVG(UserID) FROM hits").unwrap();
assert_eq!(
native_single_average_shape(&parsed),
Some(("hits", "UserID", "avg(UserID)".into()))
);
let aliased = rudb_parse::parse_ast("SELECT AVG(UserID) AS mean_user FROM hits").unwrap();
assert_eq!(
native_single_average_shape(&aliased),
Some(("hits", "UserID", "mean_user".into()))
);
for sql in [
"SELECT AVG(UserID) FROM hits WHERE UserID > 0",
"SELECT AVG(DISTINCT UserID) FROM hits",
"SELECT AVG(UserID + 1) FROM hits",
"SELECT AVG(UserID) FROM hits LIMIT 1",
"SELECT AVG(UserID) FROM hits GROUP BY RegionID",
] {
let parsed = rudb_parse::parse_ast(sql).unwrap();
assert_eq!(native_single_average_shape(&parsed), None, "{sql}");
}
}
#[test]
fn simple_average_statement_rejects_other_sql() {
assert_eq!(
native_simple_average_statement(" SELECT AVG(UserID) FROM hits; "),
Some(("hits", "UserID"))
);
for sql in [
"SELECT AVG(UserID) FROM hits WHERE UserID > 0",
"SELECT AVG(DISTINCT UserID) FROM hits",
"SELECT AVG(UserID + 1) FROM hits",
"SELECT AVG(UserID) FROM hits; SELECT 1",
"SELECT AVG(UserID) FROM hits GROUP BY RegionID",
"SELECT AVG(UserID) FROM hits -- comment",
"SELECT AVG(UserID) FROM hits; ;",
] {
assert_eq!(native_simple_average_statement(sql), None, "{sql}");
}
}
#[test]
fn cold_average_matches_regular_execution_for_bigints_and_nulls() {
let path = std::env::temp_dir().join(format!("rudb-q4-{}.rdb", std::process::id()));
let name = path.to_str().unwrap();
let database = Database::open(name).unwrap();
database.execute("CREATE TABLE hits (UserID BIGINT)").unwrap();
database
.execute("INSERT INTO hits VALUES (2414420660257356000), (2534231104689841000), (NULL)")
.unwrap();
database.execute("CREATE TABLE empty_hits (UserID BIGINT)").unwrap();
database.execute("CREATE TABLE null_hits (UserID BIGINT)").unwrap();
database.execute("INSERT INTO null_hits VALUES (NULL)").unwrap();
let cases = [
"SELECT AVG(UserID) FROM hits",
"SELECT AVG(UserID) FROM empty_hits",
"SELECT AVG(UserID) FROM null_hits",
];
let expected = cases.map(|sql| database.query(sql).unwrap().rows().collect::<Vec<_>>());
drop(database);
for (sql, expected) in cases.into_iter().zip(expected) {
let actual = Database::query_native_once(name, sql).unwrap().unwrap();
assert_eq!(actual.rows().collect::<Vec<_>>(), expected, "{sql}");
}
assert_eq!(
Database::query_native_average_value_once(name, "SELECT AVG(UserID) FROM hits")
.unwrap(),
Some((2414420660257356000_i128 + 2534231104689841000_i128) as f64 / 2.0)
);
for sql in ["SELECT AVG(UserID) FROM empty_hits", "SELECT AVG(UserID) FROM null_hits"] {
assert_eq!(Database::query_native_average_value_once(name, sql).unwrap(), None);
}
std::fs::remove_file(path).unwrap();
}
#[test]
fn cold_three_aggregate_shape_accepts_only_the_certified_query() {
let parsed = rudb_parse::parse_ast(
"SELECT SUM(AdvEngineID), COUNT(*), AVG(ResolutionWidth) FROM hits",
)
.unwrap();
assert_eq!(
native_three_aggregate_shape(&parsed),
Some((
"hits",
"AdvEngineID",
"ResolutionWidth",
["sum(AdvEngineID)".into(), "count_star()".into(), "avg(ResolutionWidth)".into()]
))
);
for sql in [
"SELECT SUM(AdvEngineID), COUNT(*), AVG(ResolutionWidth) FROM hits WHERE AdvEngineID > 0",
"SELECT SUM(AdvEngineID), COUNT(*), AVG(ResolutionWidth) FROM hits LIMIT 1",
"SELECT SUM(DISTINCT AdvEngineID), COUNT(*), AVG(ResolutionWidth) FROM hits",
"SELECT SUM(AdvEngineID), COUNT(*), AVG(ResolutionWidth) FROM hits GROUP BY RegionID",
] {
let parsed = rudb_parse::parse_ast(sql).unwrap();
assert_eq!(native_three_aggregate_shape(&parsed), None, "{sql}");
}
}
#[test]
fn simple_three_statement_uses_requested_columns_and_rejects_other_clauses() {
assert_eq!(
native_simple_three_statement(
" SELECT SUM(Points), COUNT(*), AVG(Width) FROM measurements; "
),
Some(("measurements", "Points", "Width"))
);
for sql in [
"SELECT SUM(Points), COUNT(*), AVG(Width) FROM measurements WHERE Points > 0",
"SELECT SUM(DISTINCT Points), COUNT(*), AVG(Width) FROM measurements",
"SELECT SUM(Points), COUNT(*), AVG(Width) FROM measurements GROUP BY Width",
"SELECT SUM(Points), COUNT(*), AVG(Width) FROM measurements LIMIT 1",
"SELECT SUM(Points), COUNT(*), AVG(Width) FROM measurements; SELECT 1",
"SELECT SUM(Points), COUNT(*), AVG(Width + 1) FROM measurements",
"SELECT SUM(select), COUNT(*), AVG(Width) FROM measurements",
"SELECT SUM(Points), COUNT(*), AVG(Width) FROM from",
] {
assert_eq!(native_simple_three_statement(sql), None, "{sql}");
}
}
#[test]
fn cold_three_aggregate_answers_from_native_catalog() {
let path = std::env::temp_dir().join(format!("rudb-q3-{}.rdb", std::process::id()));
let name = path.to_str().unwrap();
let database = Database::open(name).unwrap();
database
.execute("CREATE TABLE hits (AdvEngineID SMALLINT, ResolutionWidth SMALLINT)")
.unwrap();
database.execute("INSERT INTO hits VALUES (1, 100), (NULL, 200), (3, NULL)").unwrap();
database
.execute("CREATE TABLE empty_hits (AdvEngineID SMALLINT, ResolutionWidth SMALLINT)")
.unwrap();
database
.execute("CREATE TABLE null_hits (AdvEngineID SMALLINT, ResolutionWidth SMALLINT)")
.unwrap();
database.execute("INSERT INTO null_hits VALUES (NULL, NULL)").unwrap();
database.execute("CREATE TABLE measurements (Points SMALLINT, Width SMALLINT)").unwrap();
database.execute("INSERT INTO measurements VALUES (2, 10), (5, 20)").unwrap();
drop(database);
let result = Database::query_native_once(
name,
"SELECT SUM(AdvEngineID), COUNT(*), AVG(ResolutionWidth) FROM hits",
)
.unwrap()
.unwrap();
assert_eq!(
result.rows().collect::<Vec<_>>(),
vec![vec![Value::HugeInt(4), Value::BigInt(3), Value::Double(150.0)]]
);
assert_eq!(
Database::query_native_three_values_once(
name,
"SELECT SUM(AdvEngineID), COUNT(*), AVG(ResolutionWidth) FROM hits",
)
.unwrap(),
Some((4, 3, 150.0))
);
assert_eq!(
Database::query_native_three_values_once(
name,
"SELECT SUM(Points), COUNT(*), AVG(Width) FROM measurements",
)
.unwrap(),
Some((7, 2, 15.0))
);
assert_eq!(
Database::query_native_three_values_once(
name,
"SELECT SUM(AdvEngineID), COUNT(*), AVG(ResolutionWidth) FROM hits; SELECT 1",
)
.unwrap(),
None
);
for (table, count) in [("empty_hits", 0), ("null_hits", 1)] {
let sql =
format!("SELECT SUM(AdvEngineID), COUNT(*), AVG(ResolutionWidth) FROM {table}");
let result = Database::query_native_once(name, &sql).unwrap().unwrap();
assert_eq!(
result.rows().collect::<Vec<_>>(),
vec![vec![Value::Null, Value::BigInt(count), Value::Null]]
);
}
std::fs::remove_file(path).unwrap();
}
#[test]
fn cold_count_shape_accepts_only_the_certified_query() {
let parsed = rudb_parse::parse_ast("SELECT COUNT(*) FROM hits WHERE AdvEngineID <> 0")
.expect("query parses");
assert_eq!(native_nonzero_shape(&parsed), Some(("hits", "AdvEngineID", "count_star()")));
for sql in [
"SELECT COUNT(*) FROM hits WHERE AdvEngineID <> 1",
"SELECT COUNT(*) FROM hits WHERE AdvEngineID <> 0 LIMIT 1",
"SELECT COUNT(*) FROM hits WHERE AdvEngineID <> 0 AND RegionID = 1",
"SELECT COUNT(DISTINCT AdvEngineID) FROM hits WHERE AdvEngineID <> 0",
] {
let parsed = rudb_parse::parse_ast(sql).expect("query parses");
assert_eq!(native_nonzero_shape(&parsed), None, "{sql}");
}
}
#[test]
fn cold_nonzero_csv_count_uses_column_statistics() {
let path = std::env::temp_dir().join(format!("rudb-q2-csv-{}.rdb", std::process::id()));
let name = path.to_str().unwrap();
let database = Database::open(name).unwrap();
database.execute("CREATE TABLE hits (AdvEngineID SMALLINT)").unwrap();
database.execute("INSERT INTO hits VALUES (0), (0), (NULL), (2), (-3)").unwrap();
database.execute("CREATE TABLE events (engine INTEGER)").unwrap();
database.execute("INSERT INTO events VALUES (0), (8), (NULL)").unwrap();
drop(database);
let sql = "SELECT COUNT(*) FROM hits WHERE AdvEngineID <> 0";
assert_eq!(Database::query_native_nonzero_value_once(name, sql).unwrap(), Some(2));
assert_eq!(
Database::query_native_nonzero_value_once(
name,
"select count(*) from events where engine <> 0;"
)
.unwrap(),
Some(1)
);
assert_eq!(
Database::query_native_nonzero_value_once(
name,
"SELECT COUNT(*) FROM hits WHERE AdvEngineID <> 1"
)
.unwrap(),
None
);
for sql in [
"SELECT COUNT(*) FROM select WHERE AdvEngineID <> 0",
"SELECT COUNT(*) FROM hits WHERE select <> 0",
] {
assert_eq!(super::native_simple_nonzero_statement(sql), None);
assert_eq!(Database::query_native_nonzero_value_once(name, sql).unwrap(), None);
}
std::fs::remove_file(path).unwrap();
}
#[test]
fn publishing_syncs_the_directory_after_the_rename() {
let fs = SimFilesystem::new();
fs.create_dir_all(Path::new("/data")).unwrap();
let file = fs.open(Path::new("/data/db.7.tmp"), OpenMode::CreateNew).unwrap();
file.write_at(0, b"new").unwrap();
file.sync().unwrap();
fs.clear_log();
publish(&fs, Path::new("/data/db.7.tmp"), Path::new("/data/db")).unwrap();
let ops = fs.ops();
assert_eq!(ops.len(), 2, "{ops:?}");
assert!(matches!(ops[0], Op::Rename { .. }), "{ops:?}");
assert_eq!(ops[1], Op::SyncDir { path: "/data".into() });
assert_eq!(fs.contents(Path::new("/data/db")).unwrap(), b"new".to_vec());
}
#[test]
fn a_query_runs_while_a_load_into_a_new_table_does() {
let path = std::env::temp_dir().join(format!(
"rudb-load-lock-{}-{}.rdb",
std::process::id(),
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.expect("the clock advances")
.as_nanos()
));
let database = Database::open(path.to_str().expect("a UTF-8 temporary path")).unwrap();
database.execute("CREATE TABLE small AS SELECT 7 AS a").unwrap();
let (told, loading) = mpsc::channel();
*database.shared.inner.loading.lock().unwrap() = Some(told);
let connection = database.connect();
let stopper = connection.clone();
let load = std::thread::spawn(move || {
connection.execute("CREATE TABLE big AS SELECT count(*) AS n FROM range(100000000000)")
});
loading.recv_timeout(Duration::from_secs(60)).expect("the load let go of the catalog");
let reader = database.clone();
let (answered, answer) = mpsc::channel();
std::thread::spawn(move || {
let rows = reader.query("SELECT a FROM small").map(|result| result.rows().collect());
let _ = answered.send(rows);
});
let got = answer.recv_timeout(Duration::from_secs(20));
stopper.interrupt();
let loaded = load.join().expect("the load thread ran");
let rows: Vec<Vec<Value>> =
got.expect("the query finished while the load ran").expect("the query succeeded");
assert_eq!(rows, vec![vec![Value::Integer(7)]]);
assert!(loaded.is_err(), "the load was interrupted");
assert!(
database.query("SELECT * FROM big").is_err(),
"an interrupted load leaves no table"
);
drop(database);
let _ = std::fs::remove_file(&path);
}
#[test]
fn a_bare_file_name_lives_in_the_current_directory() {
assert_eq!(super::directory_of(Path::new("db")), Path::new("."));
assert_eq!(super::directory_of(Path::new("/data/db")), Path::new("/data"));
assert_eq!(super::directory_of(Path::new("a/db")), Path::new("a"));
}
#[test]
fn two_statements_over_one_catalog_plan_from_the_same_counts() {
let database = Database::new();
database.execute("CREATE TABLE t (a INTEGER)").expect("a table");
database.execute("INSERT INTO t VALUES (1), (2), (3)").expect("three rows");
let shared = &database.shared;
let first = shared.facts(&shared.read());
let again = shared.facts(&shared.read());
assert!(std::sync::Arc::ptr_eq(&first, &again), "nothing changed, so nothing was rebuilt");
database.execute("INSERT INTO t VALUES (4)").expect("a fourth row");
let after = shared.facts(&shared.read());
assert!(!std::sync::Arc::ptr_eq(&first, &after), "the catalog changed under it");
assert!(after.generation() > first.generation(), "and it says which version it is");
let rows = |facts: &rudb_opt::estimate::Facts| {
facts.get(&rudb_opt::estimate::Key::Rows {
catalog: "memory",
schema: "main",
table: "t",
})
};
assert_eq!(rows(&first), rudb_common::Stat::exact(3, rudb_common::Provenance::RowCount));
assert_eq!(rows(&after), rudb_common::Stat::exact(4, rudb_common::Provenance::RowCount));
}
#[test]
fn every_column_of_a_result_reaches_the_caller_flat() {
use rudb_vector::Form;
let database = Database::new();
database.execute("CREATE TABLE t (a INTEGER, s VARCHAR)").expect("a table");
database
.execute("INSERT INTO t VALUES (1, 'a long string that will not fit inline'), (2, 'b')")
.expect("two rows");
let result = database.query("SELECT a, s, s || 'x' AS j FROM t WHERE a > 0").expect("runs");
assert_eq!(result.len(), 2);
for chunk in result.chunk_iter() {
for (at, column) in chunk.columns().iter().enumerate() {
assert_eq!(column.form(), Form::Flat, "column {at} came out encoded");
}
}
}
#[test]
fn a_relationship_carries_both_certificates_only_when_every_child_row_found_a_parent() {
use super::Shared;
let path = std::env::temp_dir().join(format!(
"rudb-certificates-{}-{}.rdb",
std::process::id(),
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.expect("the clock advances")
.as_nanos()
));
let database =
Database::open(path.to_str().expect("a UTF-8 temporary path")).expect("a file");
database.execute("CREATE TABLE customer (c_custkey INTEGER)").unwrap();
database.execute("CREATE TABLE orders (o_custkey INTEGER)").unwrap();
database.execute("CREATE TABLE returns (r_custkey INTEGER)").unwrap();
database.execute("CREATE TABLE zones (z_key INTEGER)").unwrap();
database.execute("CREATE TABLE visits (v_zone INTEGER)").unwrap();
database.execute("INSERT INTO customer SELECT i FROM range(1, 4001) AS r(i)").unwrap();
database
.execute("INSERT INTO orders SELECT 1 + (i - 1) / 3 FROM range(1, 10001) AS r(i)")
.unwrap();
database
.execute("INSERT INTO returns SELECT 1 + (i - 1) / 3 FROM range(1, 10001) AS r(i)")
.unwrap();
database.execute("INSERT INTO returns VALUES (9999)").unwrap();
database
.execute("INSERT INTO zones SELECT 1 + i % 500 FROM range(0, 1000) AS r(i)")
.unwrap();
database
.execute("INSERT INTO visits SELECT 1 + i % 500 FROM range(0, 2000) AS r(i)")
.unwrap();
let declared = "orders(o_custkey) -> customer(c_custkey), \
returns(r_custkey) -> customer(c_custkey), \
visits(v_zone) -> zones(z_key)";
database.execute(&format!("SET graph_links = '{declared}'")).unwrap();
database.execute("CHECKPOINT").unwrap();
let shared = &database.shared;
let found = Shared::related(&shared.read(), declared);
let certificates: Vec<(&str, bool, bool)> =
found.iter().map(|link| (link.child.as_str(), link.built, link.total)).collect();
assert_eq!(
certificates,
vec![("orders", true, true), ("returns", true, false), ("visits", false, false)],
"one unmatched child row costs the relationship its totality and not its link"
);
drop(database);
std::fs::remove_file(&path).ok();
}
}