uqa_storage/sqlite/inverted_index/
block_max.rs1use 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 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 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}