uqa_storage/
block_max_index.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Per-block maximum score index for Block-Max WAND optimisation.
8//!
9//! For each `(table, field, term)` posting list, we precompute the
10//! maximum score reachable inside each fixed-size block of consecutive
11//! entries. Block-Max WAND consults these tighter per-block upper bounds
12//! during pivot resolution, achieving higher skip rates than plain WAND
13//! (Theorem 6.2.2, Paper 3).
14
15use std::collections::BTreeMap;
16
17use rusqlite::params;
18
19use uqa_core::PostingList;
20
21use crate::{StorageBackendError, StorageBackendResult};
22
23pub const DEFAULT_BLOCK_SIZE: usize = 128;
24
25/// Trait every scorer used to seed a block-max index implements. Only
26/// the inner BM25 surface is needed; `uqa_scoring::BM25Scorer` already
27/// satisfies it without extra glue.
28pub trait BlockMaxScorer {
29    fn score(&self, term_freq: u64, doc_length: u64, doc_freq: u64) -> f64;
30}
31
32#[derive(Debug, Clone)]
33pub struct BlockMaxIndex {
34    block_size: usize,
35    block_maxes: BTreeMap<(String, String, String), Vec<f64>>,
36}
37
38impl Default for BlockMaxIndex {
39    fn default() -> Self {
40        Self {
41            block_size: DEFAULT_BLOCK_SIZE,
42            block_maxes: BTreeMap::new(),
43        }
44    }
45}
46
47impl BlockMaxIndex {
48    pub fn new(block_size: usize) -> StorageBackendResult<Self> {
49        if block_size == 0 {
50            return Err(StorageBackendError::Other(
51                "block-max block size must be greater than zero".to_string(),
52            ));
53        }
54        Ok(Self {
55            block_size,
56            block_maxes: BTreeMap::new(),
57        })
58    }
59
60    pub fn block_size(&self) -> usize {
61        self.block_size
62    }
63
64    pub fn set_block_maxes(
65        &mut self,
66        table: &str,
67        field: &str,
68        term: &str,
69        scores: Vec<f64>,
70    ) -> StorageBackendResult<()> {
71        validate_scores(&scores)?;
72        self.block_maxes.insert(
73            (table.to_string(), field.to_string(), term.to_string()),
74            scores,
75        );
76        Ok(())
77    }
78
79    /// Compute and store per-block maxima for `posting_list`. Each
80    /// block's max is `max_{e in block} scorer.score(tf(e), tf(e), df)` and
81    /// uses term frequency as the document-length stand-in when none is given.
82    pub fn build<S: BlockMaxScorer + ?Sized>(
83        &mut self,
84        posting_list: &PostingList,
85        scorer: &S,
86        field: &str,
87        term: &str,
88        table: &str,
89    ) -> StorageBackendResult<()> {
90        if self.block_size == 0 {
91            return Err(StorageBackendError::Other(
92                "block-max block size must be greater than zero".to_string(),
93            ));
94        }
95        let entries = posting_list.entries();
96        let key = (table.to_string(), field.to_string(), term.to_string());
97        if entries.is_empty() {
98            self.block_maxes.insert(key, Vec::new());
99            return Ok(());
100        }
101        let df = u64::try_from(entries.len()).map_err(|_| {
102            StorageBackendError::Other("posting-list length exceeds u64".to_string())
103        })?;
104        let mut blocks = Vec::with_capacity(entries.len().div_ceil(self.block_size));
105        for chunk in entries.chunks(self.block_size) {
106            let mut max_score = 0.0_f64;
107            for entry in chunk {
108                let positions = &entry.payload.positions;
109                let tf = if positions.is_empty() {
110                    1
111                } else {
112                    u64::try_from(positions.len()).map_err(|_| {
113                        StorageBackendError::Other("term position count exceeds u64".to_string())
114                    })?
115                };
116                let s = scorer.score(tf, tf, df);
117                validate_score(s)?;
118                if s > max_score {
119                    max_score = s;
120                }
121            }
122            blocks.push(max_score);
123        }
124        self.block_maxes.insert(key, blocks);
125        Ok(())
126    }
127
128    pub fn block_max(&self, table: &str, field: &str, term: &str, block_idx: usize) -> f64 {
129        let key = (table.to_string(), field.to_string(), term.to_string());
130        self.block_maxes
131            .get(&key)
132            .and_then(|v| v.get(block_idx).copied())
133            .unwrap_or(0.0)
134    }
135
136    pub fn num_blocks(&self, table: &str, field: &str, term: &str) -> usize {
137        let key = (table.to_string(), field.to_string(), term.to_string());
138        self.block_maxes.get(&key).map_or(0, Vec::len)
139    }
140
141    /// Borrow all block scores for one posting without repeated key
142    /// construction. BMW uses this to precompute suffix bounds once per query.
143    pub fn block_maxes(&self, table: &str, field: &str, term: &str) -> Option<&[f64]> {
144        let key = (table.to_string(), field.to_string(), term.to_string());
145        self.block_maxes.get(&key).map(Vec::as_slice)
146    }
147
148    /// Block index for a given posting-list cursor position.
149    pub fn block_index_for(&self, position: usize) -> StorageBackendResult<usize> {
150        if self.block_size == 0 {
151            return Err(StorageBackendError::Other(
152                "block-max block size must be greater than zero".to_string(),
153            ));
154        }
155        Ok(position / self.block_size)
156    }
157
158    pub fn clear(&mut self) {
159        self.block_maxes.clear();
160    }
161
162    pub fn save_to_sqlite(&self, conn: &rusqlite::Connection) -> rusqlite::Result<()> {
163        for scores in self.block_maxes.values() {
164            validate_scores(scores).map_err(storage_error_to_sqlite)?;
165        }
166        ensure_global_blockmax_shape(conn)?;
167        let transaction = conn.unchecked_transaction()?;
168        transaction.execute("DELETE FROM _global_blockmax", [])?;
169        for ((table, field, term), scores) in &self.block_maxes {
170            for (block_idx, score) in scores.iter().enumerate() {
171                let block_idx = i64::try_from(block_idx)
172                    .map_err(|error| rusqlite::Error::ToSqlConversionFailure(Box::new(error)))?;
173                transaction.execute(
174                    "INSERT INTO _global_blockmax
175                        (table_name, field, term, block_idx, max_score)
176                     VALUES (?1, ?2, ?3, ?4, ?5)",
177                    params![table, field, term, block_idx, *score],
178                )?;
179            }
180        }
181        transaction.commit()
182    }
183
184    pub fn load_from_sqlite(&mut self, conn: &rusqlite::Connection) -> rusqlite::Result<()> {
185        ensure_global_blockmax_shape(conn)?;
186        let mut stmt = conn.prepare(
187            "SELECT table_name, field, term, block_idx, max_score
188             FROM _global_blockmax
189             ORDER BY table_name, field, term, block_idx",
190        )?;
191        let rows = stmt.query_map([], |row| {
192            Ok((
193                row.get::<_, String>(0)?,
194                row.get::<_, String>(1)?,
195                row.get::<_, String>(2)?,
196                row.get::<_, i64>(3)?,
197                row.get::<_, f64>(4)?,
198            ))
199        })?;
200        let mut loaded = BTreeMap::<(String, String, String), Vec<f64>>::new();
201        for row in rows {
202            let (table, field, term, block_idx, score) = row?;
203            validate_score(score).map_err(storage_error_to_sqlite)?;
204            let idx = usize::try_from(block_idx)
205                .map_err(|_| rusqlite::Error::IntegralValueOutOfRange(3, block_idx))?;
206            let entry = loaded.entry((table, field, term)).or_default();
207            if idx != entry.len() {
208                return Err(rusqlite::Error::FromSqlConversionFailure(
209                    3,
210                    rusqlite::types::Type::Integer,
211                    Box::new(std::io::Error::new(
212                        std::io::ErrorKind::InvalidData,
213                        format!(
214                            "invalid block-max ordinal sequence: expected {}, found {idx}",
215                            entry.len()
216                        ),
217                    )),
218                ));
219            }
220            entry.push(score);
221        }
222        self.block_maxes = loaded;
223        Ok(())
224    }
225}
226
227fn validate_scores(scores: &[f64]) -> StorageBackendResult<()> {
228    for &score in scores {
229        validate_score(score)?;
230    }
231    Ok(())
232}
233
234fn validate_score(score: f64) -> StorageBackendResult<()> {
235    if score.is_finite() && score >= 0.0 {
236        Ok(())
237    } else {
238        Err(StorageBackendError::Other(format!(
239            "block-max score must be finite and non-negative, got {score}"
240        )))
241    }
242}
243
244fn storage_error_to_sqlite(error: StorageBackendError) -> rusqlite::Error {
245    rusqlite::Error::ToSqlConversionFailure(Box::new(error))
246}
247
248fn ensure_global_blockmax_shape(conn: &rusqlite::Connection) -> rusqlite::Result<()> {
249    conn.execute(
250        "CREATE TABLE IF NOT EXISTS _global_blockmax (
251            table_name TEXT NOT NULL DEFAULT '',
252            field     TEXT NOT NULL,
253            term      TEXT NOT NULL,
254            block_idx INTEGER NOT NULL,
255            max_score REAL NOT NULL,
256            PRIMARY KEY (table_name, field, term, block_idx)
257        )",
258        [],
259    )?;
260    let mut stmt = conn.prepare("PRAGMA table_info(_global_blockmax)")?;
261    let cols = stmt
262        .query_map([], |row| row.get::<_, String>(1))?
263        .collect::<Result<Vec<_>, _>>()?;
264    drop(stmt);
265    if !cols.iter().any(|c| c == "table_name") {
266        conn.execute(
267            "ALTER TABLE _global_blockmax ADD COLUMN table_name TEXT NOT NULL DEFAULT ''",
268            [],
269        )?;
270    }
271    Ok(())
272}
273
274#[cfg(test)]
275mod tests {
276    use super::*;
277    use uqa_core::{Payload, PostingEntry, PostingList};
278
279    /// Trivial scorer: tf raised to a constant — strictly increasing in
280    /// tf so the per-block max equals the max tf in the block.
281    struct LinearScorer;
282    impl BlockMaxScorer for LinearScorer {
283        fn score(&self, term_freq: u64, _doc_length: u64, _doc_freq: u64) -> f64 {
284            term_freq as f64
285        }
286    }
287
288    struct InvalidScorer(f64);
289    impl BlockMaxScorer for InvalidScorer {
290        fn score(&self, _term_freq: u64, _doc_length: u64, _doc_freq: u64) -> f64 {
291            self.0
292        }
293    }
294
295    fn pl_with_tfs(tfs: &[u32]) -> PostingList {
296        let entries: Vec<PostingEntry> = tfs
297            .iter()
298            .enumerate()
299            .map(|(i, &tf)| {
300                let positions = (0..tf).collect();
301                PostingEntry::new(
302                    u64::try_from(i).unwrap() + 1,
303                    Payload {
304                        positions,
305                        score: 0.0,
306                        fields: BTreeMap::default(),
307                    },
308                )
309            })
310            .collect();
311        PostingList::from_unsorted(entries)
312    }
313
314    #[test]
315    fn block_max_records_per_block_maximum() {
316        let mut idx = BlockMaxIndex::new(2).unwrap();
317        let pl = pl_with_tfs(&[1, 5, 3, 7, 2]);
318        idx.build(&pl, &LinearScorer, "title", "rust", "articles")
319            .unwrap();
320        // Blocks of size 2: [1, 5] [3, 7] [2] -> maxes 5, 7, 2
321        assert_eq!(idx.num_blocks("articles", "title", "rust"), 3);
322        assert!((idx.block_max("articles", "title", "rust", 0) - 5.0).abs() < 1e-12);
323        assert!((idx.block_max("articles", "title", "rust", 1) - 7.0).abs() < 1e-12);
324        assert!((idx.block_max("articles", "title", "rust", 2) - 2.0).abs() < 1e-12);
325    }
326
327    #[test]
328    fn empty_posting_list_records_no_blocks() {
329        let mut idx = BlockMaxIndex::new(4).unwrap();
330        idx.build(&PostingList::new(), &LinearScorer, "title", "rust", "t")
331            .unwrap();
332        assert_eq!(idx.num_blocks("t", "title", "rust"), 0);
333        assert!((idx.block_max("t", "title", "rust", 0) - 0.0).abs() < 1e-12);
334    }
335
336    #[test]
337    fn block_index_for_position() {
338        let idx = BlockMaxIndex::new(4).unwrap();
339        assert_eq!(idx.block_index_for(0).unwrap(), 0);
340        assert_eq!(idx.block_index_for(3).unwrap(), 0);
341        assert_eq!(idx.block_index_for(4).unwrap(), 1);
342        assert_eq!(idx.block_index_for(9).unwrap(), 2);
343    }
344
345    #[test]
346    fn rejects_zero_block_size_and_invalid_scores_without_replacing_state() {
347        assert!(BlockMaxIndex::new(0).is_err());
348
349        let mut index = BlockMaxIndex::new(2).unwrap();
350        index
351            .set_block_maxes("docs", "body", "term", vec![3.0])
352            .unwrap();
353        let postings = pl_with_tfs(&[1, 2]);
354        assert!(index
355            .build(&postings, &InvalidScorer(f64::NAN), "body", "term", "docs")
356            .is_err());
357        assert_eq!(index.block_max("docs", "body", "term", 0), 3.0);
358        assert!(index
359            .set_block_maxes("docs", "body", "term", vec![-1.0])
360            .is_err());
361        assert_eq!(index.block_max("docs", "body", "term", 0), 3.0);
362    }
363
364    #[test]
365    fn corrupt_persisted_ordinal_does_not_replace_loaded_state() {
366        let connection = rusqlite::Connection::open_in_memory().unwrap();
367        ensure_global_blockmax_shape(&connection).unwrap();
368        connection
369            .execute(
370                "INSERT INTO _global_blockmax
371                    (table_name, field, term, block_idx, max_score)
372                 VALUES ('docs', 'body', 'bad', -1, 9.0)",
373                [],
374            )
375            .unwrap();
376        let mut index = BlockMaxIndex::default();
377        index
378            .set_block_maxes("old", "body", "term", vec![1.0])
379            .unwrap();
380
381        assert!(index.load_from_sqlite(&connection).is_err());
382        assert_eq!(index.block_max("old", "body", "term", 0), 1.0);
383    }
384
385    #[test]
386    fn failed_save_rolls_back_deleted_snapshot() {
387        let connection = rusqlite::Connection::open_in_memory().unwrap();
388        ensure_global_blockmax_shape(&connection).unwrap();
389        connection
390            .execute(
391                "INSERT INTO _global_blockmax
392                    (table_name, field, term, block_idx, max_score)
393                 VALUES ('old', 'body', 'term', 0, 1.0)",
394                [],
395            )
396            .unwrap();
397        connection
398            .execute_batch(
399                "CREATE TRIGGER fail_blockmax_insert
400                 BEFORE INSERT ON _global_blockmax
401                 BEGIN
402                     SELECT RAISE(ABORT, 'injected block-max failure');
403                 END;",
404            )
405            .unwrap();
406        let mut index = BlockMaxIndex::default();
407        index
408            .set_block_maxes("new", "body", "term", vec![2.0])
409            .unwrap();
410
411        assert!(index.save_to_sqlite(&connection).is_err());
412        let persisted: (String, f64) = connection
413            .query_row(
414                "SELECT table_name, max_score FROM _global_blockmax",
415                [],
416                |row| Ok((row.get(0)?, row.get(1)?)),
417            )
418            .unwrap();
419        assert_eq!(persisted, ("old".to_string(), 1.0));
420    }
421}
uqa_storage/block_max_index.rs

uqa_storage/
block_max_index.rs
