zeppelin_embed/planner/
lexical.rs

1//! Selectivity-driven exact lexical filtering.
2
3use crate::fts::bm25::Bm25Params;
4use crate::fts::index::{IndexError, LexicalIndex};
5use crate::fts::prune::{search_pruned_filtered_controlled, select_strategy};
6use crate::fts::search::{
7    ControlledSearchError, SearchResult, TermQuery, search_allow_list_driven_controlled,
8};
9use crate::meta::DocBitmap;
10
11/// Corpus divisor below which the allow-list drives lexical iteration.
12///
13/// PLACEHOLDER -- NOT YET MEASURED. Both branches are exact; a poor value is
14/// a latency defect only.
15pub const LEXICAL_ALLOW_LIST_DIVISOR: u64 = 64;
16
17/// Exact lexical filter execution branch.
18#[derive(Clone, Copy, Debug, Eq, PartialEq)]
19pub enum LexicalBranch {
20    /// Enumerate allow-listed rows and seek terms to each row.
21    AllowListDrive,
22    /// Preserve pruning traversal and post-check every scored candidate.
23    PostCheck,
24}
25
26/// Filtered lexical result and the branch that produced it.
27#[derive(Clone, Debug, PartialEq)]
28pub struct LexicalSearchOutcome {
29    /// Exact BM25 hits and counters.
30    pub result: SearchResult,
31    /// Actual exact filter branch.
32    pub branch: LexicalBranch,
33}
34
35/// Typed lexical filter planning or scoring failure.
36#[derive(Debug)]
37pub enum LexicalFilterError {
38    /// One bitmap per lexical segment is required.
39    SegmentCount {
40        /// Lexical segment count.
41        expected: usize,
42        /// Supplied bitmap count.
43        actual: usize,
44    },
45    /// An allow-list row exceeded its segment's dense row space.
46    RowOutOfRange {
47        /// Segment ordinal.
48        segment: usize,
49        /// Invalid row.
50        row: u32,
51        /// Segment row count.
52        row_count: u32,
53    },
54    /// The lexical scorer rejected the index or query.
55    Index(IndexError),
56}
57
58impl std::fmt::Display for LexicalFilterError {
59    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
60        match self {
61            Self::SegmentCount { expected, actual } => write!(
62                formatter,
63                "lexical filter supplied {actual} segment bitmaps, expected {expected}"
64            ),
65            Self::RowOutOfRange {
66                segment,
67                row,
68                row_count,
69            } => write!(
70                formatter,
71                "lexical filter row {row} is outside segment {segment} row count {row_count}"
72            ),
73            Self::Index(error) => error.fmt(formatter),
74        }
75    }
76}
77
78impl std::error::Error for LexicalFilterError {
79    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
80        match self {
81            Self::Index(error) => Some(error),
82            Self::SegmentCount { .. } | Self::RowOutOfRange { .. } => None,
83        }
84    }
85}
86
87impl From<IndexError> for LexicalFilterError {
88    fn from(error: IndexError) -> Self {
89        Self::Index(error)
90    }
91}
92
93/// Plans and executes one exact lexical query over segment-local allow-lists.
94pub fn search_lexical_filtered(
95    index: &LexicalIndex,
96    query: &TermQuery,
97    k: usize,
98    params: Bm25Params,
99    allow_lists: &[DocBitmap],
100    forced: Option<LexicalBranch>,
101) -> Result<LexicalSearchOutcome, LexicalFilterError> {
102    let references = allow_lists.iter().collect::<Vec<_>>();
103    search_lexical_filtered_refs(index, query, k, params, &references, forced)
104}
105
106pub(crate) fn search_lexical_filtered_refs(
107    index: &LexicalIndex,
108    query: &TermQuery,
109    k: usize,
110    params: Bm25Params,
111    allow_lists: &[&DocBitmap],
112    forced: Option<LexicalBranch>,
113) -> Result<LexicalSearchOutcome, LexicalFilterError> {
114    search_lexical_filtered_refs_controlled(index, query, k, params, allow_lists, forced, || Ok(()))
115}
116
117pub(crate) fn search_lexical_filtered_refs_controlled<E: From<LexicalFilterError>>(
118    index: &LexicalIndex,
119    query: &TermQuery,
120    k: usize,
121    params: Bm25Params,
122    allow_lists: &[&DocBitmap],
123    forced: Option<LexicalBranch>,
124    mut check: impl FnMut() -> Result<(), E>,
125) -> Result<LexicalSearchOutcome, E> {
126    let mut work = crate::fts::control::WorkCheck::new(&mut check);
127    work.check_now()?;
128    validate_allow_lists(index, allow_lists, &mut work)?;
129    let mut allowed = 0_u64;
130    for allow_list in allow_lists {
131        work.step()?;
132        allowed = allowed.saturating_add(allow_list.cardinality());
133    }
134    work.check_now()?;
135    let corpus = index.document_count();
136    let branch = forced.unwrap_or_else(|| {
137        if allowed.saturating_mul(LEXICAL_ALLOW_LIST_DIVISOR) <= corpus {
138            LexicalBranch::AllowListDrive
139        } else {
140            LexicalBranch::PostCheck
141        }
142    });
143    let result = match branch {
144        LexicalBranch::AllowListDrive => {
145            search_allow_list_driven_controlled(index, query, k, params, allow_lists, &mut check)
146        }
147        LexicalBranch::PostCheck => search_pruned_filtered_controlled(
148            index,
149            query,
150            k,
151            params,
152            select_strategy(query.terms.len(), k),
153            allow_lists,
154            &mut check,
155        ),
156    }
157    .map_err(|error| match error {
158        ControlledSearchError::Index(error) => E::from(LexicalFilterError::Index(error)),
159        ControlledSearchError::Control(error) => error,
160    })?;
161    check()?;
162    Ok(LexicalSearchOutcome { result, branch })
163}
164
165// Keep this loop's stack layout independent of the larger query caller.
166#[inline(never)]
167fn validate_allow_lists<E: From<LexicalFilterError>>(
168    index: &LexicalIndex,
169    allow_lists: &[&DocBitmap],
170    work: &mut crate::fts::control::WorkCheck<impl FnMut() -> Result<(), E>>,
171) -> Result<(), E> {
172    if index.segments().len() != allow_lists.len() {
173        return Err(LexicalFilterError::SegmentCount {
174            expected: index.segments().len(),
175            actual: allow_lists.len(),
176        }
177        .into());
178    }
179    for (segment, (sealed, allow_list)) in index.segments().iter().zip(allow_lists).enumerate() {
180        work.step()?;
181        if let Some(row) = allow_list.first_at_or_after(sealed.row_count()) {
182            return Err(LexicalFilterError::RowOutOfRange {
183                segment,
184                row,
185                row_count: sealed.row_count(),
186            }
187            .into());
188        }
189    }
190    Ok(())
191}
zeppelin_embed/planner/lexical.rs

zeppelin_embed/planner/
lexical.rs
