infino 0.2.0

// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The Infino Authors

//! Two-pass candidate planning for SQL `WHERE` predicates over
//! FTS-indexed columns.
//!
//! ## Why
//!
//! Without an index, `SELECT title FROM supertable WHERE title = 'rust
//! async runtime'` decodes the whole `title` column and drops the
//! non-matching rows. The inverted index already knows which rows
//! contain a term, so we resolve a small **candidate row set** from the
//! postings and decode only those rows — the row-level analog of the
//! term-bloom *superfile* prune.
//!
//! ## How (two passes)
//!
//!   1. **Candidate generation (this module).** The `WHERE` `Expr` tree
//!      is lowered to a [`CandidatePlan`] — a boolean tree whose leaves
//!      retrieve rows via [`SuperfileReader::token_match`]. Evaluated
//!      against one superfile it yields a `RoaringBitmap` of candidate
//!      `local_doc_id`s, or `None` ("no usable bound — scan the
//!      superfile").
//!   2. **Verification (DataFusion).** The provider turns the candidate
//!      set into a Parquet row selection so only those rows decode, and
//!      DataFusion's `FilterExec` (filters are reported `Inexact`)
//!      re-applies the **exact** predicate. The candidate set only has
//!      to be a *superset* of the true matches.
//!
//! ## Soundness
//!
//! A row equal to `'a b'` tokenizes to a set containing both `a` and
//! `b`, so it is in the term-AND `token_match(col, [a, b], And)`.
//! Requiring the literal's tokens can only keep a non-matching row
//! (wrong order, extra words, different spacing), never drop a matching
//! one — the exact equality is verified in pass 2. `AND` with an
//! un-boundable child drops that child (keeps more rows — still a
//! superset); `OR` with any un-boundable child is itself `Unbounded`;
//! `NOT`, non-FTS columns, range ops, and `LIKE` are `Unbounded` (a
//! word-token index can't soundly bound substring / negation).

use std::{collections::HashSet, sync::Arc};

use datafusion::{
    logical_expr::{Expr, Operator},
    scalar::ScalarValue,
};
use futures::future::BoxFuture;
use roaring::RoaringBitmap;

use crate::{
    superfile::{
        ReadError, SuperfileReader,
        fts::{reader::BoolMode, tokenize::Tokenizer},
    },
    supertable::{
        error::QueryError,
        manifest::ManifestSnapshot,
        query::prune::{PruneLeaf, select_superfiles},
    },
};

/// A superfile-independent boolean plan over FTS term retrievals, lowered
/// once from a SQL `WHERE` clause and [`evaluate`](CandidatePlan::evaluate)d
/// per superfile to a superset of the rows satisfying the FTS-resolvable
/// part of the predicate.
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) enum CandidatePlan {
    /// Rows whose `column` contains every one of `tokens` (term-AND).
    /// The candidate superset of `column = '<text tokenizing to tokens>'`;
    /// the exact predicate is re-verified by the `FilterExec` above the
    /// scan (filters are reported `Inexact`). Resolved per superfile by a
    /// single `token_match(.., And)` — postings only, no column decode —
    /// so verification + projection happen together in DataFusion's one
    /// scan pass. (An `exact_match`-per-leaf alternative was measured and
    /// rejected: it decodes the predicate column in its own pass 2, once
    /// per `OR`/`IN` branch, on top of the scan — multi-decode.)
    TermsAll { column: String, tokens: Vec<String> },
    /// Intersection of children (logical `AND`).
    And(Vec<CandidatePlan>),
    /// Union of children (logical `OR`).
    Or(Vec<CandidatePlan>),
    /// No usable bound: scan the superfile and let `FilterExec` verify.
    Unbounded,
}

impl CandidatePlan {
    /// Lower the conjunction of top-level `filters` (DataFusion ANDs the
    /// provider's filters together) into one plan. `fts_cols` is the set
    /// of FTS-indexed column names; `resolve` maps an FTS column to the
    /// tokenizer it was indexed with, so per-column analyzers lower query
    /// text the same way the column was tokenized at ingest. Empty
    /// `fts_cols` ⇒ no FTS columns ⇒ always [`Unbounded`].
    pub(crate) fn from_filters(
        filters: &[Expr],
        fts_cols: &HashSet<&str>,
        resolve: &dyn Fn(&str) -> Arc<dyn Tokenizer>,
    ) -> CandidatePlan {
        if fts_cols.is_empty() {
            return CandidatePlan::Unbounded;
        }
        and_combine(
            filters
                .iter()
                .map(|f| lower(f, fts_cols, resolve))
                .collect(),
        )
    }

    /// Evaluate against one superfile's reader. `Ok(None)` means "no bound
    /// — scan all rows"; `Ok(Some(bitmap))` is the candidate
    /// `local_doc_id` superset (possibly empty). `TermsAll` is one
    /// `token_match(.., And)`; `And`/`Or` intersect/union children.
    pub(crate) fn evaluate<'a>(
        &'a self,
        reader: &'a SuperfileReader,
    ) -> BoxFuture<'a, Result<Option<RoaringBitmap>, ReadError>> {
        Box::pin(async move {
            match self {
                CandidatePlan::Unbounded => Ok(None),
                CandidatePlan::TermsAll { column, tokens } => {
                    let refs: Vec<&str> = tokens.iter().map(String::as_str).collect();
                    let docs = reader.token_match(column, &refs, BoolMode::And).await?;
                    Ok(Some(docs.into_iter().collect()))
                }
                CandidatePlan::And(children) => {
                    let mut acc: Option<RoaringBitmap> = None;
                    for c in children {
                        if let Some(bm) = c.evaluate(reader).await? {
                            acc = Some(match acc {
                                Some(a) => a & bm,
                                None => bm,
                            });
                            if acc.as_ref().is_some_and(RoaringBitmap::is_empty) {
                                return Ok(Some(RoaringBitmap::new()));
                            }
                        }
                        // A `None` (unbounded) child adds no constraint.
                    }
                    Ok(acc)
                }
                CandidatePlan::Or(children) => {
                    let mut acc = RoaringBitmap::new();
                    for c in children {
                        match c.evaluate(reader).await? {
                            Some(bm) => acc |= bm,
                            // An unbounded branch makes the union unbounded.
                            None => return Ok(None),
                        }
                    }
                    Ok(Some(acc))
                }
            }
        })
    }
}

impl CandidatePlan {
    /// ManifestSnapshot-only superfile survival gate for this plan: the superfile
    /// ids that *could* match according to term blooms. `None` means no
    /// gate — the plan is [`Unbounded`], contains an `OR`, or otherwise
    /// cannot be expressed as a conjunction of bloom leaves.
    pub(crate) async fn surviving_superfile_ids(
        &self,
        manifest: &ManifestSnapshot,
    ) -> Result<Option<HashSet<u128>>, QueryError> {
        let mut groups = Vec::new();
        if !self.collect_survival_or_groups(&mut groups) {
            return Ok(None);
        }
        if groups.is_empty() {
            return Ok(Some(HashSet::new()));
        }
        let mut acc = HashSet::new();
        for leaves in groups {
            if leaves.is_empty() {
                continue;
            }
            acc.extend(
                select_superfiles(manifest, &leaves)
                    .await?
                    .iter()
                    .map(|e| e.superfile_id.as_u128()),
            );
        }
        Ok(Some(acc))
    }

    /// Flatten this plan into one bloom-conjunction group per `OR` branch.
    fn collect_survival_or_groups(&self, groups: &mut Vec<Vec<PruneLeaf>>) -> bool {
        match self {
            CandidatePlan::Or(children) => children
                .iter()
                .all(|child| child.collect_survival_or_branch(groups)),
            other => other.collect_survival_or_branch(groups),
        }
    }

    /// Append one `OR` branch as a single conjunctive leaf group.
    fn collect_survival_or_branch(&self, groups: &mut Vec<Vec<PruneLeaf>>) -> bool {
        match self {
            CandidatePlan::Unbounded => false,
            CandidatePlan::Or(children) => children
                .iter()
                .all(|child| child.collect_survival_or_branch(groups)),
            other => {
                let mut leaves = Vec::new();
                if !other.append_prune_leaves(&mut leaves) {
                    return false;
                }
                groups.push(leaves);
                true
            }
        }
    }

    /// Append conjunctive [`PruneLeaf`]s for this subtree. Returns `false`
    /// when the subtree contains an `OR` (not expressible as one bloom
    /// conjunction).
    fn append_prune_leaves(&self, leaves: &mut Vec<PruneLeaf>) -> bool {
        match self {
            CandidatePlan::Unbounded => true,
            CandidatePlan::TermsAll { column, tokens } => {
                if tokens.is_empty() {
                    return true;
                }
                leaves.push(PruneLeaf::TermPresence {
                    column: column.clone(),
                    terms: tokens.clone(),
                    mode: BoolMode::And,
                });
                true
            }
            CandidatePlan::And(children) => children
                .iter()
                .all(|child| child.append_prune_leaves(leaves)),
            CandidatePlan::Or(_) => false,
        }
    }

    /// Cheap upper-bound estimate of how many rows this plan would match
    /// in `reader`'s superfile, computed from per-term `df` only (no
    /// `token_match`, no posting decode). The bound follows the boolean
    /// tree: a term-`AND` can't exceed the **smallest** term's `df`
    /// (`min`); an `OR`/`IN` union can't exceed the **sum** of branch
    /// estimates (capped at `n_docs`); `Unbounded` is `n_docs` (no
    /// bound). The provider uses this to skip the index pushdown when a
    /// predicate would match a large fraction of the superfile — there the
    /// matches saturate the data pages so an index `RowSelection` can't
    /// skip any, and a plain scan is cheaper.
    pub(crate) fn estimate<'a>(
        &'a self,
        reader: &'a SuperfileReader,
    ) -> BoxFuture<'a, Result<u64, ReadError>> {
        Box::pin(async move {
            let n_docs = reader.n_docs();
            match self {
                CandidatePlan::Unbounded => Ok(n_docs),
                CandidatePlan::TermsAll { column, tokens } => {
                    if tokens.is_empty() {
                        return Ok(n_docs);
                    }
                    // Intersection ≤ the rarest token's df — resolved with
                    // one batched df lookup (single FST parse + coalesced
                    // header fetch) rather than one parse + fetch per token.
                    let refs: Vec<&str> = tokens.iter().map(String::as_str).collect();
                    let min_df = reader
                        .term_dfs(column, &refs)
                        .await?
                        .into_iter()
                        .min()
                        .unwrap_or(u64::MAX);
                    Ok(min_df.min(n_docs))
                }
                CandidatePlan::And(children) => {
                    let mut m = n_docs;
                    for c in children {
                        m = m.min(c.estimate(reader).await?);
                    }
                    Ok(m)
                }
                CandidatePlan::Or(children) => {
                    let mut sum: u64 = 0;
                    for c in children {
                        sum = sum.saturating_add(c.estimate(reader).await?);
                    }
                    Ok(sum.min(n_docs))
                }
            }
        })
    }
}

/// Lower one `Expr` node.
fn lower(
    expr: &Expr,
    fts_cols: &HashSet<&str>,
    resolve: &dyn Fn(&str) -> Arc<dyn Tokenizer>,
) -> CandidatePlan {
    match expr {
        Expr::BinaryExpr(be) => match be.op {
            Operator::And => and_combine(vec![
                lower(&be.left, fts_cols, resolve),
                lower(&be.right, fts_cols, resolve),
            ]),
            Operator::Or => or_combine(vec![
                lower(&be.left, fts_cols, resolve),
                lower(&be.right, fts_cols, resolve),
            ]),
            Operator::Eq => eq_leaf(&be.left, &be.right, fts_cols, resolve),
            // Range / inequality / arithmetic ops aren't term-bounded.
            _ => CandidatePlan::Unbounded,
        },
        // `IN (a, b, …)` on an FTS column is an OR of equalities.
        Expr::InList(il) if !il.negated => in_list_leaf(il, fts_cols, resolve),
        // NOT, LIKE, IS NULL, functions, etc. — not soundly term-bounded.
        _ => CandidatePlan::Unbounded,
    }
}

/// Lower `col = 'literal'` (either operand order) on an FTS column.
fn eq_leaf(
    left: &Expr,
    right: &Expr,
    fts_cols: &HashSet<&str>,
    resolve: &dyn Fn(&str) -> Arc<dyn Tokenizer>,
) -> CandidatePlan {
    let (column, value) = match (left, right) {
        (Expr::Column(c), Expr::Literal(v, _)) => (&c.name, v),
        (Expr::Literal(v, _), Expr::Column(c)) => (&c.name, v),
        _ => return CandidatePlan::Unbounded,
    };
    terms_all(column, value, fts_cols, resolve)
}

/// Lower `col IN ('a', 'b', …)` on an FTS column to an OR of term-ANDs.
fn in_list_leaf(
    il: &datafusion::logical_expr::expr::InList,
    fts_cols: &HashSet<&str>,
    resolve: &dyn Fn(&str) -> Arc<dyn Tokenizer>,
) -> CandidatePlan {
    let Expr::Column(c) = il.expr.as_ref() else {
        return CandidatePlan::Unbounded;
    };
    let mut branches = Vec::with_capacity(il.list.len());
    for item in &il.list {
        let Expr::Literal(v, _) = item else {
            return CandidatePlan::Unbounded;
        };
        branches.push(terms_all(&c.name, v, fts_cols, resolve));
    }
    or_combine(branches)
}

/// Build a `TermsAll` leaf for `column = value`, or `Unbounded` if the
/// column isn't FTS-indexed, the value isn't a string, or it tokenizes
/// to nothing (e.g. the empty string — no tokens to bound with).
fn terms_all(
    column: &str,
    value: &ScalarValue,
    fts_cols: &HashSet<&str>,
    resolve: &dyn Fn(&str) -> Arc<dyn Tokenizer>,
) -> CandidatePlan {
    if !fts_cols.contains(column) {
        return CandidatePlan::Unbounded;
    }
    let Some(s) = scalar_str(value) else {
        return CandidatePlan::Unbounded;
    };
    let tok = resolve(column);
    let tokens: Vec<String> = tok.tokenize(s).collect();
    if tokens.is_empty() {
        return CandidatePlan::Unbounded;
    }
    CandidatePlan::TermsAll {
        column: column.to_owned(),
        tokens,
    }
}

/// Extract a UTF-8 string from a scalar literal, if it is one.
fn scalar_str(v: &ScalarValue) -> Option<&str> {
    match v {
        ScalarValue::Utf8(Some(s))
        | ScalarValue::LargeUtf8(Some(s))
        | ScalarValue::Utf8View(Some(s)) => Some(s.as_str()),
        _ => None,
    }
}

/// Combine children under `AND`: an `Unbounded` child drops out (adds
/// no constraint), nested `And`s flatten, all-unbounded → `Unbounded`.
fn and_combine(children: Vec<CandidatePlan>) -> CandidatePlan {
    let mut flat = Vec::with_capacity(children.len());
    for c in children {
        match c {
            CandidatePlan::Unbounded => {}
            CandidatePlan::And(inner) => flat.extend(inner),
            other => flat.push(other),
        }
    }
    collapse(flat, true)
}

/// Combine children under `OR`: any `Unbounded` child makes the whole
/// union `Unbounded`; nested `Or`s flatten.
fn or_combine(children: Vec<CandidatePlan>) -> CandidatePlan {
    let mut flat = Vec::with_capacity(children.len());
    for c in children {
        match c {
            CandidatePlan::Unbounded => return CandidatePlan::Unbounded,
            CandidatePlan::Or(inner) => flat.extend(inner),
            other => flat.push(other),
        }
    }
    collapse(flat, false)
}

/// Wrap a flattened child list back into `And`/`Or`, collapsing the
/// 0- and 1-child degenerate cases.
fn collapse(mut flat: Vec<CandidatePlan>, is_and: bool) -> CandidatePlan {
    match flat.len() {
        0 => CandidatePlan::Unbounded,
        1 => flat.pop().expect("len checked == 1"),
        _ if is_and => CandidatePlan::And(flat),
        _ => CandidatePlan::Or(flat),
    }
}

#[cfg(test)]
mod tests {
    use datafusion::{
        logical_expr::expr::InList,
        prelude::{col, lit},
    };

    use super::*;
    use crate::superfile::fts::tokenize::{AsciiLowerTokenizer, StandardTokenizer};

    fn fts_cols() -> HashSet<&'static str> {
        let mut s = HashSet::new();
        s.insert("title");
        s
    }

    /// Resolver for the lowering tests: every column tokenizes with the
    /// ASCII-lower analyzer.
    fn ascii_resolver(_col: &str) -> Arc<dyn Tokenizer> {
        Arc::new(AsciiLowerTokenizer)
    }

    fn plan(expr: Expr) -> CandidatePlan {
        CandidatePlan::from_filters(&[expr], &fts_cols(), &ascii_resolver)
    }

    /// Bloom-survival flattening: an `AND` of term-alls collapses to one
    /// conjunctive group; a top-level `OR` yields one group per branch; a plan
    /// that isn't a pure conjunction (`Unbounded`, or an `OR` nested under an
    /// `AND`) is inexpressible and returns `false`.
    #[test]
    fn candidate_plan_flattens_into_bloom_survival_groups() {
        let terms = |c: &str, t: &str| CandidatePlan::TermsAll {
            column: c.to_string(),
            tokens: vec![t.to_string()],
        };

        let mut groups = Vec::new();
        assert!(
            CandidatePlan::And(vec![terms("a", "x"), terms("b", "y")])
                .collect_survival_or_groups(&mut groups)
        );
        assert_eq!(groups.len(), 1, "AND is one conjunctive group");
        assert_eq!(groups[0].len(), 2, "with a leaf per term-all");

        let mut groups = Vec::new();
        assert!(
            CandidatePlan::Or(vec![terms("a", "x"), terms("b", "y")])
                .collect_survival_or_groups(&mut groups)
        );
        assert_eq!(groups.len(), 2, "top-level OR is one group per branch");

        let mut groups = Vec::new();
        assert!(!CandidatePlan::Unbounded.collect_survival_or_groups(&mut groups));

        let mut groups = Vec::new();
        assert!(
            !CandidatePlan::And(vec![
                terms("a", "x"),
                CandidatePlan::Or(vec![terms("b", "y"), terms("c", "z")]),
            ])
            .collect_survival_or_groups(&mut groups),
            "OR nested under AND is not a single bloom conjunction"
        );
    }

    #[test]
    fn eq_on_fts_column_lowers_to_terms_all() {
        let p = plan(col("title").eq(lit("rust async")));
        assert_eq!(
            p,
            CandidatePlan::TermsAll {
                column: "title".into(),
                tokens: vec!["rust".into(), "async".into()],
            }
        );
    }

    #[test]
    fn eq_operands_reversed_still_lowers() {
        let p = plan(lit("rust").eq(col("title")));
        assert_eq!(
            p,
            CandidatePlan::TermsAll {
                column: "title".into(),
                tokens: vec!["rust".into()],
            }
        );
    }

    #[test]
    fn eq_on_non_fts_column_is_unbounded() {
        assert_eq!(
            plan(col("category").eq(lit("rust"))),
            CandidatePlan::Unbounded
        );
    }

    #[test]
    fn empty_literal_is_unbounded() {
        assert_eq!(plan(col("title").eq(lit(""))), CandidatePlan::Unbounded);
    }

    #[test]
    fn range_op_is_unbounded() {
        assert_eq!(plan(col("title").gt(lit("m"))), CandidatePlan::Unbounded);
    }

    #[test]
    fn and_of_fts_and_non_fts_keeps_only_fts_branch() {
        let p = plan(
            col("title")
                .eq(lit("rust"))
                .and(col("category").eq(lit("lang"))),
        );
        assert_eq!(
            p,
            CandidatePlan::TermsAll {
                column: "title".into(),
                tokens: vec!["rust".into()],
            }
        );
    }

    #[test]
    fn and_of_two_fts_equalities_intersects() {
        let p = plan(
            col("title")
                .eq(lit("rust"))
                .and(col("title").eq(lit("async"))),
        );
        match p {
            CandidatePlan::And(children) => assert_eq!(children.len(), 2),
            other => panic!("expected And, got {other:?}"),
        }
    }

    #[test]
    fn or_of_two_fts_equalities_unions() {
        let p = plan(
            col("title")
                .eq(lit("rust"))
                .or(col("title").eq(lit("python"))),
        );
        match p {
            CandidatePlan::Or(children) => assert_eq!(children.len(), 2),
            other => panic!("expected Or, got {other:?}"),
        }
    }

    #[test]
    fn or_with_non_fts_branch_is_unbounded() {
        let p = plan(
            col("title")
                .eq(lit("rust"))
                .or(col("category").eq(lit("lang"))),
        );
        assert_eq!(p, CandidatePlan::Unbounded);
    }

    #[test]
    fn not_is_unbounded() {
        assert_eq!(
            plan(!col("title").eq(lit("rust"))),
            CandidatePlan::Unbounded
        );
    }

    #[test]
    fn not_eq_is_unbounded() {
        // `title != 'rust'` (Operator::NotEq) can't be term-bounded.
        assert_eq!(
            plan(col("title").not_eq(lit("rust"))),
            CandidatePlan::Unbounded
        );
    }

    #[test]
    fn and_with_not_child_keeps_fts_branch() {
        // `title = 'rust' AND NOT (title = 'compiler')` — the NOT branch
        // is un-boundable and drops out of candidate generation (verified
        // in pass 2), so candidates still come from the FTS branch.
        let p = plan(
            col("title")
                .eq(lit("rust"))
                .and(!col("title").eq(lit("compiler"))),
        );
        assert_eq!(
            p,
            CandidatePlan::TermsAll {
                column: "title".into(),
                tokens: vec!["rust".into()],
            }
        );
    }

    #[test]
    fn like_is_unbounded() {
        assert_eq!(
            plan(col("title").like(lit("rust%"))),
            CandidatePlan::Unbounded
        );
    }

    #[test]
    fn in_list_on_fts_column_is_or_of_terms_all() {
        let expr = Expr::InList(InList::new(
            Box::new(col("title")),
            vec![lit("rust"), lit("python")],
            false,
        ));
        match plan(expr) {
            CandidatePlan::Or(children) => {
                assert_eq!(children.len(), 2);
                assert!(matches!(children[0], CandidatePlan::TermsAll { .. }));
            }
            other => panic!("expected Or, got {other:?}"),
        }
    }

    #[test]
    fn negated_in_list_is_unbounded() {
        let expr = Expr::InList(InList::new(Box::new(col("title")), vec![lit("rust")], true));
        assert_eq!(plan(expr), CandidatePlan::Unbounded);
    }

    #[test]
    fn no_fts_columns_is_unbounded() {
        // With no FTS columns there is nothing to term-bound. (The old
        // "FTS columns but no tokenizer" state is unrepresentable now
        // that a per-column tokenizer always exists when FTS columns do.)
        let p = CandidatePlan::from_filters(
            &[col("title").eq(lit("rust"))],
            &HashSet::new(),
            &ascii_resolver,
        );
        assert_eq!(p, CandidatePlan::Unbounded);
    }

    #[test]
    fn lowering_uses_the_per_column_tokenizer() {
        // `title` is analyzed with the Unicode-aware standard tokenizer,
        // which keeps non-ASCII letters; ascii_lower drops the whole
        // token. The lowering must pick the column's own analyzer.
        let resolve = |col: &str| -> Arc<dyn Tokenizer> {
            if col == "title" {
                Arc::new(StandardTokenizer)
            } else {
                Arc::new(AsciiLowerTokenizer)
            }
        };
        let bounded =
            CandidatePlan::from_filters(&[col("title").eq(lit("Süd"))], &fts_cols(), &resolve);
        assert_eq!(
            bounded,
            CandidatePlan::TermsAll {
                column: "title".to_owned(),
                tokens: vec!["süd".to_owned()],
            }
        );

        // The same literal under ascii_lower drops the non-ASCII token,
        // leaving nothing to bound with ⇒ Unbounded. Proves the result
        // above came from the standard tokenizer, not a table-wide default.
        let unbounded = CandidatePlan::from_filters(
            &[col("title").eq(lit("Süd"))],
            &fts_cols(),
            &ascii_resolver,
        );
        assert_eq!(unbounded, CandidatePlan::Unbounded);
    }
}