1use 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, TokenTermKey,
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 self.require_graph_format()?;
26 let term = TokenTermKey::from_text(term);
27 let table = self.skip_table_name(field);
28 let target_doc_id = encode_index_u64("target document", target_doc_id)?;
29 Ok(self.conn.with(|conn| {
30 if !table_exists(conn, &table)? {
31 return Ok((0, 0));
32 }
33 let sql = format!(
34 "SELECT skip_doc_id, skip_offset FROM {}
35 WHERE term = ?1 AND skip_doc_id <= ?2
36 ORDER BY skip_doc_id DESC LIMIT 1",
37 quote_ident(&table)
38 );
39 let row: Option<(i64, i64)> = conn
40 .query_row(&sql, params![term.as_bytes(), target_doc_id], |row| {
41 Ok((row.get(0)?, row.get(1)?))
42 })
43 .optional()?;
44 match row {
45 Some((doc_id, offset)) => Ok((
46 decode_index_u64("skip document", doc_id)?,
47 decode_index_usize("skip offset", offset)?,
48 )),
49 None => Ok((0, 0)),
50 }
51 })?)
52 }
53
54 pub fn build_block_max_scores<S: BlockMaxScorer + ?Sized>(
55 &self,
56 field: &str,
57 term: &str,
58 scorer: &S,
59 ) -> StorageBackendResult<()> {
60 self.build_block_max_scores_versioned(field, term, scorer, "")
61 }
62
63 pub fn build_block_max_scores_versioned<S: BlockMaxScorer + ?Sized>(
67 &self,
68 field: &str,
69 term: &str,
70 scorer: &S,
71 scorer_fingerprint: &str,
72 ) -> StorageBackendResult<()> {
73 self.build_block_max_scores_key(
74 field,
75 &TokenTermKey::from_text(term),
76 scorer,
77 scorer_fingerprint,
78 )
79 }
80
81 pub(super) fn build_block_max_scores_key<S: BlockMaxScorer + ?Sized>(
82 &self,
83 field: &str,
84 term: &TokenTermKey,
85 scorer: &S,
86 scorer_fingerprint: &str,
87 ) -> StorageBackendResult<()> {
88 let mut cursor = self.posting_cursor_key(field, term)?;
89 if cursor.doc_freq() == 0 && !self.has_field(field)? {
90 return Ok(());
91 }
92 let df = cursor.doc_freq();
93 let scored_capacity = usize::try_from(df)
94 .map_err(|_| SQLiteError::StorageBackend("document frequency exceeds usize".into()))?;
95 let mut scored_entries = Vec::with_capacity(scored_capacity);
96 while let Some(entry) = cursor.current() {
97 scored_entries.push((entry.term_freq, entry.doc_length));
98 cursor.advance()?;
99 }
100 self.ensure_aux_tables(field)?;
101 let table = self.blockmax_table_name(field);
102 self.conn.with_mut(|conn| {
103 let tx = conn.savepoint()?;
104 tx.execute(
105 &format!("DELETE FROM {} WHERE term = ?1", quote_ident(&table)),
106 [term.as_bytes()],
107 )?;
108 for (block_idx, chunk) in scored_entries.chunks(Self::BLOCK_SIZE).enumerate() {
109 let mut max_score = 0.0_f64;
110 for &(tf, doc_length) in chunk {
111 let score = scorer.score(tf, doc_length, df);
112 if !score.is_finite() || score < 0.0 {
113 return Err(SQLiteError::StorageBackend(format!(
114 "block-max score must be finite and non-negative, got {score}"
115 )));
116 }
117 max_score = max_score.max(score);
118 }
119 let block_idx = encode_index_usize("block index", block_idx)?;
120 tx.execute(
121 &format!(
122 "INSERT OR REPLACE INTO {}
123 (term, block_idx, max_score, scorer_fingerprint)
124 VALUES (?1, ?2, ?3, ?4)",
125 quote_ident(&table)
126 ),
127 params![term.as_bytes(), block_idx, max_score, scorer_fingerprint],
128 )?;
129 }
130 tx.commit()?;
131 Ok(())
132 })?;
133 Ok(())
134 }
135
136 pub fn build_all_block_max_scores<S: BlockMaxScorer + ?Sized>(
137 &self,
138 field: &str,
139 scorer: &S,
140 ) -> StorageBackendResult<()> {
141 let terms = self.vocabulary_keys(field)?;
142 for term in terms {
143 self.build_block_max_scores_key(field, &term, scorer, "")?;
144 }
145 Ok(())
146 }
147
148 pub fn get_block_max_score(
149 &self,
150 field: &str,
151 term: &str,
152 block_idx: usize,
153 ) -> StorageBackendResult<f64> {
154 self.require_graph_format()?;
155 let term = TokenTermKey::from_text(term);
156 let table = self.blockmax_table_name(field);
157 let block_idx = encode_index_usize("block index", block_idx)?;
158 Ok(self.conn.with(|conn| {
159 if !table_exists(conn, &table)? {
160 return Ok(0.0);
161 }
162 let sql = format!(
163 "SELECT max_score FROM {}
164 WHERE term = ?1 AND block_idx = ?2",
165 quote_ident(&table)
166 );
167 let score: Option<f64> = conn
168 .query_row(&sql, params![term.as_bytes(), block_idx], |row| row.get(0))
169 .optional()?;
170 Ok(score.unwrap_or(0.0))
171 })?)
172 }
173
174 pub fn get_all_block_max_scores(
175 &self,
176 field: &str,
177 term: &str,
178 ) -> StorageBackendResult<Vec<f64>> {
179 self.get_all_block_max_scores_key(field, &TokenTermKey::from_text(term))
180 }
181
182 pub fn get_all_block_max_scores_key(
183 &self,
184 field: &str,
185 term: &TokenTermKey,
186 ) -> StorageBackendResult<Vec<f64>> {
187 self.require_graph_format()?;
188 let table = self.blockmax_table_name(field);
189 Ok(self.conn.with(|conn| {
190 if !table_exists(conn, &table)? {
191 return Ok(Vec::new());
192 }
193 let sql = format!(
194 "SELECT block_idx, max_score FROM {}
195 WHERE term = ?1 ORDER BY block_idx",
196 quote_ident(&table)
197 );
198 let mut stmt = conn.prepare(&sql)?;
199 let rows = stmt
200 .query_map([term.as_bytes()], |row| {
201 Ok((row.get::<_, i64>(0)?, row.get::<_, f64>(1)?))
202 })?
203 .collect::<Result<Vec<_>, _>>()?;
204 let mut scores = Vec::with_capacity(rows.len());
205 for (expected, (block_idx, score)) in rows.into_iter().enumerate() {
206 let block_idx = decode_index_usize("block index", block_idx)?;
207 if block_idx != expected {
208 return Err(SQLiteError::StorageBackend(format!(
209 "corrupt inverted index: expected block index {expected}, found {block_idx}"
210 )));
211 }
212 scores.push(score);
213 }
214 Ok(scores)
215 })?)
216 }
217
218 pub fn get_versioned_block_max_scores(
219 &self,
220 field: &str,
221 term: &str,
222 scorer_fingerprint: &str,
223 ) -> StorageBackendResult<Option<Vec<f64>>> {
224 Ok(self
225 .get_versioned_block_max_scores_bulk(field, &[term.to_string()], scorer_fingerprint)?
226 .pop()
227 .flatten())
228 }
229
230 pub fn get_versioned_block_max_scores_bulk(
231 &self,
232 field: &str,
233 terms: &[String],
234 scorer_fingerprint: &str,
235 ) -> StorageBackendResult<Vec<Option<Vec<f64>>>> {
236 let keys = terms
237 .iter()
238 .map(|term| TokenTermKey::from_text(term))
239 .collect::<Vec<_>>();
240 self.get_versioned_block_max_scores_keys_bulk(field, &keys, scorer_fingerprint)
241 }
242
243 pub fn get_versioned_block_max_scores_keys_bulk(
244 &self,
245 field: &str,
246 terms: &[TokenTermKey],
247 scorer_fingerprint: &str,
248 ) -> StorageBackendResult<Vec<Option<Vec<f64>>>> {
249 if terms.is_empty() {
250 return Ok(Vec::new());
251 }
252 let table = self.blockmax_table_name(field);
253 Ok(self.conn.with(|conn| {
254 self.require_graph_format_on(conn)?;
255 if !table_exists(conn, &table)? {
256 return Ok(vec![None; terms.len()]);
257 }
258 let pragma = format!("PRAGMA table_info({})", quote_ident(&table));
259 let mut columns = conn.prepare(&pragma)?;
260 let has_fingerprint = columns
261 .query_map([], |row| row.get::<_, String>(1))?
262 .collect::<Result<Vec<_>, _>>()?
263 .iter()
264 .any(|name| name == "scorer_fingerprint");
265 drop(columns);
266 if !has_fingerprint {
267 return Ok(vec![None; terms.len()]);
268 }
269
270 let unique_terms = terms
271 .iter()
272 .cloned()
273 .collect::<std::collections::BTreeSet<_>>();
274 let unique_terms = unique_terms.into_iter().collect::<Vec<_>>();
275 let mut by_term = std::collections::BTreeMap::<Vec<u8>, Vec<(i64, f64)>>::new();
276 for chunk in unique_terms.chunks(900) {
277 let placeholders = std::iter::repeat_n("?", chunk.len())
278 .collect::<Vec<_>>()
279 .join(", ");
280 let sql = format!(
281 "SELECT term, block_idx, max_score FROM {}
282 WHERE scorer_fingerprint = ? AND term IN ({placeholders})
283 ORDER BY term, block_idx",
284 quote_ident(&table)
285 );
286 let mut values = Vec::with_capacity(chunk.len() + 1);
287 values.push(rusqlite::types::Value::Text(scorer_fingerprint.to_string()));
288 values.extend(
289 chunk
290 .iter()
291 .map(|term| rusqlite::types::Value::Blob(term.as_bytes().to_vec())),
292 );
293 let mut statement = conn.prepare(&sql)?;
294 let rows = statement.query_map(rusqlite::params_from_iter(values), |row| {
295 Ok((
296 row.get::<_, Vec<u8>>(0)?,
297 row.get::<_, i64>(1)?,
298 row.get::<_, f64>(2)?,
299 ))
300 })?;
301 for row in rows {
302 let (term, block_idx, score) = row?;
303 by_term.entry(term).or_default().push((block_idx, score));
304 }
305 }
306
307 let mut decoded = std::collections::BTreeMap::<TokenTermKey, Option<Vec<f64>>>::new();
308 for term in unique_terms {
309 let rows = by_term.remove(term.as_bytes()).unwrap_or_default();
310 if rows.is_empty() {
311 decoded.insert(term, None);
312 continue;
313 }
314 let mut scores = Vec::with_capacity(rows.len());
315 for (expected, (block_idx, score)) in rows.into_iter().enumerate() {
316 let block_idx = decode_index_usize("block index", block_idx)?;
317 if block_idx != expected || !score.is_finite() || score < 0.0 {
318 return Err(SQLiteError::StorageBackend(format!(
319 "corrupt block-max index for `{field}.{term:?}` at block {block_idx}"
320 )));
321 }
322 scores.push(score);
323 }
324 decoded.insert(term, Some(scores));
325 }
326 Ok(terms.iter().map(|term| decoded[term].clone()).collect())
327 })?)
328 }
329
330 pub fn load_block_max_into(&self, target: &mut BlockMaxIndex) -> StorageBackendResult<()> {
331 for field in self.fields_with_blockmax_tables()? {
332 for term in self.vocabulary_keys(&field)? {
333 let scores = self.get_all_block_max_scores_key(&field, &term)?;
334 if !scores.is_empty() {
335 target.set_block_maxes_key(&self.table, &field, &term, scores)?;
336 }
337 }
338 }
339 Ok(())
340 }
341}