Skip to main content

uqa_storage/sqlite/inverted_index/
block_max.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Skip-pointer and scorer-versioned block-max materialization.
8
9use super::{
10    decode_index_u64, decode_index_usize, encode_index_u64, encode_index_usize, params,
11    quote_ident, table_exists, BlockMaxIndex, BlockMaxScorer, DocId, InvertedIndex,
12    OptionalExtension, SQLiteError, SQLiteInvertedIndex, StorageBackendResult,
13};
14
15impl SQLiteInvertedIndex {
16    /// Read the nearest materialized skip pointer without mutating storage.
17    /// Call [`Self::flush_skip_pointers`] from an explicit maintenance/write
18    /// boundary after postings change.
19    pub fn skip_to(
20        &self,
21        field: &str,
22        term: &str,
23        target_doc_id: DocId,
24    ) -> StorageBackendResult<(DocId, usize)> {
25        let table = self.skip_table_name(field);
26        let target_doc_id = encode_index_u64("target document", target_doc_id)?;
27        Ok(self.conn.with(|conn| {
28            if !table_exists(conn, &table)? {
29                return Ok((0, 0));
30            }
31            let sql = format!(
32                "SELECT skip_doc_id, skip_offset FROM {}
33                     WHERE term = ?1 AND skip_doc_id <= ?2
34                     ORDER BY skip_doc_id DESC LIMIT 1",
35                quote_ident(&table)
36            );
37            let row: Option<(i64, i64)> = conn
38                .query_row(&sql, params![term, target_doc_id], |row| {
39                    Ok((row.get(0)?, row.get(1)?))
40                })
41                .optional()?;
42            match row {
43                Some((doc_id, offset)) => Ok((
44                    decode_index_u64("skip document", doc_id)?,
45                    decode_index_usize("skip offset", offset)?,
46                )),
47                None => Ok((0, 0)),
48            }
49        })?)
50    }
51
52    pub fn build_block_max_scores<S: BlockMaxScorer + ?Sized>(
53        &self,
54        field: &str,
55        term: &str,
56        scorer: &S,
57    ) -> StorageBackendResult<()> {
58        self.build_block_max_scores_versioned(field, term, scorer, "")
59    }
60
61    /// Build one scorer-versioned block-max posting. The version is checked
62    /// on load so changed BM25 parameters or field statistics can never feed
63    /// an unsafe pruning bound.
64    pub fn build_block_max_scores_versioned<S: BlockMaxScorer + ?Sized>(
65        &self,
66        field: &str,
67        term: &str,
68        scorer: &S,
69        scorer_fingerprint: &str,
70    ) -> StorageBackendResult<()> {
71        let mut cursor = self.posting_cursor(field, term)?;
72        if cursor.doc_freq() == 0 && !self.has_field(field)? {
73            return Ok(());
74        }
75        let df = cursor.doc_freq();
76        let scored_capacity = usize::try_from(df)
77            .map_err(|_| SQLiteError::StorageBackend("document frequency exceeds usize".into()))?;
78        let mut scored_entries = Vec::with_capacity(scored_capacity);
79        while let Some(entry) = cursor.current() {
80            scored_entries.push((entry.term_freq, entry.doc_length.max(entry.term_freq)));
81            cursor.advance()?;
82        }
83        self.ensure_aux_tables(field)?;
84        let table = self.blockmax_table_name(field);
85        self.conn.with_mut(|conn| {
86            let tx = conn.savepoint()?;
87            tx.execute(
88                &format!("DELETE FROM {} WHERE term = ?1", quote_ident(&table)),
89                [term],
90            )?;
91            for (block_idx, chunk) in scored_entries.chunks(Self::BLOCK_SIZE).enumerate() {
92                let mut max_score = 0.0_f64;
93                for &(tf, doc_length) in chunk {
94                    let score = scorer.score(tf, doc_length, df);
95                    if !score.is_finite() || score < 0.0 {
96                        return Err(SQLiteError::StorageBackend(format!(
97                            "block-max score must be finite and non-negative, got {score}"
98                        )));
99                    }
100                    max_score = max_score.max(score);
101                }
102                let block_idx = encode_index_usize("block index", block_idx)?;
103                tx.execute(
104                    &format!(
105                        "INSERT OR REPLACE INTO {}
106                            (term, block_idx, max_score, scorer_fingerprint)
107                         VALUES (?1, ?2, ?3, ?4)",
108                        quote_ident(&table)
109                    ),
110                    params![term, block_idx, max_score, scorer_fingerprint],
111                )?;
112            }
113            tx.commit()?;
114            Ok(())
115        })?;
116        Ok(())
117    }
118
119    pub fn build_all_block_max_scores<S: BlockMaxScorer + ?Sized>(
120        &self,
121        field: &str,
122        scorer: &S,
123    ) -> StorageBackendResult<()> {
124        let terms = self.terms_for_field(field)?;
125        for term in terms {
126            self.build_block_max_scores(field, &term, scorer)?;
127        }
128        Ok(())
129    }
130
131    pub fn get_block_max_score(
132        &self,
133        field: &str,
134        term: &str,
135        block_idx: usize,
136    ) -> StorageBackendResult<f64> {
137        let table = self.blockmax_table_name(field);
138        let block_idx = encode_index_usize("block index", block_idx)?;
139        Ok(self.conn.with(|conn| {
140            if !table_exists(conn, &table)? {
141                return Ok(0.0);
142            }
143            let sql = format!(
144                "SELECT max_score FROM {}
145                     WHERE term = ?1 AND block_idx = ?2",
146                quote_ident(&table)
147            );
148            let score: Option<f64> = conn
149                .query_row(&sql, params![term, block_idx], |row| row.get(0))
150                .optional()?;
151            Ok(score.unwrap_or(0.0))
152        })?)
153    }
154
155    pub fn get_all_block_max_scores(
156        &self,
157        field: &str,
158        term: &str,
159    ) -> StorageBackendResult<Vec<f64>> {
160        let table = self.blockmax_table_name(field);
161        Ok(self.conn.with(|conn| {
162            if !table_exists(conn, &table)? {
163                return Ok(Vec::new());
164            }
165            let sql = format!(
166                "SELECT block_idx, max_score FROM {}
167                     WHERE term = ?1 ORDER BY block_idx",
168                quote_ident(&table)
169            );
170            let mut stmt = conn.prepare(&sql)?;
171            let rows = stmt
172                .query_map([term], |row| {
173                    Ok((row.get::<_, i64>(0)?, row.get::<_, f64>(1)?))
174                })?
175                .collect::<Result<Vec<_>, _>>()?;
176            let mut scores = Vec::with_capacity(rows.len());
177            for (expected, (block_idx, score)) in rows.into_iter().enumerate() {
178                let block_idx = decode_index_usize("block index", block_idx)?;
179                if block_idx != expected {
180                    return Err(SQLiteError::StorageBackend(format!(
181                        "corrupt inverted index: expected block index {expected}, found {block_idx}"
182                    )));
183                }
184                scores.push(score);
185            }
186            Ok(scores)
187        })?)
188    }
189
190    pub fn get_versioned_block_max_scores(
191        &self,
192        field: &str,
193        term: &str,
194        scorer_fingerprint: &str,
195    ) -> StorageBackendResult<Option<Vec<f64>>> {
196        Ok(self
197            .get_versioned_block_max_scores_bulk(field, &[term.to_string()], scorer_fingerprint)?
198            .pop()
199            .flatten())
200    }
201
202    pub fn get_versioned_block_max_scores_bulk(
203        &self,
204        field: &str,
205        terms: &[String],
206        scorer_fingerprint: &str,
207    ) -> StorageBackendResult<Vec<Option<Vec<f64>>>> {
208        if terms.is_empty() {
209            return Ok(Vec::new());
210        }
211        let table = self.blockmax_table_name(field);
212        Ok(self.conn.with(|conn| {
213            if !table_exists(conn, &table)? {
214                return Ok(vec![None; terms.len()]);
215            }
216            let pragma = format!("PRAGMA table_info({})", quote_ident(&table));
217            let mut columns = conn.prepare(&pragma)?;
218            let has_fingerprint = columns
219                .query_map([], |row| row.get::<_, String>(1))?
220                .collect::<Result<Vec<_>, _>>()?
221                .iter()
222                .any(|name| name == "scorer_fingerprint");
223            drop(columns);
224            if !has_fingerprint {
225                return Ok(vec![None; terms.len()]);
226            }
227
228            let unique_terms = terms
229                .iter()
230                .cloned()
231                .collect::<std::collections::BTreeSet<_>>();
232            let unique_terms = unique_terms.into_iter().collect::<Vec<_>>();
233            let mut by_term = std::collections::BTreeMap::<String, Vec<(i64, f64)>>::new();
234            for chunk in unique_terms.chunks(900) {
235                let placeholders = std::iter::repeat_n("?", chunk.len())
236                    .collect::<Vec<_>>()
237                    .join(", ");
238                let sql = format!(
239                    "SELECT term, block_idx, max_score FROM {}
240                     WHERE scorer_fingerprint = ? AND term IN ({placeholders})
241                     ORDER BY term, block_idx",
242                    quote_ident(&table)
243                );
244                let mut values = Vec::with_capacity(chunk.len() + 1);
245                values.push(rusqlite::types::Value::Text(scorer_fingerprint.to_string()));
246                values.extend(chunk.iter().cloned().map(rusqlite::types::Value::Text));
247                let mut statement = conn.prepare(&sql)?;
248                let rows = statement.query_map(rusqlite::params_from_iter(values), |row| {
249                    Ok((
250                        row.get::<_, String>(0)?,
251                        row.get::<_, i64>(1)?,
252                        row.get::<_, f64>(2)?,
253                    ))
254                })?;
255                for row in rows {
256                    let (term, block_idx, score) = row?;
257                    by_term.entry(term).or_default().push((block_idx, score));
258                }
259            }
260
261            let mut decoded = std::collections::BTreeMap::<String, Option<Vec<f64>>>::new();
262            for term in unique_terms {
263                let rows = by_term.remove(&term).unwrap_or_default();
264                if rows.is_empty() {
265                    decoded.insert(term, None);
266                    continue;
267                }
268                let mut scores = Vec::with_capacity(rows.len());
269                for (expected, (block_idx, score)) in rows.into_iter().enumerate() {
270                    let block_idx = decode_index_usize("block index", block_idx)?;
271                    if block_idx != expected || !score.is_finite() || score < 0.0 {
272                        return Err(SQLiteError::StorageBackend(format!(
273                            "corrupt block-max index for `{field}.{term}` at block {block_idx}"
274                        )));
275                    }
276                    scores.push(score);
277                }
278                decoded.insert(term, Some(scores));
279            }
280            Ok(terms.iter().map(|term| decoded[term].clone()).collect())
281        })?)
282    }
283
284    pub fn load_block_max_into(&self, target: &mut BlockMaxIndex) -> StorageBackendResult<()> {
285        for field in self.fields_with_blockmax_tables()? {
286            for term in self.terms_for_field(&field)? {
287                let scores = self.get_all_block_max_scores(&field, &term)?;
288                if !scores.is_empty() {
289                    target.set_block_maxes(&self.table, &field, &term, scores)?;
290                }
291            }
292        }
293        Ok(())
294    }
295}