summa_core/segment/reader/
candidate_lookup.rs

1//! Candidate addressing shared by every point-scoring vertical.
2
3use super::SegmentReader;
4use crate::dsl::{Field, FieldType};
5
6use crate::{DocId, Error, Result};
7
8#[derive(Clone, Copy, Debug)]
9pub(crate) struct CandidateLocation {
10    pub doc: DocId,
11    pub ordinal: u16,
12    /// Field-local ID; chunk/BMP physical IDs and flat-vector indexes need
13    /// not agree even when they refer to the same logical document ordinal.
14    pub physical: u32,
15}
16
17impl SegmentReader {
18    /// Resolve every stored value of the selected documents in one field.
19    /// The caller controls the expansion budget; no field silently truncates
20    /// a long document or invents a body ordinal for document context.
21    pub(crate) async fn candidate_locations(
22        &self,
23        field: Field,
24        documents: &[DocId],
25        max_locations: usize,
26        budget: &mut super::SparseProbeBudget,
27    ) -> Result<Vec<CandidateLocation>> {
28        if documents.iter().any(|&doc| doc >= self.num_docs())
29            || !documents.windows(2).all(|pair| pair[0] < pair[1])
30        {
31            return Err(Error::Query(
32                "candidate documents must be valid, unique and sorted".into(),
33            ));
34        }
35        if documents.is_empty() {
36            return Ok(Vec::new());
37        }
38        let entry = self
39            .schema
40            .get_field_entry(field)
41            .ok_or_else(|| Error::FieldNotFound(field.0.to_string()))?;
42        let mut locations = Vec::with_capacity(documents.len().min(max_locations));
43        let mut push = |doc, ordinal, physical| -> Result<()> {
44            if locations.len() == max_locations {
45                return Err(Error::Query(format!(
46                    "candidate scoring expands beyond {max_locations} field values"
47                )));
48            }
49            locations.push(CandidateLocation {
50                doc,
51                ordinal,
52                physical,
53            });
54            Ok(())
55        };
56        match &entry.field_type {
57            FieldType::Text if entry.chunked => {
58                let Some(map) = self.chunk_map(field) else {
59                    return Ok(locations);
60                };
61                if !map.has_logical_addressing() {
62                    return Err(Error::Query("legacy reordered text needs explicit Reorder to upgrade its chunk map for L1".into()));
63                }
64                for &doc in documents {
65                    for (ordinal, physical) in map.slots_for_document(doc) {
66                        push(doc, ordinal, physical)?;
67                    }
68                }
69            }
70            FieldType::SparseVector => {
71                if let Some(index) = self.seismic_index(field) {
72                    for &doc in documents {
73                        for physical in index.rows_for_document(doc) {
74                            push(doc, index.key(physical).ordinal, physical)?;
75                        }
76                    }
77                    return Ok(locations);
78                }
79
80                let Some(bmp) = self.bmp_indexes.get(&field.0) else {
81                    if self.sparse_indexes.contains_key(&field.0) {
82                        return self
83                            .maxscore_candidate_locations(
84                                field,
85                                documents,
86                                None,
87                                max_locations,
88                                budget,
89                            )
90                            .await;
91                    }
92                    return Ok(locations);
93                };
94                for &doc in documents {
95                    if let Some(forward) = bmp.forward() {
96                        for (ordinal, physical) in forward.for_document(doc) {
97                            push(doc, ordinal, physical)?;
98                        }
99                    } else if bmp.logically_ordered() {
100                        for (ordinal, physical) in bmp.ordered_slots_for_document(doc) {
101                            push(doc, ordinal, physical)?;
102                        }
103                    } else {
104                        return Err(Error::Query("reordered BMP backfill requires forward values; enable bmp_forward_index and explicitly reorder/rebuild, or disable L1 backfill".into()));
105                    }
106                }
107            }
108            FieldType::DenseVector | FieldType::BinaryDenseVector => {
109                let Some(flat) = self.flat_vectors.get(&field.0) else {
110                    if self.vector_indexes.contains_key(&field.0) {
111                        return Err(Error::Query(format!(
112                            "candidate scoring needs stored vectors for field {}",
113                            entry.name
114                        )));
115                    }
116                    return Ok(locations);
117                };
118                for &doc in documents {
119                    let (start, count) = flat.flat_indexes_for_doc_range(doc);
120                    for physical in start..start + count {
121                        let (actual_doc, ordinal) = flat.get_doc_id(physical);
122                        if actual_doc != doc {
123                            return Err(Error::Corruption(
124                                "flat vector lookup points to another document".into(),
125                            ));
126                        }
127                        push(
128                            doc,
129                            ordinal,
130                            u32::try_from(physical).map_err(|_| {
131                                Error::Query("candidate flat vector address exceeds u32".into())
132                            })?,
133                        )?;
134                    }
135                }
136            }
137            FieldType::Text => {
138                if self
139                    .meta
140                    .field_stats
141                    .get(&field.0)
142                    .is_none_or(|stats| stats.total_tokens == 0)
143                {
144                    return Ok(locations);
145                }
146                let lengths = self.doc_lengths(field).ok_or_else(|| {
147                    Error::Query(format!(
148                        "candidate scoring requires field-length metadata for plain text field {}",
149                        entry.name
150                    ))
151                })?;
152                for &doc in documents {
153                    if lengths.length(doc) == 0 {
154                        continue;
155                    }
156                    push(doc, 0, doc)?;
157                }
158            }
159            other => {
160                return Err(Error::Query(format!(
161                    "candidate scoring does not support field type {other:?}"
162                )));
163            }
164        }
165        Ok(locations)
166    }
167}
168
169impl SegmentReader {
170    /// Capability diagnostics for legacy/reordered fields. Missing values are
171    /// distinct from a field whose representation cannot support backfill.
172    pub fn unprepared_candidate_fields(&self) -> Vec<String> {
173        self.schema
174            .fields()
175            .filter_map(|(field, entry)| {
176                let prepared = if let Some(map) = self.chunk_map(field) {
177                    map.has_logical_addressing()
178                } else if self.seismic_index(field).is_some() {
179                    true
180                } else if let Some(bmp) = self.bmp_indexes.get(&field.0) {
181                    bmp.forward().is_some() || bmp.logically_ordered()
182                } else {
183                    !(self.vector_indexes.contains_key(&field.0)
184                        && !self.flat_vectors.contains_key(&field.0))
185                        && !(matches!(entry.field_type, FieldType::Text)
186                            && !entry.chunked
187                            && self
188                                .meta
189                                .field_stats
190                                .get(&field.0)
191                                .is_some_and(|stats| stats.total_tokens > 0)
192                            && self.doc_lengths(field).is_none())
193                };
194                (!prepared).then(|| entry.name.clone())
195            })
196            .collect()
197    }
198}
199
200impl SegmentReader {
201    /// Resolve only nominated logical passages. Lookup costs depend on the
202    /// candidate set, not on the number of chunks in a book.
203    pub(crate) async fn candidate_passage_locations(
204        &self,
205        field: Field,
206        targets: &[crate::segment::logical_address::LogicalUnit],
207        budget: &mut super::SparseProbeBudget,
208    ) -> Result<Vec<CandidateLocation>> {
209        use crate::segment::logical_address::{LogicalUnit, ordered_slot_for_unit};
210        if self.sparse_indexes.contains_key(&field.0) {
211            let mut documents: Vec<_> = targets.iter().map(|t| t.doc).collect();
212            documents.dedup();
213            return self
214                .maxscore_candidate_locations(
215                    field,
216                    &documents,
217                    Some(targets),
218                    targets.len(),
219                    budget,
220                )
221                .await;
222        }
223        let mut locations = Vec::with_capacity(targets.len());
224        for &target in targets {
225            let physical = if let Some(map) = self.chunk_map(field) {
226                if !map.has_logical_addressing() {
227                    return Err(Error::Query("legacy reordered text needs explicit Reorder to upgrade its chunk map for L1".into()));
228                }
229                map.slot_for_unit(target)
230            } else if let Some(index) = self.seismic_index(field) {
231                index
232                    .rows_for_document(target.doc)
233                    .find(|&row| index.key(row).ordinal == target.ordinal)
234            } else if let Some(bmp) = self.bmp_indexes.get(&field.0) {
235                if let Some(forward) = bmp.forward() {
236                    forward.find(target)
237                } else if bmp.logically_ordered() {
238                    ordered_slot_for_unit(bmp.num_virtual_docs, target, |physical| {
239                        let (doc, ordinal) = bmp.virtual_to_doc(physical);
240                        (doc != u32::MAX).then_some(LogicalUnit { doc, ordinal })
241                    })
242                } else {
243                    return Err(Error::Query(
244                        "reordered BMP backfill requires forward values; enable bmp_forward_index and explicitly reorder/rebuild, or disable L1 backfill"
245                            .into(),
246                    ));
247                }
248            } else if let Some(flat) = self.flat_vectors.get(&field.0) {
249                let (start, count) = flat.flat_indexes_for_doc_range(target.doc);
250                let mut low = start;
251                let mut high = start + count;
252                while low < high {
253                    let mid = low + (high - low) / 2;
254                    if flat.get_doc_id(mid).1 < target.ordinal {
255                        low = mid + 1;
256                    } else {
257                        high = mid;
258                    }
259                }
260                if low < start + count && flat.get_doc_id(low) == (target.doc, target.ordinal) {
261                    Some(
262                        u32::try_from(low)
263                            .map_err(|_| Error::Query("flat vector address exceeds u32".into()))?,
264                    )
265                } else {
266                    None
267                }
268            } else if self.vector_indexes.contains_key(&field.0) {
269                return Err(Error::Query(
270                    "candidate scoring needs stored flat vectors".into(),
271                ));
272            } else {
273                None
274            };
275            if let Some(physical) = physical {
276                let actual = if let Some(map) = self.chunk_map(field) {
277                    map.resolve(physical)
278                } else if let Some(index) = self.seismic_index(field) {
279                    let key = index.key(physical);
280                    (key.doc, key.ordinal)
281                } else if let Some(bmp) = self.bmp_indexes.get(&field.0) {
282                    if let Some(forward) = bmp.forward() {
283                        let key = forward.key(physical);
284                        (key.doc, key.ordinal)
285                    } else {
286                        bmp.virtual_to_doc(physical)
287                    }
288                } else {
289                    self.flat_vectors[&field.0].get_doc_id(physical as usize)
290                };
291                if actual != (target.doc, target.ordinal) {
292                    return Err(Error::Corruption(
293                        "candidate lookup points to another passage".into(),
294                    ));
295                }
296                locations.push(CandidateLocation {
297                    doc: target.doc,
298                    ordinal: target.ordinal,
299                    physical,
300                });
301            }
302        }
303        Ok(locations)
304    }
305}
306
307impl SegmentReader {
308    async fn maxscore_candidate_locations(
309        &self,
310        field: Field,
311        documents: &[DocId],
312        targets: Option<&[crate::segment::logical_address::LogicalUnit]>,
313        limit: usize,
314        budget: &mut super::SparseProbeBudget,
315    ) -> Result<Vec<CandidateLocation>> {
316        use crate::segment::logical_address::LogicalUnit;
317        let index = &self.sparse_indexes[&field.0];
318        let mut units = std::collections::BTreeSet::new();
319        index
320            .probe_candidates(documents, None, budget, |doc, ordinal, _| {
321                let key = LogicalUnit { doc, ordinal };
322                if targets.is_some_and(|targets| targets.binary_search(&key).is_err()) {
323                    return Ok(());
324                }
325                if !units.contains(&key) {
326                    if units.len() == limit {
327                        return Err(Error::Query(
328                            "L1 sparse field-value expansion budget exceeded".into(),
329                        ));
330                    }
331                    units.insert(key);
332                }
333                Ok(())
334            })
335            .await?;
336        Ok(units
337            .into_iter()
338            .enumerate()
339            .map(|(i, key)| CandidateLocation {
340                doc: key.doc,
341                ordinal: key.ordinal,
342                physical: i as u32,
343            })
344            .collect())
345    }
346}
347
348impl SegmentReader {
349    /// Admit the selected BMP payload using metadata alone. Forward records
350    /// are variable-sized; counting candidates does not bound their read work.
351    pub(crate) fn reserve_candidate_bmp_reads(
352        &self,
353        field: Field,
354        targets: &[u32],
355        remaining: &mut u64,
356    ) -> Result<()> {
357        let bmp = self
358            .bmp_index(field)
359            .ok_or_else(|| Error::Corruption("L1 BMP locations lack a BMP index".into()))?;
360        let mut previous_block = None;
361        for &target in targets {
362            let bytes = if let Some(forward) = bmp.forward() {
363                forward.vector_byte_len(target)?
364            } else {
365                if target >= bmp.num_virtual_docs {
366                    return Err(Error::Corruption("BMP candidate slot out of bounds".into()));
367                }
368                let block = target / bmp.bmp_block_size;
369                if previous_block == Some(block) {
370                    continue;
371                }
372                previous_block = Some(block);
373                let (start, end) = bmp.block_data_range(block);
374                end - start
375            };
376            *remaining = remaining.checked_sub(bytes).ok_or_else(|| {
377                Error::Query("L1 text/BMP payload read budget exceeded (256 MiB)".into())
378            })?;
379        }
380        Ok(())
381    }
382
383    /// Existing text readers return zero-copy views on mmap/RAM but materialize
384    /// ranges on lazy backends. Admit those ranges before invoking the reader.
385    pub(crate) async fn reserve_candidate_text_reads(
386        &self,
387        field: Field,
388        term: &[u8],
389        positions: bool,
390        remaining: &mut u64,
391    ) -> Result<()> {
392        let lazy_postings = !self.postings.file().is_sync();
393        let lazy_positions =
394            positions && self.postings.positions_file().is_some_and(|h| !h.is_sync());
395        if !lazy_postings && !lazy_positions {
396            return Ok(());
397        }
398        let mut key = Vec::with_capacity(4 + term.len());
399        key.extend_from_slice(&field.0.to_le_bytes());
400        key.extend_from_slice(term);
401        let Some(info) = self.term_dict.get(&key).await? else {
402            return Ok(());
403        };
404        let posting_bytes = if lazy_postings {
405            info.external_info().map_or(0, |(_, bytes)| bytes)
406        } else {
407            0
408        };
409        let position_bytes = if lazy_positions {
410            info.position_info().map_or(0, |(_, bytes)| bytes)
411        } else {
412            0
413        };
414        *remaining = posting_bytes
415            .checked_add(position_bytes)
416            .and_then(|bytes| remaining.checked_sub(bytes))
417            .ok_or_else(|| Error::Query("L1 lazy text read budget exceeded (256 MiB)".into()))?;
418        Ok(())
419    }
420}
421
422#[cfg(all(test, feature = "native"))]
423mod tests {
424    use super::*;
425    #[tokio::test]
426    async fn lazy_text_metadata_counts_and_backfill_admission_avoid_payload_reads() {
427        use crate::directories::{FileHandle, RamDirectory};
428        use crate::{Document, Index, IndexConfig, IndexWriter, Schema};
429        let mut schema = Schema::builder();
430        let field = schema.add_text_field_with_tokenizer("body", true, false, "simple");
431        let schema = std::sync::Arc::new(schema.build());
432        let dir = RamDirectory::new();
433        let config = IndexConfig::default();
434        let mut writer = IndexWriter::create(dir.clone(), (*schema).clone(), config.clone())
435            .await
436            .unwrap();
437        for _ in 0..256 {
438            let mut doc = Document::new();
439            doc.add_text(field, "common term");
440            writer.add_document(doc).unwrap();
441        }
442        writer.commit().await.unwrap();
443        let index = Index::open(dir.clone(), config).await.unwrap();
444        let searcher = index.reader().await.unwrap().searcher().await.unwrap();
445        let id = crate::segment::SegmentId(searcher.segment_readers()[0].meta().id);
446        let mut reader = SegmentReader::open(&dir, id, schema, 4).await.unwrap();
447        reader.postings = crate::structures::postings::PostingListReader::new(
448            FileHandle::lazy(
449                reader.postings.file().len(),
450                std::sync::Arc::new(|_| Box::pin(async { panic!("payload I/O before admission") })),
451            ),
452            reader.postings.positions_file().cloned(),
453        );
454        let error = reader
455            .reserve_candidate_text_reads(field, b"common", false, &mut 0)
456            .await
457            .unwrap_err();
458        assert!(error.to_string().contains("text read budget"));
459        reader
460            .reserve_candidate_text_reads(field, b"absent", false, &mut 0)
461            .await
462            .unwrap();
463        // Statistics precede scoring admission, so they must read only the
464        // dictionary even when a common term has an external posting list.
465        assert_eq!(reader.text_doc_freq(field, b"common").await.unwrap(), 256);
466        assert_eq!(reader.text_doc_freq(field, b"absent").await.unwrap(), 0);
467        use crate::query::{CountCollector, Query, TermQuery, collect_segment};
468        let required = crate::query::BooleanQuery::new()
469            .must(TermQuery::text(field, "common"))
470            .should(TermQuery::text(field, "term"))
471            .should(TermQuery::text(field, "absent"));
472        let mut count = CountCollector::new();
473        collect_segment(&reader, &required, &mut count)
474            .await
475            .unwrap();
476        assert_eq!(count.count(), 256);
477        for (term, expected) in [("common", 256), ("absent", 0)] {
478            let query = TermQuery::text(field, term);
479            let mut count = CountCollector::new();
480            collect_segment(&reader, &query, &mut count).await.unwrap();
481            assert_eq!(count.count(), expected);
482            assert_eq!(
483                u64::from(query.count_estimate(&reader).await.unwrap()),
484                expected
485            );
486        }
487    }
488}
summa_core/segment/reader/candidate_lookup.rs

summa_core/segment/reader/
candidate_lookup.rs
