use std::collections::BTreeMap;
use rudb_common::stat::{Class, Source, Stat};
use rudb_plan::{ConjunctionOp, Expr, ExprRef, JoinKind, Node, NodeRef, Plan, SetOpKind};
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>,
}
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()
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.tables.is_empty()
}
}
const GUESSED: Class = Class::Estimated { source: Source::Constant };
const CEILING: Class = Class::Certified { bound: 1.0 };
#[must_use]
pub fn rows(plan: &Plan, node: NodeRef, stats: &Statistics) -> Option<u64> {
rows_stat(plan, node, stats).value().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),
Node::Get { catalog, schema, table, .. } => {
match stats.rows_in(plan.string(catalog), plan.string(schema), plan.string(table)) {
Some(rows) => Stat::exact(rows),
None => Stat::Unknown,
}
}
Node::Values { rows: list, .. } => {
u64::try_from(plan.row_list(list).len()).map_or(Stat::Unknown, Stat::exact)
}
Node::TableFunction { .. } | 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);
}
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 },
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 },
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 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) }
}
}
}
fn ceiling(stat: Stat<u64>) -> Stat<u64> {
match stat {
Stat::Unknown => Stat::Unknown,
Stat::Known { value, class } => Stat::Known { value, class: class.combine(CEILING) },
}
}
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 },
Stat::Known { value: right, class: right_class },
) = (left, right)
else {
return Stat::Unknown;
};
let both = left_class.combine(right_class);
if conditions == 0 {
return Stat::Known { value: left.saturating_mul(right), class: both };
}
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) }
}
}
}
#[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, Source, 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 { source: Source::Constant };
#[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 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 }));
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);
}
#[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);
}
}