diff --git a/sstable/benches/stream_bench.rs b/sstable/benches/stream_bench.rs index 70dcdd8e3..f8235ceea 100644 --- a/sstable/benches/stream_bench.rs +++ b/sstable/benches/stream_bench.rs @@ -1,13 +1,68 @@ use std::collections::BTreeSet; +use std::hint::black_box; use std::io; use common::file_slice::FileSlice; use criterion::{Criterion, criterion_group, criterion_main}; use rand::rngs::StdRng; use rand::{Rng, SeedableRng}; +use tantivy_fst::Automaton; use tantivy_sstable::{Dictionary, MonotonicU64SSTable}; const CHARSET: &[u8] = b"abcdefghij"; +const AUTOMATON_PREFIX: &[u8] = b"ab"; +const NUM_AUTOMATON_MATCHES: usize = 1_017; + +// Matches `prefix.*`, but only implement can_match/will_always_match if configured to +// +// this allow comparing effects of optimisations depending on these functions +struct HintedPrefixAutomaton<'a> { + prefix: &'a [u8], + can_match_hint: bool, + always_match_hint: bool, +} + +impl<'a> HintedPrefixAutomaton<'a> { + fn new(prefix: &'a [u8], can_match_hint: bool, always_match_hint: bool) -> Self { + Self { + prefix, + can_match_hint, + always_match_hint, + } + } +} + +impl Automaton for HintedPrefixAutomaton<'_> { + type State = Option; + + fn start(&self) -> Self::State { + Some(0) + } + + fn is_match(&self, state: &Self::State) -> bool { + *state == Some(self.prefix.len()) + } + + fn can_match(&self, state: &Self::State) -> bool { + !self.can_match_hint || state.is_some() + } + + fn will_always_match(&self, state: &Self::State) -> bool { + self.always_match_hint && self.is_match(state) + } + + fn accept(&self, state: &Self::State, byte: u8) -> Self::State { + let Some(pos) = *state else { return None }; + if pos == self.prefix.len() { + return Some(pos); + } + if self.prefix[pos] == byte { + Some(pos + 1) + } else { + None + } + } +} fn generate_key(rng: &mut impl Rng) -> String { let len = rng.random_range(3..12); @@ -56,6 +111,26 @@ fn stream_bench( count } +fn automaton_bench( + dictionary: &Dictionary, + can_match_hint: bool, + always_match_hint: bool, +) -> usize { + let mut stream = dictionary + .search(HintedPrefixAutomaton::new( + AUTOMATON_PREFIX, + black_box(can_match_hint), + black_box(always_match_hint), + )) + .into_stream() + .unwrap(); + let mut count = 0; + while stream.advance() { + count += 1; + } + count +} + pub fn criterion_benchmark(c: &mut Criterion) { let dict = prepare_sstable().unwrap(); c.bench_function("short_scan_init", |b| { @@ -63,7 +138,7 @@ pub fn criterion_benchmark(c: &mut Criterion) { }); c.bench_function("short_scan_init_and_scan", |b| { b.iter(|| { - assert_eq!(stream_bench(&dict, b"fa", b"faz", true), 971); + assert_eq!(stream_bench(&dict, b"fa", b"faz", true), 1051); }) }); c.bench_function("full_scan_init_and_scan_full_with_bound", |b| { @@ -81,6 +156,18 @@ pub fn criterion_benchmark(c: &mut Criterion) { count }) }); + c.bench_function("full_scan_prefix_automaton_no_hints", |b| { + b.iter(|| assert_eq!(automaton_bench(&dict, false, false), NUM_AUTOMATON_MATCHES)) + }); + c.bench_function("full_scan_prefix_automaton_can_match_hint_only", |b| { + b.iter(|| assert_eq!(automaton_bench(&dict, true, false), NUM_AUTOMATON_MATCHES)) + }); + c.bench_function("full_scan_prefix_automaton_always_match_hint_only", |b| { + b.iter(|| assert_eq!(automaton_bench(&dict, false, true), NUM_AUTOMATON_MATCHES)) + }); + c.bench_function("full_scan_prefix_automaton_both_hints", |b| { + b.iter(|| assert_eq!(automaton_bench(&dict, true, true), NUM_AUTOMATON_MATCHES)) + }); } criterion_group!(benches, criterion_benchmark); diff --git a/sstable/src/delta.rs b/sstable/src/delta.rs index 97d868e4e..986700fef 100644 --- a/sstable/src/delta.rs +++ b/sstable/src/delta.rs @@ -1,3 +1,4 @@ +use std::cmp::Ordering; use std::io::{self, BufWriter, Write}; use std::ops::Range; @@ -12,6 +13,62 @@ const FOUR_BIT_LIMITS: usize = 1 << 4; const VINT_MODE: u8 = 1u8; const BLOCK_LEN: usize = 4_000; +/// Incrementally compares delta-encoded keys with a fixed target key. +pub(crate) struct DeltaKeyComparator { + num_matching_bytes: usize, +} + +impl DeltaKeyComparator { + pub(crate) fn new() -> Self { + DeltaKeyComparator { + num_matching_bytes: 0, + } + } + + #[inline(always)] + pub(crate) fn compare( + &mut self, + target: &[u8], + common_prefix_len: usize, + suffix: &[u8], + ) -> Ordering { + match common_prefix_len.cmp(&self.num_matching_bytes) { + // popped bytes already matched => too far + Ordering::Less => return Ordering::Greater, + Ordering::Equal => (), + // the ok prefix is less than current entry prefix => continue to next element + Ordering::Greater => return Ordering::Less, + } + + for (key_byte, target_byte) in suffix.iter().zip(&target[self.num_matching_bytes..]) { + match key_byte.cmp(target_byte) { + Ordering::Equal => self.num_matching_bytes += 1, + ordering => return ordering, + } + } + + (common_prefix_len + suffix.len()).cmp(&target.len()) + } + + #[inline(always)] + pub(crate) fn compare_across_blocks( + &mut self, + target: &[u8], + common_prefix_len: usize, + suffix: &[u8], + ) -> Ordering { + // blocks are independent. On each new block we get a common_prefix_len=0 entry. + // reset our state with it + if common_prefix_len == 0 { + let num_matching_bytes = crate::common_prefix_len(target, suffix); + self.num_matching_bytes = num_matching_bytes; + // cannot panicm at worth we might compare empty slices if num_matching_bytes==len() + return suffix[num_matching_bytes..].cmp(&target[num_matching_bytes..]); + } + self.compare(target, common_prefix_len, suffix) + } +} + pub struct DeltaWriter where W: io::Write { @@ -241,7 +298,9 @@ where TValueReader: value::ValueReader #[cfg(test)] mod tests { - use super::DeltaReader; + use std::cmp::Ordering; + + use super::{DeltaKeyComparator, DeltaReader}; use crate::value::U64MonotonicValueReader; #[test] @@ -249,4 +308,24 @@ mod tests { let mut delta_reader: DeltaReader = DeltaReader::empty(); assert!(!delta_reader.advance().unwrap()); } + + #[test] + fn test_delta_key_comparator_across_block_reset() { + let mut comparator = DeltaKeyComparator::new(); + let target = b"bba"; + + assert_eq!( + comparator.compare_across_blocks(target, 0, b"baaaaa"), + Ordering::Less + ); + assert_eq!( + comparator.compare_across_blocks(target, 2, b"baaa"), + Ordering::Less + ); + // A zero-length common prefix marks a block reset, so the suffix is a complete key. + assert_eq!( + comparator.compare_across_blocks(target, 0, b"bbbaaa"), + Ordering::Greater + ); + } } diff --git a/sstable/src/dictionary.rs b/sstable/src/dictionary.rs index 5de411467..69b57053f 100644 --- a/sstable/src/dictionary.rs +++ b/sstable/src/dictionary.rs @@ -14,6 +14,7 @@ use itertools::Itertools; use tantivy_fst::Automaton; use tantivy_fst::automaton::AlwaysMatch; +use crate::delta::DeltaKeyComparator; use crate::streamer::{Streamer, StreamerBuilder}; use crate::{BlockAddr, DeltaReader, Reader, SSTable, SSTableIndex, TermOrdinal, VoidSSTable}; @@ -356,41 +357,19 @@ impl Dictionary { ) -> io::Result { let mut term_ord = 0; let key_bytes = key.as_ref(); - let mut ok_bytes = 0; + let mut key_comparator = DeltaKeyComparator::new(); while sstable_delta_reader.advance()? { - let prefix_len = sstable_delta_reader.common_prefix_len(); - let suffix = sstable_delta_reader.suffix(); - - match prefix_len.cmp(&ok_bytes) { - Ordering::Less => return Ok(TermOrdHit::Next(term_ord)), /* popped bytes already matched => too far */ - Ordering::Equal => (), - Ordering::Greater => { - // the ok prefix is less than current entry prefix => continue to next elem + match key_comparator.compare( + key_bytes, + sstable_delta_reader.common_prefix_len(), + sstable_delta_reader.suffix(), + ) { + Ordering::Less => { term_ord += 1; - continue; } + Ordering::Equal => return Ok(TermOrdHit::Exact(term_ord)), + Ordering::Greater => return Ok(TermOrdHit::Next(term_ord)), } - - // we have ok_bytes byte of common prefix, check if this key adds more - for (key_byte, suffix_byte) in key_bytes[ok_bytes..].iter().zip(suffix) { - match suffix_byte.cmp(key_byte) { - Ordering::Less => break, // byte too small - Ordering::Equal => ok_bytes += 1, // new matching - // byte - Ordering::Greater => return Ok(TermOrdHit::Next(term_ord)), // too far - } - } - - if ok_bytes == key_bytes.len() { - if prefix_len + suffix.len() == ok_bytes { - return Ok(TermOrdHit::Exact(term_ord)); - } else { - // current key is a prefix of current element, not a match - return Ok(TermOrdHit::Next(term_ord)); - } - } - - term_ord += 1; } Ok(TermOrdHit::Next(term_ord)) diff --git a/sstable/src/lib.rs b/sstable/src/lib.rs index 1f6bd14c7..d925ad540 100644 --- a/sstable/src/lib.rs +++ b/sstable/src/lib.rs @@ -70,7 +70,7 @@ const SSTABLE_VERSION: u32 = 3; /// Given two byte string returns the length of /// the longest common prefix. -fn common_prefix_len(left: &[u8], right: &[u8]) -> usize { +pub(crate) fn common_prefix_len(left: &[u8], right: &[u8]) -> usize { left.iter() .cloned() .zip(right.iter().cloned()) diff --git a/sstable/src/streamer.rs b/sstable/src/streamer.rs index 9203b3d0a..9e5d11541 100644 --- a/sstable/src/streamer.rs +++ b/sstable/src/streamer.rs @@ -1,9 +1,11 @@ +use std::cmp::Ordering; use std::io; use std::ops::Bound; use tantivy_fst::Automaton; use tantivy_fst::automaton::AlwaysMatch; +use crate::delta::DeltaKeyComparator; use crate::dictionary::Dictionary; use crate::{DeltaReader, SSTable, TermOrdinal}; @@ -30,6 +32,22 @@ fn bound_as_byte_slice(bound: &Bound>) -> Bound<&[u8]> { } } +#[inline(always)] +fn matches_upper_bound( + comparator: &mut DeltaKeyComparator, + upper_bound: &Bound>, + common_prefix_len: usize, + suffix: &[u8], +) -> bool { + let (upper_bound_key, inclusive) = match upper_bound { + Bound::Unbounded => return true, + Bound::Included(upper_bound_key) => (upper_bound_key, true), + Bound::Excluded(upper_bound_key) => (upper_bound_key, false), + }; + let ordering = comparator.compare_across_blocks(upper_bound_key, common_prefix_len, suffix); + ordering == Ordering::Less || inclusive && ordering == Ordering::Equal +} + impl<'a, TSSTable, A> StreamerBuilder<'a, TSSTable, A> where A: Automaton, @@ -124,14 +142,23 @@ where Bound::Unbounded => 0, }; + let always_match_at = if self.automaton.will_always_match(&start_state) { + Some(0) + } else { + None + }; + Ok(Streamer { automaton: self.automaton, states: vec![start_state], + always_match_at, delta_reader, key: Vec::new(), term_ord: first_term.checked_sub(1), + lower_bound_reached: self.lower == Bound::Unbounded, lower_bound: self.lower, upper_bound: self.upper, + upper_bound_comparator: DeltaKeyComparator::new(), _lifetime: std::marker::PhantomData, }) } @@ -174,8 +201,11 @@ where term_ord: Option, lower_bound: Bound>, upper_bound: Bound>, + upper_bound_comparator: DeltaKeyComparator, // this field is used to please the type-interface of a dictionary in tantivy _lifetime: std::marker::PhantomData<&'a ()>, + lower_bound_reached: bool, + always_match_at: Option, } impl Streamer<'_, TSSTable, AlwaysMatch> @@ -184,12 +214,15 @@ where TSSTable: SSTable pub fn empty() -> Self { Streamer { automaton: AlwaysMatch, - states: Vec::new(), + states: vec![AlwaysMatch.start()], + always_match_at: Some(0), delta_reader: DeltaReader::empty(), key: Vec::new(), term_ord: None, + lower_bound_reached: true, lower_bound: Bound::Unbounded, upper_bound: Bound::Unbounded, + upper_bound_comparator: DeltaKeyComparator::new(), _lifetime: std::marker::PhantomData, } } @@ -201,53 +234,173 @@ where A::State: Clone, TSSTable: SSTable, { + #[inline(always)] + fn advance_delta_reader(&mut self) -> bool { + if !self.delta_reader.advance().unwrap() { + return false; + } + // 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), + }); + true + } + + /// Make progress up to the lower bound + /// + /// Returns whether the reader was positioned on a key matching the lower bound. + /// If false, the delta_reader has been exhausted without finding such a key. + fn initialize(&mut self) -> bool { + debug_assert!(!self.lower_bound_reached); + let mut lower_bound_comparator = DeltaKeyComparator::new(); + while self.advance_delta_reader() { + let common_prefix_len = self.delta_reader.common_prefix_len(); + let suffix = self.delta_reader.suffix(); + let (lower_bound_key, inclusive) = match &self.lower_bound { + Bound::Unbounded => unreachable!("unbounded streamers do not need initialization"), + Bound::Included(lower_bound_key) => (lower_bound_key, true), + Bound::Excluded(lower_bound_key) => (lower_bound_key, false), + }; + let ordering = lower_bound_comparator.compare_across_blocks( + lower_bound_key, + common_prefix_len, + suffix, + ); + let match_lower_bound = + ordering == Ordering::Greater || inclusive && ordering == Ordering::Equal; + if match_lower_bound { + self.key.clear(); + self.key + .extend_from_slice(&lower_bound_key[..common_prefix_len]); + self.key.extend_from_slice(suffix); + let mut state: A::State = self.states.last().unwrap().clone(); + for b in &self.key { + state = self.automaton.accept(&state, *b); + self.states.push(state.clone()); + } + self.lower_bound_reached = true; + return true; + } + } + self.lower_bound_reached = true; + false + } + /// Advance position the stream on the next item. /// Before the first call to `.advance()`, the stream /// is an uninitialized state. pub fn advance(&mut self) -> bool { - while self.delta_reader.advance().unwrap() { - // 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), - }); + if !self.lower_bound_reached { + if !self.initialize() { + // no key higher than lower-bound at all + return false; + } + if !matches_upper_bound( + &mut self.upper_bound_comparator, + &self.upper_bound, + 0, + &self.key, + ) { + return false; + } + if self.automaton.is_match(self.states.last().unwrap()) { + return true; + } + } + + match ( + // we could check always_match_at == Some(0), but this actually gets + // inlined into `true` with AlwaysMatch, which is even faster + self.automaton + .will_always_match(&self.states.first().unwrap()), + self.upper_bound == Bound::Unbounded, + ) { + (true, true) => self.advance_always_match::(), + (true, false) => self.advance_always_match::(), + (false, true) => self.advance_with_automaton::(), + (false, false) => self.advance_with_automaton::(), + } + } + + fn advance_always_match(&mut self) -> bool { + if !self.advance_delta_reader() { + return false; + } + self.reconstruct_key_and_check_upper_bound::() + } + + fn advance_with_automaton(&mut self) -> bool { + // fast path, check if prefix always match and we can skip Vec management + if let Some(always_match_at) = self.always_match_at.take() { + if !self.advance_delta_reader() { + return false; + } + let common_prefix_len = self.delta_reader.common_prefix_len(); + if always_match_at <= common_prefix_len { + self.always_match_at = Some(always_match_at); + return self.reconstruct_key_and_check_upper_bound::(); + } + } else if !self.advance_delta_reader() { + return false; + } + + loop { let common_prefix_len = self.delta_reader.common_prefix_len(); self.states.truncate(common_prefix_len + 1); - self.key.truncate(common_prefix_len); + // TODO we could detect when we reach a !can_match, and skip both state and key + // computation until we truncate that can_t_match out of our state. it's already + // done at the block layer, so not as important let mut state: A::State = self.states.last().unwrap().clone(); for &b in self.delta_reader.suffix() { state = self.automaton.accept(&state, b); self.states.push(state.clone()); } - self.key.extend_from_slice(self.delta_reader.suffix()); - let match_lower_bound = match &self.lower_bound { - Bound::Unbounded => true, - Bound::Included(lower_bound_key) => lower_bound_key[..] <= self.key[..], - Bound::Excluded(lower_bound_key) => lower_bound_key[..] < self.key[..], - }; - if !match_lower_bound { - continue; + let matches = self.automaton.is_match(&state); + if matches { + self.always_match_at = self + .states + .iter() + .enumerate() + .rev() + .take_while(|(_i, state)| self.automaton.will_always_match(state)) + .last() + .map(|(i, _state)| i); } - // We match the lower key once. All subsequent keys will pass that bar. - self.lower_bound = Bound::Unbounded; - let match_upper_bound = match &self.upper_bound { - Bound::Unbounded => true, - Bound::Included(upper_bound_key) => upper_bound_key[..] >= self.key[..], - Bound::Excluded(upper_bound_key) => upper_bound_key[..] > self.key[..], - }; - if !match_upper_bound { + + if !self.reconstruct_key_and_check_upper_bound::() { return false; } - if self.automaton.is_match(&state) { + if matches { return true; } + if !self.advance_delta_reader() { + return false; + } } - false + } + + #[inline(always)] + fn reconstruct_key_and_check_upper_bound(&mut self) -> bool { + let common_prefix_len = self.delta_reader.common_prefix_len(); + self.key.truncate(common_prefix_len); + self.key.extend_from_slice(self.delta_reader.suffix()); + + // TODO there is an idea where we only look at the upper bound when our delta_reader + // reached the last block (if we pruned blocks beforehand (do we always?) we cannot + // find that key before that block) + NO_BOUND + || matches_upper_bound( + &mut self.upper_bound_comparator, + &self.upper_bound, + common_prefix_len, + self.delta_reader.suffix(), + ) } /// Returns the `TermOrdinal` of the given term.