rudb-opt 0.4.13

//! How large a grouped aggregate's hash table should be before its first row arrives.
//!
//! A hash aggregate starts with sixty four buckets and doubles. Every doubling walks the groups it
//! already holds and writes each one into a new bucket array, so a group by that ends with a
//! million groups has written about two million buckets on the way to holding one million, and the
//! last few doublings are the ones that move nearly all of it. None of that work produces a row.
//! `spec/stats/05-every-query.md` section 5.4 says the number that removes it is the distinct count
//! of the grouping key, and that the count is exact for the keys this matters most on.
//!
//! This pass writes one number per aggregate and moves nothing. A plan it has run over produces the
//! rows it produced before, in the order it produced them, because the operator it is talking to
//! reads the number as a size and never as a count of anything.
//!
//! # The number is a ceiling and not an estimate
//!
//! Two facts bound how many groups an aggregate can make, and both point the same way.
//!
//! A grouping key that is one column of one table cannot make more groups than that column has
//! distinct values. Two key columns cannot make more than the product of theirs. Neither is affected
//! by anything between the scan and the aggregate, because a filter removes rows and removing rows
//! never adds a group.
//!
//! The rows arriving are the other bound, because a group needs a row in it.
//!
//! What it deliberately does not take is [`estimate::rows`] of the aggregate itself, which is the
//! modelled group count: that number is a guess about how the keys landed, it is
//! [`rudb_common::Class::Estimated`] wherever it comes from, and sizing an allocation from a guess is
//! the mistake section 5.1 names.
//!
//! # Only where the key is the smaller of the two bounds
//!
//! The smaller of two ceilings is a ceiling, but the two are not equally good sizes, and taking the
//! smaller one regardless was part of a three percent regression on TPC-H. The pass takes the key's
//! ceiling and refuses when the rows arriving are the smaller bound.
//!
//! The reason is that the two numbers say different things. A counted key column's distinct count is
//! how many groups there will be, exactly so where nothing filtered underneath and near enough where
//! something did. The rows arriving say nothing at all about how many groups there are: a `GROUP BY`
//! that turns six million rows into fifty seven of them has six million as its row bound, and the
//! room that asks for is six million buckets for fifty seven groups.
//!
//! So where the product of the key's counts is above the rows, the honest reading is that this pass
//! has learned nothing about this key and the table should start where it always started. That is
//! also what happens to a key of several columns, because the product of five distinct counts is a
//! ceiling nothing ever approaches. TPC-H q18 groups on five columns of a join of several million
//! rows and produces fifty seven of them, and sizing that aggregate from the rows was 0.34 G of
//! instructions on its own, most of it the kernel clearing pages for buckets no row was ever written
//! into. The larger part of what q18 lost was the other aggregate in it, and that one was the
//! executor reading the number for tables that never hold the groups, which `rudb_exec`'s `Share`
//! is about.
//!
//! # Why a ceiling is the safe end of the range here
//!
//! Section 5.1 says to pick the end of the range whose failure you can afford, and the two failures
//! are not symmetric. A table sized under the truth grows the way it grew before this pass existed,
//! which costs the rehashing this pass was written to remove and nothing else. A table sized over
//! the truth holds bucket memory it never fills, and bucket memory is charged against the query's
//! budget, so a large enough overshoot turns a query that ran into a query that reports being out of
//! memory. That is worse than a slow query, and it is why `MOST` exists.
//!
//! # What it leaves alone
//!
//! An ungrouped aggregate, which produces exactly one row and has no table.
//!
//! A key that is not columns. `GROUP BY lower(c)` has as many groups as `c` has values at most, and
//! reading through the expression to say so is a separate piece of work that is not this one.
//!
//! A key column nobody counted, which is a column of a table with no dictionary, no sketch and no
//! footer entry. There is no ceiling to take and the honest answer is the size the table always
//! started at.
//!
//! An input whose rows nobody can say, since the test above is against that number and a test that
//! cannot be run is a test that has not passed.
//!
//! A ceiling that is already inside the first bucket array, because a table that was never going to
//! grow has nothing to save.

use rudb_common::rules::Rule;
use rudb_common::{Class, Direction, Result, Stat};
use rudb_plan::{Node, Plan};

use crate::estimate::{self, DISTINCT, Facts};
use crate::pass::{Context, Pass, top_down};

/// The most groups this pass will ask for room for.
///
/// Sixteen million buckets at eight bytes each, which is the hundred and twenty eight megabytes an
/// aggregate instance would be holding before it had folded a row in. Every instance of a
/// partitioned aggregate would hold its own, so the number that matters is this one times the thread
/// count, and that is the reasoning behind it being a ceiling rather than a size: an overshoot here
/// is charged against the query's memory budget and can fail a query that used to run. Beyond this
/// the doubling can have the rest, which by then is the last two or three of them.
const MOST: u64 = 8 << 20;

/// The most groups a table that has not grown yet can hold.
///
/// `rudb_exec`'s table starts at sixty four buckets and grows when it is half full, so it holds
/// thirty two groups before the first doubling. A ceiling at or under this is a table that was never
/// going to grow, and asking for the size it already has would put an entry in the plan that changes
/// nothing.
const ALREADY: u64 = 32;

/// Writes the group count onto every aggregate that has a ceiling worth acting on.
///
/// A rudb name rather than a DuckDB one, because DuckDB has no pass that does this and [`crate::UPSTREAM`]
/// is the list of names it does have.
///
/// Two settings turn it off and they mean different things. `SET disabled_optimizers =
/// 'aggregate_presize'` is the pass, which is the door DuckDB's name for a pass goes through. `SET
/// stats_presize = 'off'` is [`Rule::Presize`], which is the rule, and that is the door
/// `spec/stats/09-measurement.md` section 9.2 asks for so that a report can say what this rule on its
/// own earned. `SET statistics = 'off'` is the master over the second of them.
#[derive(Debug)]
pub struct AggregatePresize;

impl Pass for AggregatePresize {
    fn name(&self) -> &'static str {
        "aggregate_presize"
    }

    fn run(&self, plan: &mut Plan, context: &Context) -> Result<()> {
        if context.allows(Rule::Presize) {
            size(plan, context.facts());
        }
        Ok(())
    }
}

/// Records a size for every aggregate in `plan` that has one.
///
/// Idempotent, because the answer is a function of the plan's shape and the facts, and this pass
/// changes neither. Running it twice writes the same entries over the same entries.
fn size(plan: &mut Plan, stats: &Facts) {
    let mut found = Vec::new();
    for node in top_down(plan) {
        let Node::Aggregate { input, index, groups, .. } = *plan.node(node) else {
            continue;
        };
        let keys = plan.expr_list(groups);
        if keys.is_empty() {
            continue;
        }
        let Some(ceiling) = ceiling(plan, input, keys, stats) else {
            continue;
        };
        if ceiling <= ALREADY {
            continue;
        }
        found.push((index, ceiling.min(MOST)));
    }
    for (index, ceiling) in found {
        plan.presize(index, ceiling);
    }
}

/// The most groups the aggregate over `input` keyed on `keys` could possibly produce, when that
/// number is the key's own and not the row count wearing a ceiling's clothes.
///
/// `None` when nothing bounds it, which is the ordinary answer for a key that is an expression and
/// for a column of a table nobody counted, and `None` again when the rows arriving are the smaller
/// of the two bounds. The second one is the whole of the rule and the reason is in the module
/// documentation above.
fn ceiling(
    plan: &Plan,
    input: rudb_plan::NodeRef,
    keys: &[rudb_plan::ExprRef],
    stats: &Facts,
) -> Option<u64> {
    let bindings = estimate::keyed(plan, keys)?;
    let mut values: u64 = 1;
    for binding in bindings {
        values = values.saturating_mul(bounded(estimate::stated(plan, binding, stats))?);
    }
    // A group needs a row in it, so the rows arriving bound the groups too, and where they are the
    // smaller bound this pass has learned nothing about the key and asks for nothing.
    (values <= estimate::rows(plan, input, stats)?).then_some(values)
}

/// The distinct count read as a ceiling, or `None` when it is not one.
///
/// [`Class::Exact`] is a ceiling because it is the number. [`Direction::AtMost`] is a ceiling by what
/// the direction means, whatever the bound is. The other two directions are not ceilings at all, and
/// they are taken anyway when the bound is inside a factor of two, because the table doubles: a size
/// wrong by less than that is one doubling from right in whichever direction it is wrong, and one
/// doubling is what this pass is trying to save fifteen of.
///
/// [`Class::Estimated`] is refused. The number a guess would size the allocation from is the thing
/// section 5.1 says not to allocate from.
fn bounded(stat: Stat<u64>) -> Option<u64> {
    let value = *stat.read(DISTINCT)?;
    match stat.class()? {
        Class::Exact => Some(value),
        Class::Certified { direction: Direction::AtMost, .. } => Some(value),
        Class::Certified { bound, .. } if bound <= 1.0 => Some(value),
        Class::Certified { .. } | Class::Estimated => None,
    }
}

#[cfg(test)]
mod tests {
    use rudb_common::Provenance;
    use rudb_plan::Plan;

    use super::{AggregatePresize, MOST, Rule};
    use crate::estimate::Facts;
    use crate::pass::{Context, Pass};

    /// One table of one column, named the way the printer names one.
    const SCAN: &str = "Get memory.main.t AS t #0 [a::INTEGER]";

    /// A plan over one table with one column, grouped on that column.
    fn grouped(filter: Option<&str>) -> Plan {
        let scan = SCAN;
        let text = match filter {
            None => format!("Aggregate #1 groups=[#0.0::INTEGER] aggregates=[]\n  {scan}\n"),
            Some(predicate) => format!(
                "Aggregate #1 groups=[#0.0::INTEGER] aggregates=[]\n  Filter {predicate}\n    \
                 {scan}\n"
            ),
        };
        Plan::parse(&text).expect("a plan that parses")
    }

    fn counted(rows: u64, distinct: u64) -> Context {
        let mut facts = Facts::new();
        facts.record("memory", "main", "t", rows);
        facts.record_distinct("memory", "main", "t", "a", distinct, Provenance::Dictionary);
        let mut context = Context::new();
        context.measure(std::sync::Arc::new(facts));
        context
    }

    fn run(plan: &mut Plan, context: &Context) {
        AggregatePresize.run(plan, context).expect("a pass that cannot fail");
    }

    /// The same facts with one rule turned off, which is what an ablation run does.
    fn without(rule: Rule) -> Context {
        let mut context = counted(1_000_000, 50_000);
        let mut rules = rudb_common::rules::Rules::default();
        rules.set(rule, false);
        context.govern(rules);
        context
    }

    #[test]
    fn the_rule_s_own_setting_turns_it_off() {
        let mut plan = grouped(None);
        run(&mut plan, &without(Rule::Presize));
        assert_eq!(plan.presized_count(), 0, "stats_presize = off asked for no room");
    }

    #[test]
    fn the_master_setting_turns_it_off_too() {
        // `statistics = off` reaches this without naming it, which is the point of a master: the
        // ablation of section 9.3 is one statement and it has to cover a rule written after it.
        let mut plan = grouped(None);
        run(&mut plan, &without(Rule::StatsAll));
        assert_eq!(plan.presized_count(), 0, "statistics = off asked for no room");
    }

    #[test]
    fn a_counted_key_column_sizes_the_table() {
        let mut plan = grouped(None);
        run(&mut plan, &counted(1_000_000, 50_000));
        assert_eq!(plan.presized(1), Some(50_000));
    }

    #[test]
    fn a_key_the_rows_bound_more_tightly_than_its_count_does_is_left_alone() {
        // More distinct values than there are rows, which a fold across parts can say. The rows
        // are the smaller of the two bounds, so this pass has learned nothing about the key and
        // asks for nothing. Taking the rows instead is how a group by that makes fifty seven
        // groups out of millions of rows asked for room for millions of them.
        let mut plan = grouped(None);
        run(&mut plan, &counted(4_000, 50_000));
        assert_eq!(plan.presized_count(), 0);
    }

    #[test]
    fn a_filter_underneath_does_not_raise_the_ceiling() {
        // The filter's own estimate is below the row count and the key ceiling is unchanged by it,
        // because removing rows never adds a group. Either the key is still the smaller bound and
        // is what gets taken, or the filter has cut the rows below it and nothing gets taken, and
        // neither one can name a number above the key's own count.
        let mut plan = grouped(Some("(#0.0::INTEGER > 10::INTEGER)::BOOLEAN"));
        run(&mut plan, &counted(1_000_000, 50_000));
        let sized = plan.presized(1).unwrap_or(0);
        assert!(sized <= 50_000, "{sized} is above the key's own ceiling");
    }

    #[test]
    fn a_key_of_several_columns_is_left_alone_because_the_product_is_not_a_size() {
        // Two counted columns of a table with far fewer rows than their product. The product is a
        // real ceiling and a useless size, which is the shape every wide group by in TPC-H has.
        let text = "Aggregate #1 groups=[#0.0::INTEGER, #0.1::INTEGER] aggregates=[]\n  Get \
                    memory.main.t AS t #0 [a::INTEGER, b::INTEGER]\n";
        let mut plan = Plan::parse(text).expect("a plan that parses");
        let mut facts = Facts::new();
        facts.record("memory", "main", "t", 1_000_000);
        facts.record_distinct("memory", "main", "t", "a", 50_000, Provenance::Dictionary);
        facts.record_distinct("memory", "main", "t", "b", 50_000, Provenance::Dictionary);
        let mut context = Context::new();
        context.measure(std::sync::Arc::new(facts));
        run(&mut plan, &context);
        assert_eq!(plan.presized_count(), 0);
    }

    #[test]
    fn a_column_nobody_counted_is_left_alone() {
        let mut plan = grouped(None);
        let mut facts = Facts::new();
        facts.record("memory", "main", "t", 1_000_000);
        let mut context = Context::new();
        context.measure(std::sync::Arc::new(facts));
        run(&mut plan, &context);
        assert_eq!(plan.presized(1), None);
        assert_eq!(plan.presized_count(), 0);
    }

    #[test]
    fn an_ungrouped_aggregate_has_no_table_to_size() {
        let text = format!("Aggregate #1 groups=[] aggregates=[count_star()::BIGINT]\n  {SCAN}\n");
        let mut plan = Plan::parse(&text).expect("a plan that parses");
        run(&mut plan, &counted(1_000_000, 50_000));
        assert_eq!(plan.presized_count(), 0);
    }

    #[test]
    fn a_key_that_never_fills_the_first_buckets_is_left_alone() {
        let mut plan = grouped(None);
        run(&mut plan, &counted(1_000_000, 8));
        assert_eq!(plan.presized_count(), 0);
    }

    #[test]
    fn a_ceiling_past_the_cap_is_capped() {
        let mut plan = grouped(None);
        run(&mut plan, &counted(u64::MAX, u64::MAX));
        assert_eq!(plan.presized(1), Some(MOST));
    }

    #[test]
    fn an_input_whose_rows_nobody_can_say_is_left_alone() {
        // The table's count is not on record, so the test the key has to pass cannot be run.
        let mut plan = grouped(None);
        let mut facts = Facts::new();
        facts.record_distinct("memory", "main", "t", "a", 50_000, Provenance::Dictionary);
        let mut context = Context::new();
        context.measure(std::sync::Arc::new(facts));
        run(&mut plan, &context);
        assert_eq!(plan.presized_count(), 0);
    }

    #[test]
    fn a_second_run_writes_what_the_first_one_wrote() {
        let mut plan = grouped(None);
        let context = counted(1_000_000, 50_000);
        run(&mut plan, &context);
        let once = plan.presized(1);
        run(&mut plan, &context);
        assert_eq!(plan.presized(1), once);
        assert_eq!(plan.presized_count(), 1);
    }

    #[test]
    fn the_pass_is_off_when_it_is_named() {
        let context = Context::without("aggregate_presize").expect("a name that is a pass");
        assert!(context.is_disabled("aggregate_presize"));
    }
}