writerslogic_warm/core/documents/
lexical.rs

1// Copyright 2024-2026 WritersLogic Contributors
2// SPDX-License-Identifier: Apache-2.0
3
4use anyhow::{ensure, Context, Result};
5use serde::{Deserialize, Serialize};
6use std::borrow::Cow;
7use std::collections::BTreeSet;
8
9use super::packed_terms::WordId;
10use super::positions::tokens;
11use super::{Bm25, DocumentStore, SearchOptions, StoredChunk, MAX_DOCUMENT_BYTES};
12
13#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
14pub enum LexicalExecution {
15    #[default]
16    Auto,
17    /// Unpruned posting-cursor execution.
18    DocumentAtATime,
19    /// Dynamic pruning for term disjunctions (block-max MaxScore, or a strongest-blocks-first
20    /// scan of a single list); other queries use cursors.
21    BlockMaxWand,
22    /// Full scan and full sort, useful for verifying another execution strategy.
23    Exhaustive,
24}
25
26#[derive(Clone, Debug)]
27pub struct LexicalOptions {
28    pub k: usize,
29    pub count: bool,
30    pub bm25_k1: f64,
31    pub bm25_b: f64,
32    pub execution: LexicalExecution,
33}
34
35impl Default for LexicalOptions {
36    fn default() -> Self {
37        Self {
38            k: 10,
39            count: false,
40            bm25_k1: 1.2,
41            bm25_b: 0.75,
42            execution: LexicalExecution::Auto,
43        }
44    }
45}
46
47#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
48#[serde(rename_all = "camelCase")]
49pub struct LexicalHit {
50    pub document_id: String,
51    pub chunk_index: usize,
52    pub score: f64,
53}
54
55#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
56pub struct LexicalResults {
57    pub hits: Vec<LexicalHit>,
58    pub count: Option<usize>,
59}
60
61#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
62enum Occur {
63    Must,
64    Should,
65    MustNot,
66}
67
68#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord)]
69struct Clause {
70    occur: Occur,
71    words: [Vec<String>; 2],
72    phrase: bool,
73}
74
75/// A flat Boolean query: bare optional clauses, +required clauses, -excluded clauses,
76/// and quoted adjacent phrases. Purely negative queries match no documents.
77#[derive(Clone, Debug)]
78pub struct LexicalQuery {
79    clauses: Vec<Clause>,
80}
81
82impl LexicalQuery {
83    pub fn parse(text: &str) -> Result<Self> {
84        ensure!(text.len() <= MAX_DOCUMENT_BYTES, "query exceeds 8 MiB");
85        let mut rest = text.trim();
86        let mut clauses = Vec::new();
87        while !rest.is_empty() {
88            ensure!(clauses.len() < 1024, "lexical query exceeds 1024 clauses");
89            let occur = match rest.as_bytes()[0] {
90                b'+' => {
91                    rest = &rest[1..];
92                    Occur::Must
93                }
94                b'-' => {
95                    rest = &rest[1..];
96                    Occur::MustNot
97                }
98                _ => Occur::Should,
99            };
100            let phrase = rest.starts_with('"');
101            let value;
102            if phrase {
103                rest = &rest[1..];
104                let end = rest
105                    .find('"')
106                    .ok_or_else(|| anyhow::anyhow!("unterminated lexical phrase"))?;
107                value = &rest[..end];
108                rest = &rest[end + 1..];
109                ensure!(
110                    rest.is_empty() || rest.starts_with(char::is_whitespace),
111                    "unsupported lexical phrase suffix (only zero slop is supported)"
112                );
113                ensure!(
114                    value
115                        .chars()
116                        .all(|c| c.is_alphanumeric() || c.is_whitespace()),
117                    "unsupported character in lexical phrase"
118                );
119            } else {
120                let end = rest.find(char::is_whitespace).unwrap_or(rest.len());
121                value = &rest[..end];
122                rest = &rest[end..];
123                ensure!(
124                    !value.is_empty() && value.chars().all(char::is_alphanumeric),
125                    "unsupported lexical query syntax"
126                );
127            }
128            let words = [tokens(value, false), tokens(value, true)];
129            ensure!(!words[0].is_empty(), "lexical clause requires a token");
130            ensure!(words[0].len() <= 1024, "lexical phrase exceeds 1024 tokens");
131            clauses.push(Clause {
132                occur,
133                phrase: phrase && words[0].len() > 1,
134                words,
135            });
136            rest = rest.trim_start();
137        }
138        clauses.sort_unstable();
139        clauses.dedup();
140        Ok(Self { clauses })
141    }
142
143    fn has_positive(&self) -> bool {
144        self.clauses.iter().any(|c| c.occur != Occur::MustNot)
145    }
146}
147
148/// The clause's occurrences in `chunk`: a word's frequency, or a phrase's when `exact`,
149/// otherwise 1 at its first occurrence. `ids` holds one cache per phrase word and form.
150#[cfg(test)]
151fn frequency(clause: &Clause, chunk: &StoredChunk, exact: bool, ids: &mut [Vec<WordId>; 2]) -> u32 {
152    frequency_with_positions(clause, chunk, exact, ids, chunk.terms.positions.as_deref())
153}
154
155fn frequency_with_positions(
156    clause: &Clause,
157    chunk: &StoredChunk,
158    exact: bool,
159    ids: &mut [Vec<WordId>; 2],
160    positions: Option<&super::positions::Positions>,
161) -> u32 {
162    let form = usize::from(chunk.stemmed);
163    let words = &clause.words[form];
164    if !clause.phrase {
165        return chunk.terms.get(&words[0]).copied().unwrap_or(0);
166    }
167    let mut indices = Vec::with_capacity(words.len());
168    let (mut lead, mut least) = (0, u32::MAX);
169    for (at, (word, id)) in words.iter().zip(&mut ids[form]).enumerate() {
170        // A repeated word of a short phrase reuses its earlier index.
171        let earlier = (words.len() <= 16)
172            .then(|| words[..at].iter().position(|w| w == word))
173            .flatten();
174        let found = match earlier {
175            Some(earlier) => Some(indices[earlier]),
176            None => chunk.terms.index_cached(word, id),
177        };
178        let Some(index) = found else {
179            return 0;
180        };
181        let frequency = chunk.terms.frequency_at(index);
182        if frequency == 0 {
183            return 0;
184        }
185        if frequency < least {
186            lead = indices.len();
187            least = frequency;
188        }
189        indices.push(index);
190    }
191    let positions = positions.expect("phrase capability checked");
192    let mut cursors: Vec<_> = indices
193        .iter()
194        .map(|&i| positions.iter(i).peekable())
195        .collect();
196    let mut count = 0;
197    for position in positions.iter(indices[lead]) {
198        let Some(start) = position.checked_sub(lead as u32) else {
199            continue;
200        };
201        let matched = cursors.iter_mut().enumerate().all(|(offset, iter)| {
202            let Some(target) = start.checked_add(offset as u32) else {
203                return false;
204            };
205            while iter.peek().is_some_and(|&p| p < target) {
206                iter.next();
207            }
208            iter.peek() == Some(&target)
209        });
210        count += u32::from(matched);
211        if matched && !exact {
212            break;
213        }
214    }
215    count
216}
217
218fn positions_for_chunk<'a>(
219    store: &DocumentStore,
220    chunk: &'a StoredChunk,
221) -> Result<Cow<'a, super::positions::Positions>> {
222    match chunk.terms.positions.as_deref() {
223        Some(positions) => Ok(Cow::Borrowed(positions)),
224        None => {
225            let text = store.chunk_text(chunk).context("unsupported index: a candidate phrase needs stored positions or retained chunk text")?;
226            Ok(Cow::Owned(super::positions::Positions::from_text(
227                &text,
228                chunk.stemmed,
229                &chunk.terms,
230                chunk.terms.is_whole_document(),
231            )))
232        }
233    }
234}
235
236/// One dictionary-id cache per word and form of each clause: a dictionary may hold chunks
237/// of both forms, whose words differ.
238fn word_ids(query: &LexicalQuery) -> Vec<[Vec<WordId>; 2]> {
239    query
240        .clauses
241        .iter()
242        .map(|clause| {
243            std::array::from_fn(|form| {
244                clause.words[form]
245                    .iter()
246                    .map(|_| WordId::default())
247                    .collect()
248            })
249        })
250        .collect()
251}
252
253fn prepare(options: &LexicalOptions) -> Result<Bm25> {
254    ensure!(options.k <= 1_000_000, "lexical k exceeds 1000000");
255    let bm25 = Bm25::from_options(&SearchOptions {
256        bm25_k1: Some(options.bm25_k1),
257        bm25_b: Some(options.bm25_b),
258        ..Default::default()
259    })?;
260    Ok(bm25)
261}
262
263fn clause_idfs(store: &DocumentStore, query: &LexicalQuery) -> Vec<[f64; 2]> {
264    let n = store.chunk_count as f64;
265    query
266        .clauses
267        .iter()
268        .map(|clause| {
269            std::array::from_fn(|form| {
270                clause.words[form]
271                    .iter()
272                    .map(|word| {
273                        let df = store.postings[form].get(word).map_or(0, |l| l.len()) as f64;
274                        Bm25::idf(n, df)
275                    })
276                    .sum()
277            })
278        })
279        .collect()
280}
281
282fn matches(query: &LexicalQuery, frequencies: &[u32]) -> bool {
283    let mut required = false;
284    let mut optional = false;
285    for (clause, &frequency) in query.clauses.iter().zip(frequencies) {
286        match clause.occur {
287            Occur::Must => {
288                required = true;
289                if frequency == 0 {
290                    return false;
291                }
292            }
293            Occur::MustNot if frequency != 0 => return false,
294            Occur::Should => optional |= frequency != 0,
295            _ => {}
296        }
297    }
298    required || optional
299}
300
301fn score(
302    query: &LexicalQuery,
303    frequencies: &[u32],
304    idfs: &[[f64; 2]],
305    form: usize,
306    norm: f64,
307    bm25: Bm25,
308) -> Result<f64> {
309    let mut score = 0.0;
310    for ((clause, &frequency), idf) in query.clauses.iter().zip(frequencies).zip(idfs) {
311        if clause.occur != Occur::MustNot && frequency > 0 {
312            score += bm25.weight(idf[form], f64::from(frequency), norm);
313        }
314    }
315    ensure!(
316        score.is_finite(),
317        "lexical score is nonfinite; reduce BM25 parameters"
318    );
319    Ok(score)
320}
321
322fn order_hits(a: &LexicalHit, b: &LexicalHit) -> std::cmp::Ordering {
323    b.score
324        .total_cmp(&a.score)
325        .then_with(|| a.document_id.cmp(&b.document_id))
326        .then_with(|| a.chunk_index.cmp(&b.chunk_index))
327}
328
329pub(crate) fn execute(
330    store: &DocumentStore,
331    query: &LexicalQuery,
332    options: &LexicalOptions,
333) -> Result<LexicalResults> {
334    let bm25 = prepare(options)?;
335    if options.execution != LexicalExecution::Exhaustive
336        && options.k == 0
337        && options.count
338        && query
339            .clauses
340            .iter()
341            .all(|clause| clause.occur == Occur::Should && !clause.phrase)
342    {
343        return Ok(LexicalResults {
344            hits: Vec::new(),
345            count: Some(execution::count_union(store, query)),
346        });
347    }
348    match options.execution {
349        LexicalExecution::Exhaustive => exhaustive(store, query, options, bm25),
350        LexicalExecution::Auto | LexicalExecution::BlockMaxWand
351            if options.k > 0
352                && (options.execution == LexicalExecution::BlockMaxWand || options.k <= 1000) =>
353        {
354            match wand::execute(store, query, options, bm25)? {
355                Some(result) => Ok(result),
356                None => document_at_a_time(store, query, options, bm25),
357            }
358        }
359        _ => document_at_a_time(store, query, options, bm25),
360    }
361}
362
363fn exhaustive(
364    store: &DocumentStore,
365    query: &LexicalQuery,
366    options: &LexicalOptions,
367    bm25: Bm25,
368) -> Result<LexicalResults> {
369    let mut hits = Vec::new();
370    let mut matched = BTreeSet::new();
371    if !query.has_positive() {
372        return Ok(LexicalResults {
373            hits,
374            count: options.count.then_some(0),
375        });
376    }
377    let idfs = if options.k > 0 {
378        clause_idfs(store, query)
379    } else {
380        Vec::new()
381    };
382    let average = (store.word_count as f64 / (store.chunk_count as f64).max(1.0)).max(1.0);
383    let mut ids = word_ids(query);
384    for (id, document) in &store.documents {
385        for (chunk_index, chunk) in document.chunks.iter().enumerate() {
386            let possible = query.clauses.iter().any(|clause| {
387                clause.phrase
388                    && clause.words[usize::from(chunk.stemmed)].iter().all(|word| {
389                        chunk
390                            .terms
391                            .get(word)
392                            .is_some_and(|&frequency| frequency > 0)
393                    })
394            });
395            let positions = possible
396                .then(|| positions_for_chunk(store, chunk))
397                .transpose()?;
398            let frequencies: Vec<_> = query
399                .clauses
400                .iter()
401                .zip(&mut ids)
402                .map(|(clause, ids)| {
403                    frequency_with_positions(
404                        clause,
405                        chunk,
406                        options.k > 0,
407                        ids,
408                        positions.as_deref(),
409                    )
410                })
411                .collect();
412            if !matches(query, &frequencies) {
413                continue;
414            }
415            if options.count {
416                matched.insert(id);
417            }
418            if options.k > 0 {
419                hits.push(LexicalHit {
420                    document_id: id.clone(),
421                    chunk_index,
422                    score: score(
423                        query,
424                        &frequencies,
425                        &idfs,
426                        usize::from(chunk.stemmed),
427                        bm25.norm(chunk.word_count as f64, average),
428                        bm25,
429                    )?,
430                });
431            }
432        }
433    }
434    hits.sort_unstable_by(order_hits);
435    hits.truncate(options.k);
436    Ok(LexicalResults {
437        hits,
438        count: options.count.then_some(matched.len()),
439    })
440}
441
442pub(super) mod execution;
443mod wand;
444use execution::document_at_a_time;
445
446#[cfg(test)]
447mod tests;
writerslogic_warm/core/documents/lexical.rs

writerslogic_warm/core/documents/
lexical.rs
