Merge pull request #3030 from foundational-io/fix/streamer-term-ord-skipped-blocks

fix(sstable): report the real term ordinal when an automaton prunes blocks

sstable/src/block_reader.rs

@@ -8,8 +8,15 @@ use zstd::bulk::Decompressor;
 pub struct BlockReader {
     buffer: Vec<u8>,
     reader: OwnedBytes,
-    next_readers: std::vec::IntoIter<OwnedBytes>,
+    next_readers: std::vec::IntoIter<(OwnedBytes, u64)>,
     offset: usize,
+    /// First term ordinal of the slice we just moved to, taken once by the caller.
+    ///
+    /// When an automaton prunes blocks, the slices handed to us are not contiguous, so a
+    /// caller counting terms cannot know the ordinal it is at. Each slice therefore carries
+    /// the ordinal of its first term; consecutive blocks merged into one slice stay
+    /// contiguous, so counting within a slice remains correct.
+    pending_first_ordinal: Option<u64>,
 }
 
 impl BlockReader {
@@ -19,20 +26,34 @@ impl BlockReader {
             reader,
             next_readers: Vec::new().into_iter(),
             offset: 0,
+            pending_first_ordinal: None,
         }
     }
 
-    pub fn from_multiple_blocks(readers: Vec<OwnedBytes>) -> BlockReader {
+    /// Build a reader over non-contiguous slices, each labelled with the term ordinal of its
+    /// first term. See [`BlockReader::take_first_ordinal`].
+    pub fn from_multiple_blocks(readers: Vec<(OwnedBytes, u64)>) -> BlockReader {
         let mut next_readers = readers.into_iter();
-        let reader = next_readers.next().unwrap_or_else(OwnedBytes::empty);
+        let (reader, first_ordinal) = next_readers
+            .next()
+            .unwrap_or_else(|| (OwnedBytes::empty(), 0));
         BlockReader {
             buffer: Vec::new(),
             reader,
             next_readers,
             offset: 0,
+            pending_first_ordinal: Some(first_ordinal),
         }
     }
 
+    /// The first term ordinal of the slice most recently moved to, if it has not been taken yet.
+    ///
+    /// Returns `Some` exactly once per slice, on the first term read from it, so a caller can
+    /// reset its ordinal counter instead of incrementing across the gap left by pruned blocks.
+    pub fn take_first_ordinal(&mut self) -> Option<u64> {
+        self.pending_first_ordinal.take()
+    }
+
     pub fn deserialize_u64(&mut self) -> u64 {
         let (num_bytes, val) = super::vint::deserialize_read(self.buffer());
         self.advance(num_bytes);
@@ -53,8 +74,9 @@ impl BlockReader {
                 0 => {
                     // we are out of data for this block. Check if we have another block after
                     match self.next_readers.next() {
-                        Some(new_reader) => {
+                        Some((new_reader, first_ordinal)) => {
                             self.reader = new_reader;
+                            self.pending_first_ordinal = Some(first_ordinal);
                             continue;
                         }
                         _ => {

sstable/src/delta.rs

@@ -147,7 +147,9 @@ where TValueReader: value::ValueReader
         }
     }
 
-    pub fn from_multiple_blocks(reader: Vec<OwnedBytes>) -> Self {
+    /// Build a reader over slices that may not be contiguous, each labelled with the term
+    /// ordinal of its first term. See [`DeltaReader::take_first_ordinal`].
+    pub fn from_multiple_blocks(reader: Vec<(OwnedBytes, u64)>) -> Self {
         DeltaReader {
             idx: 0,
             common_prefix_len: 0,
@@ -157,6 +159,14 @@ where TValueReader: value::ValueReader
         }
     }
 
+    /// The first term ordinal of the slice just moved to, returned once per slice.
+    ///
+    /// A caller tracking term ordinals must consult this after every [`DeltaReader::advance`]:
+    /// when an automaton has pruned blocks the ordinal jumps, and only the slice knows where to.
+    pub fn take_first_ordinal(&mut self) -> Option<u64> {
+        self.block_reader.take_first_ordinal()
+    }
+
     pub fn empty() -> Self {
         DeltaReader::new(OwnedBytes::empty())
     }

sstable/src/dictionary.rs

@@ -116,8 +116,14 @@ impl<TSSTable: SSTable> Dictionary<TSSTable> {
             ));
             let data = blocks
                 .map(|block_addr| {
-                    self.sstable_slice
-                        .read_bytes_slice_async(block_addr.byte_range)
+                    let first_ordinal = block_addr.first_ordinal;
+                    async move {
+                        let bytes = self
+                            .sstable_slice
+                            .read_bytes_slice_async(block_addr.byte_range)
+                            .await?;
+                        io::Result::Ok((bytes, first_ordinal))
+                    }
                 })
                 .buffered(5)
                 .try_collect::<Vec<_>>()
@@ -142,7 +148,12 @@ impl<TSSTable: SSTable> Dictionary<TSSTable> {
             // merging across holes
             let blocks = self.get_block_iterator_for_range_and_automaton(key_range, automaton, 0);
             let data = blocks
-                .map(|block_addr| self.sstable_slice.read_bytes_slice(block_addr.byte_range))
+                .map(|block_addr| {
+                    let first_ordinal = block_addr.first_ordinal;
+                    self.sstable_slice
+                        .read_bytes_slice(block_addr.byte_range)
+                        .map(|bytes| (bytes, first_ordinal))
+                })
                 .collect::<Result<Vec<_>, _>>()?;
             Ok(DeltaReader::from_multiple_blocks(data))
         }
@@ -1125,4 +1136,38 @@ mod tests {
         assert_eq!(stream.key(), &[0, 255, 12]);
         assert!(!stream.advance());
     }
+
+    #[test]
+    fn test_search_term_ord_is_the_dictionary_ordinal_not_a_scan_counter() {
+        // An automaton prunes the blocks it cannot match, so the stream never reads the terms
+        // in them. `term_ord()` must still report the term's ordinal in the whole dictionary;
+        // counting advances alone reports its position among the blocks actually scanned, which
+        // silently mislabels terms for any caller that resolves ordinals back to terms.
+        let (dict, _) = make_test_sstable();
+
+        // Matches exactly one key, far enough in that everything before it is pruned.
+        let late_key = b"3FFFE";
+        let expected_ord: TermOrdinal = 0x3FFFE;
+        assert_eq!(dict.term_ord(late_key).unwrap(), Some(expected_ord));
+
+        let pattern = tantivy_fst::Regex::new("3FFFE").unwrap();
+        let mut stream = dict.search(pattern).into_stream().unwrap();
+        assert!(stream.advance());
+        assert_eq!(stream.key(), late_key);
+        assert_eq!(stream.term_ord(), expected_ord);
+        assert!(!stream.advance());
+
+        // The ordinal a match reports must round-trip back to that same match, for every match
+        // of a pattern whose hits are spread across several pruned-apart blocks.
+        let pattern = tantivy_fst::Regex::new("[0-3]FFFF").unwrap();
+        let mut stream = dict.search(pattern).into_stream().unwrap();
+        let mut seen = 0;
+        while stream.advance() {
+            let key = stream.key().to_vec();
+            let ord = stream.term_ord();
+            assert_eq!(dict.term_ord(&key).unwrap(), Some(ord), "key {key:?}");
+            seen += 1;
+        }
+        assert_eq!(seen, 3); // 0FFFF, 1FFFF, 2FFFF — the dictionary stops below 3FFFF
+    }
 }

sstable/src/streamer.rs

@@ -206,11 +206,16 @@ where
     /// is an uninitialized state.
     pub fn advance(&mut self) -> bool {
         while self.delta_reader.advance().unwrap() {
-            self.term_ord = Some(
-                self.term_ord
+            // An automaton prunes whole blocks, so the ordinal is not simply the previous one
+            // plus one: on entering a new slice it jumps to that slice's first term ordinal.
+            // Counting alone would report a term's position among the blocks actually scanned.
+            self.term_ord = Some(match self.delta_reader.take_first_ordinal() {
+                Some(first_ordinal) => first_ordinal,
+                None => self
+                    .term_ord
                     .map(|term_ord| term_ord + 1u64)
                     .unwrap_or(0u64),
-            );
+            });
             let common_prefix_len = self.delta_reader.common_prefix_len();
             self.states.truncate(common_prefix_len + 1);
             self.key.truncate(common_prefix_len);