rudb-exec 0.4.32

//! The operators that pass over their input one chunk at a time and never hold it.
//!
//! These three are what a pipeline is made of. None of them allocates anything proportional to the
//! input, none of them can block, and each one either hands its chunk on, narrows it or replaces
//! its columns. That is the property the morsel driven scheduler in section 7.2 needs, because a
//! morsel is a run of a scan pushed through every streaming operator above it by one thread.
//!
//! All three are [`Stream`] implementations, which means they take `&self` and are handed the
//! mutable part separately. A filter's mutable part is the scratch space its predicate evaluates
//! into and a limit's is the two counters, and naming them is what lets one of these be
//! instantiated on thirty two threads later without copying the predicate thirty two times. What
//! drives them is the serial driver in `rudb-pipeline`, which pushes one chunk through every stream
//! of a pipeline in the order `build.rs` stacked them.

use rudb_common::{Field, Result, Session};
use rudb_kernels::row_count;
use rudb_pipeline::{Compaction, Gauge, Progress, Stream, narrow};
use rudb_plan::{ExprRef, Node, NodeRef, Plan, Slice};
use rudb_seam::{Context, SeamId, Settings};
use rudb_vector::{Chunk, Selection};

use crate::prepared::{Prepared, Scratch};
use crate::register::compaction;
use crate::schema::Schema;

/// Keeps the rows where a predicate is true.
///
/// True, not "not false". A null predicate drops the row, which is what makes `WHERE x <> 5` leave
/// out the rows where `x` is null.
///
/// The predicate is evaluated with [`Prepared::evaluate_filter`] rather than as an expression, so a
/// top level `AND` runs a conjunct at a time over the rows the conjuncts before it left and stops
/// the moment nothing is left. The difference on a four conjunct predicate is the difference between
/// reading every row four times and reading it once.
///
/// What happens to the kept rows is the `chunk.compaction` seam. They become a selection over the
/// chunk, which is section 7.1's rule and costs the selection and not the payload, or they are
/// copied out into a chunk of their own, which costs the copy and saves every later read a
/// redirection. The implementation that never copies is the reference and the default, so the
/// engine does what it did before this seam existed until a sweep says otherwise. A chunk that
/// keeps nothing is left empty rather than passed on with rows in it, and whoever is driving skips
/// it, because an empty chunk travelling up a deep pipeline is work every operator above does for
/// no rows.
#[derive(Debug)]
pub(crate) struct Filter {
    predicate: Prepared,
    compaction: &'static dyn Compaction,
    passes: u32,
    /// Whether the aggregate this feeds reads a marked chunk. See [`Chunk::marked`].
    marks: bool,
}

/// Everything one instance of a filter mutates, which is the predicate's scratch and the seam's.
#[derive(Debug)]
pub(crate) struct Filtering {
    scratch: Scratch,
    gauge: Gauge,
}

impl Filter {
    /// Applies the session semantics to the predicate's prepared casts.
    #[must_use]
    pub(crate) fn in_session(mut self, session: &Session) -> Self {
        self.predicate = self.predicate.in_session(session);
        self
    }

    /// Marks the rows it keeps rather than cutting them, for a filter that feeds an aggregate that
    /// reads a marked chunk.
    #[must_use]
    pub(crate) fn marking(mut self, marks: bool) -> Self {
        self.marks = marks;
        self
    }

    /// # Errors
    ///
    /// If the predicate does not resolve against the input's schema, which is a failure of the plan
    /// and is found when the operator is built rather than on the first chunk. Also if the session
    /// has pinned the compaction seam to something that cannot run over these columns, which is an
    /// error rather than a quiet fall back, because a run that did not do what the setting asked
    /// for is a run whose number says something other than what it means.
    pub(crate) fn new(
        plan: &Plan,
        node: NodeRef,
        predicate: ExprRef,
        input: &Schema,
        seams: &Settings,
    ) -> Result<Self> {
        let types = input.types();
        let context = Context::new(SeamId::ChunkCompaction, seams).with_types(&types);
        let compaction = compaction().choose(&context)?.strategy();
        Ok(Self {
            predicate: Prepared::one(plan, predicate, input)?,
            compaction,
            passes: later_passes(plan, node),
            marks: false,
        })
    }
}

impl Stream for Filter {
    type Local = Filtering;

    fn local(&self) -> Filtering {
        Filtering { scratch: self.predicate.scratch(), gauge: Gauge::new(self.passes) }
    }

    fn push(&self, chunk: &mut Chunk, local: &mut Filtering) -> Result<Progress> {
        let kept = self.predicate.evaluate_filter(chunk, &mut local.scratch)?;
        if kept.len() != chunk.len() {
            if self.marks && marking_pays(kept.len(), chunk.len()) {
                let whole = std::mem::replace(chunk, Chunk::empty(&[]));
                *chunk = whole.marked(kept);
            } else {
                narrow(self.compaction, chunk, &kept, &mut local.gauge)?;
            }
        }
        Ok(Progress::More)
    }

    /// A pass over the rows for every step of the predicate that does anything to one.
    ///
    /// A comparison against a literal is one, and a filter of five of them under an `AND` is six.
    /// See [`Prepared::passes`].
    fn weight(&self) -> usize {
        self.predicate.passes()
    }
}

/// Whether a filter that feeds an aggregate should mark the rows it kept rather than cut them out.
///
/// Marking saves the copy of every column and costs the aggregate a pass over the dropped rows,
/// which it counts into no group. On q01 a row the aggregate reads costs about what copying it
/// costs, so marking wins while the filter keeps most rows and loses when it drops most. Three
/// quarters is the line, which leaves the dropped rows at most a third of the kept ones.
pub(crate) fn marking_pays(kept: usize, rows: usize) -> bool {
    kept.saturating_mul(4) >= rows.saturating_mul(3)
}

/// How many more times the rows a filter keeps will be read, counted from the plan above it.
///
/// This is the number the gain function is written in terms of, and the measurement in
/// [`Chunk::compact`] is why: compacting loses at every selectivity when there is one later pass
/// over the kept rows and wins at every selectivity when there are sixteen. A filter that cannot
/// find itself in the plan is treated as having one pass above it, which is the answer that makes
/// the gain function say no.
pub(crate) fn later_passes(plan: &Plan, filter: NodeRef) -> u32 {
    passes_between(plan, plan.root(), filter).unwrap_or(1)
}

/// The reads on the path from `node` down to `filter`, or `None` when the filter is not under it.
fn passes_between(plan: &Plan, node: NodeRef, filter: NodeRef) -> Option<u32> {
    if node == filter {
        return Some(0);
    }
    let here = reads(plan.node(node));
    for child in plan.node(node).children().into_iter().flatten() {
        if let Some(below) = passes_between(plan, child, filter) {
            return Some(below + here);
        }
    }
    None
}

/// How many times an operator reads the rows that reach it.
///
/// Coarse on purpose. What the gain function needs is the difference between an operator that hands
/// its chunk on and one that holds it, since the second kind reads the rows again after storing
/// them and is where the paper's ten percent comes from. The numbers themselves are first estimates
/// and the sweep over this seam is what turns them into measurements, which is the same status the
/// nanosecond constants in the gain function have.
fn reads(node: &Node) -> u32 {
    match node {
        // A limit hands rows on without looking at them, and a source is never above a filter.
        Node::Limit { .. }
        | Node::Get { .. }
        | Node::Dummy
        | Node::Values { .. }
        | Node::TableFunction { .. }
        | Node::CteScan { .. } => 0,
        // The rows of a materialisation are written once and read once per read of it, and the
        // filter this is counting for is either in the definition, where the write is the pass, or
        // in the body, where this node is not on the path at all.
        Node::MaterializedCte { .. } => 1,
        // A lateral call reads its chunk once to work out the calls and once more to copy the rows
        // out beside what they produced, which is the same order of work as a fetch and is counted
        // the same way.
        Node::Filter { .. }
        | Node::Project { .. }
        | Node::Fetch { .. }
        | Node::TableFetch { .. }
        | Node::LateralFunction { .. } => 1,
        // A share of the input holds every chunk until the input has ended, writes them once and
        // reads them back once, which is the same shape as the breakers beside it here.
        Node::Aggregate { .. }
        | Node::Window { .. }
        | Node::Distinct { .. }
        | Node::LimitPercent { .. }
        | Node::SetOp { .. } => 2,
        Node::TopN { .. } | Node::CrossProduct { .. } => 2,
        Node::Join { .. } | Node::DependentJoin { .. } => 3,
        // Once, which is the whole claim of section 5.2. There is no build side to write and read
        // back, no hash table to probe and no parent row copied anywhere: the child is scanned, the
        // link is read beside its columns, and the parent's columns are gathered rather than
        // materialised. A link join reads its rows the number of times a filter does.
        Node::LinkJoin { .. } => 1,
        Node::Sort { .. } => 3,
    }
}

/// Replaces the input's columns with a list of expressions.
#[derive(Debug)]
pub(crate) struct Project {
    exprs: Prepared,
    schema: Schema,
}

impl Project {
    /// Applies the session semantics to the projection's prepared casts.
    #[must_use]
    pub(crate) fn in_session(mut self, session: &Session) -> Self {
        self.exprs = self.exprs.in_session(session);
        self
    }

    /// A projection producing the plan's expressions under the plan's names.
    ///
    /// # Errors
    ///
    /// If there are not as many names as expressions, which [`Plan::validate`] already rejects and
    /// which is checked again here because this operator would otherwise produce a schema that is
    /// silently short.
    pub(crate) fn new(
        plan: &Plan,
        input: &Schema,
        index: u32,
        exprs: Slice,
        names: Slice,
    ) -> Result<Self> {
        let exprs: Vec<ExprRef> = plan.expr_list(exprs).to_vec();
        let names = plan.name_list(names);
        if names.len() != exprs.len() {
            return Err(rudb_common::Error::internal(format!(
                "a projection of {} expressions under {} names",
                exprs.len(),
                names.len()
            )));
        }
        let fields = exprs
            .iter()
            .zip(names)
            .map(|(&expr, &name)| Field::new(plan.string(name), plan.expr_type(expr).clone()))
            .collect();
        let exprs = Prepared::new(plan, &exprs, input)?;
        Ok(Self { exprs, schema: Schema::numbered(fields, index) })
    }

    /// The columns this projection produces.
    pub(crate) fn schema(&self) -> &Schema {
        &self.schema
    }
}

impl Stream for Project {
    type Local = Scratch;

    fn local(&self) -> Scratch {
        self.exprs.scratch()
    }

    fn push(&self, chunk: &mut Chunk, scratch: &mut Scratch) -> Result<Progress> {
        let mut columns = Vec::with_capacity(self.exprs.len());
        let rows = chunk.len();
        let taken = std::mem::replace(chunk, Chunk::with_rows(Vec::new(), 0)?);
        self.exprs.evaluate_taking(taken, scratch, &mut columns)?;
        *chunk = Chunk::with_rows(columns, rows)?;
        Ok(Progress::More)
    }

    /// A pass over the rows for every step that does anything to one. See [`Prepared::passes`].
    ///
    /// A projection that only names columns answers zero, which is what most of them are and what
    /// the rule above the scan was calibrated against.
    fn weight(&self) -> usize {
        self.exprs.passes()
    }
}

/// Skips `offset` rows and then emits at most `count` of them.
///
/// The offset is consumed a row at a time rather than a chunk at a time, because an offset that
/// falls in the middle of a chunk is the ordinary case and rounding it to a chunk boundary is a
/// wrong answer. The chunk that reaches the count comes back with [`Progress::Done`] on it, which
/// is how `LIMIT 10` over a large table stops the scan instead of reading rows in order to throw
/// them away.
///
/// Either end can be a number the query did not write out, which is what `LIMIT (SELECT 3)` and
/// `LIMIT RANDOM() * 10` are. Those arrive as an expression over the input, because the binder put
/// the value in a column of every row, and the number is read off the first chunk and kept for the
/// rest of the query. Read once and not per chunk: a volatile call would otherwise answer a
/// different limit every chunk, and the pin reads it once too.
#[derive(Debug)]
pub(crate) struct Limit {
    count: Edge,
    offset: Edge,
}

/// One end of a limit while the query runs.
#[derive(Debug)]
pub(crate) enum Edge {
    /// Every row, which is only ever a count.
    All,
    /// A number the binder worked out.
    Rows(u64),
    /// An expression over the input, holding the number in every row.
    Read(Prepared),
}

/// How much of the limit one instance has used up, and what the limit turned out to be.
#[derive(Debug, Default)]
pub(crate) struct Taken {
    skipped: u64,
    emitted: u64,
    /// The count and the offset, once the first chunk has been looked at. `None` until then.
    settled: Option<(Option<u64>, u64)>,
    /// Working space for whichever ends are read off the rows. Empty for the ends that are not.
    counting: Scratch,
    skipping: Scratch,
}

impl Limit {
    pub(crate) fn new(count: Edge, offset: Edge) -> Self {
        Self { count, offset }
    }

    /// Applies the session semantics to whatever casts the two ends hold.
    #[must_use]
    pub(crate) fn in_session(mut self, session: &Session) -> Self {
        self.count = self.count.in_session(session);
        self.offset = self.offset.in_session(session);
        self
    }

    /// How many rows are still wanted, given what has already been emitted.
    fn room(count: Option<u64>, taken: &Taken) -> Option<u64> {
        count.map(|count| count.saturating_sub(taken.emitted))
    }
}

impl Edge {
    #[must_use]
    pub(crate) fn in_session(self, session: &Session) -> Self {
        match self {
            Self::Read(prepared) => Self::Read(prepared.in_session(session)),
            settled => settled,
        }
    }

    /// Working space sized for this end, which is nothing at all unless it is read off the rows.
    pub(crate) fn scratch(&self) -> Scratch {
        match self {
            Self::Read(prepared) => prepared.scratch(),
            Self::All | Self::Rows(_) => Scratch::default(),
        }
    }

    /// The number this end asks for, reading the chunk when that is where the number is.
    ///
    /// `None` is every row, which only a count answers. A null reads as every row too, because
    /// `LIMIT (SELECT NULL)` is every row on the pin, the same as `LIMIT NULL` written out.
    pub(crate) fn rows(
        &self,
        chunk: &Chunk,
        scratch: &mut Scratch,
        clause: &str,
    ) -> Result<Option<u64>> {
        match self {
            Self::All => Ok(None),
            Self::Rows(rows) => Ok(Some(*rows)),
            Self::Read(prepared) => {
                let value = prepared.evaluate_one(chunk, scratch)?.value_at(0);
                if value.is_null() {
                    return Ok(None);
                }
                row_count(&value, clause).map(Some)
            }
        }
    }
}

impl Stream for Limit {
    type Local = Taken;

    fn local(&self) -> Taken {
        Taken {
            counting: self.count.scratch(),
            skipping: self.offset.scratch(),
            ..Taken::default()
        }
    }

    /// Never, and this is the one operator in the tree that says so.
    ///
    /// How much of the limit has been used up is local state, and four instances each allowed to
    /// emit ten rows emit forty. Sharing a counter between them would fix the count and not the
    /// answer, because `LIMIT 10` with no `ORDER BY` would then return whichever ten rows the
    /// threads got to first, which is a different ten on every run. So the pipeline a limit is in
    /// runs on one thread, and the way a large query with a limit on it gets parallelism back is a
    /// top n, which sorts and is a sink and combines.
    fn parallel(&self) -> bool {
        false
    }

    fn push(&self, chunk: &mut Chunk, taken: &mut Taken) -> Result<Progress> {
        let rows = chunk.len() as u64;
        // Nothing to take from and, more to the point, no row to read a bound out of. An empty
        // chunk is skipped by whoever is driving, so this is belt and braces rather than the way
        // the first chunk usually arrives.
        if rows == 0 {
            return Ok(Progress::More);
        }
        let (count, offset) = match taken.settled {
            Some(settled) => settled,
            None => {
                let count = self.count.rows(chunk, &mut taken.counting, "LIMIT")?;
                // No offset written and a null offset are the same thing, which is none skipped.
                let offset = self.offset.rows(chunk, &mut taken.skipping, "OFFSET")?.unwrap_or(0);
                *taken.settled.insert((count, offset))
            }
        };
        let skipping = (offset - taken.skipped).min(rows);
        taken.skipped += skipping;
        let available = rows - skipping;
        let taking = match Self::room(count, taken) {
            Some(room) => room.min(available),
            None => available,
        };
        taken.emitted += taking;
        if skipping != 0 || taking != rows {
            let mut kept = Selection::with_capacity(taking as usize);
            for row in skipping..skipping + taking {
                kept.push(row as usize);
            }
            keep(chunk, &kept)?;
        }
        match Self::room(count, taken) {
            Some(0) => Ok(Progress::Done),
            _ => Ok(Progress::More),
        }
    }
}

/// Narrow a chunk to the rows a selection kept, in place.
///
/// [`Chunk::select`] takes the chunk by value, because taking it by value is what lets it move the
/// payload into the new vectors rather than copy it, and a push operator has a `&mut` and not a
/// value. So the chunk is swapped out for an empty one, narrowed, and put back. The empty one is
/// never observed, since a failure here fails the query.
fn keep(chunk: &mut Chunk, kept: &Selection) -> Result<()> {
    let whole = std::mem::replace(chunk, Chunk::empty(&[]));
    *chunk = whole.select(kept)?;
    Ok(())
}

/// The limit driven through the trait rather than through a plan.
///
/// The plan level tests are in `tests.rs` and they go through `build`, which is the right place to
/// check that `LIMIT 3 OFFSET 1` answers with the rows it should. These check the part only the
/// trait has, which is when [`Progress::Done`] comes back, because that is the signal that stops a
/// scan and nothing above the operator can see it once the answer has been assembled.
#[cfg(test)]
mod tests {
    use rudb_common::{LogicalType, Value};
    use rudb_vector::{Data, Vector};

    use super::{Edge, Limit, Progress, Stream};
    use rudb_vector::Chunk;

    fn chunk(values: &[i32]) -> Chunk {
        let column = Vector::flat(LogicalType::Integer, Data::Int32(values.to_vec().into()))
            .expect("integers are an i32 layout");
        Chunk::new(vec![column]).expect("one column is one length")
    }

    fn rows(chunk: &Chunk) -> Vec<Value> {
        (0..chunk.len()).map(|row| chunk.value_at(row, 0)).collect()
    }

    #[test]
    fn the_chunk_that_fills_the_count_is_the_one_that_says_done() {
        let limit = Limit::new(Edge::Rows(3), Edge::Rows(0));
        let mut taken = limit.local();

        let mut first = chunk(&[1, 2]);
        assert_eq!(limit.push(&mut first, &mut taken).expect("two rows fit"), Progress::More);
        assert_eq!(rows(&first), vec![Value::Integer(1), Value::Integer(2)]);

        let mut second = chunk(&[3, 4]);
        assert_eq!(limit.push(&mut second, &mut taken).expect("one row fits"), Progress::Done);
        assert_eq!(rows(&second), vec![Value::Integer(3)]);
    }

    #[test]
    fn an_offset_that_falls_inside_a_chunk_is_counted_in_rows() {
        let limit = Limit::new(Edge::All, Edge::Rows(3));
        let mut taken = limit.local();

        let mut first = chunk(&[1, 2]);
        assert_eq!(limit.push(&mut first, &mut taken).expect("all skipped"), Progress::More);
        assert!(first.is_empty());

        let mut second = chunk(&[3, 4, 5]);
        assert_eq!(limit.push(&mut second, &mut taken).expect("one more skipped"), Progress::More);
        assert_eq!(rows(&second), vec![Value::Integer(4), Value::Integer(5)]);
    }

    /// `LIMIT 0` is a query that reads nothing, and it is worth its own test because the count is
    /// reached before a row has been seen, which is the one path where the operator is finished on
    /// the call that starts it.
    #[test]
    fn a_limit_of_nothing_is_done_on_the_first_chunk() {
        let limit = Limit::new(Edge::Rows(0), Edge::Rows(0));
        let mut taken = limit.local();
        let mut first = chunk(&[1, 2]);
        assert_eq!(limit.push(&mut first, &mut taken).expect("nothing wanted"), Progress::Done);
        assert!(first.is_empty());
    }

    /// An instance's counters are its own, which is what makes the operator shareable. Two locals
    /// off one limit both get their own three rows.
    #[test]
    fn two_instances_of_one_limit_do_not_share_a_count() {
        let limit = Limit::new(Edge::Rows(3), Edge::Rows(0));
        let mut one = limit.local();
        let mut two = limit.local();

        let mut first = chunk(&[1, 2, 3]);
        assert_eq!(limit.push(&mut first, &mut one).expect("three rows"), Progress::Done);
        let mut second = chunk(&[4, 5, 6]);
        assert_eq!(limit.push(&mut second, &mut two).expect("three rows"), Progress::Done);
        assert_eq!(rows(&second), vec![Value::Integer(4), Value::Integer(5), Value::Integer(6)]);
    }
}