diff --git a/sstable/src/block_reader.rs b/sstable/src/block_reader.rs index 0de2b44cd..29623cf7a 100644 --- a/sstable/src/block_reader.rs +++ b/sstable/src/block_reader.rs @@ -8,8 +8,15 @@ use zstd::bulk::Decompressor; pub struct BlockReader { buffer: Vec, reader: OwnedBytes, - next_readers: std::vec::IntoIter, + 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, } 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) -> 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 { + 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; } _ => { diff --git a/sstable/src/delta.rs b/sstable/src/delta.rs index 4d4317ea8..bca56317c 100644 --- a/sstable/src/delta.rs +++ b/sstable/src/delta.rs @@ -147,7 +147,9 @@ where TValueReader: value::ValueReader } } - pub fn from_multiple_blocks(reader: Vec) -> 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 { + self.block_reader.take_first_ordinal() + } + pub fn empty() -> Self { DeltaReader::new(OwnedBytes::empty()) } diff --git a/sstable/src/dictionary.rs b/sstable/src/dictionary.rs index 134e69274..5de411467 100644 --- a/sstable/src/dictionary.rs +++ b/sstable/src/dictionary.rs @@ -116,8 +116,14 @@ impl Dictionary { )); 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::>() @@ -142,7 +148,12 @@ impl Dictionary { // 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::, _>>()?; 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 + } } diff --git a/sstable/src/streamer.rs b/sstable/src/streamer.rs index 8477f5fd2..9203b3d0a 100644 --- a/sstable/src/streamer.rs +++ b/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);