use super::*;
use crate::{
diskann_index::format::{DiskANNCanonicalOrigin, DiskANNVectorVersion},
key_value::KeyValueDiskANNPopulationRecords,
mvcc::{
overlay::run::{read_counts, SpilledRunWriter},
PrivateRecordRevision, StorageTransactionId,
},
};
fn field(name: &str) -> Vec<u8> {
let mut key = vec![b'v'];
for value in ["table", name] {
key.extend_from_slice(&(value.len() as u32).to_be_bytes());
key.extend_from_slice(value.as_bytes());
}
key
}
#[test]
fn population_visits_do_not_reopen_spill_after_a_terminal_page() {
use crate::mvcc::MemoryVersionStore;
use std::sync::Arc;
for count in [1_usize, 63, 64, 65, 127, 128, 129] {
let control = StorageReadControl::with_limit(128 << 10);
let revision = PrivateRecordRevision::for_tests();
let mut writer =
SpilledRunWriter::new(count as u64 + 1, count as u64 * 8, control.memory()).unwrap();
for id in 0..count {
writer
.push(
format!("h{id:06}").as_bytes(),
None,
RecordWriteKind::Canonical,
revision,
(!id.is_multiple_of(7)).then_some(b"header".as_slice()),
&control,
)
.unwrap();
}
writer
.push(
b"z-unrelated",
None,
RecordWriteKind::Canonical,
revision,
Some(b"unrelated"),
&control,
)
.unwrap();
let changes = PrivateRecordChanges::from_spilled_run(
writer.finish().unwrap(),
revision,
None,
&control,
)
.unwrap();
let store = MemoryVersionStore::new(control.memory());
let view = MergedRecordSnapshot::new(
Arc::new(store.snapshot().unwrap()),
changes.snapshot().unwrap(),
);
read_counts::take();
let mut found = Vec::new();
visit(&view, b"h", &control, &mut |key, value| {
assert_eq!(value, b"header");
found.push(key.to_vec());
Ok(())
})
.unwrap();
assert_eq!(
found,
(0..count)
.filter(|id| !id.is_multiple_of(7))
.map(|id| format!("h{id:06}").into_bytes())
.collect::<Vec<_>>()
);
let reads = read_counts::take();
assert_eq!(reads.blocks, count / 64 + 1, "count={count}: {reads:?}");
assert_eq!(reads.values, found.len());
control.cancellation().cancel();
assert!(visit(&view, b"h", &control, &mut |_, _| Ok(())).is_err());
drop((view, changes, store));
assert_eq!(control.memory().used(), 0);
}
}
#[test]
fn population_deletion_preview_does_not_admit_an_unused_origin_lookup() {
use crate::mvcc::MemoryVersionStore;
use std::sync::Arc;
let control = StorageReadControl::with_limit(128 << 10);
let store = MemoryVersionStore::new(control.memory());
let private = PrivateRecordChanges::new(control.memory());
let view = MergedRecordSnapshot::new(
Arc::new(store.snapshot().unwrap()),
private.snapshot().unwrap(),
);
let layout = KeyValueDiskANNPopulationRecords;
let prefix = layout.origin_prefix(&field("deleted"), &control).unwrap();
let mut writes = BudgetedVec::new(control.memory());
for document in 0..512_u64 {
let mut key = prefix.to_vec();
key.extend_from_slice(&document.to_be_bytes());
writes
.push(
PreparedRecordWrite::copy_bytes(&key, None, None, &control)
.unwrap()
.with_kind(RecordWriteKind::DiskANNOrigin),
)
.unwrap();
}
let prepared = PreparedRecordCommit::from_unique_owned(writes, &control).unwrap();
let _pressure = control
.memory()
.reserve(control.memory().limit() - control.memory().used() - 4096)
.unwrap();
assert!(PreparedLookup::new(&prepared, &control).is_err());
let generated = Reconciliation {
before: &view,
after: &view,
layout: &layout,
history: DatabaseId::from_bytes([9; 16]),
control: &control,
structural: None,
}
.run(&prepared, &[])
.unwrap();
assert!(generated.is_empty());
control.cancellation().cancel();
assert!(Reconciliation {
before: &view,
after: &view,
layout: &layout,
history: DatabaseId::from_bytes([9; 16]),
control: &control,
structural: None,
}
.run(&prepared, &[])
.is_err());
}
#[test]
fn spilled_population_field_reads_seek_and_load_only_the_selected_origins() {
const FIELDS: usize = 32;
const DOCUMENTS: usize = 64;
let control = StorageReadControl::with_limit(128 << 10);
let layout = KeyValueDiskANNPopulationRecords;
let identity = StorageTransactionId::new(DatabaseId::from_bytes([7; 16]), 1).unwrap();
let origin =
DiskANNCanonicalOrigin::new(DiskANNVectorVersion::new(identity, 1).unwrap(), 128, 1)
.unwrap();
let fields: Vec<_> = (0..FIELDS)
.map(|index| field(&format!("field{index:02}")))
.collect();
let mut writer = SpilledRunWriter::new(
(FIELDS * DOCUMENTS) as u64,
(FIELDS * DOCUMENTS * 64) as u64,
control.memory(),
)
.unwrap();
for field in &fields {
let prefix = layout.origin_prefix(field, &control).unwrap();
for document in 0..=DOCUMENTS as u64 {
let mut key = prefix.to_vec();
key.extend_from_slice(&document.to_be_bytes());
writer
.push(
&key,
None,
if document == DOCUMENTS as u64 {
RecordWriteKind::Canonical
} else {
RecordWriteKind::DiskANNOrigin
},
PrivateRecordRevision::for_tests(),
Some(&origin.encode()),
&control,
)
.unwrap();
}
}
let prepared =
PreparedRecordCommit::from_spilled_run(writer.finish().unwrap().unwrap(), &control)
.unwrap();
assert!(prepared.resident().is_none());
let index = PreparedLookup::new(&prepared, &control).unwrap();
read_counts::take();
for field in &fields {
for _ in 0..3 {
let mut documents = Vec::new();
visit_field_origins(&layout, &index, field, &control, &mut |change| {
assert_eq!(&*change.field, field);
assert_eq!(change.origin, origin);
documents.push(change.document);
Ok(())
})
.unwrap();
assert_eq!(documents, (0..DOCUMENTS as u64).collect::<Vec<_>>());
}
}
let reads = read_counts::take();
assert_eq!(reads.values, 3 * FIELDS * DOCUMENTS, "{reads:?}");
assert!(reads.entries < 3 * FIELDS * (DOCUMENTS + 256), "{reads:?}");
control.cancellation().cancel();
assert!(visit_field_origins(&layout, &index, &fields[0], &control, &mut |_| Ok(())).is_err());
drop(index);
drop(prepared);
assert_eq!(control.memory().used(), 0);
}
#[test]
fn resident_population_ranges_distinguish_fields_with_shared_name_prefixes() {
let control = StorageReadControl::with_limit(128 << 10);
let layout = KeyValueDiskANNPopulationRecords;
let identity = StorageTransactionId::new(DatabaseId::from_bytes([8; 16]), 1).unwrap();
let origin =
DiskANNCanonicalOrigin::new(DiskANNVectorVersion::new(identity, 1).unwrap(), 3, 1).unwrap();
let fields = ["a", "aa", "a\0"].map(field);
let mut writes = BudgetedVec::new(control.memory());
for field in fields.iter().rev() {
let mut key = layout.origin_prefix(field, &control).unwrap();
key.extend_from_slice(&7_u64.to_be_bytes()).unwrap();
writes
.push(
PreparedRecordWrite::copy_bytes(&key, None, Some(&origin.encode()), &control)
.unwrap()
.with_kind(RecordWriteKind::DiskANNOrigin),
)
.unwrap();
}
let prepared = PreparedRecordCommit::from_unique_owned(writes, &control).unwrap();
let index = PreparedLookup::new(&prepared, &control).unwrap();
for field in fields {
let mut count = 0;
visit_field_origins(&layout, &index, &field, &control, &mut |change| {
assert_eq!(&*change.field, &field);
count += 1;
Ok(())
})
.unwrap();
assert_eq!(count, 1);
}
}