summa_core/segment/reader/
term_expansion.rs

1//! Bounded dictionary expansion and canonical posting reads.
2use super::{MAX_PREFIX_POSTINGS, MAX_PREFIX_TERMS, SegmentReader, checked_file_range};
3use crate::dsl::Field;
4use crate::structures::BlockPostingList;
5use crate::{Error, Result};
6
7// Inline metadata stays allocation-free; external lists retain borrowed views.
8#[allow(clippy::large_enum_variant)]
9#[derive(Debug)]
10pub(crate) enum ExpandedPosting {
11    Inline(crate::structures::DecodedInlinePostings),
12    External(crate::structures::postings::DeferredPosting),
13}
14
15impl ExpandedPosting {
16    pub(crate) fn doc_count(&self) -> u32 {
17        match self {
18            Self::Inline(postings) => postings.docs().len() as u32,
19            Self::External(postings) => postings.doc_count(),
20        }
21    }
22
23    fn into_block(self) -> Result<BlockPostingList> {
24        match self {
25            Self::External(postings) => Ok(postings.into_list()),
26            Self::Inline(postings) => {
27                let mut list = crate::structures::PostingList::with_capacity(postings.docs().len());
28                for (&doc, &tf) in postings.docs().iter().zip(postings.frequencies()) {
29                    list.push(doc, tf);
30                }
31                Ok(BlockPostingList::from_posting_list(&list)?)
32            }
33        }
34    }
35}
36
37impl SegmentReader {
38    /// Read the bounded union inputs for a nonempty literal prefix.
39    pub async fn get_prefix_postings(
40        &self,
41        field: Field,
42        prefix: &[u8],
43    ) -> Result<Vec<BlockPostingList>> {
44        self.get_prefix_expansion(field, prefix)
45            .await?
46            .into_iter()
47            .map(ExpandedPosting::into_block)
48            .collect()
49    }
50
51    pub(crate) async fn get_prefix_expansion(
52        &self,
53        field: Field,
54        prefix: &[u8],
55    ) -> Result<Vec<ExpandedPosting>> {
56        if prefix.is_empty() {
57            return Err(Error::Query("prefix must not be empty".into()));
58        }
59        self.get_matching_postings(field, &[prefix.to_vec()], "prefix", usize::MAX, |_| true)
60            .await
61    }
62
63    pub(crate) async fn get_matching_postings(
64        &self,
65        field: Field,
66        prefixes: &[Vec<u8>],
67        label: &str,
68        max_scanned: usize,
69        mut accepts: impl FnMut(&[u8]) -> bool + Send,
70    ) -> Result<Vec<ExpandedPosting>> {
71        let mut entries = Vec::new();
72        let mut scanned = 0usize;
73        let mut truncated = false;
74        for prefix in prefixes {
75            let mut key_prefix = Vec::with_capacity(4 + prefix.len());
76            key_prefix.extend_from_slice(&field.0.to_le_bytes());
77            key_prefix.extend_from_slice(prefix);
78            let remaining = max_scanned.saturating_sub(scanned);
79            let (mut range, more) = self
80                .term_dict
81                .prefix_scan_values(
82                    &key_prefix,
83                    MAX_PREFIX_TERMS - entries.len(),
84                    remaining,
85                    |key| {
86                        scanned += 1;
87                        accepts(&key[4..])
88                    },
89                )
90                .await?;
91            entries.append(&mut range);
92            if more {
93                truncated = true;
94                break;
95            }
96        }
97        if truncated {
98            return Err(Error::Query(format!(
99                "{label} expands to more than {MAX_PREFIX_TERMS} terms"
100            )));
101        }
102        let posting_count: u64 = entries
103            .iter()
104            .map(|term_info| term_info.doc_freq() as u64)
105            .sum();
106        if posting_count > MAX_PREFIX_POSTINGS {
107            return Err(Error::Query(format!(
108                "{label} expands to {posting_count} postings (maximum {MAX_PREFIX_POSTINGS})"
109            )));
110        }
111        let mut results = Vec::with_capacity(entries.len());
112        let postings = self.postings.for_expansion();
113
114        for term_info in entries {
115            if term_info.is_inline() {
116                let inline = term_info
117                    .decode_inline_fixed()
118                    .ok_or_else(|| Error::Corruption("invalid expanded inline postings".into()))?;
119                results.push(ExpandedPosting::Inline(inline));
120            } else if let Some((posting_offset, posting_len)) = term_info.external_info() {
121                let range = checked_file_range(
122                    posting_offset,
123                    posting_len,
124                    self.postings.file().len(),
125                    "expanded term posting",
126                )?;
127                results.push(ExpandedPosting::External(
128                    postings.read_deferred(range).await?,
129                ));
130            }
131        }
132
133        Ok(results)
134    }
135
136    #[cfg(feature = "sync")]
137    /// Read the bounded union inputs for a nonempty literal prefix.
138    pub fn get_prefix_postings_sync(
139        &self,
140        field: Field,
141        prefix: &[u8],
142    ) -> Result<Vec<BlockPostingList>> {
143        self.get_prefix_expansion_sync(field, prefix)?
144            .into_iter()
145            .map(ExpandedPosting::into_block)
146            .collect()
147    }
148
149    #[cfg(feature = "sync")]
150    pub(crate) fn get_prefix_expansion_sync(
151        &self,
152        field: Field,
153        prefix: &[u8],
154    ) -> Result<Vec<ExpandedPosting>> {
155        if prefix.is_empty() {
156            return Err(Error::Query("prefix must not be empty".into()));
157        }
158        self.get_matching_postings_sync(field, &[prefix.to_vec()], "prefix", usize::MAX, |_| true)
159    }
160
161    #[cfg(feature = "sync")]
162    pub(crate) fn get_matching_postings_sync(
163        &self,
164        field: Field,
165        prefixes: &[Vec<u8>],
166        label: &str,
167        max_scanned: usize,
168        mut accepts: impl FnMut(&[u8]) -> bool + Send,
169    ) -> Result<Vec<ExpandedPosting>> {
170        let mut entries = Vec::new();
171        let mut scanned = 0usize;
172        let mut truncated = false;
173        for prefix in prefixes {
174            let mut key_prefix = Vec::with_capacity(4 + prefix.len());
175            key_prefix.extend_from_slice(&field.0.to_le_bytes());
176            key_prefix.extend_from_slice(prefix);
177            let remaining = max_scanned.saturating_sub(scanned);
178            let (mut range, more) = self.term_dict.prefix_scan_values_sync(
179                &key_prefix,
180                MAX_PREFIX_TERMS - entries.len(),
181                remaining,
182                |key| {
183                    scanned += 1;
184                    accepts(&key[4..])
185                },
186            )?;
187            entries.append(&mut range);
188            if more {
189                truncated = true;
190                break;
191            }
192        }
193        if truncated {
194            return Err(Error::Query(format!(
195                "{label} expands to more than {MAX_PREFIX_TERMS} terms"
196            )));
197        }
198        let posting_count: u64 = entries
199            .iter()
200            .map(|term_info| term_info.doc_freq() as u64)
201            .sum();
202        if posting_count > MAX_PREFIX_POSTINGS {
203            return Err(Error::Query(format!(
204                "{label} expands to {posting_count} postings (maximum {MAX_PREFIX_POSTINGS})"
205            )));
206        }
207        let mut results = Vec::with_capacity(entries.len());
208        let postings = self.postings.for_expansion();
209
210        for term_info in entries {
211            if term_info.is_inline() {
212                let inline = term_info
213                    .decode_inline_fixed()
214                    .ok_or_else(|| Error::Corruption("invalid expanded inline postings".into()))?;
215                results.push(ExpandedPosting::Inline(inline));
216            } else if let Some((posting_offset, posting_len)) = term_info.external_info() {
217                let range = checked_file_range(
218                    posting_offset,
219                    posting_len,
220                    self.postings.file().len(),
221                    "expanded term posting",
222                )?;
223                results.push(ExpandedPosting::External(
224                    postings.read_deferred_sync(range)?,
225                ));
226            }
227        }
228
229        Ok(results)
230    }
231}
232
233#[cfg(test)]
234mod tests {
235    use super::*;
236    use crate::segment::{SegmentBuilder, SegmentBuilderConfig, SegmentId};
237    use crate::{Document, RamDirectory, SchemaBuilder};
238    use std::sync::Arc;
239
240    #[tokio::test]
241    async fn malformed_inline_expansion_is_an_error_in_both_execution_modes() {
242        use crate::directories::{FileHandle, OwnedBytes};
243        use crate::structures::{AsyncSSTableReader, SSTableWriter, TermInfo};
244        let mut schema = SchemaBuilder::default();
245        let field = schema.add_text_field_with_tokenizer("tag", true, false, "raw");
246        let schema = Arc::new(schema.build());
247        let directory = RamDirectory::new();
248        let mut builder =
249            SegmentBuilder::new(schema.clone(), SegmentBuilderConfig::default()).unwrap();
250        let mut document = Document::new();
251        document.add_text(field, "alpha");
252        builder.add_document(document).unwrap();
253        let id = SegmentId::new();
254        builder.build(&directory, id, None).await.unwrap();
255        let mut reader = SegmentReader::open(&directory, id, schema, 16)
256            .await
257            .unwrap();
258        let mut dictionary = SSTableWriter::<_, TermInfo>::new(Vec::new());
259        let mut key = field.0.to_le_bytes().to_vec();
260        key.extend_from_slice(b"alpha");
261        // Valid metadata envelope, truncated posting payload: doc ID but no TF.
262        dictionary
263            .insert(
264                &key,
265                &TermInfo::Inline {
266                    doc_freq: 1,
267                    data: [0; 16],
268                    data_len: 1,
269                },
270            )
271            .unwrap();
272        reader.term_dict = Arc::new(
273            AsyncSSTableReader::open(
274                FileHandle::from_bytes(OwnedBytes::new(dictionary.finish().unwrap())),
275                8,
276            )
277            .await
278            .unwrap(),
279        );
280        assert!(
281            reader
282                .get_prefix_expansion(field, b"a")
283                .await
284                .unwrap_err()
285                .to_string()
286                .contains("invalid expanded inline postings")
287        );
288        #[cfg(feature = "sync")]
289        assert!(
290            reader
291                .get_prefix_expansion_sync(field, b"a")
292                .unwrap_err()
293                .to_string()
294                .contains("invalid expanded inline postings")
295        );
296    }
297
298    #[tokio::test]
299    async fn disjoint_dictionary_ranges_share_one_scan_and_expansion_budget() {
300        let mut schema = SchemaBuilder::default();
301        let field = schema.add_text_field_with_tokenizer("tag", true, false, "raw");
302        let schema = Arc::new(schema.build());
303        let directory = RamDirectory::new();
304        let mut builder =
305            SegmentBuilder::new(schema.clone(), SegmentBuilderConfig::default()).unwrap();
306        let mut document = Document::new();
307        for prefix in ["alpha", "beta"] {
308            for n in 0..600 {
309                document.add_text(field, format!("{prefix}{n:03}"));
310            }
311        }
312        builder.add_document(document).unwrap();
313        let id = SegmentId::new();
314        builder.build(&directory, id, None).await.unwrap();
315        let reader = SegmentReader::open(&directory, id, schema, 16)
316            .await
317            .unwrap();
318        let prefixes = vec![b"alpha".to_vec(), b"beta".to_vec()];
319        for (scan_budget, accepts, expected) in [
320            (1200, false, "ok"),
321            (1199, false, "scan exceeds"),
322            (1200, true, "more than 1024"),
323        ] {
324            let outcome = reader
325                .get_matching_postings(field, &prefixes, "regex", scan_budget, |_| accepts)
326                .await;
327            if expected == "ok" {
328                assert!(outcome.unwrap().is_empty());
329            } else {
330                assert!(outcome.unwrap_err().to_string().contains(expected));
331            }
332            #[cfg(feature = "sync")]
333            {
334                let outcome = reader.get_matching_postings_sync(
335                    field,
336                    &prefixes,
337                    "regex",
338                    scan_budget,
339                    |_| accepts,
340                );
341                if expected == "ok" {
342                    assert!(outcome.unwrap().is_empty());
343                } else {
344                    assert!(outcome.unwrap_err().to_string().contains(expected));
345                }
346            }
347        }
348    }
349}
summa_core/segment/reader/term_expansion.rs

summa_core/segment/reader/
term_expansion.rs
