use std::collections::BTreeMap;
use rudb_common::Value;
use rudb_common::stat::{Class, Direction, Provenance, Stat};
use rudb_plan::{
ConjunctionOp, Expr, ExprRef, JoinKind, Node, NodeRef, Plan, SetOpKind, Slice, StrRef,
};
use crate::walk;
const KEPT_BY_A_CONDITION: f64 = 0.2;
const KEPT_BY_A_GROUP_BY: f64 = 0.1;
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct Statistics {
tables: BTreeMap<(String, String, String), u64>,
files: BTreeMap<String, u64>,
}
impl Statistics {
#[must_use]
pub fn new() -> Self {
Self::default()
}
pub fn record(&mut self, catalog: &str, schema: &str, table: &str, rows: u64) {
self.tables.insert((catalog.to_owned(), schema.to_owned(), table.to_owned()), rows);
}
#[must_use]
pub fn rows_in(&self, catalog: &str, schema: &str, table: &str) -> Option<u64> {
self.tables.get(&(catalog.to_owned(), schema.to_owned(), table.to_owned())).copied()
}
pub fn record_file(&mut self, pattern: &str, rows: u64) {
self.files.insert(pattern.to_owned(), rows);
}
#[must_use]
pub fn rows_in_file(&self, pattern: &str) -> Option<u64> {
self.files.get(pattern).copied()
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.tables.is_empty() && self.files.is_empty()
}
}
const GUESSED: Class = Class::Estimated;
const FROM_A_CONSTANT: Provenance = Provenance::Default;
const CEILING: Class = Class::Certified { bound: 1.0, direction: Direction::AtMost };
#[must_use]
pub fn rows(plan: &Plan, node: NodeRef, stats: &Statistics) -> Option<u64> {
rows_stat(plan, node, stats).decide().copied()
}
#[must_use]
pub fn rows_stat(plan: &Plan, node: NodeRef, stats: &Statistics) -> Stat<u64> {
let of = |child: NodeRef| rows_stat(plan, child, stats);
match *plan.node(node) {
Node::Dummy => Stat::exact(1, Provenance::RowCount),
Node::Get { catalog, schema, table, .. } => {
match stats.rows_in(plan.string(catalog), plan.string(schema), plan.string(table)) {
Some(rows) => Stat::exact(rows, Provenance::RowCount),
None => Stat::Unknown,
}
}
Node::Values { rows: list, .. } => u64::try_from(plan.row_list(list).len())
.map_or(Stat::Unknown, |rows| Stat::exact(rows, Provenance::RowCount)),
Node::TableFunction { function, args, .. } => in_file(plan, function, args, stats),
Node::LateralFunction { .. } => Stat::Unknown,
Node::Filter { input, predicate } => {
let kept = KEPT_BY_A_CONDITION.powi(conjuncts(plan, predicate));
guess(of(input), kept)
}
Node::Project { input, .. }
| Node::Window { input, .. }
| Node::Sort { input, .. }
| Node::Fetch { input, .. }
| Node::TableFetch { input, .. } => of(input),
Node::Aggregate { input, groups, .. } => {
if plan.expr_list(groups).is_empty() {
return Stat::exact(1, Provenance::RowCount);
}
guess(of(input), KEPT_BY_A_GROUP_BY)
}
Node::Distinct { input, .. } => guess(of(input), KEPT_BY_A_GROUP_BY),
Node::Limit { input, count, offset } => {
let input = of(input);
match count {
None => input.map(|n| n.saturating_sub(offset)),
Some(count) => match input {
Stat::Unknown => {
Stat::Known { value: count, class: CEILING, provenance: FROM_A_CONSTANT }
}
known => known.map(|n| n.saturating_sub(offset).min(count)),
},
}
}
Node::TopN { input, count, offset, .. } => match of(input) {
Stat::Unknown => {
Stat::Known { value: count, class: CEILING, provenance: FROM_A_CONSTANT }
}
known => known.map(|n| n.saturating_sub(offset).min(count)),
},
Node::Join { left, right, kind, conditions, .. } => {
join(of(left), of(right), kind, plan.expr_list(conditions).len())
}
Node::DependentJoin { .. } => Stat::Unknown,
Node::CrossProduct { left, right } => of(left).zip(of(right), u64::saturating_mul),
Node::MaterializedCte { body, .. } => of(body),
Node::CteScan { .. } => Stat::Unknown,
Node::SetOp { left, right, kind, all, .. } => {
let total = of(left).zip(of(right), u64::saturating_add);
match (kind, all) {
(SetOpKind::Union, true) => total,
_ => ceiling(total),
}
}
}
}
fn in_file(plan: &Plan, function: StrRef, args: Slice, stats: &Statistics) -> Stat<u64> {
if plan.string(function) != "read_parquet" {
return Stat::Unknown;
}
let [path] = plan.expr_list(args) else {
return Stat::Unknown;
};
let Expr::Constant(value) = *plan.expr(*path) else {
return Stat::Unknown;
};
let Value::Varchar(pattern) = plan.value(value) else {
return Stat::Unknown;
};
match stats.rows_in_file(pattern) {
Some(rows) => Stat::exact(rows, Provenance::RowCount),
None => Stat::Unknown,
}
}
fn guess(input: Stat<u64>, kept: f64) -> Stat<u64> {
match input {
Stat::Unknown => Stat::Unknown,
Stat::Known { value, class, .. } => Stat::Known {
value: scale(value, kept).max(1),
class: class.combine(GUESSED),
provenance: FROM_A_CONSTANT,
},
}
}
fn ceiling(stat: Stat<u64>) -> Stat<u64> {
match stat {
Stat::Unknown => Stat::Unknown,
Stat::Known { value, class, provenance } => {
Stat::Known { value, class: class.combine(CEILING), provenance }
}
}
}
fn conjuncts(plan: &Plan, predicate: ExprRef) -> i32 {
let counted = match *plan.expr(predicate) {
Expr::Conjunction { op: ConjunctionOp::And, children } => {
plan.expr_list(children).iter().filter(|&&part| !walk::constant(plan, part)).count()
}
_ => usize::from(!walk::constant(plan, predicate)),
};
i32::try_from(counted.min(8)).unwrap_or(8)
}
fn join(left: Stat<u64>, right: Stat<u64>, kind: JoinKind, conditions: usize) -> Stat<u64> {
match kind {
JoinKind::Semi => guess(left, KEPT_BY_A_CONDITION),
JoinKind::Anti => guess(left, 1.0 - KEPT_BY_A_CONDITION),
JoinKind::Single | JoinKind::Mark => left,
JoinKind::Positional => left.zip(right, u64::min),
_ => {
let (
Stat::Known { value: left, class: left_class, provenance: left_from },
Stat::Known { value: right, class: right_class, provenance: right_from },
) = (left, right)
else {
return Stat::Unknown;
};
let both = left_class.combine(right_class);
let from = if left_from == right_from { left_from } else { Provenance::Propagation };
if conditions == 0 {
return Stat::Known {
value: left.saturating_mul(right),
class: both,
provenance: from,
};
}
let matched = left.max(right);
let value = match kind {
JoinKind::Left => matched.max(left),
JoinKind::Right => matched.max(right),
JoinKind::Full => matched.max(left).max(right),
_ => matched,
};
Stat::Known { value, class: both.combine(GUESSED), provenance: FROM_A_CONSTANT }
}
}
}
#[expect(
clippy::cast_precision_loss,
clippy::cast_possible_truncation,
clippy::cast_sign_loss,
reason = "an estimate going through f64 is the point, and the result is clamped"
)]
fn scale(rows: u64, by: f64) -> u64 {
let scaled = rows as f64 * by;
if scaled.is_finite() && scaled >= 0.0 { scaled.min(u64::MAX as f64) as u64 } else { 0 }
}
#[cfg(test)]
mod tests {
use rudb_common::stat::{Class, Direction, Provenance, Stat};
use rudb_plan::Plan;
use super::{Statistics, rows, rows_stat};
fn scan(table: &str, index: u32) -> String {
format!("Get memory.main.{table} AS {table} #{index} [a::INTEGER]\n")
}
fn statistics(tables: &[(&str, u64)]) -> Statistics {
let mut stats = Statistics::new();
for (table, count) in tables {
stats.record("memory", "main", table, *count);
}
stats
}
fn estimate(text: &str, tables: &[(&str, u64)]) -> Option<u64> {
let stats = statistics(tables);
let plan =
Plan::parse(text).unwrap_or_else(|error| panic!("{text} did not parse: {error}"));
rows(&plan, plan.root(), &stats)
}
fn stat(text: &str, tables: &[(&str, u64)]) -> Stat<u64> {
let stats = statistics(tables);
let plan =
Plan::parse(text).unwrap_or_else(|error| panic!("{text} did not parse: {error}"));
rows_stat(&plan, plan.root(), &stats)
}
const GUESSED: Class = Class::Estimated;
#[test]
fn a_scan_is_what_the_catalog_said_and_nothing_when_nobody_said() {
let text = scan("t", 0);
assert_eq!(estimate(&text, &[("t", 5000)]), Some(5000));
assert_eq!(estimate(&text, &[]), None);
assert_eq!(estimate(&text, &[("t", 0)]), Some(0));
}
#[test]
fn not_knowing_travels_up_rather_than_being_rounded_away() {
let text = format!("Filter (#0.0::INTEGER > 1::INTEGER)::BOOLEAN\n {}", scan("t", 0));
assert_eq!(estimate(&text, &[]), None);
assert!(estimate(&text, &[("t", 1000)]).is_some());
}
#[test]
fn an_ungrouped_aggregate_is_one_row_whatever_is_under_it() {
let text = format!("Aggregate #1 groups=[] aggregates=[]\n {}", scan("t", 0));
assert_eq!(estimate(&text, &[]), Some(1));
assert_eq!(estimate(&text, &[("t", 9_000_000)]), Some(1));
}
#[test]
fn a_group_by_collapses_its_input_and_a_scan_under_it_still_decides_whether_it_can() {
let text = format!("Aggregate #1 groups=[#0.0::INTEGER] aggregates=[]\n {}", scan("t", 0));
assert_eq!(estimate(&text, &[("t", 1000)]), Some(100));
assert_eq!(estimate(&text, &[]), None);
}
#[test]
fn a_limit_is_a_ceiling_even_over_an_input_nobody_measured() {
let text = format!("Limit 10 offset 0\n {}", scan("t", 0));
assert_eq!(estimate(&text, &[]), Some(10));
assert_eq!(estimate(&text, &[("t", 3)]), Some(3));
assert_eq!(estimate(&text, &[("t", 3_000_000)]), Some(10));
}
#[test]
fn an_offset_with_no_limit_takes_rows_away_and_cannot_add_any() {
let text = format!("Limit ALL offset 5\n {}", scan("t", 0));
assert_eq!(estimate(&text, &[("t", 12)]), Some(7));
assert_eq!(estimate(&text, &[("t", 2)]), Some(0));
assert_eq!(estimate(&text, &[]), None);
}
#[test]
fn a_filter_never_estimates_a_relation_away_entirely() {
let and = "(#0.0::INTEGER > 1::INTEGER)::BOOLEAN AND (#0.0::INTEGER > 2::INTEGER)::BOOLEAN \
AND (#0.0::INTEGER > 3::INTEGER)::BOOLEAN AND \
(#0.0::INTEGER > 4::INTEGER)::BOOLEAN AND (#0.0::INTEGER > 5::INTEGER)::BOOLEAN \
AND (#0.0::INTEGER > 6::INTEGER)::BOOLEAN";
let text = format!("Filter ({and})::BOOLEAN\n {}", scan("t", 0));
assert_eq!(estimate(&text, &[("t", 10)]), Some(1));
let one = format!("Filter (#0.0::INTEGER > 1::INTEGER)::BOOLEAN\n {}", scan("t", 0));
assert_eq!(estimate(&one, &[("t", 1_000_000)]), Some(200_000));
}
#[test]
fn a_condition_that_reads_no_column_is_not_counted_as_a_condition() {
let both = format!(
"Filter ((#0.0::INTEGER > 1::INTEGER)::BOOLEAN AND (1::INTEGER > 2::INTEGER)::BOOLEAN)::BOOLEAN\n {}",
scan("t", 0)
);
assert_eq!(estimate(&both, &[("t", 1_000_000)]), Some(200_000));
let alone = format!("Filter (1::INTEGER > 2::INTEGER)::BOOLEAN\n {}", scan("t", 0));
assert_eq!(estimate(&alone, &[("t", 1_000_000)]), Some(1_000_000));
}
#[test]
fn an_inner_join_comes_out_the_size_of_its_larger_side() {
let text = format!(
"Join INNER on=[(#0.0::INTEGER = #1.0::INTEGER)::BOOLEAN]\n {} {}",
scan("small", 0),
scan("big", 1)
);
assert_eq!(estimate(&text, &[("small", 10_000), ("big", 50_000)]), Some(50_000));
assert_eq!(estimate(&text, &[("small", 10_000)]), None);
}
#[test]
fn a_join_with_no_condition_is_the_product_and_says_so() {
let text = format!("Join INNER on=[]\n {} {}", scan("small", 0), scan("big", 1));
assert_eq!(estimate(&text, &[("small", 1000), ("big", 1000)]), Some(1_000_000));
}
#[test]
fn an_outer_join_never_estimates_below_the_side_it_preserves() {
let text = format!(
"Join LEFT on=[(#0.0::INTEGER = #1.0::INTEGER)::BOOLEAN]\n {} {}",
scan("big", 0),
scan("small", 1)
);
assert_eq!(estimate(&text, &[("big", 50_000), ("small", 10)]), Some(50_000));
}
#[test]
fn a_semi_join_is_bounded_by_its_left_side_and_ignores_the_right() {
let text = format!(
"Join SEMI on=[(#0.0::INTEGER = #1.0::INTEGER)::BOOLEAN]\n {} {}",
scan("small", 0),
scan("big", 1)
);
let estimated =
estimate(&text, &[("small", 1000), ("big", 9_000_000)]).expect("both sides known");
assert!(estimated <= 1000, "a semi join produced {estimated} out of 1000 left rows");
}
#[test]
fn a_cross_product_of_two_enormous_sides_saturates_rather_than_wrapping() {
let text = format!("CrossProduct\n {} {}", scan("a", 0), scan("b", 1));
assert_eq!(estimate(&text, &[("a", u64::MAX), ("b", 2)]), Some(u64::MAX));
}
#[test]
fn a_union_all_is_both_sides_and_so_is_the_bound_on_the_rest_of_them() {
let text = format!("SetOp UNION ALL #2\n {} {}", scan("a", 0), scan("b", 1));
assert_eq!(estimate(&text, &[("a", 30), ("b", 12)]), Some(42));
}
#[test]
fn a_count_that_came_from_the_catalog_says_it_is_exact() {
assert_eq!(stat(&scan("t", 0), &[("t", 5000)]).class(), Some(Class::Exact));
assert_eq!(stat(&scan("t", 0), &[]).class(), None);
}
#[test]
fn every_number_here_says_where_it_came_from() {
assert_eq!(stat(&scan("t", 0), &[("t", 5000)]).provenance(), Some(Provenance::RowCount));
let text = format!("Filter (#0.0::INTEGER > 1::INTEGER)::BOOLEAN\n {}", scan("t", 0));
assert_eq!(stat(&text, &[("t", 1000)]).provenance(), Some(Provenance::Default));
}
#[test]
fn a_cardinality_is_for_deciding_and_answers_nothing() {
let text = format!("Filter (#0.0::INTEGER > 1::INTEGER)::BOOLEAN\n {}", scan("t", 0));
let guessed = stat(&text, &[("t", 1000)]);
assert_eq!(guessed.decide(), Some(&200));
assert_eq!(guessed.answer(), None);
assert_eq!(guessed.enable(), None);
let counted = stat(&scan("t", 0), &[("t", 5000)]);
assert_eq!(counted.answer(), Some(&5000));
assert_eq!(counted.enable(), Some(&5000));
let nothing = stat(&scan("t", 0), &[]);
assert_eq!(nothing.decide(), None);
assert_eq!(nothing.answer(), None);
assert_eq!(nothing.enable(), None);
}
#[test]
fn one_guess_anywhere_under_a_node_makes_the_node_a_guess() {
let text = format!("Filter (#0.0::INTEGER > 1::INTEGER)::BOOLEAN\n {}", scan("t", 0));
assert_eq!(stat(&text, &[("t", 1000)]).class(), Some(GUESSED));
let twice = format!(
"Aggregate #1 groups=[#0.0::INTEGER] aggregates=[]\n Filter (#0.0::INTEGER > \
1::INTEGER)::BOOLEAN\n {}",
scan("t", 0)
);
assert_eq!(stat(&twice, &[("t", 1000)]).class(), Some(GUESSED));
}
#[test]
fn an_ungrouped_aggregate_is_exact_because_one_row_is_a_fact() {
let text = format!("Aggregate #1 groups=[] aggregates=[]\n {}", scan("t", 0));
assert_eq!(stat(&text, &[]).class(), Some(Class::Exact));
}
#[test]
fn a_limit_over_an_unmeasured_input_is_certified_rather_than_estimated() {
let text = format!("Limit 10 offset 0\n {}", scan("t", 0));
assert_eq!(
stat(&text, &[]).class(),
Some(Class::Certified { bound: 1.0, direction: Direction::AtMost })
);
assert_eq!(stat(&text, &[("t", 3)]).class(), Some(Class::Exact));
}
#[test]
fn a_join_with_no_condition_is_a_product_and_the_product_is_exact() {
let product = format!("Join INNER on=[]\n {} {}", scan("a", 0), scan("b", 1));
assert_eq!(stat(&product, &[("a", 1000), ("b", 1000)]).class(), Some(Class::Exact));
let equi = format!(
"Join INNER on=[(#0.0::INTEGER = #1.0::INTEGER)::BOOLEAN]\n {} {}",
scan("a", 0),
scan("b", 1)
);
assert_eq!(stat(&equi, &[("a", 1000), ("b", 1000)]).class(), Some(GUESSED));
}
#[test]
fn a_table_function_is_unknown_and_stays_unknown_over_it() {
let text = "TableFunction range args=[] #0 [a::BIGINT]\n";
assert_eq!(stat(text, &[]), Stat::Unknown);
}
fn over_files(text: &str, files: &[(&str, u64)]) -> Stat<u64> {
let mut stats = Statistics::new();
for (path, count) in files {
stats.record_file(path, *count);
}
let plan =
Plan::parse(text).unwrap_or_else(|error| panic!("{text} did not parse: {error}"));
rows_stat(&plan, plan.root(), &stats)
}
fn read_parquet(path: &str) -> String {
format!("TableFunction read_parquet args=['{path}'::VARCHAR] #0 [a::BIGINT]\n")
}
#[test]
fn a_parquet_file_somebody_counted_is_exact_because_the_footer_counted_it() {
let text = read_parquet("lineitem.parquet");
let stat = over_files(&text, &[("lineitem.parquet", 6_001_215)]);
assert_eq!(stat.value().copied(), Some(6_001_215));
assert_eq!(stat.class(), Some(Class::Exact));
}
#[test]
fn a_parquet_file_nobody_counted_is_unknown_rather_than_zero() {
assert_eq!(over_files(&read_parquet("gone.parquet"), &[]), Stat::Unknown);
assert_eq!(over_files(&read_parquet("a.parquet"), &[("b.parquet", 10)]), Stat::Unknown);
}
#[test]
fn a_reader_that_is_not_parquet_is_unknown_whatever_was_counted() {
let text = "TableFunction read_csv args=['t.csv'::VARCHAR] #0 [a::BIGINT]\n";
assert_eq!(over_files(text, &[("t.csv", 500)]), Stat::Unknown);
}
#[test]
fn statistics_that_nobody_filled_in_say_so() {
let mut stats = Statistics::new();
assert!(stats.is_empty());
stats.record("memory", "main", "t", 7);
assert!(!stats.is_empty());
assert_eq!(stats.rows_in("memory", "main", "t"), Some(7));
assert_eq!(stats.rows_in("memory", "other", "t"), None);
}
#[test]
fn a_file_count_on_its_own_is_something_known() {
let mut stats = Statistics::new();
stats.record_file("t.parquet", 12);
assert!(!stats.is_empty(), "a database with no tables and one file knows something");
assert_eq!(stats.rows_in_file("t.parquet"), Some(12));
assert_eq!(stats.rows_in_file("other.parquet"), None);
}
}