use crate::config::{BloomConstructionPolicy, FilterAdvisor, FilterPolicy, FilterPolicyEntry};
use crate::{AbstractTree, AnyTree, Config, SeqNo, SequenceNumberCounter, Table};
const KEYS: u32 = 1_000;
fn key(i: u32) -> String {
format!("key{i:06}")
}
fn tree_of_even_keys(
folder: &std::path::Path,
bits_per_key: f32,
counted: bool,
) -> crate::Result<AnyTree> {
let tree = Config::new(
folder,
SequenceNumberCounter::default(),
SequenceNumberCounter::default(),
)
.filter_policy(FilterPolicy::all(FilterPolicyEntry::Bloom(
BloomConstructionPolicy::BitsPerKey(bits_per_key),
)))
.filter_advisor(counted.then(|| FilterAdvisor::new(u64::MAX)))
.open()?;
for i in (0..=KEYS).step_by(2) {
tree.insert(key(i), "value", u64::from(i));
}
tree.flush_active_memtable(0)?;
Ok(tree)
}
fn only_table(tree: &AnyTree) -> Table {
let version = tree.current_version();
let mut tables = version.iter_tables();
let Some(table) = tables.next().cloned() else {
panic!("the flush wrote a table");
};
assert!(tables.next().is_none(), "one table");
table
}
fn counts(table: &Table) -> (u64, u64) {
let Some(stats) = table.probe_stats() else {
panic!("the table counts its probes");
};
(stats.probes(), stats.negatives())
}
#[test]
fn absent_keys_count_as_negative_probes_and_present_ones_do_not() -> crate::Result<()> {
let folder = tempfile::tempdir()?;
let tree = tree_of_even_keys(folder.path(), 10.0, true)?;
for i in 0..KEYS {
let found = tree.get(key(i), SeqNo::MAX)?;
assert_eq!(found.is_some(), i % 2 == 0, "key {i}");
}
assert_eq!(
counts(&only_table(&tree)),
(u64::from(KEYS), u64::from(KEYS / 2))
);
Ok(())
}
#[test]
fn keys_outside_the_table_range_are_not_probes() -> crate::Result<()> {
let folder = tempfile::tempdir()?;
let tree = tree_of_even_keys(folder.path(), 10.0, true)?;
assert!(tree.get("a", SeqNo::MAX)?.is_none());
assert!(tree.get(key(KEYS + 1), SeqNo::MAX)?.is_none());
assert_eq!(counts(&only_table(&tree)), (0, 0));
Ok(())
}
#[test]
fn false_positives_count_as_negative_probes() -> crate::Result<()> {
let folder = tempfile::tempdir()?;
let tree = tree_of_even_keys(folder.path(), 1.0, true)?;
#[cfg(feature = "metrics")]
let skipped = tree.metrics().io_skipped_by_filter();
for i in (1..KEYS).step_by(2) {
assert!(tree.get(key(i), SeqNo::MAX)?.is_none());
}
#[cfg(feature = "metrics")]
{
let rejected = tree.metrics().io_skipped_by_filter() - skipped;
assert!(
rejected < usize::try_from(KEYS / 4).unwrap_or(usize::MAX),
"the filter must let most absent keys through, rejected {rejected}",
);
}
assert_eq!(
counts(&only_table(&tree)),
(u64::from(KEYS / 2), u64::from(KEYS / 2))
);
Ok(())
}
#[test]
fn invisible_and_deleted_keys_are_not_negative_probes() -> crate::Result<()> {
let folder = tempfile::tempdir()?;
let tree = Config::new(
folder.path(),
SequenceNumberCounter::default(),
SequenceNumberCounter::default(),
)
.filter_advisor(Some(FilterAdvisor::new(u64::MAX)))
.open()?;
tree.insert("early", "value", 1);
tree.insert("late", "value", 20);
tree.insert("deleted", "value", 2);
tree.remove("deleted", 3);
tree.flush_active_memtable(0)?;
let table = only_table(&tree);
assert!(tree.get("late", 10)?.is_none());
assert!(tree.get("deleted", 10)?.is_none());
assert_eq!(counts(&table), (2, 0));
let mut absent = 0;
for i in 0..100u32 {
let key = format!("e-absent{i}");
let (_, before) = counts(&table);
assert!(tree.get(&key, SeqNo::MAX)?.is_none());
absent += counts(&table).1 - before;
}
assert_eq!(absent, 100, "every absent key is a negative probe");
Ok(())
}
#[test]
fn multi_get_counts_each_key_once() -> crate::Result<()> {
let folder = tempfile::tempdir()?;
let tree = tree_of_even_keys(folder.path(), 10.0, true)?;
let keys: Vec<String> = (0..KEYS).map(key).collect();
let found = tree.multi_get(&keys, SeqNo::MAX)?;
assert_eq!(
found.iter().filter(|found| found.is_some()).count(),
usize::try_from(KEYS / 2).unwrap_or(usize::MAX)
);
assert_eq!(
counts(&only_table(&tree)),
(u64::from(KEYS), u64::from(KEYS / 2))
);
Ok(())
}
#[test]
fn a_level_every_key_is_filtered_out_of_is_answered_and_counted_once() -> crate::Result<()> {
let folder = tempfile::tempdir()?;
let any = tree_of_even_keys(folder.path(), 24.0, true)?;
let AnyTree::Standard(tree) = &any else {
panic!("a standard tree");
};
let keys: Vec<String> = (1..KEYS).step_by(2).map(key).collect();
let mut remaining: Vec<(usize, u64)> = keys
.iter()
.enumerate()
.map(|(index, key)| (index, crate::hash::hash64(key.as_bytes())))
.collect();
let mut results: Vec<Option<crate::value::InternalValue>> = alloc::vec![None; keys.len()];
let resolved = tree.resolve_level_staged(0, &mut remaining, keys.clone(), &mut results)?;
assert!(resolved, "the level is answered: nothing in it");
assert!(results.iter().all(Option::is_none));
let count = keys.len() as u64;
assert_eq!(
counts(&only_table(&any)),
(count, count),
"(probes, negatives)"
);
Ok(())
}
#[test]
fn nothing_is_counted_without_an_advisor() -> crate::Result<()> {
let folder = tempfile::tempdir()?;
let tree = tree_of_even_keys(folder.path(), 10.0, false)?;
for i in 0..KEYS {
tree.get(key(i), SeqNo::MAX)?;
}
assert!(only_table(&tree).probe_stats().is_none());
Ok(())
}
#[test]
fn a_chunked_resolve_counts_a_passed_key_with_no_block() -> crate::Result<()> {
let folder = tempfile::tempdir()?;
let any = Config::new(
folder.path(),
SequenceNumberCounter::default(),
SequenceNumberCounter::default(),
)
.filter_policy(FilterPolicy::all(FilterPolicyEntry::Bloom(
BloomConstructionPolicy::BitsPerKey(1.0),
)))
.filter_advisor(Some(FilterAdvisor::new(u64::MAX)))
.open()?;
for i in 0..KEYS {
any.insert(format!("a{i:06}"), "value", u64::from(i));
}
any.remove_range("b", "c", u64::from(KEYS));
any.flush_active_memtable(0)?;
for i in (0..KEYS).step_by(2) {
any.insert(format!("b{i:06}"), "value", u64::from(KEYS + 1 + i));
}
any.flush_active_memtable(0)?;
let AnyTree::Standard(tree) = &any else {
panic!("a standard tree");
};
let version = tree.current_version();
let widened = version
.iter_tables()
.find(|table| table.metadata.key_range.min().starts_with(b"a"))
.cloned()
.unwrap_or_else(|| panic!("the table the range tombstone widens"));
let keys: Vec<String> = (1..KEYS).step_by(2).map(|i| format!("b{i:06}")).collect();
let mut remaining: Vec<(usize, u64)> = keys
.iter()
.enumerate()
.map(|(index, key)| (index, crate::hash::hash64(key.as_bytes())))
.collect();
let mut results: Vec<Option<crate::value::InternalValue>> = alloc::vec![None; keys.len()];
let resolved = tree.resolve_level_staged(0, &mut remaining, keys.clone(), &mut results)?;
assert!(resolved, "the other table has blocks to read");
let count = keys.len() as u64;
assert_eq!(counts(&widened), (count, count), "(probes, negatives)");
Ok(())
}
#[test]
fn pinned_and_batch_reads_count_absent_keys_once() -> crate::Result<()> {
let folder = tempfile::tempdir()?;
let tree = tree_of_even_keys(folder.path(), 1.0, true)?;
let table = only_table(&tree);
let absent: Vec<String> = (1..KEYS).step_by(2).map(key).collect();
let count = absent.len() as u64;
for key in &absent {
assert!(tree.get_pinned(key, SeqNo::MAX)?.is_none());
}
assert_eq!(counts(&table), (count, count), "pinned reads");
let hashed = |keys: &[String]| -> Vec<(Vec<u8>, u64)> {
keys.iter()
.map(|key| (key.as_bytes().to_vec(), crate::hash::hash64(key.as_bytes())))
.collect()
};
let absent_hashed = hashed(&absent);
let sorted: Vec<(&[u8], u64)> = absent_hashed
.iter()
.map(|(key, hash)| (key.as_slice(), *hash))
.collect();
let found = table.batch_get(&sorted, SeqNo::MAX)?;
assert!(found.iter().all(Option::is_none));
assert_eq!(counts(&table), (2 * count, 2 * count), "batch read");
let beyond: Vec<String> = (0..100).map(|i| format!("zz{i:04}")).collect();
let beyond_hashed = hashed(&beyond);
let sorted: Vec<(&[u8], u64)> = beyond_hashed
.iter()
.map(|(key, hash)| (key.as_slice(), *hash))
.collect();
let found = table.batch_get(&sorted, SeqNo::MAX)?;
assert!(found.iter().all(Option::is_none));
assert_eq!(
counts(&table),
(2 * count + 100, 2 * count + 100),
"batch read past every block"
);
Ok(())
}
#[test]
fn a_chunked_resolve_counts_false_positives_once() -> crate::Result<()> {
let folder = tempfile::tempdir()?;
let any = tree_of_even_keys(folder.path(), 1.0, true)?;
let AnyTree::Standard(tree) = &any else {
panic!("a standard tree");
};
let keys: Vec<String> = (0..=KEYS).map(key).collect();
let mut remaining: Vec<(usize, u64)> = keys
.iter()
.enumerate()
.map(|(index, key)| (index, crate::hash::hash64(key.as_bytes())))
.collect();
let mut results: Vec<Option<crate::value::InternalValue>> = alloc::vec![None; keys.len()];
let resolved = tree.resolve_level_staged(0, &mut remaining, keys.clone(), &mut results)?;
assert!(resolved, "the level has blocks to read");
let present = results.iter().filter(|result| result.is_some()).count();
let absent = keys.len() - present;
assert_eq!(present, keys.len() / 2 + 1, "every even key is found");
assert_eq!(
counts(&only_table(&any)),
(keys.len() as u64, absent as u64),
"(probes, negatives)"
);
Ok(())
}
fn hot_narrow_and_cold_wide(folder: &std::path::Path) -> crate::Result<(AnyTree, Table, Table)> {
use crate::config::BlockSizePolicy;
use crate::table::probe_stats::ProbeCounts;
let tree = Config::new(
folder,
SequenceNumberCounter::default(),
SequenceNumberCounter::default(),
)
.data_block_size_policy(BlockSizePolicy::all(256))
.filter_advisor(Some(FilterAdvisor::new(u64::MAX)))
.open()?;
for i in 100..200 {
tree.insert(key(i), "value", 1);
}
tree.flush_active_memtable(0)?;
for i in 0..1_000 {
tree.insert(key(i), "value", 2);
}
tree.flush_active_memtable(0)?;
let version = tree.current_version();
let mut tables: Vec<Table> = version.iter_tables().cloned().collect();
tables.sort_by_key(|table| table.metadata.item_count);
let [hot, cold] = <[Table; 2]>::try_from(tables).unwrap_or_else(|_| panic!("two tables"));
let seed = |table: &Table, negatives| {
let Some(stats) = table.probe_stats() else {
panic!("the table counts its probes");
};
stats.add(ProbeCounts {
probes: negatives,
negatives,
});
};
seed(&hot, 10_000);
seed(&cold, 1_000);
Ok((tree, hot, cold))
}
#[test]
fn an_output_inherits_each_input_by_the_range_it_covers() -> crate::Result<()> {
use crate::table::probe_stats::inherited_counts;
let folder = tempfile::tempdir()?;
let (_tree, hot, cold) = hot_narrow_and_cold_wide(folder.path())?;
let inputs = [hot, cold];
let narrow = inherited_counts(key(100).as_bytes(), key(199).as_bytes(), &inputs)?;
assert_eq!(narrow.negatives, 10_000);
assert_eq!(narrow.probes, narrow.negatives);
let whole = inherited_counts(key(0).as_bytes(), key(999).as_bytes(), &inputs)?;
let cold_share = whole.negatives - 10_000;
assert!((870..=910).contains(&cold_share), "cold share {cold_share}");
let density = whole.negatives / 1_000;
assert!(density > 1 && density < 100, "density {density}");
let low = inherited_counts(key(0).as_bytes(), key(499).as_bytes(), &inputs)?;
let high = inherited_counts(key(500).as_bytes(), key(999).as_bytes(), &inputs)?;
assert!(low.negatives >= 10_000 + 350, "low {low:?}");
assert!(high.negatives <= 550, "high {high:?}");
let halves = low.negatives + high.negatives;
assert!(
halves <= whole.negatives + 4 && halves + 4 >= whole.negatives,
"halves {halves}, whole {}",
whole.negatives
);
Ok(())
}
#[test]
fn a_compaction_hands_the_inputs_counts_to_its_outputs() -> crate::Result<()> {
let folder = tempfile::tempdir()?;
let (tree, hot, cold) = hot_narrow_and_cold_wide(folder.path())?;
let inherited = crate::table::probe_stats::inherited_counts(
key(0).as_bytes(),
key(999).as_bytes(),
&[hot.clone(), cold.clone()],
)?
.negatives;
drop((hot, cold));
tree.major_compact(u64::MAX, 0)?;
assert_eq!(counts(&only_table(&tree)), (inherited, inherited));
assert!(
(10_870..=10_910).contains(&inherited),
"inherited {inherited}"
);
let folder = tempfile::tempdir()?;
let (tree, hot, cold) = hot_narrow_and_cold_wide(folder.path())?;
drop((hot, cold));
tree.major_compact(4 * 1_024, 0)?;
let version = tree.current_version();
let outputs: Vec<Table> = version.iter_tables().cloned().collect();
assert!(outputs.len() > 2, "{} outputs", outputs.len());
let total: u64 = outputs.iter().map(|table| counts(table).1).sum();
let rounding = 4 * outputs.len() as u64;
assert!(
(inherited - rounding..=inherited + rounding).contains(&total),
"total {total} over {} outputs",
outputs.len()
);
let hot_range: u64 = outputs
.iter()
.filter(|table| {
let range = &table.metadata.key_range;
range.min().as_ref() <= key(199).as_bytes()
&& range.max().as_ref() >= key(100).as_bytes()
})
.map(|table| counts(table).1)
.sum();
assert!(hot_range >= 10_000, "hot range {hot_range}");
Ok(())
}
#[test]
fn a_restricted_view_keeps_the_suffix_share_of_the_counts() -> crate::Result<()> {
use crate::table::probe_stats::share_of;
use core::ops::Bound::{Excluded, Unbounded};
let folder = tempfile::tempdir()?;
let (_tree, _hot, cold) = hot_narrow_and_cold_wide(folder.path())?;
let boundary = key(500);
let restricted = cold.reopen_restricted(boundary.as_bytes().to_vec().into())?;
let (_, suffix) = counts(&restricted);
let prefix = share_of(&cold, (Unbounded, Excluded(boundary.as_bytes())))?.negatives;
assert!((450..=550).contains(&suffix), "suffix {suffix}");
assert!(
(1_000 - 2..=1_000).contains(&(prefix + suffix)),
"prefix {prefix} + suffix {suffix}"
);
Ok(())
}
#[test]
fn a_restricted_view_keeps_the_trees_read_budget() -> crate::Result<()> {
use crate::config::ReadBudget;
let budget = ReadBudget::new(4_096, 2);
assert_ne!(budget, ReadBudget::default());
let folder = tempfile::tempdir()?;
let tree = Config::new(
folder.path(),
SequenceNumberCounter::default(),
SequenceNumberCounter::default(),
)
.columnar_read_budget(budget)
.open()?;
for i in 0..100 {
tree.insert(key(i), "value", 1);
}
tree.flush_active_memtable(0)?;
let table = only_table(&tree);
assert_eq!(table.read_budget(), budget);
let restricted = table.reopen_restricted(key(50).into_bytes().into())?;
assert_eq!(restricted.read_budget(), budget);
Ok(())
}
#[test]
fn a_filterless_input_leaves_the_counts_to_a_filtered_one() -> crate::Result<()> {
use crate::table::probe_stats::{ProbeCounts, inherited_counts};
let folder = tempfile::tempdir()?;
let any = Config::new(
folder.path(),
SequenceNumberCounter::default(),
SequenceNumberCounter::default(),
)
.filter_policy(FilterPolicy::new([
FilterPolicyEntry::Bloom(BloomConstructionPolicy::BitsPerKey(10.0)),
FilterPolicyEntry::None,
]))
.filter_advisor(Some(FilterAdvisor::new(u64::MAX)))
.open()?;
for i in 0..KEYS {
any.insert(key(i), "value", u64::from(i));
}
any.flush_active_memtable(0)?;
any.major_compact(u64::MAX, 0)?;
for i in 0..KEYS {
any.insert(key(i), "newer", u64::from(KEYS + i));
}
any.flush_active_memtable(0)?;
let version = any.current_version();
let (filtered, filterless): (Vec<Table>, Vec<Table>) = version
.iter_tables()
.cloned()
.partition(|table| table.filter_size() > 0);
let ([newer], [older]) = (filtered.as_slice(), filterless.as_slice()) else {
panic!("one table with a filter and one without");
};
let Some(stats) = newer.probe_stats() else {
panic!("the table counts its probes");
};
stats.add(ProbeCounts {
probes: 1_000,
negatives: 1_000,
});
let inherited = inherited_counts(
key(0).as_bytes(),
key(KEYS - 1).as_bytes(),
&[older.clone(), newer.clone()],
)?;
assert_eq!(inherited.negatives, 1_000);
Ok(())
}
#[test]
fn an_output_outside_an_input_inherits_nothing_from_it() -> crate::Result<()> {
use crate::table::probe_stats::inherited_counts;
let folder = tempfile::tempdir()?;
let (_tree, hot, cold) = hot_narrow_and_cold_wide(folder.path())?;
let beyond = inherited_counts(key(2_000).as_bytes(), key(3_000).as_bytes(), &[hot, cold])?;
assert_eq!(beyond, crate::table::probe_stats::ProbeCounts::default());
Ok(())
}
#[test]
fn recovered_tables_count_from_zero() -> crate::Result<()> {
let folder = tempfile::tempdir()?;
{
let tree = tree_of_even_keys(folder.path(), 10.0, true)?;
tree.get(key(1), SeqNo::MAX)?;
}
let tree = Config::new(
folder.path(),
SequenceNumberCounter::default(),
SequenceNumberCounter::default(),
)
.filter_advisor(Some(FilterAdvisor::new(u64::MAX)))
.open()?;
assert_eq!(counts(&only_table(&tree)), (0, 0));
tree.get(key(1), SeqNo::MAX)?;
assert_eq!(counts(&only_table(&tree)), (1, 1));
Ok(())
}