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