Merge pull request #3126 from quickwit-oss/trinity.pointard/sstable-perf

improve Streamer::advance perf

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<usize>;
+
+    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<MonotonicU64SSTable>,
+    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);

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<W, TValueWriter>
 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<U64MonotonicValueReader> = 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
+        );
+    }
 }

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<TSSTable: SSTable> Dictionary<TSSTable> {
     ) -> io::Result<TermOrdHit> {
         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))

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())

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<Vec<u8>>) -> Bound<&[u8]> {
     }
 }
 
+#[inline(always)]
+fn matches_upper_bound(
+    comparator: &mut DeltaKeyComparator,
+    upper_bound: &Bound<Vec<u8>>,
+    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<TermOrdinal>,
     lower_bound: Bound<Vec<u8>>,
     upper_bound: Bound<Vec<u8>>,
+    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<usize>,
 }
 
 impl<TSSTable> 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>(),
+            (true, false) => self.advance_always_match::<false>(),
+            (false, true) => self.advance_with_automaton::<true>(),
+            (false, false) => self.advance_with_automaton::<false>(),
+        }
+    }
+
+    fn advance_always_match<const NO_BOUND: bool>(&mut self) -> bool {
+        if !self.advance_delta_reader() {
+            return false;
+        }
+        self.reconstruct_key_and_check_upper_bound::<NO_BOUND>()
+    }
+
+    fn advance_with_automaton<const NO_BOUND: bool>(&mut self) -> bool {
+        // fast path, check if prefix always match and we can skip Vec<state> 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::<NO_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::<NO_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<const NO_BOUND: bool>(&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.