mod budget;
mod cancel;
mod cascade;
mod clock;
mod close;
pub mod durability;
mod expansion;
pub(crate) mod graph_cache;
mod group_count;
mod hybrid;
pub mod lock;
pub(crate) mod materialize;
pub(crate) mod namespace_batch;
pub use cascade::CascadeRule;
#[cfg(any(test, feature = "test-seams"))]
pub use namespace_batch::namespace_delete_cascade_with_steps;
pub use namespace_batch::{namespace_declare_cascade, namespace_delete_cascade};
#[cfg(feature = "graph-cypher")]
pub(crate) mod native_graph;
mod pool;
pub(crate) mod prepared;
mod prepared_lexical;
mod reindex;
pub(crate) mod rescored_scan;
#[cfg(test)]
mod shared_bound_tests;
mod snapshot;
mod snapshot_copy;
mod snapshot_view;
pub use namespace_batch::{
LiveNamespaceMutation, NamespaceMutation, namespace_batch, namespace_batch_live,
};
#[cfg(any(test, feature = "test-seams"))]
#[doc(hidden)]
pub use namespace_batch::{
namespace_batch_live_on_vfs, namespace_batch_live_with_steps, namespace_batch_with_steps,
namespace_reclaim, namespace_reclaim_on_vfs,
};
pub(crate) mod stats;
use crate::diag::{timing_elapsed, timing_start};
pub use cancel::{
CancelToken, Deadline, DeadlineError, QueryCancellation, QueryControl, QueryError,
};
#[cfg(any(test, feature = "test-seams"))]
pub use clock::ManualMonotonicClock;
pub use clock::{MonotonicClock, SystemMonotonicClock};
#[cfg(any(test, feature = "test-seams"))]
pub use hybrid::{
HybridScoreTestObservations, begin_hybrid_fresh_round_test_observations,
begin_hybrid_score_test_observations, take_hybrid_score_test_observations,
};
#[cfg(any(test, feature = "test-seams"))]
pub use materialize::MaterializationTestCounters;
pub use materialize::{
HybridPreparationError, MaterializationError, MaterializedRow, QueryMaterializer,
};
#[cfg(test)]
mod hybrid_overlap_tests;
pub use group_count::{
DocumentGroup, DocumentGroupCounts, DocumentGroupValue, MAX_DOCUMENT_GROUP_LIMIT,
};
pub use rescored_scan::ScanRescoreOptions;
pub use snapshot::{
InMemorySegment, InMemorySegmentFactors, PreparedSegment, PublishedSnapshot, SnapshotLease,
};
pub use snapshot_copy::SnapshotTargetReason;
pub use stats::Stats;
use std::collections::HashSet;
#[cfg(any(test, feature = "test-seams"))]
use std::io::IoSlice;
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Condvar, Mutex, RwLock};
use std::time::Duration;
use durability::{
CommitTier, DurabilityMode, DurabilityPolicy, DurabilityPolicyError, SyncRequirement,
};
use lock::{StoreLock, StoreLockError};
use self::close::BackgroundThread;
#[cfg(any(test, feature = "test-seams"))]
#[derive(Clone)]
pub struct StoreTestDependencies {
pub(crate) vfs: Arc<dyn crate::vfs::Vfs>,
pub(crate) clock: Arc<dyn MonotonicClock>,
hybrid_leg_fault: Option<HybridLegTestFault>,
storage_fault_controller: Option<StorageFaultController>,
ingest_retention_fault_controller: Option<crate::ingest::IngestRetentionFaultController>,
metadata_test_controller: Option<Arc<crate::planner::MetadataTestController>>,
pub(crate) vector_fault_controller: Option<crate::scan::vector_fault::VectorFaultController>,
pub(crate) kernel_fault_controller: Option<crate::kernels::vector_fault::KernelFaultController>,
vector_seal_scheme: Option<crate::quant::QuantScheme>,
}
#[cfg(any(test, feature = "test-seams"))]
impl StoreTestDependencies {
#[must_use]
pub fn new(vfs: Arc<dyn crate::vfs::Vfs>, clock: Arc<dyn MonotonicClock>) -> Self {
Self {
vfs,
clock,
hybrid_leg_fault: None,
storage_fault_controller: None,
ingest_retention_fault_controller: None,
metadata_test_controller: None,
vector_fault_controller: None,
kernel_fault_controller: None,
vector_seal_scheme: None,
}
}
#[must_use]
pub const fn with_hybrid_leg_fault(mut self, fault: HybridLegTestFault) -> Self {
self.hybrid_leg_fault = Some(fault);
self
}
#[must_use]
pub fn with_storage_fault_controller(mut self, controller: StorageFaultController) -> Self {
self.storage_fault_controller = Some(controller);
self
}
#[must_use]
pub fn with_ingest_retention_fault_controller(
mut self,
controller: crate::ingest::IngestRetentionFaultController,
) -> Self {
self.ingest_retention_fault_controller = Some(controller);
self
}
#[must_use]
pub fn with_metadata_test_controller(
mut self,
controller: Arc<crate::planner::MetadataTestController>,
) -> Self {
self.metadata_test_controller = Some(controller);
self
}
#[must_use]
pub fn with_vector_fault_controller(
mut self,
controller: crate::scan::vector_fault::VectorFaultController,
) -> Self {
self.vector_fault_controller = Some(controller);
self
}
#[must_use]
pub fn with_kernel_fault_controller(
mut self,
controller: crate::kernels::vector_fault::KernelFaultController,
) -> Self {
self.kernel_fault_controller = Some(controller);
self
}
#[must_use]
pub const fn with_vector_seal_scheme(mut self, scheme: crate::quant::QuantScheme) -> Self {
self.vector_seal_scheme = Some(scheme);
self
}
}
#[cfg(any(test, feature = "test-seams"))]
#[derive(Clone, Debug, Eq, PartialEq)]
#[doc(hidden)]
pub enum StorageTestFault {
TornWalHeader,
TornWalBody,
TornWalChecksum,
PostCommitError,
ManifestPreRename,
ManifestPostRename,
CorruptSegmentRegion,
WrongManifestObject,
WrongSegmentObject,
ListOmission { file_name: String },
DeleteOmission { file_name: String },
}
#[cfg(any(test, feature = "test-seams"))]
impl StorageTestFault {
const fn operation(&self) -> &'static str {
match self {
Self::TornWalHeader | Self::TornWalBody | Self::TornWalChecksum => "wal-prefix",
Self::PostCommitError => "retry",
Self::ManifestPreRename | Self::ManifestPostRename => "publication",
Self::CorruptSegmentRegion | Self::WrongManifestObject | Self::WrongSegmentObject => {
"format-check"
}
Self::ListOmission { .. } | Self::DeleteOmission { .. } => "orphan-cleanup",
}
}
const fn key(&self) -> &'static str {
match self {
Self::TornWalHeader => "torn-wal-header",
Self::TornWalBody => "torn-wal-body",
Self::TornWalChecksum => "torn-wal-checksum",
Self::PostCommitError => "post-commit-error",
Self::ManifestPreRename => "manifest-pre-rename-crash",
Self::ManifestPostRename => "manifest-post-rename-crash",
Self::CorruptSegmentRegion => "corrupt-segment-region",
Self::WrongManifestObject => "wrong-manifest-object",
Self::WrongSegmentObject => "wrong-segment-object",
Self::ListOmission { .. } | Self::DeleteOmission { .. } => "list-delete-omission",
}
}
}
#[cfg(any(test, feature = "test-seams"))]
#[derive(Clone, Debug, Eq, PartialEq)]
#[doc(hidden)]
pub struct StorageFaultPlan {
op_index: u32,
artifact: String,
offset: Option<u64>,
segment: Option<crate::segment::SegmentId>,
region_kind: Option<u16>,
chunk: Option<u32>,
}
#[cfg(any(test, feature = "test-seams"))]
impl StorageFaultPlan {
#[must_use]
pub fn new(op_index: u32, artifact: impl Into<String>) -> Self {
Self {
op_index,
artifact: artifact.into(),
offset: None,
segment: None,
region_kind: None,
chunk: None,
}
}
#[must_use]
pub const fn with_offset(mut self, offset: u64) -> Self {
self.offset = Some(offset);
self
}
#[must_use]
pub const fn with_segment_region(
mut self,
segment: crate::segment::SegmentId,
region_kind: u16,
chunk: u32,
) -> Self {
self.segment = Some(segment);
self.region_kind = Some(region_kind);
self.chunk = Some(chunk);
self
}
#[must_use]
pub const fn with_segment(mut self, segment: crate::segment::SegmentId) -> Self {
self.segment = Some(segment);
self
}
#[must_use]
pub const fn op_index(&self) -> u32 {
self.op_index
}
#[must_use]
pub fn artifact(&self) -> &str {
&self.artifact
}
#[must_use]
pub const fn offset(&self) -> Option<u64> {
self.offset
}
#[must_use]
pub const fn segment(&self) -> Option<crate::segment::SegmentId> {
self.segment
}
#[must_use]
pub const fn region_kind(&self) -> Option<u16> {
self.region_kind
}
#[must_use]
pub const fn chunk(&self) -> Option<u32> {
self.chunk
}
}
#[cfg(any(test, feature = "test-seams"))]
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
#[doc(hidden)]
pub enum StorageReceiptSite {
WalOpenHeaderValidation,
WalOpenRecordValidation,
WalOpenRecordChecksum,
WalCommitAppendAfterInnerSuccess,
ManifestCommitBeforeRename,
ManifestCommitAfterRename,
SegmentReadRegionChecksum,
ManifestOpenFamilyValidation,
SegmentOpenFamilyValidation,
SegmentOpenObjectIdentity,
OrphanCleanupList,
OrphanCleanupDelete,
}
#[cfg(any(test, feature = "test-seams"))]
impl StorageReceiptSite {
const fn key(self) -> &'static str {
match self {
Self::WalOpenHeaderValidation => "WalOpen.HeaderValidation",
Self::WalOpenRecordValidation => "WalOpen.RecordValidation",
Self::WalOpenRecordChecksum => "WalOpen.RecordChecksum",
Self::WalCommitAppendAfterInnerSuccess => "WalCommit.AppendAfterInnerSuccess",
Self::ManifestCommitBeforeRename => "ManifestCommit.BeforeRename",
Self::ManifestCommitAfterRename => "ManifestCommit.AfterRename",
Self::SegmentReadRegionChecksum => "SegmentRead.RegionChecksum",
Self::ManifestOpenFamilyValidation => "ManifestOpen.FamilyValidation",
Self::SegmentOpenFamilyValidation => "SegmentOpen.FamilyValidation",
Self::SegmentOpenObjectIdentity => "SegmentOpen.ObjectIdentity",
Self::OrphanCleanupList => "OrphanCleanup.List",
Self::OrphanCleanupDelete => "OrphanCleanup.Delete",
}
}
}
#[cfg(any(test, feature = "test-seams"))]
#[derive(Clone, Debug, Eq, PartialEq)]
#[doc(hidden)]
pub enum StorageReceiptObserved {
WalHeader {
artifact: String,
reason: crate::wal::header::WalHeaderError,
},
WalRecord {
artifact: String,
offset: u64,
reason: crate::wal::replay::CorruptionReason,
},
WalAppend {
artifact: String,
encoded_len: u64,
first_seq: u64,
last_seq: u64,
inner_append_completed: bool,
caller_saw_error: bool,
},
ManifestRename {
temporary: String,
committed: String,
rename_performed: bool,
new_segment_final: bool,
directory_sync_returned: bool,
},
SegmentChecksum {
artifact: String,
segment: crate::segment::SegmentId,
region_kind: u16,
chunk: u32,
expected_checksum: u64,
actual_checksum: u64,
},
Format {
artifact: String,
check: crate::format::frame::FormatCheck,
expected_family: Option<u16>,
actual_family: Option<u16>,
expected_id: Option<crate::segment::SegmentId>,
actual_id: Option<crate::segment::SegmentId>,
},
Omission {
artifact: String,
deletion_observed: bool,
},
}
#[cfg(any(test, feature = "test-seams"))]
#[derive(Clone, Debug, Eq, PartialEq)]
#[doc(hidden)]
pub struct StorageFaultReceipt {
operation: &'static str,
fault: &'static str,
site: StorageReceiptSite,
plan: StorageFaultPlan,
observed: StorageReceiptObserved,
cardinality: u32,
}
#[cfg(any(test, feature = "test-seams"))]
impl StorageFaultReceipt {
#[must_use]
pub const fn campaign(&self) -> &'static str {
"storage-durability"
}
#[must_use]
pub const fn operation(&self) -> &'static str {
self.operation
}
#[must_use]
pub const fn fault(&self) -> &'static str {
self.fault
}
#[must_use]
pub const fn site(&self) -> &'static str {
self.site.key()
}
#[must_use]
pub const fn typed_site(&self) -> StorageReceiptSite {
self.site
}
#[must_use]
pub const fn cardinality(&self) -> u32 {
self.cardinality
}
#[must_use]
pub const fn plan(&self) -> &StorageFaultPlan {
&self.plan
}
#[must_use]
pub const fn observed(&self) -> &StorageReceiptObserved {
&self.observed
}
pub fn from_manifest_abort_acknowledgment(acknowledgment: &str) -> Result<Self, String> {
let body = acknowledgment
.strip_suffix('\n')
.filter(|body| !body.contains('\n'))
.ok_or_else(|| {
"manifest abort acknowledgment must end in exactly one newline".to_owned()
})?;
let fields = body
.split('|')
.map(|field| {
field.split_once('=').ok_or_else(|| {
format!("manifest abort acknowledgment field lacks '=': {field}")
})
})
.collect::<Result<Vec<_>, _>>()?;
let expected_keys = [
"campaign",
"operation",
"fault",
"site",
"op_index",
"cardinality",
"artifact",
"temporary",
"committed",
"rename_performed",
"new_segment_final",
"directory_sync_returned",
];
if fields.len() != expected_keys.len()
|| fields
.iter()
.map(|(key, _)| *key)
.ne(expected_keys.iter().copied())
{
return Err("manifest abort acknowledgment field inventory differs".to_owned());
}
let mut values = fields.into_iter().map(|(_, value)| value);
let campaign = values
.next()
.ok_or_else(|| "manifest abort acknowledgment lacks campaign".to_owned())?;
let operation = values
.next()
.ok_or_else(|| "manifest abort acknowledgment lacks operation".to_owned())?;
let fault = values
.next()
.ok_or_else(|| "manifest abort acknowledgment lacks fault".to_owned())?;
let site = values
.next()
.ok_or_else(|| "manifest abort acknowledgment lacks site".to_owned())?;
let op_index = values
.next()
.ok_or_else(|| "manifest abort acknowledgment lacks op index".to_owned())?
.parse::<u32>()
.map_err(|error| format!("manifest abort op index is invalid: {error}"))?;
let cardinality = values
.next()
.ok_or_else(|| "manifest abort acknowledgment lacks cardinality".to_owned())?
.parse::<u32>()
.map_err(|error| format!("manifest abort cardinality is invalid: {error}"))?;
let artifact = values
.next()
.ok_or_else(|| "manifest abort acknowledgment lacks artifact".to_owned())?;
let temporary = values
.next()
.ok_or_else(|| "manifest abort acknowledgment lacks temporary path".to_owned())?;
let committed = values
.next()
.ok_or_else(|| "manifest abort acknowledgment lacks committed path".to_owned())?;
let parse_bool = |field: &str, value: &str| match value {
"true" => Ok(true),
"false" => Ok(false),
_ => Err(format!("manifest abort {field} is not a bool: {value}")),
};
let rename_performed = parse_bool(
"rename_performed",
values
.next()
.ok_or_else(|| "manifest abort acknowledgment lacks rename result".to_owned())?,
)?;
let new_segment_final = parse_bool(
"new_segment_final",
values.next().ok_or_else(|| {
"manifest abort acknowledgment lacks final-segment result".to_owned()
})?,
)?;
let directory_sync_returned = parse_bool(
"directory_sync_returned",
values.next().ok_or_else(|| {
"manifest abort acknowledgment lacks directory-sync result".to_owned()
})?,
)?;
if campaign != "storage-durability"
|| operation != "publication"
|| cardinality != 1
|| artifact != crate::manifest::io::MANIFEST_FILE
|| temporary != crate::manifest::io::MANIFEST_TEMP_FILE
|| committed != crate::manifest::io::MANIFEST_FILE
|| !new_segment_final
|| directory_sync_returned
{
return Err("manifest abort acknowledgment fixed facts differ".to_owned());
}
let (fault, site, expected_rename) = match (fault, site) {
("manifest-pre-rename-crash", "ManifestCommit.BeforeRename") => (
"manifest-pre-rename-crash",
StorageReceiptSite::ManifestCommitBeforeRename,
false,
),
("manifest-post-rename-crash", "ManifestCommit.AfterRename") => (
"manifest-post-rename-crash",
StorageReceiptSite::ManifestCommitAfterRename,
true,
),
_ => return Err("manifest abort fault and production site differ".to_owned()),
};
if rename_performed != expected_rename {
return Err("manifest abort rename result differs from production site".to_owned());
}
Ok(Self {
operation: "publication",
fault,
site,
plan: StorageFaultPlan::new(op_index, artifact),
observed: StorageReceiptObserved::ManifestRename {
temporary: temporary.to_owned(),
committed: committed.to_owned(),
rename_performed,
new_segment_final,
directory_sync_returned,
},
cardinality,
})
}
}
#[cfg(any(test, feature = "test-seams"))]
#[derive(Debug)]
struct StorageFaultState {
fault: StorageTestFault,
plan: StorageFaultPlan,
consumed: bool,
receipt: Option<StorageFaultReceipt>,
cleanup_report: Option<StorageCleanupReport>,
#[cfg(unix)]
child_abort_ack: Option<std::os::unix::net::UnixStream>,
}
#[cfg(any(test, feature = "test-seams"))]
#[derive(Clone, Debug, Eq, PartialEq)]
#[doc(hidden)]
pub struct StorageCleanupReport {
reclaimed_bytes: u64,
deleted_paths: Vec<String>,
retained_eligible_paths: Vec<String>,
directory_synced: bool,
}
#[cfg(any(test, feature = "test-seams"))]
impl StorageCleanupReport {
#[must_use]
pub const fn reclaimed_bytes(&self) -> u64 {
self.reclaimed_bytes
}
#[must_use]
pub fn deleted_paths(&self) -> &[String] {
&self.deleted_paths
}
#[must_use]
pub fn retained_eligible_paths(&self) -> &[String] {
&self.retained_eligible_paths
}
#[must_use]
pub const fn directory_synced(&self) -> bool {
self.directory_synced
}
}
#[cfg(any(test, feature = "test-seams"))]
#[derive(Clone, Debug)]
#[doc(hidden)]
pub struct StorageFaultController {
state: Arc<Mutex<StorageFaultState>>,
}
#[cfg(any(test, feature = "test-seams"))]
impl StorageFaultController {
#[must_use]
pub fn new(fault: StorageTestFault, plan: StorageFaultPlan) -> Self {
Self {
state: Arc::new(Mutex::new(StorageFaultState {
fault,
plan,
consumed: false,
receipt: None,
cleanup_report: None,
#[cfg(unix)]
child_abort_ack: None,
})),
}
}
#[cfg(unix)]
#[must_use]
pub fn with_child_abort_ack(self, acknowledgment: std::os::unix::net::UnixStream) -> Self {
if let Ok(mut state) = self.state.lock() {
state.child_abort_ack = Some(acknowledgment);
}
self
}
pub fn take_receipt(&self) -> Option<StorageFaultReceipt> {
self.state
.lock()
.ok()
.and_then(|mut state| state.receipt.take())
}
pub fn take_cleanup_report(&self) -> Option<StorageCleanupReport> {
self.state
.lock()
.ok()
.and_then(|mut state| state.cleanup_report.take())
}
fn record_cleanup_report(&self, report: crate::manifest::io::OrphanCleanupReport) {
let normalize = |paths: Vec<PathBuf>| {
paths
.into_iter()
.filter_map(|path| {
path.file_name()
.and_then(|name| name.to_str())
.map(str::to_owned)
})
.collect::<Vec<_>>()
};
if let Ok(mut state) = self.state.lock() {
state.cleanup_report = Some(StorageCleanupReport {
reclaimed_bytes: report.reclaimed_bytes,
deleted_paths: normalize(report.deleted_paths),
retained_eligible_paths: normalize(report.retained_eligible_paths),
directory_synced: report.directory_synced,
});
}
}
fn armed_fault(&self) -> Option<StorageTestFault> {
self.state
.lock()
.ok()
.and_then(|state| (!state.consumed).then(|| state.fault.clone()))
}
fn planned_segment(&self) -> Option<crate::segment::SegmentId> {
self.state
.lock()
.ok()
.and_then(|state| state.plan.segment())
}
fn emit(&self, site: StorageReceiptSite, observed: StorageReceiptObserved) -> bool {
let Ok(mut state) = self.state.lock() else {
return false;
};
if state.consumed {
return false;
}
state.consumed = true;
state.receipt = Some(StorageFaultReceipt {
operation: state.fault.operation(),
fault: state.fault.key(),
site,
plan: state.plan.clone(),
observed,
cardinality: 1,
});
true
}
fn abort_after_manifest_receipt(&self) -> ! {
#[cfg(unix)]
if let Ok(mut state) = self.state.lock() {
let wire_receipt = state
.receipt
.as_ref()
.and_then(storage_manifest_abort_wire_receipt);
if let (Some(wire_receipt), Some(acknowledgment)) =
(wire_receipt, state.child_abort_ack.as_mut())
{
use std::io::Write as _;
let _ = acknowledgment.write_all(wire_receipt.as_bytes());
let _ = acknowledgment.flush();
}
}
std::process::abort()
}
}
#[cfg(any(test, feature = "test-seams"))]
thread_local! {
static STORAGE_OPEN_CONTROLLER: std::cell::RefCell<
Option<std::sync::Weak<Mutex<StorageFaultState>>>
> = const { std::cell::RefCell::new(None) };
}
#[cfg(any(test, feature = "test-seams"))]
fn with_storage_open_controller<R>(
controller: Option<&StorageFaultController>,
operation: impl FnOnce() -> R,
) -> R {
STORAGE_OPEN_CONTROLLER.with(|slot| {
let previous = slot.replace(controller.map(|controller| Arc::downgrade(&controller.state)));
let result = operation();
let _ = slot.replace(previous);
result
})
}
#[cfg(any(test, feature = "test-seams"))]
fn storage_open_controller() -> Option<StorageFaultController> {
STORAGE_OPEN_CONTROLLER.with(|slot| {
slot.borrow()
.as_ref()
.and_then(std::sync::Weak::upgrade)
.map(|state| StorageFaultController { state })
})
}
#[cfg(any(test, feature = "test-seams"))]
pub(crate) fn record_storage_wal_recovery_fault(error: &crate::wal::WalRecoveryError) {
let Some(controller) = storage_open_controller() else {
return;
};
match (controller.armed_fault(), error) {
(
Some(StorageTestFault::TornWalHeader),
crate::wal::WalRecoveryError::InvalidHeader(reason),
) => {
let _ = controller.emit(
StorageReceiptSite::WalOpenHeaderValidation,
StorageReceiptObserved::WalHeader {
artifact: "wal.ze".to_owned(),
reason: *reason,
},
);
}
(
Some(StorageTestFault::TornWalBody),
crate::wal::WalRecoveryError::CorruptAt {
offset,
reason:
reason @ crate::wal::replay::CorruptionReason::Record {
error:
crate::wal::record::RecordError::HeaderTruncated { .. }
| crate::wal::record::RecordError::BodyTruncated { .. }
| crate::wal::record::RecordError::LengthOverflow { .. },
..
},
},
) => record_storage_wal_record_fault(
&controller,
StorageReceiptSite::WalOpenRecordValidation,
*offset,
*reason,
),
(
Some(StorageTestFault::TornWalChecksum),
crate::wal::WalRecoveryError::CorruptAt {
offset,
reason:
reason @ crate::wal::replay::CorruptionReason::Record {
error: crate::wal::record::RecordError::ChecksumMismatch { .. },
..
},
},
) => record_storage_wal_record_fault(
&controller,
StorageReceiptSite::WalOpenRecordChecksum,
*offset,
*reason,
),
_ => {}
}
}
#[cfg(any(test, feature = "test-seams"))]
fn record_storage_wal_record_fault(
controller: &StorageFaultController,
site: StorageReceiptSite,
offset: usize,
reason: crate::wal::replay::CorruptionReason,
) {
let Ok(offset) = u64::try_from(offset) else {
return;
};
let _ = controller.emit(
site,
StorageReceiptObserved::WalRecord {
artifact: "wal.ze".to_owned(),
offset,
reason,
},
);
}
#[cfg(any(test, feature = "test-seams"))]
pub(crate) fn record_storage_manifest_format_fault(
error: &crate::format::frame::FormatError,
actual_family: Option<u16>,
) {
let Some(controller) = storage_open_controller() else {
return;
};
if matches!(
controller.armed_fault(),
Some(StorageTestFault::WrongManifestObject)
) && error.check() == crate::format::frame::FormatCheck::Family
{
let _ = controller.emit(
StorageReceiptSite::ManifestOpenFamilyValidation,
StorageReceiptObserved::Format {
artifact: storage_artifact_name(error.artifact()).to_owned(),
check: error.check(),
expected_family: Some(crate::format::FormatFamily::Manifest.id()),
actual_family,
expected_id: None,
actual_id: None,
},
);
}
}
#[cfg(any(test, feature = "test-seams"))]
pub(crate) fn record_storage_segment_format_fault(
error: &crate::format::frame::FormatError,
actual_family: Option<u16>,
expected_id: crate::segment::SegmentId,
) {
let Some(controller) = storage_open_controller() else {
return;
};
if matches!(
controller.armed_fault(),
Some(StorageTestFault::WrongSegmentObject)
) && error.check() == crate::format::frame::FormatCheck::Family
{
let _ = controller.emit(
StorageReceiptSite::SegmentOpenFamilyValidation,
StorageReceiptObserved::Format {
artifact: storage_artifact_name(error.artifact()).to_owned(),
check: error.check(),
expected_family: Some(crate::format::FormatFamily::Segment.id()),
actual_family,
expected_id: Some(expected_id),
actual_id: None,
},
);
}
}
#[cfg(any(test, feature = "test-seams"))]
pub(crate) fn record_storage_segment_identity_fault(
artifact: &str,
expected: crate::segment::SegmentId,
actual: crate::segment::SegmentId,
) {
let Some(controller) = storage_open_controller() else {
return;
};
if matches!(
controller.armed_fault(),
Some(StorageTestFault::WrongSegmentObject)
) {
let _ = controller.emit(
StorageReceiptSite::SegmentOpenObjectIdentity,
StorageReceiptObserved::Format {
artifact: storage_artifact_name(artifact).to_owned(),
check: crate::format::frame::FormatCheck::ObjectIdentity,
expected_family: None,
actual_family: None,
expected_id: Some(expected),
actual_id: Some(actual),
},
);
}
}
#[cfg(unix)]
#[cfg(any(test, feature = "test-seams"))]
fn storage_manifest_abort_wire_receipt(receipt: &StorageFaultReceipt) -> Option<String> {
let StorageReceiptObserved::ManifestRename {
temporary,
committed,
rename_performed,
new_segment_final,
directory_sync_returned,
} = &receipt.observed
else {
return None;
};
Some(format!(
"campaign={}|operation={}|fault={}|site={}|op_index={}|cardinality={}|artifact={}|temporary={temporary}|committed={committed}|rename_performed={rename_performed}|new_segment_final={new_segment_final}|directory_sync_returned={directory_sync_returned}\n",
receipt.campaign(),
receipt.operation,
receipt.fault,
receipt.site.key(),
receipt.plan.op_index,
receipt.cardinality,
receipt.plan.artifact,
))
}
#[cfg(any(test, feature = "test-seams"))]
fn storage_artifact_name(artifact: &str) -> &str {
Path::new(artifact)
.file_name()
.and_then(|name| name.to_str())
.unwrap_or(artifact)
}
#[cfg(any(test, feature = "test-seams"))]
type StorageSegmentControllerRegistry = Mutex<
Vec<(
PathBuf,
crate::segment::SegmentId,
std::sync::Weak<Mutex<StorageFaultState>>,
)>,
>;
#[cfg(any(test, feature = "test-seams"))]
static STORAGE_SEGMENT_CONTROLLERS: std::sync::OnceLock<StorageSegmentControllerRegistry> =
std::sync::OnceLock::new();
#[cfg(any(test, feature = "test-seams"))]
fn register_storage_segment_controller(
store_directory: &Path,
segment: crate::segment::SegmentId,
controller: &StorageFaultController,
) {
let registry = STORAGE_SEGMENT_CONTROLLERS.get_or_init(|| Mutex::new(Vec::new()));
if let Ok(mut entries) = registry.lock() {
entries.retain(|(_, _, state)| state.strong_count() != 0);
entries.push((
store_directory.to_path_buf(),
segment,
Arc::downgrade(&controller.state),
));
}
}
#[cfg(any(test, feature = "test-seams"))]
pub(crate) fn record_storage_segment_checksum_fault(
store_directory: &Path,
segment: crate::segment::SegmentId,
kind: crate::segment::layout::RegionKind,
chunk: u32,
expected: u64,
actual: u64,
) {
let Some(registry) = STORAGE_SEGMENT_CONTROLLERS.get() else {
return;
};
let states = match registry.lock() {
Ok(mut entries) => {
entries.retain(|(_, _, state)| state.strong_count() != 0);
entries
.iter()
.filter(|(directory, registered, _)| {
directory == store_directory && *registered == segment
})
.filter_map(|(_, _, state)| state.upgrade())
.collect::<Vec<_>>()
}
Err(_) => return,
};
for state in states {
let controller = StorageFaultController { state };
if matches!(
controller.armed_fault(),
Some(StorageTestFault::CorruptSegmentRegion)
) {
let _ = controller.emit(
StorageReceiptSite::SegmentReadRegionChecksum,
StorageReceiptObserved::SegmentChecksum {
artifact: segment.file_name(),
segment,
region_kind: kind.id(),
chunk,
expected_checksum: expected,
actual_checksum: actual,
},
);
}
}
}
#[cfg(any(test, feature = "test-seams"))]
struct StorageFaultVfs {
inner: Arc<dyn crate::vfs::Vfs>,
controller: StorageFaultController,
}
#[cfg(any(test, feature = "test-seams"))]
struct StorageFaultFile {
inner: Box<dyn crate::vfs::VfsFile>,
path: PathBuf,
controller: StorageFaultController,
}
#[cfg(any(test, feature = "test-seams"))]
impl StorageFaultFile {
fn after_append(&self, encoded_len: u64, first_seq: u64, last_seq: u64) -> std::io::Result<()> {
if matches!(
self.controller.armed_fault(),
Some(StorageTestFault::PostCommitError)
) && self.path.file_name().is_some_and(|name| name == "wal.ze")
&& self.controller.emit(
StorageReceiptSite::WalCommitAppendAfterInnerSuccess,
StorageReceiptObserved::WalAppend {
artifact: "wal.ze".to_owned(),
encoded_len,
first_seq,
last_seq,
inner_append_completed: true,
caller_saw_error: true,
},
)
{
return Err(std::io::Error::other(
"scheduled post-commit error after inner append",
));
}
Ok(())
}
}
#[cfg(any(test, feature = "test-seams"))]
impl crate::vfs::VfsFile for StorageFaultFile {
fn append(&mut self, bytes: &[u8]) -> std::io::Result<()> {
let record_facts = storage_wal_append_facts(bytes)?;
self.inner.append(bytes)?;
record_facts.map_or(Ok(()), |(encoded_len, first_seq, last_seq)| {
self.after_append(encoded_len, first_seq, last_seq)
})
}
fn append_vectored(&mut self, buffers: &mut [IoSlice<'_>]) -> std::io::Result<()> {
let mut aggregate = None::<(u64, u64, u64)>;
for bytes in buffers.iter() {
if let Some((encoded_len, first_seq, last_seq)) = storage_wal_append_facts(bytes)? {
aggregate = Some(match aggregate {
None => (encoded_len, first_seq, last_seq),
Some((total, first, _)) => (
total.checked_add(encoded_len).ok_or_else(|| {
std::io::Error::other("WAL record append length overflow")
})?,
first,
last_seq,
),
});
}
}
self.inner.append_vectored(buffers)?;
aggregate.map_or(Ok(()), |(encoded_len, first_seq, last_seq)| {
self.after_append(encoded_len, first_seq, last_seq)
})
}
fn sync(&self, kind: crate::vfs::SyncKind) -> std::io::Result<()> {
self.inner.sync(kind)
}
}
#[cfg(any(test, feature = "test-seams"))]
impl crate::vfs::Vfs for StorageFaultVfs {
fn segment_data_read_counter(&self) -> Option<Arc<AtomicU64>> {
self.inner.segment_data_read_counter()
}
fn ensure_directory(&self, path: &Path, create: bool) -> std::io::Result<bool> {
self.inner.ensure_directory(path, create)
}
fn open(&self, path: &Path) -> std::io::Result<u64> {
self.inner.open(path)
}
fn open_for_map(&self, path: &Path) -> std::io::Result<std::fs::File> {
self.inner.open_for_map(path)
}
fn read(&self, path: &Path) -> std::io::Result<Vec<u8>> {
self.inner.read(path)
}
fn read_range(&self, path: &Path, offset: u64, length: usize) -> std::io::Result<Vec<u8>> {
self.inner.read_range(path, offset, length)
}
fn write(&self, path: &Path, bytes: &[u8]) -> std::io::Result<()> {
self.inner.write(path, bytes)
}
fn open_append(&self, path: &Path) -> std::io::Result<Box<dyn crate::vfs::VfsFile>> {
Ok(Box::new(StorageFaultFile {
inner: self.inner.open_append(path)?,
path: path.to_path_buf(),
controller: self.controller.clone(),
}))
}
fn rename(&self, from: &Path, to: &Path) -> std::io::Result<()> {
let is_manifest = from
.file_name()
.is_some_and(|name| name == crate::manifest::io::MANIFEST_TEMP_FILE)
&& to
.file_name()
.is_some_and(|name| name == crate::manifest::io::MANIFEST_FILE);
if is_manifest
&& matches!(
self.controller.armed_fault(),
Some(StorageTestFault::ManifestPreRename)
)
&& self.controller.emit(
StorageReceiptSite::ManifestCommitBeforeRename,
StorageReceiptObserved::ManifestRename {
temporary: crate::manifest::io::MANIFEST_TEMP_FILE.to_owned(),
committed: crate::manifest::io::MANIFEST_FILE.to_owned(),
rename_performed: false,
new_segment_final: storage_planned_segment_is_final(
&*self.inner,
from,
self.controller.planned_segment(),
),
directory_sync_returned: false,
},
)
{
self.controller.abort_after_manifest_receipt();
}
self.inner.rename(from, to)?;
if is_manifest
&& matches!(
self.controller.armed_fault(),
Some(StorageTestFault::ManifestPostRename)
)
&& self.controller.emit(
StorageReceiptSite::ManifestCommitAfterRename,
StorageReceiptObserved::ManifestRename {
temporary: crate::manifest::io::MANIFEST_TEMP_FILE.to_owned(),
committed: crate::manifest::io::MANIFEST_FILE.to_owned(),
rename_performed: true,
new_segment_final: storage_planned_segment_is_final(
&*self.inner,
to,
self.controller.planned_segment(),
),
directory_sync_returned: false,
},
)
{
self.controller.abort_after_manifest_receipt();
}
Ok(())
}
fn sync(&self, path: &Path, kind: crate::vfs::SyncKind) -> std::io::Result<()> {
self.inner.sync(path, kind)
}
fn list(&self, directory: &Path) -> std::io::Result<Vec<PathBuf>> {
let mut paths = self.inner.list(directory)?;
if let Some(StorageTestFault::ListOmission { file_name }) = self.controller.armed_fault()
&& let Some(position) = paths.iter().position(|path| {
path.file_name()
.is_some_and(|name| name == file_name.as_str())
})
{
let _ = paths.remove(position);
let _ = self.controller.emit(
StorageReceiptSite::OrphanCleanupList,
StorageReceiptObserved::Omission {
artifact: file_name,
deletion_observed: false,
},
);
}
Ok(paths)
}
fn for_each_direct_child(
&self,
directory: &Path,
visitor: &mut dyn FnMut(&Path) -> std::io::Result<()>,
) -> std::io::Result<()> {
self.inner.for_each_direct_child(directory, visitor)
}
fn delete(&self, path: &Path) -> std::io::Result<()> {
if let Some(StorageTestFault::DeleteOmission { file_name }) = self.controller.armed_fault()
&& path
.file_name()
.is_some_and(|name| name == file_name.as_str())
&& self.controller.emit(
StorageReceiptSite::OrphanCleanupDelete,
StorageReceiptObserved::Omission {
artifact: file_name,
deletion_observed: false,
},
)
{
return Ok(());
}
self.inner.delete(path)
}
}
#[cfg(any(test, feature = "test-seams"))]
fn storage_wal_append_facts(bytes: &[u8]) -> std::io::Result<Option<(u64, u64, u64)>> {
const RECORD_HEADER_LEN: usize = 14;
const RECORD_CHECKSUM_LEN: usize = 8;
let has_wal_header = bytes
.get(..8)
.is_some_and(|magic| magic == crate::format::frame::FILE_MAGIC)
&& bytes
.get(8..10)
.is_some_and(|family| family == crate::format::FormatFamily::Wal.id().to_le_bytes());
let mut offset = if has_wal_header {
crate::wal::header::WAL_HEADER_LEN
} else {
0
};
if offset == bytes.len() {
return Ok(None);
}
let mut encoded_len = 0_u64;
let mut first_seq = None;
let mut last_seq = None;
while offset < bytes.len() {
let header_end = offset
.checked_add(RECORD_HEADER_LEN)
.ok_or_else(|| std::io::Error::other("WAL record header offset overflow"))?;
let header = bytes
.get(offset..header_end)
.ok_or_else(|| std::io::Error::other("WAL append has a torn record header"))?;
let payload_len = usize::try_from(u32::from_le_bytes(
header
.get(..4)
.ok_or_else(|| std::io::Error::other("WAL payload length is absent"))?
.try_into()
.map_err(|_| std::io::Error::other("WAL payload length width changed"))?,
))
.map_err(|_| std::io::Error::other("WAL payload length does not fit usize"))?;
let seq = u64::from_le_bytes(
header
.get(4..12)
.ok_or_else(|| std::io::Error::other("WAL sequence is absent"))?
.try_into()
.map_err(|_| std::io::Error::other("WAL sequence width changed"))?,
);
let record_len = RECORD_HEADER_LEN
.checked_add(payload_len)
.and_then(|length| length.checked_add(RECORD_CHECKSUM_LEN))
.ok_or_else(|| std::io::Error::other("WAL record length overflow"))?;
offset = offset
.checked_add(record_len)
.ok_or_else(|| std::io::Error::other("WAL record end overflow"))?;
if offset > bytes.len() {
return Err(std::io::Error::other("WAL append has a torn record body"));
}
encoded_len = encoded_len
.checked_add(
u64::try_from(record_len)
.map_err(|_| std::io::Error::other("WAL record length does not fit u64"))?,
)
.ok_or_else(|| std::io::Error::other("WAL record append length overflow"))?;
first_seq.get_or_insert(seq);
last_seq = Some(seq);
}
match (first_seq, last_seq) {
(Some(first), Some(last)) => Ok(Some((encoded_len, first, last))),
_ => Ok(None),
}
}
#[cfg(any(test, feature = "test-seams"))]
fn storage_planned_segment_is_final(
vfs: &dyn crate::vfs::Vfs,
artifact: &Path,
planned: Option<crate::segment::SegmentId>,
) -> bool {
let Some(file_name) = planned.map(crate::segment::SegmentId::file_name) else {
return false;
};
artifact.parent().is_some_and(|directory| {
vfs.list(directory).is_ok_and(|paths| {
paths.iter().any(|path| {
path.file_name()
.and_then(|name| name.to_str())
.is_some_and(|name| name == file_name)
})
})
})
}
#[cfg(any(test, feature = "test-seams"))]
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
#[doc(hidden)]
pub enum HybridLegTestFault {
Panic(crate::fusion::FusionLeg),
}
#[allow(clippy::panic)]
fn maybe_trigger_hybrid_leg_panic(armed: bool, detail: &'static str) {
#[cfg(any(test, feature = "test-seams"))]
if armed {
std::panic::panic_any(detail);
}
#[cfg(not(any(test, feature = "test-seams")))]
let _ = (armed, detail);
}
#[cfg(any(test, feature = "test-seams"))]
#[derive(Clone, Debug, Eq, PartialEq)]
#[doc(hidden)]
pub struct HybridExecutionReceipt {
pub vector_thread: std::thread::ThreadId,
pub lexical_thread_name: Option<String>,
pub vector_completed: bool,
pub lexical_completed: bool,
}
pub const DEFAULT_READER_DRAIN_TIMEOUT: Duration = Duration::from_millis(250);
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum StoreState {
Open,
Closing,
Closed,
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct OpenOptions {
inspection: bool,
access_mode: AccessMode,
durability_mode: DurabilityMode,
commit_tier: CommitTier,
reader_drain_timeout: Duration,
max_resident_bytes: u64,
max_temp_bytes: u64,
epoch: Option<crate::epoch::StoreEpoch>,
tokenizer: Option<crate::fts::tokenizer::TokenizerConfig>,
schema: Option<crate::meta::Schema>,
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum AccessMode {
ReadWrite,
ReadOnly,
}
impl OpenOptions {
#[must_use]
pub const fn new() -> Self {
Self {
inspection: false,
access_mode: AccessMode::ReadWrite,
durability_mode: DurabilityMode::Derived,
commit_tier: CommitTier::Ordered,
reader_drain_timeout: DEFAULT_READER_DRAIN_TIMEOUT,
max_resident_bytes: u64::MAX,
max_temp_bytes: u64::MAX,
epoch: None,
tokenizer: None,
schema: None,
}
}
#[must_use]
pub const fn read_only() -> Self {
Self {
inspection: false,
access_mode: AccessMode::ReadOnly,
durability_mode: DurabilityMode::Derived,
commit_tier: CommitTier::Ordered,
reader_drain_timeout: DEFAULT_READER_DRAIN_TIMEOUT,
max_resident_bytes: u64::MAX,
max_temp_bytes: u64::MAX,
epoch: None,
tokenizer: None,
schema: None,
}
}
#[must_use]
pub const fn with_durability(mut self, mode: DurabilityMode, tier: CommitTier) -> Self {
self.durability_mode = mode;
self.commit_tier = tier;
self
}
#[cfg(feature = "graph-cypher")]
#[must_use]
pub(crate) const fn with_access_mode(mut self, access_mode: AccessMode) -> Self {
self.access_mode = access_mode;
self
}
#[must_use]
pub const fn with_reader_drain_timeout(mut self, timeout: Duration) -> Self {
self.reader_drain_timeout = timeout;
self
}
#[must_use]
pub const fn with_max_resident_bytes(mut self, bytes: u64) -> Self {
self.max_resident_bytes = bytes;
self
}
#[must_use]
pub fn with_schema(mut self, schema: crate::meta::Schema) -> Self {
self.schema = Some(schema);
self
}
#[must_use]
pub const fn with_max_temp_bytes(mut self, bytes: u64) -> Self {
self.max_temp_bytes = bytes;
self
}
#[must_use]
pub fn with_epoch(mut self, epoch: crate::epoch::StoreEpoch) -> Self {
self.epoch = Some(epoch);
self
}
#[must_use]
pub fn with_tokenizer(mut self, tokenizer: crate::fts::tokenizer::TokenizerConfig) -> Self {
self.tokenizer = Some(tokenizer);
self
}
}
impl Default for OpenOptions {
fn default() -> Self {
Self::new()
}
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct GraphSearchOptions {
profile: crate::graph::search::GraphSearchProfile,
ef: Option<usize>,
seed: u64,
#[cfg(any(test, feature = "test-seams"))]
cancel_after_hops: Option<usize>,
}
impl GraphSearchOptions {
#[must_use]
pub const fn new(profile: crate::graph::search::GraphSearchProfile) -> Self {
Self {
profile,
ef: None,
seed: 0,
#[cfg(any(test, feature = "test-seams"))]
cancel_after_hops: None,
}
}
#[must_use]
pub const fn with_ef(mut self, ef: usize) -> Self {
self.ef = Some(ef);
self
}
#[must_use]
pub const fn with_seed(mut self, seed: u64) -> Self {
self.seed = seed;
self
}
#[cfg(any(test, feature = "test-seams"))]
#[must_use]
pub const fn with_cancel_after_hops(mut self, hops: usize) -> Self {
self.cancel_after_hops = Some(hops);
self
}
pub(crate) const fn profile(self) -> crate::graph::search::GraphSearchProfile {
self.profile
}
pub(crate) const fn ef(self) -> Option<usize> {
self.ef
}
pub(crate) const fn seed(self) -> u64 {
self.seed
}
#[cfg(any(test, feature = "test-seams"))]
pub(crate) const fn cancel_after_hops(self) -> Option<usize> {
self.cancel_after_hops
}
}
impl Default for GraphSearchOptions {
fn default() -> Self {
Self::new(crate::graph::search::GraphSearchProfile::SiftClass)
}
}
pub(crate) fn auto_graph_search_options(
snapshot: &PublishedSnapshot,
) -> Result<GraphSearchOptions, QueryError> {
snapshot
.graph_profile()
.map(|profile| GraphSearchOptions::new(profile.search_profile()))
.map_err(crate::graph::search::GraphSearchError::Profile)
.map_err(QueryError::Graph)
}
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
pub enum SearchTier {
#[default]
Auto,
Exact,
Scan,
Graph(GraphSearchOptions),
}
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
pub struct SearchOptions {
scan: crate::scan::ScanOptions,
tier: Option<SearchTier>,
scan_rescore: Option<ScanRescoreOptions>,
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum GraphBoundMode {
Shared,
#[cfg(test)]
Independent,
}
impl SearchOptions {
#[must_use]
pub const fn new(scan: crate::scan::ScanOptions) -> Self {
Self {
scan,
tier: None,
scan_rescore: None,
}
}
#[must_use]
pub const fn with_tier(mut self, tier: SearchTier) -> Self {
self.tier = Some(tier);
self.scan_rescore = None;
self
}
#[must_use]
pub const fn with_scan_rescore(mut self, options: ScanRescoreOptions) -> Self {
self.tier = Some(SearchTier::Scan);
self.scan_rescore = Some(options);
self
}
#[must_use]
pub const fn scan_rescore(self) -> Option<ScanRescoreOptions> {
self.scan_rescore
}
#[must_use]
pub const fn scan(self) -> crate::scan::ScanOptions {
self.scan
}
#[must_use]
pub const fn tier(self) -> SearchTier {
match self.tier {
Some(tier) => tier,
None => SearchTier::Auto,
}
}
pub(crate) const fn explicit_tier(self) -> Option<SearchTier> {
self.tier
}
}
impl From<crate::scan::ScanOptions> for SearchOptions {
fn from(scan: crate::scan::ScanOptions) -> Self {
Self::new(scan)
}
}
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
pub struct DocumentFields(u8);
impl DocumentFields {
pub const NONE: Self = Self(0);
pub const VECTOR: Self = Self(1 << 0);
pub const TEXT: Self = Self(1 << 1);
pub const METADATA: Self = Self(1 << 2);
pub const ATTRIBUTES: Self = Self(1 << 3);
pub const ALL: Self =
Self(Self::VECTOR.0 | Self::TEXT.0 | Self::METADATA.0 | Self::ATTRIBUTES.0);
#[must_use]
pub const fn contains(self, fields: Self) -> bool {
self.0 & fields.0 == fields.0
}
}
impl std::ops::BitOr for DocumentFields {
type Output = Self;
fn bitor(self, rhs: Self) -> Self::Output {
Self(self.0 | rhs.0)
}
}
#[derive(Clone, Debug, PartialEq)]
pub struct StoredDocument {
pub doc_id: crate::ingest::DocId,
pub revision: crate::ingest::Revision,
pub timestamp: i64,
pub vector: Option<Vec<f32>>,
pub text: Option<String>,
pub metadata: Option<Vec<u8>>,
pub attributes: Option<Vec<(crate::meta::ColumnId, crate::meta::PredicateValue)>>,
}
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
pub enum ScanOrder {
#[default]
Storage,
TimestampAscending,
TimestampDescending,
Attribute {
column: crate::meta::ColumnId,
direction: ScanDirection,
},
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum ScanDirection {
Ascending,
Descending,
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum DocumentScanSource {
Sealed(crate::segment::SegmentId),
Active,
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct DocumentScanCursor {
pub generation: u64,
pub order: ScanOrder,
pub source: DocumentScanSource,
pub next_row: u32,
}
#[derive(Clone, Debug)]
pub struct DocumentScanRequest<'a> {
limit: usize,
fields: DocumentFields,
order: ScanOrder,
cursor: Option<DocumentScanCursor>,
timestamp_range: Option<(i64, i64)>,
predicate: Option<&'a crate::meta::Predicate>,
control: QueryControl,
}
impl<'a> DocumentScanRequest<'a> {
#[must_use]
pub fn new(limit: usize, fields: DocumentFields, control: QueryControl) -> Self {
Self {
limit,
fields,
order: ScanOrder::Storage,
cursor: None,
timestamp_range: None,
predicate: None,
control,
}
}
#[must_use]
pub const fn with_order(mut self, order: ScanOrder) -> Self {
self.order = order;
self
}
#[must_use]
pub const fn with_cursor(mut self, cursor: DocumentScanCursor) -> Self {
self.cursor = Some(cursor);
self
}
#[must_use]
pub const fn with_timestamp_range(mut self, start: i64, end: i64) -> Self {
self.timestamp_range = Some((start, end));
self
}
#[must_use]
pub const fn with_predicate(mut self, predicate: &'a crate::meta::Predicate) -> Self {
self.predicate = Some(predicate);
self
}
}
#[derive(Clone, Debug, PartialEq)]
pub struct DocumentScanPage {
pub documents: Vec<StoredDocument>,
pub generation: u64,
pub continuation: Option<DocumentScanCursor>,
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct DocumentCount {
pub count: u64,
pub generation: u64,
}
pub const MAX_DOCUMENT_SCAN_LIMIT: usize = 1 << 20;
#[derive(Debug)]
pub enum StoreError {
CascadeCycle {
cycle: Vec<String>,
},
Io {
path: PathBuf,
source: std::io::Error,
},
NotDirectory {
path: PathBuf,
},
StoreBusy {
path: PathBuf,
},
Lock(StoreLockError),
Durability(DurabilityPolicyError),
Kernel(crate::kernels::KernelInitError),
Manifest(crate::manifest::ManifestError),
Tokenizer(crate::fts::tokenizer::TokenizerError),
EpochMismatch(crate::epoch::EpochMismatch),
EpochUndeclared,
EpochUnstamped,
SchemaMismatch {
persisted: crate::meta::Schema,
declared: crate::meta::Schema,
},
ScanStale {
cursor_generation: u64,
current_generation: u64,
},
InvalidScan {
detail: String,
},
GroupLimitExceeded {
limit: usize,
},
Segment(crate::segment::SegmentError),
DuplicateLiveDocument {
doc_id: crate::ingest::DocId,
},
GraphUnavailable {
segment_id: crate::segment::SegmentId,
},
Wal(crate::wal::WalReadError),
WalRecovery(crate::wal::WalRecoveryError),
WalRecord {
seq: crate::wal::LogSeq,
source: crate::wal::VisibleRecordError,
},
WalMutation {
seq: crate::wal::LogSeq,
op: u16,
source: crate::ingest::wal_payload::PayloadError,
},
WalRevisionOrder {
seq: crate::wal::LogSeq,
doc_id: crate::ingest::DocId,
current: crate::ingest::Revision,
attempted: crate::ingest::Revision,
},
UnsupportedWalMutation {
seq: crate::wal::LogSeq,
op: u16,
},
WalVector {
seq: crate::wal::LogSeq,
source: crate::quant::QuantError,
},
WalWrite(crate::wal::WalWriteError),
WalRetire(crate::wal::WalRetireError),
Statistics {
component: &'static str,
source: std::io::Error,
},
BudgetExceeded {
needed: u64,
budget: u64,
component: &'static str,
},
AllocationFailed {
needed: u64,
component: &'static str,
},
DimensionMismatch {
expected: usize,
actual: usize,
},
ActiveRowOverflow,
EmptyActiveSegment,
SealCancelled,
ReadOnly,
SealedTombstoneRecoveryRequired,
ForeignPreparedSegment,
GenerationOverflow,
ConcurrentActiveMutation {
expected_generation: u64,
actual_generation: u64,
},
PartitionBytesOverflow,
PurgeRecovery {
detail: String,
},
Closing,
Closed,
ReadCancelled,
BackgroundStart {
source: std::io::Error,
},
BackgroundHandshake,
BackgroundThreadPanicked,
QueryPoolStart {
source: std::io::Error,
},
QueryPoolHandshake,
QueryPoolThreadPanicked,
QueryPoolCapacity {
source: String,
},
Synchronization {
component: &'static str,
},
NativeGraphDirectory {
path: PathBuf,
},
SnapshotTarget {
path: PathBuf,
reason: SnapshotTargetReason,
},
SnapshotPin {
detail: &'static str,
},
SnapshotPurgePending,
SnapshotViewsOpen,
}
impl std::fmt::Display for StoreError {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::CascadeCycle { cycle } => {
write!(formatter, "cascade cycle: {}", cycle.join(" -> "))
}
Self::Io { path, source } => {
write!(formatter, "store I/O {}: {source}", path.display())
}
Self::NotDirectory { path } => {
write!(
formatter,
"store path is not a directory: {}",
path.display()
)
}
Self::SnapshotViewsOpen => {
formatter.write_str("close open snapshots before physical removal")
}
Self::StoreBusy { path } => {
write!(formatter, "store already has a writer: {}", path.display())
}
Self::Lock(error) => error.fmt(formatter),
Self::Durability(error) => error.fmt(formatter),
Self::Kernel(error) => error.fmt(formatter),
Self::Manifest(error) => error.fmt(formatter),
Self::Tokenizer(error) => write!(formatter, "store tokenizer: {error}"),
Self::EpochMismatch(error) => error.fmt(formatter),
Self::EpochUndeclared => {
formatter.write_str("store epoch is persisted but the caller declared none")
}
Self::EpochUnstamped => formatter
.write_str("store manifest predates epoch identity and cannot adopt a declaration"),
Self::SchemaMismatch {
persisted,
declared,
} => match persisted.additive_evolution(declared) {
Err(conflict) => write!(
formatter,
"declared store schema does not match the persisted schema: {conflict}"
),
Ok(Some(_)) => {
formatter.write_str(
"declared store schema adds attributes that a read-only open \
cannot persist:",
)?;
for added in declared
.columns()
.iter()
.filter(|column| persisted.column(column.id()).is_none())
{
write!(formatter, " '{}' (id {})", added.name(), added.id().get())?;
}
Ok(())
}
Ok(None) => write!(
formatter,
"declared store schema {declared:?} does not match persisted schema \
{persisted:?}"
),
},
Self::ScanStale {
cursor_generation,
current_generation,
} => write!(
formatter,
"scan cursor generation {cursor_generation} differs from current generation {current_generation}"
),
Self::InvalidScan { detail } => write!(formatter, "invalid document scan: {detail}"),
Self::GroupLimitExceeded { limit } => write!(
formatter,
"grouped count found more than {limit} distinct values; raise the group limit"
),
Self::Segment(error) => error.fmt(formatter),
Self::DuplicateLiveDocument { doc_id } => {
write!(
formatter,
"document {} has multiple live rows",
doc_id.get()
)
}
Self::GraphUnavailable { segment_id } => {
write!(formatter, "sealed segment {segment_id} has no graph region")
}
Self::Wal(error) => error.fmt(formatter),
Self::WalRecovery(error) => error.fmt(formatter),
Self::WalRecord { seq, source } => {
write!(
formatter,
"WAL sequence {} retained record: {source}",
seq.get()
)
}
Self::WalMutation { seq, op, source } => write!(
formatter,
"WAL sequence {} operation {op} payload: {source}",
seq.get()
),
Self::WalRevisionOrder {
seq,
doc_id,
current,
attempted,
} => write!(
formatter,
"WAL sequence {} document {} revision {} does not follow {}",
seq.get(),
doc_id.get(),
attempted.get(),
current.get()
),
Self::UnsupportedWalMutation { seq, op } => write!(
formatter,
"WAL sequence {} operation {op} has no active-segment recovery path",
seq.get()
),
Self::WalVector { seq, source } => {
write!(formatter, "WAL sequence {} vector: {source}", seq.get())
}
Self::WalWrite(error) => error.fmt(formatter),
Self::WalRetire(error) => error.fmt(formatter),
Self::Statistics { component, source } => {
write!(formatter, "store statistics {component}: {source}")
}
Self::BudgetExceeded {
needed,
budget,
component,
} => write!(
formatter,
"store {component} allocation needs {needed} bytes, budget is {budget} bytes"
),
Self::AllocationFailed { needed, component } => write!(
formatter,
"store {component} allocator rejected {needed} bytes"
),
Self::DimensionMismatch { expected, actual } => write!(
formatter,
"active vector dimension {actual} does not match {expected}"
),
Self::ActiveRowOverflow => {
formatter.write_str("active segment row or byte geometry overflow")
}
Self::EmptyActiveSegment => formatter.write_str("active segment is empty"),
Self::SealCancelled => formatter.write_str("seal was cancelled before commit"),
Self::ReadOnly => formatter.write_str("store handle is read-only"),
Self::SealedTombstoneRecoveryRequired => formatter.write_str(
"store requires writable sealed-tombstone recovery before reads are safe",
),
Self::ForeignPreparedSegment => {
formatter.write_str("prepared segment belongs to another store")
}
Self::GenerationOverflow => formatter.write_str("store snapshot generation overflow"),
Self::ConcurrentActiveMutation {
expected_generation,
actual_generation,
} => write!(
formatter,
"active generation changed from {expected_generation} to {actual_generation} while preparing a commit"
),
Self::PartitionBytesOverflow => {
formatter.write_str("partition reclaimed-byte count overflow")
}
Self::PurgeRecovery { detail } => {
write!(formatter, "physical purge recovery failed: {detail}")
}
Self::Closing => formatter.write_str("store is closing"),
Self::Closed => formatter.write_str("store is closed"),
Self::ReadCancelled => formatter.write_str("store close cancelled the admitted read"),
Self::BackgroundStart { source } => {
write!(
formatter,
"store lifecycle thread could not start: {source}"
)
}
Self::BackgroundHandshake => {
formatter.write_str("store lifecycle thread startup handshake failed")
}
Self::BackgroundThreadPanicked => {
formatter.write_str("store lifecycle thread panicked")
}
Self::QueryPoolStart { source } => {
write!(formatter, "store query worker could not start: {source}")
}
Self::QueryPoolHandshake => {
formatter.write_str("store query worker startup handshake failed")
}
Self::QueryPoolThreadPanicked => formatter.write_str("store query worker panicked"),
Self::QueryPoolCapacity { source } => {
write!(formatter, "store query worker capacity failed: {source}")
}
Self::Synchronization { component } => {
write!(
formatter,
"store lifecycle synchronization poisoned: {component}"
)
}
Self::NativeGraphDirectory { path } => write!(
formatter,
"native graph directory cannot be opened as a legacy store: {}",
path.display()
),
Self::SnapshotTarget { path, reason } => {
write!(formatter, "snapshot target {} {reason}", path.display())
}
Self::SnapshotPin { detail } => {
write!(formatter, "snapshot pin is inconsistent: {detail}")
}
Self::SnapshotPurgePending => formatter
.write_str("a physical purge is pending; await it before taking a snapshot"),
}
}
}
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub enum StoreErrorKind {
Io,
InvalidArgument,
StoreBusy,
Unsupported,
Corrupt,
BudgetExceeded,
OutOfMemory,
DimensionMismatch,
EpochMismatch,
EpochUndeclared,
EpochUnstamped,
Internal,
EmptyBatch,
Cancelled,
ReadOnly,
Closing,
Closed,
Panic,
Synchronization,
}
impl StoreError {
#[must_use]
pub fn kind(&self) -> StoreErrorKind {
match self {
Self::Io { .. }
| Self::Lock(_)
| Self::Statistics { .. }
| Self::BackgroundStart { .. }
| Self::QueryPoolStart { .. }
| Self::WalWrite(_)
| Self::WalRetire(_) => StoreErrorKind::Io,
Self::CascadeCycle { .. }
| Self::NotDirectory { .. }
| Self::Tokenizer(_)
| Self::SchemaMismatch { .. }
| Self::ScanStale { .. }
| Self::InvalidScan { .. }
| Self::NativeGraphDirectory { .. }
| Self::SnapshotTarget { .. } => StoreErrorKind::InvalidArgument,
Self::StoreBusy { .. } | Self::SnapshotViewsOpen => StoreErrorKind::StoreBusy,
Self::Durability(_)
| Self::Kernel(_)
| Self::GraphUnavailable { .. }
| Self::UnsupportedWalMutation { .. }
| Self::SnapshotPurgePending => StoreErrorKind::Unsupported,
Self::Manifest(_)
| Self::Segment(_)
| Self::DuplicateLiveDocument { .. }
| Self::Wal(_)
| Self::WalRecovery(_)
| Self::WalRecord { .. }
| Self::WalMutation { .. }
| Self::WalRevisionOrder { .. }
| Self::WalVector { .. }
| Self::PurgeRecovery { .. } => StoreErrorKind::Corrupt,
Self::BudgetExceeded { .. } | Self::GroupLimitExceeded { .. } => {
StoreErrorKind::BudgetExceeded
}
Self::AllocationFailed { .. } => StoreErrorKind::OutOfMemory,
Self::DimensionMismatch { .. } => StoreErrorKind::DimensionMismatch,
Self::EpochMismatch(_) => StoreErrorKind::EpochMismatch,
Self::EpochUndeclared => StoreErrorKind::EpochUndeclared,
Self::EpochUnstamped => StoreErrorKind::EpochUnstamped,
Self::ActiveRowOverflow
| Self::GenerationOverflow
| Self::ConcurrentActiveMutation { .. }
| Self::PartitionBytesOverflow
| Self::ForeignPreparedSegment
| Self::BackgroundHandshake
| Self::QueryPoolHandshake
| Self::QueryPoolCapacity { .. }
| Self::SnapshotPin { .. } => StoreErrorKind::Internal,
Self::EmptyActiveSegment => StoreErrorKind::EmptyBatch,
Self::SealCancelled | Self::ReadCancelled => StoreErrorKind::Cancelled,
Self::ReadOnly | Self::SealedTombstoneRecoveryRequired => StoreErrorKind::ReadOnly,
Self::Closing => StoreErrorKind::Closing,
Self::Closed => StoreErrorKind::Closed,
Self::BackgroundThreadPanicked | Self::QueryPoolThreadPanicked => StoreErrorKind::Panic,
Self::Synchronization { .. } => StoreErrorKind::Synchronization,
}
}
}
impl std::error::Error for StoreError {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
match self {
Self::Io { source, .. } => Some(source),
Self::Lock(error) => Some(error),
Self::Durability(error) => Some(error),
Self::Kernel(error) => Some(error),
Self::Manifest(error) => Some(error),
Self::Tokenizer(error) => Some(error),
Self::EpochMismatch(error) => Some(error),
Self::Segment(error) => Some(error),
Self::Wal(error) => Some(error),
Self::WalRecovery(error) => Some(error),
Self::WalRecord { source, .. } => Some(source),
Self::WalMutation { source, .. } => Some(source),
Self::WalVector { source, .. } => Some(source),
Self::WalWrite(error) => Some(error),
Self::WalRetire(error) => Some(error),
Self::Statistics { source, .. } => Some(source),
Self::BackgroundStart { source } => Some(source),
Self::QueryPoolStart { source } => Some(source),
Self::CascadeCycle { .. }
| Self::NotDirectory { .. }
| Self::StoreBusy { .. }
| Self::EpochUndeclared
| Self::EpochUnstamped
| Self::SchemaMismatch { .. }
| Self::ScanStale { .. }
| Self::InvalidScan { .. }
| Self::GraphUnavailable { .. }
| Self::DuplicateLiveDocument { .. }
| Self::WalRevisionOrder { .. }
| Self::UnsupportedWalMutation { .. }
| Self::BudgetExceeded { .. }
| Self::GroupLimitExceeded { .. }
| Self::AllocationFailed { .. }
| Self::DimensionMismatch { .. }
| Self::ActiveRowOverflow
| Self::EmptyActiveSegment
| Self::SealCancelled
| Self::ReadOnly
| Self::SealedTombstoneRecoveryRequired
| Self::ForeignPreparedSegment
| Self::GenerationOverflow
| Self::ConcurrentActiveMutation { .. }
| Self::PartitionBytesOverflow
| Self::PurgeRecovery { .. }
| Self::Closing
| Self::Closed
| Self::ReadCancelled
| Self::BackgroundHandshake
| Self::BackgroundThreadPanicked
| Self::QueryPoolHandshake
| Self::QueryPoolThreadPanicked
| Self::QueryPoolCapacity { .. }
| Self::Synchronization { .. }
| Self::NativeGraphDirectory { .. }
| Self::SnapshotTarget { .. }
| Self::SnapshotPin { .. }
| Self::SnapshotPurgePending
| Self::SnapshotViewsOpen => None,
}
}
}
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
pub struct OpenMigrations {
pub schema_added: bool,
pub wal_tail_cut: bool,
pub generation: u64,
}
pub struct Store {
open_migrations: OpenMigrations,
pub(crate) directory: PathBuf,
pub(crate) vfs: Arc<dyn crate::vfs::Vfs>,
pub(crate) clock: Arc<dyn MonotonicClock>,
pub(crate) state: Mutex<StoreState>,
pub(crate) state_changed: Condvar,
pub(crate) background: Mutex<Option<BackgroundThread>>,
pub(crate) query_pool: Mutex<Option<Arc<pool::QueryPool>>>,
pub(crate) lexical_worker: Mutex<Option<Arc<pool::LexicalWorker>>>,
#[cfg(feature = "graph-cypher")]
pub(crate) native_graph: Arc<native_graph::NativeGraphPublication>,
pub(crate) snapshot: RwLock<Option<Arc<PublishedSnapshot>>>,
pub(crate) active: Mutex<Option<crate::ingest::ActiveState>>,
pub(crate) wal_writer: Mutex<Option<crate::ingest::StoreWal>>,
pub(crate) writer_lock: Mutex<Option<Arc<StoreLock>>>,
snapshot_pins: Arc<AtomicU64>,
snapshot_pin: Mutex<Option<snapshot_view::SnapshotPin>>,
logical_writer_lock: Mutex<Option<StoreLock>>,
reclamation_pin: Mutex<Option<Arc<StoreLock>>>,
pub(crate) maintenance: Mutex<()>,
pub(crate) health_state: Mutex<crate::diag::HealthState>,
pub(crate) durability_policy: DurabilityPolicy,
pub(crate) reader_drain_timeout: Duration,
pub(crate) accounting: Arc<stats::Accounting>,
pub(crate) active_queries: AtomicU64,
#[cfg(any(test, feature = "test-seams"))]
text_materialization_work: materialize::TestMaterializationWork,
pub(crate) epoch: Option<crate::epoch::StoreEpoch>,
pub(crate) epoch_alias: crate::epoch::EpochAliasCell,
pub(crate) tokenizer: crate::fts::tokenizer::Analyzer,
pub(crate) schema: crate::meta::Schema,
lexical_index_cache: LexicalIndexCache,
#[cfg(any(test, feature = "test-seams"))]
pub(crate) ingest_retention_fault_controller:
Option<crate::ingest::IngestRetentionFaultController>,
#[cfg(any(test, feature = "test-seams"))]
hybrid_leg_fault: Mutex<Option<HybridLegTestFault>>,
#[cfg(any(test, feature = "test-seams"))]
hybrid_execution_receipt: Mutex<Option<HybridExecutionReceipt>>,
#[cfg(any(test, feature = "test-seams"))]
pub(crate) metadata_test_controller: Option<Arc<crate::planner::MetadataTestController>>,
#[cfg(any(test, feature = "test-seams"))]
pub(crate) vector_fault_controller: Option<crate::scan::vector_fault::VectorFaultController>,
#[cfg(any(test, feature = "test-seams"))]
pub(crate) kernel_fault_controller: Option<crate::kernels::vector_fault::KernelFaultController>,
#[cfg(any(test, feature = "test-seams"))]
pub(crate) vector_seal_scheme: Option<crate::quant::QuantScheme>,
#[cfg(test)]
pub(crate) teardown_probe: Arc<close::TeardownProbe>,
}
fn acquire_writer_lock(
path: &Path,
access_mode: AccessMode,
) -> Result<Option<StoreLock>, StoreError> {
Ok(match access_mode {
AccessMode::ReadWrite => Some(StoreLock::acquire(path).map_err(|error| match error {
StoreLockError::Io { path, source }
if source.kind() == std::io::ErrorKind::WouldBlock =>
{
StoreError::StoreBusy { path }
}
StoreLockError::Io { path, source } => {
StoreError::Lock(StoreLockError::Io { path, source })
}
})?),
AccessMode::ReadOnly => None,
})
}
fn refuse_native_graph_directory(vfs: &dyn crate::vfs::Vfs, path: &Path) -> Result<(), StoreError> {
match vfs.ensure_directory(path, false) {
Ok(true) => {}
Ok(false) => return Ok(()),
Err(source) if source.kind() == std::io::ErrorKind::NotFound => return Ok(()),
Err(source) => {
return Err(StoreError::Io {
path: path.to_path_buf(),
source,
});
}
}
let mut native = false;
vfs.for_each_direct_child(path, &mut |child| {
let Some(name) = child.file_name().and_then(|name| name.to_str()) else {
return Ok(());
};
if name == "graph-root.ze"
|| (name.starts_with("graph-wal-") && name.ends_with(".ze"))
|| (name.starts_with("graph-") && name.ends_with(".zgraph"))
|| name.starts_with("graph-root.ze.tmp-")
{
native = true;
}
Ok(())
})
.map_err(|source| StoreError::Io {
path: path.to_path_buf(),
source,
})?;
if native {
return Err(StoreError::NativeGraphDirectory {
path: path.to_path_buf(),
});
}
Ok(())
}
#[cfg(feature = "graph-cypher")]
pub(crate) fn is_legacy_store_directory(
vfs: &dyn crate::vfs::Vfs,
path: &Path,
) -> Result<bool, StoreError> {
for name in [
crate::manifest::io::MANIFEST_FILE,
crate::manifest::io::MANIFEST_TEMP_FILE,
"wal.ze",
] {
let candidate = path.join(name);
match vfs.open(&candidate) {
Ok(_) => return Ok(true),
Err(source)
if matches!(
source.kind(),
std::io::ErrorKind::NotFound | std::io::ErrorKind::NotADirectory
) => {}
Err(source) => {
return Err(StoreError::Io {
path: candidate,
source,
});
}
}
}
Ok(false)
}
#[cfg(feature = "graph-cypher")]
fn acquire_native_graph_lock(
path: &Path,
access_mode: AccessMode,
existing: bool,
) -> Result<Option<StoreLock>, StoreError> {
let acquired = match access_mode {
AccessMode::ReadWrite if existing => StoreLock::acquire_existing(path),
AccessMode::ReadWrite => StoreLock::acquire(path),
AccessMode::ReadOnly => StoreLock::acquire_shared(path),
};
acquired.map(Some).map_err(|error| match error {
StoreLockError::Io { path, source } if source.kind() == std::io::ErrorKind::WouldBlock => {
StoreError::StoreBusy { path }
}
StoreLockError::Io { path, source } => {
StoreError::Lock(StoreLockError::Io { path, source })
}
})
}
fn resolve_open_schema(
manifest_exists: bool,
snapshot: &PublishedSnapshot,
declared: Option<&crate::meta::Schema>,
access_mode: AccessMode,
) -> Result<(crate::meta::Schema, bool), StoreError> {
if !manifest_exists {
return Ok((
declared
.cloned()
.unwrap_or_else(crate::meta::Schema::timestamp_only),
false,
));
}
let persisted = snapshot.schema();
let Some(declared) = declared else {
return Ok((persisted.clone(), false));
};
let mismatch = || StoreError::SchemaMismatch {
persisted: persisted.clone(),
declared: declared.clone(),
};
match persisted.additive_evolution(declared) {
Ok(None) => Ok((persisted.clone(), false)),
Ok(Some(evolved)) if access_mode == AccessMode::ReadWrite => Ok((evolved, true)),
Ok(Some(_)) | Err(_) => Err(mismatch()),
}
}
fn validate_persisted_epoch_before_open(
vfs: &dyn crate::vfs::Vfs,
path: &Path,
options: &OpenOptions,
) -> Result<(), StoreError> {
let manifest_path = path.join(crate::manifest::io::MANIFEST_FILE);
let bytes = match vfs.read(&manifest_path) {
Ok(bytes) => bytes,
Err(source) if source.kind() == std::io::ErrorKind::NotFound => return Ok(()),
Err(source) => {
return Err(StoreError::Io {
path: manifest_path,
source,
});
}
};
let manifest = crate::manifest::decode_manifest(&manifest_path.display().to_string(), &bytes)
.map_err(StoreError::Manifest)?;
let declared = if options.inspection {
manifest.epoch_alias
} else {
options
.epoch
.as_ref()
.map(crate::epoch::StoreEpoch::identity)
};
validate_epoch_identity(manifest.epoch_alias, declared, true, options.access_mode)
}
fn validate_epoch_identity(
persisted: Option<crate::epoch::EpochIdentity>,
declared: Option<crate::epoch::EpochIdentity>,
manifest_exists: bool,
access_mode: AccessMode,
) -> Result<(), StoreError> {
match (persisted, declared) {
(Some(expected), Some(declared)) if expected != declared => {
Err(StoreError::EpochMismatch(crate::epoch::EpochMismatch {
expected,
declared,
}))
}
(Some(_), None) => Err(StoreError::EpochUndeclared),
(None, Some(_)) if manifest_exists || access_mode == AccessMode::ReadOnly => {
Err(StoreError::EpochUnstamped)
}
(Some(_), Some(_)) | (None, Some(_)) | (None, None) => Ok(()),
}
}
fn validate_tokenizer_epoch(
expected: Option<crate::epoch::EpochIdentity>,
tokenizer: &crate::fts::tokenizer::Analyzer,
) -> Result<(), StoreError> {
if let Some(expected) = expected
&& tokenizer.epoch() != expected.tokenizer
{
return Err(StoreError::EpochMismatch(crate::epoch::EpochMismatch {
expected,
declared: crate::epoch::EpochIdentity {
embedding: expected.embedding,
tokenizer: tokenizer.epoch(),
},
}));
}
Ok(())
}
fn cleanup_open_orphans(
store: &Store,
) -> Result<crate::manifest::io::OrphanCleanupReport, StoreError> {
let reachable_segments = store
.snapshot
.read()
.map_err(|_| StoreError::Synchronization {
component: "published snapshot",
})?
.as_ref()
.map(|snapshot| {
snapshot
.all_segments()
.iter()
.map(|segment| store.directory.join(segment.meta().id.file_name()))
.collect::<HashSet<_>>()
})
.ok_or(StoreError::Closed)?;
crate::manifest::io::cleanup_store_orphans(
store.vfs.as_ref(),
&store.directory,
&reachable_segments,
store.durability_policy,
)
.map_err(StoreError::Manifest)
}
impl Store {
#[cfg(feature = "graph-cypher")]
pub(crate) fn new_native_graph_owner(
path: &Path,
options: OpenOptions,
vfs: Arc<dyn crate::vfs::Vfs>,
clock: Arc<dyn MonotonicClock>,
) -> Result<Self, StoreError> {
Self::new_native_graph_owner_inner(path, options, vfs, clock, false)
}
#[cfg(feature = "graph-cypher")]
pub(crate) fn new_native_graph_recovery_owner(
path: &Path,
options: OpenOptions,
vfs: Arc<dyn crate::vfs::Vfs>,
clock: Arc<dyn MonotonicClock>,
) -> Result<Self, StoreError> {
Self::new_native_graph_owner_inner(path, options, vfs, clock, true)
}
#[cfg(feature = "graph-cypher")]
fn new_native_graph_owner_inner(
path: &Path,
options: OpenOptions,
vfs: Arc<dyn crate::vfs::Vfs>,
clock: Arc<dyn MonotonicClock>,
existing: bool,
) -> Result<Self, StoreError> {
crate::kernels::initialize().map_err(StoreError::Kernel)?;
let durability_policy = DurabilityPolicy::new(options.durability_mode, options.commit_tier)
.map_err(StoreError::Durability)?;
let accounting = Arc::new(stats::Accounting::new(
options.max_resident_bytes,
options.max_temp_bytes,
));
let writer_lock = acquire_native_graph_lock(path, options.access_mode, existing)?;
let tokenizer = crate::fts::tokenizer::Analyzer::new(
options
.tokenizer
.clone()
.unwrap_or_else(crate::fts::tokenizer::TokenizerConfig::text_default),
)
.map_err(StoreError::Tokenizer)?;
let native_graph = native_graph::NativeGraphPublication::new(&accounting)?;
#[cfg(test)]
let teardown_probe = Arc::new(close::TeardownProbe::new());
Ok(Self {
open_migrations: OpenMigrations::default(),
directory: path.to_path_buf(),
vfs,
clock,
state: Mutex::new(StoreState::Open),
state_changed: Condvar::new(),
background: Mutex::new(None),
query_pool: Mutex::new(None),
lexical_worker: Mutex::new(None),
native_graph,
snapshot: RwLock::new(None),
active: Mutex::new(None),
wal_writer: Mutex::new(None),
writer_lock: Mutex::new(writer_lock.map(Arc::new)),
snapshot_pins: Arc::new(AtomicU64::new(0)),
snapshot_pin: Mutex::new(None),
logical_writer_lock: Mutex::new(None),
reclamation_pin: Mutex::new(None),
maintenance: Mutex::new(()),
health_state: Mutex::new(crate::diag::HealthState::default()),
durability_policy,
reader_drain_timeout: options.reader_drain_timeout,
accounting,
active_queries: AtomicU64::new(0),
#[cfg(any(test, feature = "test-seams"))]
text_materialization_work: materialize::TestMaterializationWork::default(),
epoch: options.epoch,
epoch_alias: crate::epoch::EpochAliasCell::new(None),
tokenizer,
schema: options
.schema
.unwrap_or_else(crate::meta::Schema::timestamp_only),
lexical_index_cache: LexicalIndexCache::new(),
#[cfg(any(test, feature = "test-seams"))]
ingest_retention_fault_controller: None,
#[cfg(any(test, feature = "test-seams"))]
hybrid_leg_fault: Mutex::new(None),
#[cfg(any(test, feature = "test-seams"))]
hybrid_execution_receipt: Mutex::new(None),
#[cfg(any(test, feature = "test-seams"))]
metadata_test_controller: None,
#[cfg(any(test, feature = "test-seams"))]
vector_fault_controller: None,
#[cfg(any(test, feature = "test-seams"))]
kernel_fault_controller: None,
#[cfg(any(test, feature = "test-seams"))]
vector_seal_scheme: None,
#[cfg(test)]
teardown_probe,
})
}
pub fn open_for_inspection(path: impl AsRef<Path>) -> Result<Self, StoreError> {
let mut options = OpenOptions::read_only();
options.inspection = true;
Self::open(path, options)
}
pub fn open(path: impl AsRef<Path>, options: OpenOptions) -> Result<Self, StoreError> {
let path = path.as_ref();
if options.access_mode == AccessMode::ReadWrite {
namespace_batch::reclaim_for_open(path)?;
}
let _admission =
namespace_batch::reader_admission(path, options.access_mode == AccessMode::ReadWrite)?;
let selected = namespace_batch::resolve(path)?;
let logical_lock = if selected != path {
acquire_writer_lock(path, options.access_mode)?
} else {
None
};
if logical_lock.is_some() && namespace_batch::resolve(path)? != selected {
return Err(StoreError::StoreBusy {
path: path.to_path_buf(),
});
}
if logical_lock.is_some() {
if let Some(root) = path.parent() {
crate::vfs::Vfs::sync(&crate::vfs::StdVfs, root, crate::vfs::SyncKind::Full)
.map_err(|source| StoreError::Io {
path: root.to_path_buf(),
source,
})?;
}
}
let store = Self::open_with_infrastructure(
&selected,
options,
Arc::new(crate::vfs::StdVfs),
Arc::new(SystemMonotonicClock),
#[cfg(any(test, feature = "test-seams"))]
None,
#[cfg(any(test, feature = "test-seams"))]
None,
#[cfg(any(test, feature = "test-seams"))]
None,
#[cfg(any(test, feature = "test-seams"))]
None,
#[cfg(any(test, feature = "test-seams"))]
None,
#[cfg(any(test, feature = "test-seams"))]
None,
#[cfg(any(test, feature = "test-seams"))]
None,
)?;
if store
.writer_lock
.lock()
.map_err(|_| StoreError::Synchronization {
component: "writer lock",
})?
.is_some()
&& namespace_batch::resolve(path)? != selected
{
return Err(StoreError::StoreBusy {
path: path.to_path_buf(),
});
}
*store
.logical_writer_lock
.lock()
.map_err(|_| StoreError::Synchronization {
component: "logical writer lock",
})? = logical_lock;
Ok(store)
}
#[allow(
clippy::too_many_arguments,
reason = "test-support controllers are explicit optional infrastructure dependencies"
)]
fn open_with_infrastructure(
path: impl AsRef<Path>,
options: OpenOptions,
vfs: Arc<dyn crate::vfs::Vfs>,
clock: Arc<dyn MonotonicClock>,
#[cfg(any(test, feature = "test-seams"))] hybrid_leg_fault: Option<HybridLegTestFault>,
#[cfg(any(test, feature = "test-seams"))] storage_fault_controller: Option<
StorageFaultController,
>,
#[cfg(any(test, feature = "test-seams"))] ingest_retention_fault_controller: Option<
crate::ingest::IngestRetentionFaultController,
>,
#[cfg(any(test, feature = "test-seams"))] metadata_test_controller: Option<
Arc<crate::planner::MetadataTestController>,
>,
#[cfg(any(test, feature = "test-seams"))] vector_fault_controller: Option<
crate::scan::vector_fault::VectorFaultController,
>,
#[cfg(any(test, feature = "test-seams"))] kernel_fault_controller: Option<
crate::kernels::vector_fault::KernelFaultController,
>,
#[cfg(any(test, feature = "test-seams"))] vector_seal_scheme: Option<
crate::quant::QuantScheme,
>,
) -> Result<Self, StoreError> {
let path = path.as_ref();
#[cfg(any(test, feature = "test-seams"))]
match kernel_fault_controller.as_ref() {
Some(controller) => {
controller
.initialize_for_store()
.map_err(StoreError::Kernel)?;
}
None => {
crate::kernels::initialize().map_err(StoreError::Kernel)?;
}
}
#[cfg(not(any(test, feature = "test-seams")))]
crate::kernels::initialize().map_err(StoreError::Kernel)?;
#[cfg(any(test, feature = "test-seams"))]
let vfs = match storage_fault_controller.as_ref() {
Some(controller) => Arc::new(StorageFaultVfs {
inner: vfs,
controller: controller.clone(),
}) as Arc<dyn crate::vfs::Vfs>,
None => vfs,
};
let durability_policy = DurabilityPolicy::new(options.durability_mode, options.commit_tier)
.map_err(StoreError::Durability)?;
refuse_native_graph_directory(vfs.as_ref(), path)?;
let accounting = Arc::new(stats::Accounting::new(
options.max_resident_bytes,
options.max_temp_bytes,
));
let is_directory = vfs
.ensure_directory(path, options.access_mode == AccessMode::ReadWrite)
.map_err(|source| StoreError::Io {
path: path.to_path_buf(),
source,
})?;
if !is_directory {
return Err(StoreError::NotDirectory {
path: path.to_path_buf(),
});
}
validate_persisted_epoch_before_open(vfs.as_ref(), path, &options)?;
let _reclamation_admission =
namespace_batch::reader_admission(path, options.access_mode == AccessMode::ReadWrite)?;
let reclamation_pin =
namespace_batch::reader_lease(path, options.access_mode == AccessMode::ReadWrite)?
.map(Arc::new);
namespace_batch::refuse_retired(vfs.as_ref(), path)?;
let writer_lock = acquire_writer_lock(path, options.access_mode)?;
validate_persisted_epoch_before_open(vfs.as_ref(), path, &options)?;
let mut wal_tail_cut = false;
if options.access_mode == AccessMode::ReadWrite {
wal_tail_cut =
crate::ingest::cut_interrupted_append(vfs.as_ref(), path, durability_policy)?;
}
if options.access_mode == AccessMode::ReadWrite {
namespace_batch::adopt_for_open(vfs.as_ref(), path)?;
}
let manifest_path = path.join(crate::manifest::io::MANIFEST_FILE);
let manifest_exists = match vfs.open(&manifest_path) {
Ok(_) => true,
Err(source) if source.kind() == std::io::ErrorKind::NotFound => false,
Err(source) => {
return Err(StoreError::Io {
path: manifest_path,
source,
});
}
};
let mut snapshot =
PublishedSnapshot::load_for_open_on_vfs(path, &accounting, vfs.as_ref())?;
let (schema, schema_evolved) = resolve_open_schema(
manifest_exists,
&snapshot,
options.schema.as_ref(),
options.access_mode,
)?;
let persisted_epoch = snapshot.epoch_alias();
let declared_epoch = if options.inspection {
persisted_epoch
} else {
options
.epoch
.as_ref()
.map(crate::epoch::StoreEpoch::identity)
};
validate_epoch_identity(
persisted_epoch,
declared_epoch,
manifest_exists,
options.access_mode,
)?;
let tokenizer = crate::fts::tokenizer::Analyzer::new(
options
.tokenizer
.clone()
.unwrap_or_else(crate::fts::tokenizer::TokenizerConfig::text_default),
)
.map_err(StoreError::Tokenizer)?;
validate_tokenizer_epoch(persisted_epoch.or(declared_epoch), &tokenizer)?;
let absorbed_through = snapshot.absorbed_through();
let wal_path = path.join("wal.ze");
let (mut active, recovered_wal, sealed_tombstones) = crate::ingest::ActiveState::recover(
vfs.as_ref(),
&wal_path,
snapshot.generation(),
absorbed_through,
&accounting,
&schema,
&tokenizer,
)?;
if options.access_mode == AccessMode::ReadWrite && manifest_exists {
match durability_policy.directory_sync() {
SyncRequirement::Skip => {}
SyncRequirement::Sync(kind) => {
vfs.sync(path, kind).map_err(|source| StoreError::Io {
path: path.to_path_buf(),
source,
})?;
}
}
}
if schema_evolved {
let generation = active
.generation
.checked_add(1)
.ok_or(StoreError::GenerationOverflow)?;
let committed =
crate::manifest::io::load_manifest(vfs.as_ref(), &manifest_path, absorbed_through)
.map_err(StoreError::Manifest)?;
crate::manifest::io::commit_manifest(
vfs.as_ref(),
path,
&crate::manifest::Manifest {
generation,
schema: schema.clone(),
..committed
},
durability_policy,
)
.map_err(StoreError::Manifest)?;
active.generation = generation;
snapshot = PublishedSnapshot::load_on_vfs(path, &accounting, vfs.as_ref())?;
}
if options.access_mode == AccessMode::ReadWrite
&& !manifest_exists
&& (options.epoch.is_some() || options.schema.is_some())
{
crate::manifest::io::commit_manifest(
vfs.as_ref(),
path,
&crate::manifest::Manifest {
generation: active.generation,
log_seq: 0,
segments: Vec::new(),
epochs: options
.epoch
.as_ref()
.cloned()
.map(crate::manifest::EpochMeta::from)
.into_iter()
.collect(),
epoch_alias: options
.epoch
.as_ref()
.map(crate::epoch::StoreEpoch::identity),
schema: schema.clone(),
},
durability_policy,
)
.map_err(StoreError::Manifest)?;
}
let wal_writer = match options.access_mode {
AccessMode::ReadWrite => Some(match recovered_wal {
Some(recovered) => crate::ingest::StoreWal::resume(
vfs.as_ref(),
&wal_path,
recovered,
durability_policy,
absorbed_through,
&accounting,
)?,
None => crate::ingest::StoreWal::create(
Arc::clone(&vfs),
path,
&wal_path,
durability_policy,
&accounting,
)?,
}),
AccessMode::ReadOnly => {
drop(recovered_wal);
None
}
};
let background = match options.access_mode {
AccessMode::ReadWrite => Some(BackgroundThread::start()?),
AccessMode::ReadOnly => None,
};
#[cfg(feature = "graph-cypher")]
let native_graph = native_graph::NativeGraphPublication::new(&accounting)?;
#[cfg(test)]
let (snapshot, background, teardown_probe) = {
let mut snapshot = snapshot;
let mut background = background;
let teardown_probe = Arc::new(close::TeardownProbe::new());
snapshot.set_teardown_probe(Arc::clone(&teardown_probe));
if let Some(background) = background.as_mut() {
background.set_teardown_probe(Arc::clone(&teardown_probe));
}
(snapshot, background, teardown_probe)
};
let mut store = Self {
open_migrations: OpenMigrations {
schema_added: schema_evolved,
wal_tail_cut,
generation: active.generation,
},
directory: path.to_path_buf(),
vfs,
clock,
state: Mutex::new(StoreState::Open),
state_changed: Condvar::new(),
background: Mutex::new(background),
query_pool: Mutex::new(None),
lexical_worker: Mutex::new(None),
#[cfg(feature = "graph-cypher")]
native_graph,
snapshot: RwLock::new(None),
active: Mutex::new(Some(active)),
wal_writer: Mutex::new(wal_writer),
writer_lock: Mutex::new(writer_lock.map(Arc::new)),
snapshot_pins: Arc::new(AtomicU64::new(0)),
snapshot_pin: Mutex::new(None),
logical_writer_lock: Mutex::new(None),
reclamation_pin: Mutex::new(reclamation_pin),
maintenance: Mutex::new(()),
health_state: Mutex::new(crate::diag::HealthState::default()),
durability_policy,
reader_drain_timeout: options.reader_drain_timeout,
accounting,
active_queries: AtomicU64::new(0),
#[cfg(any(test, feature = "test-seams"))]
text_materialization_work: materialize::TestMaterializationWork::default(),
epoch_alias: crate::epoch::EpochAliasCell::new(persisted_epoch.or(declared_epoch)),
epoch: options.epoch,
tokenizer,
schema,
lexical_index_cache: LexicalIndexCache::new(),
#[cfg(any(test, feature = "test-seams"))]
ingest_retention_fault_controller,
#[cfg(any(test, feature = "test-seams"))]
hybrid_leg_fault: Mutex::new(hybrid_leg_fault),
#[cfg(any(test, feature = "test-seams"))]
hybrid_execution_receipt: Mutex::new(None),
#[cfg(any(test, feature = "test-seams"))]
metadata_test_controller,
#[cfg(any(test, feature = "test-seams"))]
vector_fault_controller,
#[cfg(any(test, feature = "test-seams"))]
kernel_fault_controller,
#[cfg(any(test, feature = "test-seams"))]
vector_seal_scheme,
#[cfg(test)]
teardown_probe,
};
store.publish_snapshot(snapshot)?;
store
.recover_pending_physical_purge()
.map_err(|error| StoreError::PurgeRecovery {
detail: error.to_string(),
})?;
store.recover_sealed_tombstones(&sealed_tombstones)?;
if options.access_mode == AccessMode::ReadWrite {
let cleanup_report = cleanup_open_orphans(&store)?;
#[cfg(any(test, feature = "test-seams"))]
if let Some(controller) = storage_fault_controller.as_ref() {
controller.record_cleanup_report(cleanup_report);
}
#[cfg(not(any(test, feature = "test-seams")))]
let _ = cleanup_report;
}
#[cfg(any(test, feature = "test-seams"))]
if let Some(controller) = storage_fault_controller.as_ref() {
let published = store
.snapshot
.read()
.map_err(|_| StoreError::Synchronization {
component: "published snapshot",
})?;
let snapshot = published.as_ref().ok_or(StoreError::Closed)?;
for segment in snapshot.all_segments() {
register_storage_segment_controller(
&store.directory,
segment.meta().id,
controller,
);
}
}
store.open_migrations.generation = store
.active
.lock()
.map_err(|_| StoreError::Synchronization {
component: "active segment",
})?
.as_ref()
.ok_or(StoreError::Closed)?
.generation;
Ok(store)
}
#[cfg(any(test, feature = "test-seams"))]
#[doc(hidden)]
pub fn open_with_test_dependencies(
path: impl AsRef<Path>,
options: OpenOptions,
dependencies: StoreTestDependencies,
) -> Result<Self, StoreError> {
let controller = dependencies.storage_fault_controller.clone();
with_storage_open_controller(controller.as_ref(), || {
Self::open_with_infrastructure(
path,
options,
dependencies.vfs,
dependencies.clock,
dependencies.hybrid_leg_fault,
dependencies.storage_fault_controller,
dependencies.ingest_retention_fault_controller,
dependencies.metadata_test_controller,
dependencies.vector_fault_controller,
dependencies.kernel_fault_controller,
dependencies.vector_seal_scheme,
)
})
}
#[cfg(any(test, feature = "test-seams"))]
#[doc(hidden)]
pub fn take_hybrid_execution_receipt(&self) -> Option<HybridExecutionReceipt> {
self.hybrid_execution_receipt
.lock()
.ok()
.and_then(|mut receipt| receipt.take())
}
pub(crate) fn epoch_registry(
&self,
prior: &[crate::manifest::EpochMeta],
) -> Vec<crate::manifest::EpochMeta> {
if prior.is_empty() {
self.epoch
.as_ref()
.map(crate::manifest::EpochMeta::from)
.into_iter()
.collect()
} else {
prior.to_vec()
}
}
pub fn open_migrations(&self) -> OpenMigrations {
self.open_migrations
}
#[must_use]
pub fn epoch_identity(&self) -> Option<crate::epoch::EpochIdentity> {
self.epoch_alias.load()
}
#[must_use]
pub const fn schema(&self) -> &crate::meta::Schema {
&self.schema
}
pub fn state(&self) -> Result<StoreState, StoreError> {
self.state
.lock()
.map(|state| *state)
.map_err(|_| StoreError::Synchronization { component: "state" })
}
pub fn snapshot(&self) -> Result<SnapshotLease, StoreError> {
let state = self
.state
.lock()
.map_err(|_| StoreError::Synchronization { component: "state" })?;
match *state {
StoreState::Open => {}
StoreState::Closing => return Err(StoreError::Closing),
StoreState::Closed => return Err(StoreError::Closed),
}
let active = self
.active
.lock()
.map_err(|_| StoreError::Synchronization {
component: "active segment",
})?;
let generation = active.as_ref().ok_or(StoreError::Closed)?.generation;
let snapshot = self
.snapshot
.read()
.map_err(|_| StoreError::Synchronization {
component: "published snapshot",
})?
.as_ref()
.cloned()
.ok_or(StoreError::Closed)?;
drop(active);
drop(state);
Ok(SnapshotLease::new_at(snapshot, generation))
}
#[cfg(any(test, feature = "test-seams"))]
#[doc(hidden)]
pub fn test_metadata_row_values(
&self,
row_id: crate::ingest::GlobalRowId,
) -> Result<Vec<(crate::meta::ColumnId, Option<crate::meta::PredicateValue>)>, StoreError> {
match row_id.source() {
crate::ingest::RowSource::Active => {
let row = usize::try_from(row_id.local_row())
.map_err(|_| StoreError::ActiveRowOverflow)?;
let active = self
.active
.lock()
.map_err(|_| StoreError::Synchronization {
component: "active segment",
})?;
let segment = &active.as_ref().ok_or(StoreError::Closed)?.segment;
let timestamp = segment
.timestamps()
.get(row)
.copied()
.ok_or(StoreError::ActiveRowOverflow)?;
let mut values = segment.column_values(row)?;
values.push((
crate::meta::TIMESTAMP_COLUMN,
crate::meta::PredicateValue::I64(timestamp),
));
self.schema
.columns()
.iter()
.map(|definition| {
let value = values
.iter()
.find(|(column, _)| *column == definition.id())
.map(|(_, value)| value.clone());
Ok((definition.id(), value))
})
.collect()
}
crate::ingest::RowSource::Sealed(segment_id) => {
let snapshot = self.snapshot()?;
let segment = snapshot
.segments()
.iter()
.find(|segment| segment.meta().id == segment_id)
.ok_or_else(|| {
StoreError::Segment(crate::segment::SegmentError::Geometry(format!(
"public metadata row names absent segment {segment_id}"
)))
})?;
let columns = segment.columns().map_err(StoreError::Segment)?;
if row_id.local_row() >= columns.row_count() {
return Err(StoreError::Segment(crate::segment::SegmentError::Geometry(
format!(
"public metadata row {} is outside {} rows",
row_id.local_row(),
columns.row_count()
),
)));
}
columns
.schema()
.columns()
.iter()
.map(|definition| {
let column = columns.column(definition.id()).ok_or_else(|| {
StoreError::Segment(crate::segment::SegmentError::Geometry(format!(
"public metadata row is missing column {}",
definition.id().get()
)))
})?;
let value = match column {
crate::meta::Column::Id128(column) => column
.get(row_id.local_row())
.map(crate::meta::PredicateValue::Id128),
crate::meta::Column::U64(column) => column
.get(row_id.local_row())
.map(crate::meta::PredicateValue::U64),
crate::meta::Column::I64(column) => column
.get(row_id.local_row())
.map(crate::meta::PredicateValue::I64),
crate::meta::Column::F64(column) => column
.get(row_id.local_row())
.map(crate::meta::PredicateValue::F64),
crate::meta::Column::Bool(column) => column
.get(row_id.local_row())
.map(crate::meta::PredicateValue::Bool),
crate::meta::Column::DictionaryString(column) => column
.get(row_id.local_row())
.map(|value| crate::meta::PredicateValue::String(value.to_owned())),
crate::meta::Column::RawString(column) => column
.get(row_id.local_row())
.map(|value| crate::meta::PredicateValue::String(value.to_owned())),
};
Ok((definition.id(), value))
})
.collect()
}
}
}
#[cfg(any(test, feature = "test-seams"))]
#[doc(hidden)]
pub fn test_metadata_evaluate_source(
&self,
source: crate::ingest::RowSource,
predicate: &crate::meta::Predicate,
) -> Result<(u32, Vec<u32>, Vec<u32>), StoreError> {
let evaluate = |columns: &crate::meta::ColumnStore,
alive: &crate::meta::AliveSet|
-> Result<(u32, Vec<u32>, Vec<u32>), StoreError> {
let matches = crate::meta::evaluate(predicate, columns, alive).map_err(|error| {
StoreError::Segment(crate::segment::SegmentError::Columns(format!(
"test metadata source evaluation: {error}"
)))
})?;
Ok((
columns.row_count(),
alive.iter_alive().collect(),
matches.iter().collect(),
))
};
match source {
crate::ingest::RowSource::Active => {
let active = self
.active
.lock()
.map_err(|_| StoreError::Synchronization {
component: "active segment",
})?;
let segment = &active.as_ref().ok_or(StoreError::Closed)?.segment;
let mut builder = crate::meta::ColumnStoreBuilder::new(self.schema.clone());
for (row, timestamp) in segment.timestamps().iter().copied().enumerate() {
let values = segment.column_values(row)?;
builder
.push_row(timestamp, &crate::ingest::column_inputs(&values))
.map_err(|error| {
StoreError::Segment(crate::segment::SegmentError::Columns(
error.to_string(),
))
})?;
}
let columns = builder.finish().map_err(|error| {
StoreError::Segment(crate::segment::SegmentError::Columns(error.to_string()))
})?;
let alive = segment.alive()?;
evaluate(&columns, &alive)
}
crate::ingest::RowSource::Sealed(segment_id) => {
let snapshot = self.snapshot()?;
let segment = snapshot
.segments()
.iter()
.find(|segment| segment.meta().id == segment_id)
.ok_or_else(|| {
StoreError::Segment(crate::segment::SegmentError::Geometry(format!(
"metadata source evaluation names absent segment {segment_id}"
)))
})?;
let columns = segment.columns().map_err(StoreError::Segment)?;
let alive = segment.alive().map_err(StoreError::Segment)?;
evaluate(&columns, &alive)
}
}
}
pub(crate) fn publish_snapshot(&self, snapshot: PublishedSnapshot) -> Result<(), StoreError> {
let state = self
.state
.lock()
.map_err(|_| StoreError::Synchronization { component: "state" })?;
match *state {
StoreState::Open => {}
StoreState::Closing => return Err(StoreError::Closing),
StoreState::Closed => return Err(StoreError::Closed),
}
let generation = snapshot.generation();
let mut active = self
.active
.lock()
.map_err(|_| StoreError::Synchronization {
component: "active segment",
})?;
let active_generation = &mut active.as_mut().ok_or(StoreError::Closed)?.generation;
*active_generation = (*active_generation).max(generation);
let mut published = self
.snapshot
.write()
.map_err(|_| StoreError::Synchronization {
component: "published snapshot",
})?;
*published = Some(Arc::new(snapshot));
drop(active);
drop(state);
Ok(())
}
pub fn top_k_with_options(
&self,
request: crate::scan::ScanRequest<'_>,
k: usize,
options: crate::scan::ScanOptions,
control: QueryControl,
) -> Result<crate::scan::ScanOutcome, QueryError> {
let control = control.with_clock(Arc::clone(&self.clock));
let state = self
.state
.lock()
.map_err(|_| QueryError::Store(StoreError::Synchronization { component: "state" }))?;
match *state {
StoreState::Open => {}
StoreState::Closing => return Err(QueryError::Store(StoreError::Closing)),
StoreState::Closed => return Err(QueryError::Store(StoreError::Closed)),
}
let snapshot = self
.snapshot
.read()
.map_err(|_| {
QueryError::Store(StoreError::Synchronization {
component: "published snapshot",
})
})?
.as_ref()
.cloned()
.ok_or(QueryError::Store(StoreError::Closed))?;
let mut pool_slot = self.query_pool.lock().map_err(|_| {
QueryError::Store(StoreError::Synchronization {
component: "query pool",
})
})?;
if pool_slot.is_none() {
let capacity = crate::scan::physical_thread_capacity().map_err(|error| {
QueryError::Store(StoreError::QueryPoolCapacity {
source: error.to_string(),
})
})?;
let pool =
pool::QueryPool::start(capacity, &self.accounting).map_err(QueryError::Store)?;
*pool_slot = Some(Arc::new(pool));
}
let pool = pool_slot.as_ref().cloned().ok_or({
QueryError::Store(StoreError::Synchronization {
component: "query pool initialization",
})
})?;
self.active_queries.fetch_add(1, Ordering::Relaxed);
let active = ActiveQuery {
count: &self.active_queries,
};
drop(pool_slot);
drop(state);
let result = pool
.execute(request, k, options, control, SnapshotLease::new(snapshot))
.map(|mut outcome| {
outcome.candidates.truncate(k);
outcome
});
drop(active);
result
}
#[cfg(any(test, feature = "test-seams"))]
pub fn panic_query_worker_for_test(&self) -> Result<(), StoreError> {
let pool = self
.query_pool
.lock()
.map_err(|_| StoreError::Synchronization {
component: "query pool",
})?
.as_ref()
.cloned()
.ok_or(StoreError::Synchronization {
component: "query pool initialization",
})?;
pool.panic_one_and_join()
}
pub fn search(
&self,
request: crate::ingest::SearchRequest<'_>,
k: usize,
options: impl Into<SearchOptions>,
control: QueryControl,
) -> Result<crate::ingest::SearchOutcome, QueryError> {
self.search_with_graph_bound_mode(
request,
k,
options.into(),
control,
GraphBoundMode::Shared,
)
}
fn admit_lexical_query(&self) -> Result<AdmittedLexicalQuery<'_>, QueryError> {
#[cfg(any(test, feature = "test-seams"))]
self.text_materialization_work
.admissions
.fetch_add(1, Ordering::Relaxed);
#[cfg(all(feature = "query-timing", any(test, feature = "test-seams")))]
let lock_started = std::time::Instant::now();
let state = self
.state
.lock()
.map_err(|_| QueryError::Store(StoreError::Synchronization { component: "state" }))?;
#[cfg(all(feature = "query-timing", any(test, feature = "test-seams")))]
self.text_materialization_work
.observe_lexical_lock(lock_started);
match *state {
StoreState::Open => {}
StoreState::Closing => return Err(QueryError::Store(StoreError::Closing)),
StoreState::Closed => return Err(QueryError::Store(StoreError::Closed)),
}
#[cfg(all(feature = "query-timing", any(test, feature = "test-seams")))]
let lock_started = std::time::Instant::now();
let active_guard = self.active.lock().map_err(|_| {
QueryError::Store(StoreError::Synchronization {
component: "active segment",
})
})?;
#[cfg(all(feature = "query-timing", any(test, feature = "test-seams")))]
self.text_materialization_work
.observe_lexical_lock(lock_started);
let active_state = active_guard
.as_ref()
.ok_or(QueryError::Store(StoreError::Closed))?;
let generation = active_state.generation;
let active = Arc::clone(&active_state.segment);
#[cfg(all(feature = "query-timing", any(test, feature = "test-seams")))]
let lock_started = std::time::Instant::now();
let snapshot = self
.snapshot
.read()
.map_err(|_| {
QueryError::Store(StoreError::Synchronization {
component: "published snapshot",
})
})
.inspect(|_| {
#[cfg(all(feature = "query-timing", any(test, feature = "test-seams")))]
self.text_materialization_work
.observe_lexical_lock(lock_started);
})?
.as_ref()
.cloned()
.ok_or(QueryError::Store(StoreError::Closed))?;
self.active_queries.fetch_add(1, Ordering::Relaxed);
let active_query = ActiveQuery {
count: &self.active_queries,
};
drop(active_guard);
drop(state);
Ok(AdmittedLexicalQuery {
generation,
active,
snapshot: snapshot::ReadSnapshot::new(snapshot),
active_query,
})
}
fn admit_document_read(&self) -> Result<AdmittedDocumentRead<'_>, StoreError> {
let state = self
.state
.lock()
.map_err(|_| StoreError::Synchronization { component: "state" })?;
match *state {
StoreState::Open => {}
StoreState::Closing => return Err(StoreError::Closing),
StoreState::Closed => return Err(StoreError::Closed),
}
let active_guard = self
.active
.lock()
.map_err(|_| StoreError::Synchronization {
component: "active segment",
})?;
let active_state = active_guard.as_ref().ok_or(StoreError::Closed)?;
let generation = active_state.generation;
let active = Arc::clone(&active_state.segment);
let snapshot = self
.snapshot
.read()
.map_err(|_| StoreError::Synchronization {
component: "published snapshot",
})?
.as_ref()
.cloned()
.ok_or(StoreError::Closed)?;
self.active_queries.fetch_add(1, Ordering::Relaxed);
let active_query = ActiveQuery {
count: &self.active_queries,
};
drop(active_guard);
drop(state);
Ok(AdmittedDocumentRead {
generation,
active,
snapshot: snapshot::ReadSnapshot::new(snapshot),
active_query,
})
}
pub fn warm_lexical(
&self,
control: QueryControl,
) -> Result<(), crate::ingest::StoreLexicalError> {
let control = control.with_clock(Arc::clone(&self.clock));
let admission = self.admit_lexical_query()?;
let lease = SnapshotLease::new_at(Arc::clone(&admission.snapshot), admission.generation);
let cancellation = QueryCancellation::new(&control, &lease);
cancellation.check_graph().map_err(QueryError::Scan)?;
let receipt = assemble_unfiltered_lexical_index(
LexicalInputs {
filter: None,
generation: admission.generation,
cache: &self.lexical_index_cache,
snapshot: &admission.snapshot,
active: &admission.active,
accounting: &self.accounting,
},
true,
Some(&cancellation),
)
.map_err(map_store_lexical_assembly_error)?;
receipt.vocabulary(&self.accounting, &cancellation)?;
Ok(())
}
pub fn search_lexical(
&self,
query: &crate::fts::search::TermQuery,
k: usize,
control: QueryControl,
) -> Result<crate::ingest::StoreLexicalSearchOutcome, crate::ingest::StoreLexicalError> {
self.search_lexical_then(query, k, control, None, |outcome, _, _, _, _, _| {
Ok(outcome)
})
}
pub fn search_lexical_filtered(
&self,
query: &crate::fts::search::TermQuery,
k: usize,
control: QueryControl,
filter: Option<&QueryFilter>,
) -> Result<crate::ingest::StoreLexicalSearchOutcome, crate::ingest::StoreLexicalError> {
self.search_lexical_then(query, k, control, filter, |outcome, _, _, _, _, _| {
Ok(outcome)
})
}
fn search_lexical_then<R>(
&self,
query: &crate::fts::search::TermQuery,
k: usize,
control: QueryControl,
filter: Option<&QueryFilter>,
finish: impl FnOnce(
crate::ingest::StoreLexicalSearchOutcome,
&PublishedSnapshot,
&crate::ingest::ActiveSegment,
&LexicalAssembly,
&[crate::fts::search::ScoredDoc],
&QueryCancellation<'_>,
) -> Result<R, crate::ingest::StoreLexicalError>,
) -> Result<R, crate::ingest::StoreLexicalError> {
let control = control.with_clock(Arc::clone(&self.clock));
let started = self.clock.now();
let admission_started = timing_start(self.clock.as_ref());
let AdmittedLexicalQuery {
generation,
active,
snapshot,
active_query,
} = self.admit_lexical_query()?;
let admission_time = timing_elapsed(self.clock.as_ref(), admission_started);
let lexical_started = timing_start(self.clock.as_ref());
let lease = SnapshotLease::new_at(Arc::clone(&snapshot), generation);
let cancellation = QueryCancellation::new(&control, &lease);
cancellation.check_graph().map_err(QueryError::Scan)?;
let assembly = assemble_lexical_index(
LexicalInputs {
filter,
generation,
cache: &self.lexical_index_cache,
snapshot: &snapshot,
active: &active,
accounting: &self.accounting,
},
true,
Some(&cancellation),
)
.map_err(map_store_lexical_assembly_error)?;
let mut work = crate::fts::control::WorkCheck::new(|| {
cancellation
.check_graph()
.map_err(QueryError::Scan)
.map_err(crate::ingest::StoreLexicalError::from)
});
work.check_now()?;
let mut allow_lists = Vec::with_capacity(assembly.alive_sets.len());
for alive in &assembly.alive_sets {
work.step()?;
allow_lists.push(alive.alive_bitmap());
}
let lexical = crate::planner::search_lexical_filtered_refs_controlled(
&assembly.index,
query,
k,
crate::fts::bm25::Bm25Params::beir(),
&allow_lists,
None,
|| {
cancellation
.check_graph()
.map_err(QueryError::Scan)
.map_err(crate::ingest::StoreLexicalError::from)
},
)?;
cancellation.check_graph().map_err(QueryError::Scan)?;
let mut candidates = Vec::with_capacity(lexical.result.hits.len());
for hit in &lexical.result.hits {
work.step()?;
let document =
structured_lexical_document(&snapshot, &active, &assembly.sources, hit.doc, true)
.map_err(map_store_lexical_document_error)?
.ok_or(QueryError::Store(StoreError::ActiveRowOverflow))?;
candidates.push(crate::ingest::LexicalCandidate {
document,
score: hit.score,
});
}
let mut diagnostics =
crate::diag::QueryDiagnostics::lexical(crate::diag::LexicalDiagnostics {
snapshot_generation: generation,
indexed_through_seq: active
.indexed_through_seq()
.max(crate::wal::LogSeq::new(snapshot.absorbed_through())),
requested_k: k,
returned: candidates.len(),
counters: lexical.result.counters,
elapsed: self.clock.now().saturating_duration_since(started),
});
diagnostics.counters.lexical_cache_hits = usize::from(assembly.cache_hit);
diagnostics.counters.lexical_cache_builds = usize::from(!assembly.cache_hit);
diagnostics.timings =
crate::diag::QUERY_TIMING_ENABLED.then(|| crate::diag::QueryTimings {
admission: admission_time,
lexical: timing_elapsed(self.clock.as_ref(), lexical_started),
..Default::default()
});
let outcome = crate::ingest::StoreLexicalSearchOutcome {
candidates,
generation,
diagnostics,
};
work.check_now()?;
let result = finish(
outcome,
&snapshot,
&active,
&assembly,
&lexical.result.hits,
&cancellation,
);
drop(active_query);
result
}
pub fn stored_text(
&self,
version: crate::ingest::DocumentVersion,
) -> Result<Option<String>, crate::ingest::StoreLexicalError> {
let AdmittedLexicalQuery {
active,
snapshot,
active_query,
..
} = self.admit_lexical_query()?;
#[cfg(any(test, feature = "test-seams"))]
self.text_materialization_work
.reverse_version_lookups
.fetch_add(1, Ordering::Relaxed);
if let Some((row, active_version, _)) = active.existing(version.doc_id())
&& active_version == version
{
let text = active
.text(row)
.map_err(crate::ingest::StoreLexicalError::from)?;
drop(active_query);
return Ok(text.map(str::to_owned));
}
for segment in snapshot.segments() {
#[cfg(any(test, feature = "test-seams"))]
self.text_materialization_work
.reverse_version_lookups
.fetch_add(1, Ordering::Relaxed);
let Some(row) = segment
.query_row_for_document_version(version)
.map_err(crate::ingest::StoreLexicalError::from)?
else {
continue;
};
let text = segment
.query_stored_text()
.map_err(StoreError::Segment)
.map_err(crate::ingest::StoreLexicalError::from)?
.and_then(|rows| rows.row(row).flatten())
.map(str::to_owned);
drop(active_query);
return Ok(text);
}
drop(active_query);
Ok(None)
}
pub fn get_documents(
&self,
ids: &[crate::ingest::DocId],
fields: DocumentFields,
) -> Result<Vec<Option<StoredDocument>>, StoreError> {
self.get_documents_with_generation(ids, fields)
.map(|(_, documents)| documents)
}
#[doc(hidden)]
pub fn get_documents_with_generation(
&self,
ids: &[crate::ingest::DocId],
fields: DocumentFields,
) -> Result<(u64, Vec<Option<StoredDocument>>), StoreError> {
let AdmittedDocumentRead {
generation,
active,
snapshot,
active_query,
} = self.admit_document_read()?;
let mut documents = Vec::new();
documents
.try_reserve_exact(ids.len())
.map_err(|_| StoreError::AllocationFailed {
needed: allocation_bytes::<Option<StoredDocument>>(ids.len()),
component: "get documents",
})?;
for id in ids {
let mut document = active
.existing(*id)
.filter(|(row, _, _)| !active.is_tombstoned(*row))
.map(|(row, _, _)| materialize_active_document(&active, row, fields))
.transpose()?;
for segment in snapshot.segments() {
for row in segment.query_rows_for_doc_id(*id)? {
if document.is_some() {
return Err(StoreError::DuplicateLiveDocument { doc_id: *id });
}
document = Some(materialize_sealed_document(segment, row, fields)?);
}
}
documents.push(document);
}
drop(active_query);
Ok((generation, documents))
}
pub fn scan_documents(
&self,
request: DocumentScanRequest<'_>,
) -> Result<DocumentScanPage, QueryError> {
let DocumentScanRequest {
limit,
fields,
order,
cursor,
timestamp_range,
predicate,
control,
} = request;
if limit == 0 || limit > MAX_DOCUMENT_SCAN_LIMIT {
return Err(QueryError::Store(StoreError::InvalidScan {
detail: format!("limit must be in 1..={MAX_DOCUMENT_SCAN_LIMIT}, received {limit}"),
}));
}
if let ScanOrder::Attribute { column, .. } = order {
validate_scan_order_attribute(&self.schema, column)?;
}
if let Some(cursor) = cursor.filter(|cursor| cursor.order != order) {
return Err(QueryError::Store(StoreError::InvalidScan {
detail: format!(
"scan cursor was issued for order {:?}, but the request orders by {order:?}",
cursor.order
),
}));
}
let predicate = document_scan_predicate(predicate, timestamp_range)?;
if let Some(predicate) = &predicate {
crate::planner::validate_predicate(predicate, &self.schema).map_err(|error| {
QueryError::Store(StoreError::InvalidScan {
detail: error.to_string(),
})
})?;
}
let control = control.with_clock(Arc::clone(&self.clock));
control.checkpoint()?;
let AdmittedDocumentRead {
generation,
active,
snapshot,
active_query,
} = self.admit_document_read().map_err(QueryError::Store)?;
if let Some(cursor) = cursor.filter(|cursor| cursor.generation != generation) {
return Err(QueryError::Store(StoreError::ScanStale {
cursor_generation: cursor.generation,
current_generation: generation,
}));
}
let lease = SnapshotLease::new_at(Arc::clone(&snapshot), generation);
let mut work = 0_usize;
let rows = match order {
ScanOrder::Storage => collect_storage_scan_rows(
&snapshot,
&active,
&self.schema,
predicate.as_ref(),
cursor,
limit,
&control,
&lease,
&mut work,
)?,
ScanOrder::TimestampAscending
| ScanOrder::TimestampDescending
| ScanOrder::Attribute { .. } => collect_sorted_scan_rows(
&snapshot,
&active,
&self.schema,
predicate.as_ref(),
cursor,
limit,
order,
&control,
&lease,
&mut work,
)?,
};
let continuation = rows.get(limit).map(|row| DocumentScanCursor {
generation,
order,
source: row.source,
next_row: row.row,
});
let result_len = rows.len().min(limit);
let mut documents = Vec::new();
documents.try_reserve_exact(result_len).map_err(|_| {
QueryError::Store(StoreError::AllocationFailed {
needed: allocation_bytes::<StoredDocument>(result_len),
component: "scan documents",
})
})?;
for row in rows.into_iter().take(result_len) {
scan_control_checkpoint(&control, &lease, &mut work)?;
let document = match row.source {
DocumentScanSource::Active => {
materialize_active_document(&active, row.row as usize, fields)
.map_err(QueryError::Store)?
}
DocumentScanSource::Sealed(segment_id) => {
let segment = snapshot
.segments()
.iter()
.find(|segment| segment.meta().id == segment_id)
.ok_or_else(|| {
QueryError::Store(StoreError::InvalidScan {
detail: format!("scan row names absent segment {segment_id}"),
})
})?;
materialize_sealed_document(segment, row.row as usize, fields)
.map_err(QueryError::Store)?
}
};
documents.push(document);
}
drop(active_query);
Ok(DocumentScanPage {
documents,
generation,
continuation,
})
}
pub fn count_documents(
&self,
predicate: Option<&crate::meta::Predicate>,
timestamp_range: Option<(i64, i64)>,
) -> Result<DocumentCount, QueryError> {
let predicate = document_scan_predicate(predicate, timestamp_range)?;
if let Some(predicate) = &predicate {
crate::planner::validate_predicate(predicate, &self.schema).map_err(|error| {
QueryError::Store(StoreError::InvalidScan {
detail: error.to_string(),
})
})?;
}
let AdmittedDocumentRead {
generation,
active,
snapshot,
active_query,
} = self.admit_document_read().map_err(QueryError::Store)?;
let mut count = 0_u64;
for segment in snapshot.segments() {
count = count
.checked_add(sealed_scan_rows(segment, predicate.as_ref())?.cardinality())
.ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
}
count = count
.checked_add(active_scan_rows(&active, &self.schema, predicate.as_ref())?.cardinality())
.ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
drop(active_query);
Ok(DocumentCount { count, generation })
}
pub(crate) fn live_document_ids_matching(
&self,
active: &crate::ingest::ActiveSegment,
segments: &[crate::segment::reader::SegmentReader],
predicate: &crate::meta::Predicate,
) -> Result<Vec<crate::ingest::DocId>, QueryError> {
let mut ids = Vec::new();
for segment in segments {
for row in sealed_scan_rows(segment, Some(predicate))?.iter() {
ids.push(sealed_document_scan_row(segment, row)?.doc_id);
}
}
for row in active_scan_rows(active, &self.schema, Some(predicate))?.iter() {
ids.push(active_document_scan_row(active, row)?.doc_id);
}
ids.sort_unstable();
ids.dedup();
Ok(ids)
}
pub fn search_lexical_structured(
&self,
query: &crate::fts::query::LexicalQuery,
k: usize,
snippet_bytes: usize,
control: QueryControl,
) -> Result<crate::ingest::StoreStructuredLexicalSearchOutcome, crate::ingest::StoreLexicalError>
{
self.search_lexical_structured_filtered(query, k, snippet_bytes, control, None)
}
#[allow(clippy::too_many_arguments)]
pub fn search_lexical_structured_filtered(
&self,
query: &crate::fts::query::LexicalQuery,
k: usize,
snippet_bytes: usize,
control: QueryControl,
filter: Option<&QueryFilter>,
) -> Result<crate::ingest::StoreStructuredLexicalSearchOutcome, crate::ingest::StoreLexicalError>
{
let control = control.with_clock(Arc::clone(&self.clock));
let started = self.clock.now();
let admission_started = timing_start(self.clock.as_ref());
let AdmittedLexicalQuery {
generation,
active,
snapshot,
active_query,
} = self.admit_lexical_query()?;
let admission_time = timing_elapsed(self.clock.as_ref(), admission_started);
let lexical_started = timing_start(self.clock.as_ref());
let lease = SnapshotLease::new_at(Arc::clone(&snapshot), generation);
let cancellation = QueryCancellation::new(&control, &lease);
cancellation.check_graph().map_err(QueryError::Scan)?;
let assembly = assemble_lexical_index(
LexicalInputs {
filter,
generation,
cache: &self.lexical_index_cache,
snapshot: &snapshot,
active: &active,
accounting: &self.accounting,
},
true,
Some(&cancellation),
)
.map_err(map_store_lexical_assembly_error)?;
let index = &assembly.index;
let sources = &assembly.sources;
let mut expansions = assembly
.expand(query, &self.accounting, &cancellation)
.map_err(|error| match error {
expansion::ExpansionError::Query(error) => {
crate::ingest::StoreLexicalError::Query(error)
}
expansion::ExpansionError::Index(error) => crate::ingest::StoreLexicalError::from(
crate::planner::LexicalFilterError::from(error),
),
expansion::ExpansionError::Shape(error) => {
crate::ingest::StoreLexicalError::Structured(error)
}
})?;
let allow_lists = assembly
.alive_sets
.iter()
.map(|alive| alive.alive_bitmap())
.collect::<Vec<_>>();
let mut prepared_memory =
stats::AccountedCounter::new(&self.accounting, stats::AllocationComponent::Temporary)
.map_err(QueryError::Store)?;
let mut work = crate::fts::control::WorkCheck::new(|| cancellation.check_graph());
let fields = query
.fields_controlled(&mut work)
.map_err(QueryError::Scan)?;
prepared_memory
.set(
crate::fts::query::PreparedWeightedQuery::allocation_bytes_controlled(
&expansions,
&fields,
&mut work,
)
.map_err(QueryError::Scan)?
.ok_or(QueryError::Store(StoreError::ActiveRowOverflow))?,
)
.map_err(QueryError::Store)?;
let mut prepared = if index.segments().is_empty() || expansions.is_empty() {
None
} else {
Some(
crate::fts::query::PreparedWeightedQuery::new_controlled(
index,
std::mem::take(&mut expansions),
fields,
|| cancellation.check_graph(),
)
.map_err(|error| match error {
crate::fts::search::ControlledSearchError::Index(error) => {
crate::ingest::StoreLexicalError::from(
crate::planner::LexicalFilterError::from(error),
)
}
crate::fts::search::ControlledSearchError::Control(error) => {
crate::ingest::StoreLexicalError::from(QueryError::Scan(error))
}
})?,
)
};
let mut scratch_memory =
stats::AccountedCounter::new(&self.accounting, stats::AllocationComponent::Temporary)
.map_err(QueryError::Store)?;
let mut phrase = prepared_lexical::PhraseEligibility::new(
index,
Some(query),
&self.accounting,
&cancellation,
)?;
let result = match &prepared {
None => crate::fts::prune::weighted::WeightedResult::default(),
Some(prepared) => crate::fts::prune::weighted::search(
index,
prepared,
k,
&allow_lists,
|| {
cancellation
.check_graph()
.map_err(QueryError::Scan)
.map_err(crate::ingest::StoreLexicalError::from)
},
|doc| {
phrase
.matches(doc, &cancellation)
.map_err(crate::ingest::StoreLexicalError::from)
},
|bytes| {
prepared_lexical::reserve_weighted_scratch(&mut scratch_memory, bytes)
.map_err(crate::ingest::StoreLexicalError::from)
},
)
.map_err(|error| match error {
crate::fts::search::ControlledSearchError::Index(error) => {
crate::ingest::StoreLexicalError::from(
crate::planner::LexicalFilterError::from(error),
)
}
crate::fts::search::ControlledSearchError::Control(error) => error,
})?,
};
let mut counters = result.counters;
let scored = result.hits;
let mut provenance_memory =
stats::AccountedCounter::new(&self.accounting, stats::AllocationComponent::Temporary)
.map_err(QueryError::Store)?;
let mut candidates = Vec::with_capacity(scored.len());
let mut provenance_work = crate::fts::control::WorkCheck::new(|| {
cancellation.check_graph().map_err(QueryError::Scan)
});
for (doc, score) in scored {
let mut provenance = Vec::new();
let segment = index
.segments()
.get(doc.segment as usize)
.ok_or(QueryError::Store(StoreError::ActiveRowOverflow))?;
let scoring = prepared
.as_ref()
.ok_or(QueryError::Store(StoreError::ActiveRowOverflow))?;
for expansion in scoring.expansions() {
cancellation.check_graph().map_err(QueryError::Scan)?;
let fields = &scoring.scoring().query().fields;
prepared_lexical::reserve_weighted_scratch(
&mut provenance_memory,
crate::fts::sealed::TermStream::allocation_bytes(
segment,
&expansion.term,
fields,
),
)?;
if let Some(mut stream) = crate::fts::sealed::TermStream::open_sized_controlled(
segment,
&expansion.term,
fields,
&mut provenance_work,
)? {
stream.seek_controlled(doc.row, &mut provenance_work)?;
if stream.current_row() == Some(doc.row) {
provenance.push(expansion.clone());
counters.postings_decoded = counters.postings_decoded.saturating_add(1);
}
counters.blocks_decoded = counters
.blocks_decoded
.saturating_add(stream.blocks_decoded());
counters.blocks_skipped = counters
.blocks_skipped
.saturating_add(stream.blocks_skipped());
}
provenance_memory.set(0).map_err(QueryError::Store)?;
}
cancellation.check_graph().map_err(QueryError::Scan)?;
let (text, document) = structured_lexical_row(&snapshot, &active, sources, doc)?;
let terms = provenance
.iter()
.map(|entry| entry.term.clone())
.collect::<Vec<_>>();
let snippet = crate::fts::snippet::best_window(
&self.tokenizer,
text,
&terms,
snippet_bytes,
true,
)?
.ok_or(crate::ingest::StoreLexicalError::MissingSnippetMatch {
segment: doc.segment,
row: doc.row,
})?;
let snippet_text = snippet
.text(text)
.ok_or(QueryError::Store(StoreError::ActiveRowOverflow))?;
candidates.push(crate::ingest::ExplainedLexicalCandidate {
document,
score,
provenance,
snippet: crate::fts::query::OwnedLexicalSnippet {
text: snippet_text.to_owned(),
source: snippet.window,
highlights: snippet.highlights,
source_len: text.len(),
},
});
}
cancellation.check_graph().map_err(QueryError::Scan)?;
let mut diagnostics =
crate::diag::QueryDiagnostics::lexical(crate::diag::LexicalDiagnostics {
snapshot_generation: generation,
indexed_through_seq: active
.indexed_through_seq()
.max(crate::wal::LogSeq::new(snapshot.absorbed_through())),
requested_k: k,
returned: candidates.len(),
counters,
elapsed: self.clock.now().saturating_duration_since(started),
});
diagnostics.counters.lexical_cache_hits = usize::from(assembly.cache_hit);
diagnostics.counters.lexical_cache_builds = usize::from(!assembly.cache_hit);
diagnostics.timings =
crate::diag::QUERY_TIMING_ENABLED.then(|| crate::diag::QueryTimings {
admission: admission_time,
lexical: timing_elapsed(self.clock.as_ref(), lexical_started),
..Default::default()
});
drop(active_query);
Ok(crate::ingest::StoreStructuredLexicalSearchOutcome {
candidates,
expansions: prepared
.as_mut()
.map_or(expansions, |query| query.take_expansions()),
generation,
diagnostics,
})
}
pub fn search_hybrid(
&self,
vector_query: crate::ingest::SearchRequest<'_>,
lexical_query: &crate::fts::search::TermQuery,
hybrid_query: &crate::fusion::HybridQuery,
options: impl Into<SearchOptions>,
control: QueryControl,
) -> Result<crate::ingest::StoreHybridSearchOutcome, crate::fusion::FusionError> {
self.search_hybrid_inner(
vector_query,
PinnedLexicalQuery::Term(lexical_query),
hybrid_query,
options.into(),
control,
)
}
pub fn search_hybrid_structured(
&self,
vector_query: crate::ingest::SearchRequest<'_>,
lexical_query: &crate::fts::query::LexicalQuery,
hybrid_query: &crate::fusion::HybridQuery,
options: impl Into<SearchOptions>,
control: QueryControl,
) -> Result<crate::ingest::StoreHybridSearchOutcome, crate::fusion::FusionError> {
self.search_hybrid_inner(
vector_query,
PinnedLexicalQuery::Structured(lexical_query),
hybrid_query,
options.into(),
control,
)
}
fn search_hybrid_inner(
&self,
vector_query: crate::ingest::SearchRequest<'_>,
lexical_query: PinnedLexicalQuery<'_>,
hybrid_query: &crate::fusion::HybridQuery,
options: SearchOptions,
control: QueryControl,
) -> Result<crate::ingest::StoreHybridSearchOutcome, crate::fusion::FusionError> {
self.search_hybrid_then(
vector_query,
lexical_query,
hybrid_query,
options,
control,
false,
|outcome, _, _, _, _| Ok(outcome),
)
}
#[allow(clippy::too_many_arguments)] fn search_hybrid_then<R>(
&self,
vector_query: crate::ingest::SearchRequest<'_>,
lexical_query: PinnedLexicalQuery<'_>,
hybrid_query: &crate::fusion::HybridQuery,
options: SearchOptions,
control: QueryControl,
capture_rows: bool,
finish: impl FnOnce(
crate::ingest::StoreHybridSearchOutcome,
&PublishedSnapshot,
&crate::ingest::ActiveSegment,
&QueryCancellation<'_>,
Option<&materialize::HybridAddresses>,
) -> Result<R, crate::fusion::FusionError>,
) -> Result<R, crate::fusion::FusionError> {
self.search_hybrid_prepared_then(
|| Ok::<_, std::convert::Infallible>(vector_query),
vector_query.filter,
lexical_query,
hybrid_query,
options,
control,
capture_rows,
finish,
)
.map_err(|error| match error {
HybridPreparationError::Preparation(never) => match never {},
HybridPreparationError::Search(error) => error,
})
}
#[allow(clippy::too_many_arguments)] fn search_hybrid_prepared_then<'vector, E, R>(
&self,
prepare_vector: impl FnOnce() -> Result<crate::ingest::SearchRequest<'vector>, E>,
filter: Option<&QueryFilter>,
lexical_query: PinnedLexicalQuery<'_>,
hybrid_query: &crate::fusion::HybridQuery,
options: SearchOptions,
control: QueryControl,
capture_rows: bool,
finish: impl FnOnce(
crate::ingest::StoreHybridSearchOutcome,
&PublishedSnapshot,
&crate::ingest::ActiveSegment,
&QueryCancellation<'_>,
Option<&materialize::HybridAddresses>,
) -> Result<R, crate::fusion::FusionError>,
) -> Result<R, HybridPreparationError<E>> {
let control = control.with_clock(Arc::clone(&self.clock));
let started = self.clock.now();
let admission_started = timing_start(self.clock.as_ref());
let mut options = options;
let requested_tier = options.explicit_tier();
let admitted = self
.admit_vector_search_for(options, requested_tier.is_none())
.map_err(crate::fusion::FusionError::from)?;
options = admitted.execution_options;
let tier_resolution = if requested_tier.is_none() {
if snapshot_has_graph(&admitted.snapshot) {
Some(crate::diag::HybridTierResolution::Auto)
} else {
options = options.with_tier(SearchTier::Exact);
Some(crate::diag::HybridTierResolution::Exact)
}
} else {
None
};
let corpus_rows = hybrid::corpus_rows(&admitted.snapshot, &admitted.active_segment)?;
let mut width = if hybrid_query.max_rounds == 0 {
corpus_rows
} else {
hybrid::hybrid_window(hybrid_query.k, corpus_rows)?.width
};
let panic_vector = self.consume_hybrid_test_fault(crate::fusion::FusionLeg::Vector);
let panic_lexical = self.consume_hybrid_test_fault(crate::fusion::FusionLeg::Lexical);
let lexical_worker = self.ensure_lexical_worker()?;
let mut timings = crate::diag::QueryTimings {
admission: timing_elapsed(self.clock.as_ref(), admission_started),
..Default::default()
};
let lexical_inputs = LexicalInputs {
filter,
generation: admitted.generation,
cache: &self.lexical_index_cache,
snapshot: &admitted.snapshot,
active: &admitted.active_segment,
accounting: &self.accounting,
};
#[cfg(any(test, feature = "test-seams"))]
let preparation_observer = crate::fts::preparation_observer::current();
let mut lexical_preparation = prepared_lexical::PreparedLexicalQuery::new(lexical_query);
#[cfg(any(test, feature = "test-seams"))]
let fresh_preparation = hybrid::fresh_round_test_control();
let run_lexical_leg = |preparation: &mut prepared_lexical::PreparedLexicalQuery<'_>,
bound,
queued| {
#[cfg(any(test, feature = "test-seams"))]
if fresh_preparation {
preparation.reset_for_test();
}
#[cfg(any(test, feature = "test-seams"))]
let _preparation_scope =
crate::fts::preparation_observer::enter(preparation_observer.clone());
let queue_time = timing_elapsed(self.clock.as_ref(), queued);
let lexical_started = timing_start(self.clock.as_ref());
let name = std::thread::current().name().map(str::to_owned);
let result = (|| {
maybe_trigger_hybrid_leg_panic(panic_lexical, "injected lexical hybrid leg panic");
let lease =
SnapshotLease::new_at(Arc::clone(&admitted.snapshot), admitted.generation);
let cancellation = QueryCancellation::new(&control, &lease);
cancellation
.check_graph()
.map_err(QueryError::Scan)
.map_err(crate::fusion::FusionError::from)?;
match lexical_query {
PinnedLexicalQuery::Term(_) => {
exact_lexical_leg(lexical_inputs, bound, &cancellation, preparation)
}
PinnedLexicalQuery::Structured(query) => exact_structured_lexical_leg(
lexical_inputs,
&self.tokenizer,
query,
bound,
&cancellation,
preparation,
),
}
})();
(
name,
result,
queue_time,
timing_elapsed(self.clock.as_ref(), lexical_started),
)
};
let caller_chose_tier = requested_tier.is_some();
let run_vector_leg = |vector_query: crate::ingest::SearchRequest<'_>,
vector_preparation: &mut prepared::PreparedVectorQuery<'_>,
k: usize| {
let graph_available = snapshot_has_graph(&admitted.snapshot);
let scan_reason_override = if caller_chose_tier || !graph_available {
None
} else if k > corpus_rows {
Some(crate::planner::ScanReason::FullMaterialization)
} else if graph_round_is_worth_it(k, corpus_rows) {
None
} else {
Some(crate::planner::ScanReason::WideningCap {
ef: graph_round_ef(k),
rows: corpus_rows,
})
};
search_pinned(
admitted.pool.as_deref(),
&admitted.snapshot,
&admitted.active_segment,
&self.accounting,
admitted.generation,
self.epoch_identity(),
vector_query,
k,
if caller_chose_tier || graph_round_is_worth_it(k, corpus_rows) {
options
} else {
options.with_tier(SearchTier::Exact)
},
control.clone(),
GraphBoundMode::Shared,
requested_tier,
scan_reason_override,
started,
self.clock.as_ref(),
vector_preparation,
#[cfg(any(test, feature = "test-seams"))]
self.vector_fault_controller.as_ref(),
)
.map_err(crate::fusion::FusionError::from)
};
let lexical_queued = timing_start(self.clock.as_ref());
let (prepared_vector_result, lexical_result) = lexical_worker.run_scoped(
|| {
run_lexical_leg(
&mut lexical_preparation,
width.saturating_add(1),
lexical_queued,
)
},
|| {
std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
let vector_query =
prepare_vector().map_err(HybridPreparationError::Preparation)?;
if vector_query.filter.is_some() && filter.is_none() {
return Err(HybridPreparationError::Search(
QueryError::Store(StoreError::InvalidScan {
detail:
"deferred vector preparation cannot introduce a query filter"
.to_owned(),
})
.into(),
));
}
let mut vector_preparation = prepared::PreparedVectorQuery::new(
vector_query.vector(),
self.epoch_identity(),
&self.accounting,
);
maybe_trigger_hybrid_leg_panic(
panic_vector,
"injected vector hybrid leg panic",
);
let result = run_vector_leg(
vector_query,
&mut vector_preparation,
width.saturating_add(1),
);
Ok((vector_query, vector_preparation, result))
}))
.unwrap_or(Err(HybridPreparationError::Search(
crate::fusion::FusionError::LegPanic {
leg: crate::fusion::FusionLeg::Vector,
detail: "vector hybrid leg panicked",
},
)))
},
)?;
let (lexical_thread_name, lexical_result, queue_time, lexical_time) = lexical_result
.unwrap_or_else(|error| {
(
Some("zeppelin-fts".to_owned()),
Err(error),
std::time::Duration::ZERO,
std::time::Duration::ZERO,
)
});
#[cfg(any(test, feature = "test-seams"))]
if let Ok(mut receipt) = self.hybrid_execution_receipt.lock() {
*receipt = Some(HybridExecutionReceipt {
vector_thread: std::thread::current().id(),
lexical_thread_name,
vector_completed: true,
lexical_completed: true,
});
}
#[cfg(not(any(test, feature = "test-seams")))]
let _ = lexical_thread_name;
let (vector_query, mut vector_preparation, vector_result) = match prepared_vector_result {
Ok(prepared) => prepared,
Err(error) => {
let lease =
SnapshotLease::new_at(Arc::clone(&admitted.snapshot), admitted.generation);
QueryCancellation::new(&control, &lease)
.check_graph()
.map_err(map_scan_error)?;
if let Err(lexical_error) = lexical_result
&& matches!(
lexical_error,
crate::fusion::FusionError::ReadCancelled { .. }
| crate::fusion::FusionError::Timeout { .. }
| crate::fusion::FusionError::Cancelled { .. }
)
{
return Err(HybridPreparationError::Search(lexical_error));
}
return Err(error);
}
};
let (
mut vector_outcome,
(mut lexical_hits, mut lexical_assembly, mut lexical_counters, mut lexical_cache_hit),
) = resolve_hybrid_leg_results(vector_result, lexical_result)?;
timings.lexical_queue = queue_time;
timings.lexical = lexical_time;
let anchors = hybrid::HybridAnchors {
vector_ceiling: vector_outcome.vector_ceiling,
lexical_maximum: lexical_hits.first().map_or(0.0, |hit| hit.bm25),
};
let fusion_lease =
SnapshotLease::new_at(Arc::clone(&admitted.snapshot), admitted.generation);
let fusion_cancellation = QueryCancellation::new(&control, &fusion_lease);
let mut work = hybrid::HybridWork::new();
let mut cross_scores = hybrid::CrossScoreCache::new(vector_outcome.epoch);
let mut round_buffers = hybrid::HybridRoundBuffers::new();
let mut fusion_scratch =
crate::fusion::StoreFusionScratch::new(Arc::clone(&self.accounting));
let mut rescore_counters = crate::diag::ScanRescoreCounters::default();
let mut rounds = 1_usize;
let mut budget_exhausted = false;
let (fused, report_window, addresses) = loop {
if let Some(vector_timings) = vector_outcome.diagnostics.timings {
timings.vector += vector_timings.vector;
}
let fusion_started = timing_start(self.clock.as_ref());
let round = hybrid::build_round(
&admitted.snapshot,
&admitted.active_segment,
&lexical_assembly,
&lexical_preparation,
&self.tokenizer,
lexical_query,
vector_query.vector(),
&vector_outcome.candidates,
anchors,
&lexical_hits,
width,
&fusion_cancellation,
capture_rows,
&self.accounting,
&mut cross_scores,
&mut round_buffers,
)?;
let cross_fill_dims = u64::try_from(round.vector_scores_computed)
.ok()
.and_then(|rows| rows.checked_mul(vector_query.vector().len() as u64))
.ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
let cross_fill_bytes = cross_fill_dims
.checked_mul(4)
.ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
vector_outcome.stats.dims_touched = vector_outcome
.stats
.dims_touched
.checked_add(cross_fill_dims)
.ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
vector_outcome.stats.bytes_read = vector_outcome
.stats
.bytes_read
.checked_add(cross_fill_bytes)
.ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
if cross_fill_dims != 0 {
let caller = std::thread::current().id();
if !vector_outcome.stats.worker_thread_ids.contains(&caller) {
vector_outcome.stats.worker_thread_ids.push(caller);
}
vector_outcome.stats.threads_used = vector_outcome.stats.worker_thread_ids.len();
}
accumulate_search_counters(&mut lexical_counters, &round.lexical_counters);
if let Some(counters) = vector_outcome.diagnostics.scan_rescore {
rescored_scan::accumulate(&mut rescore_counters, counters)?;
}
work.complete_round(
&vector_outcome.stats,
vector_outcome.graph_stats,
lexical_counters,
)?;
work.lexical_cache_hits += usize::from(lexical_cache_hit);
work.lexical_cache_builds += usize::from(!lexical_cache_hit);
work.plans.append(&mut vector_outcome.diagnostics.plan);
work.cross_filled_vector = work
.cross_filled_vector
.checked_add(round.vector_scores_computed)
.ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
work.cross_filled_lexical = work
.cross_filled_lexical
.checked_add(round.lexical_scores_computed)
.ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
work.vector_candidates_produced = work
.vector_candidates_produced
.checked_add(vector_outcome.candidates.len())
.ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
work.lexical_candidates_produced = work
.lexical_candidates_produced
.checked_add(lexical_hits.len())
.ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
let provenance = hybrid::round_provenance(&vector_outcome, lexical_hits.len(), width);
let fused = hybrid::fuse_round(
hybrid_query,
&round,
provenance,
requested_tier == Some(SearchTier::Exact),
&mut fusion_scratch,
)?;
timings.fusion_cross_fill += timing_elapsed(self.clock.as_ref(), fusion_started);
if fused.report.termination != crate::fusion::FusionTermination::WindowUnproven
|| width >= corpus_rows
{
break (
fused,
crate::diag::HybridReport {
normalization_policy_version:
crate::fusion::HYBRID_NORMALIZATION_POLICY_VERSION,
lexical_full_materializations: 0,
provenance,
window: width,
vector_returned: round.vector.len(),
lexical_returned: round.lexical.len(),
cross_filled_vector: round.cross_filled_vector,
cross_filled_lexical: round.cross_filled_lexical,
total_cross_filled_vector: work.cross_filled_vector,
total_cross_filled_lexical: work.cross_filled_lexical,
vector_candidates_produced: work.vector_candidates_produced,
lexical_candidates_produced: work.lexical_candidates_produced,
},
round.addresses,
);
}
width = if rounds >= hybrid_query.max_rounds {
budget_exhausted = true;
corpus_rows
} else {
width.saturating_mul(2).min(corpus_rows)
};
rounds = rounds.saturating_add(1);
let lexical_queued = timing_start(self.clock.as_ref());
let (vector_result, lexical_result) = lexical_worker.run_scoped(
|| {
run_lexical_leg(
&mut lexical_preparation,
width.saturating_add(1),
lexical_queued,
)
},
|| {
run_vector_leg(
vector_query,
&mut vector_preparation,
width.saturating_add(1),
)
},
)?;
let (_, lexical_result, queue_time, lexical_time) =
lexical_result.unwrap_or_else(|error| {
(
Some("zeppelin-fts".to_owned()),
Err(error),
std::time::Duration::ZERO,
std::time::Duration::ZERO,
)
});
let (next_vector, (next_hits, next_sources, next_counters, next_cache_hit)) =
resolve_hybrid_leg_results(vector_result, lexical_result)?;
timings.lexical_queue += queue_time;
timings.lexical += lexical_time;
vector_outcome = next_vector;
lexical_hits = next_hits;
lexical_assembly = next_sources;
lexical_counters = next_counters;
lexical_cache_hit = next_cache_hit;
};
let crate::ingest::SearchOutcome {
candidates: _,
vector_ceiling: _,
stats: _,
graph_stats: _,
generation,
epoch,
diagnostics: vector_diagnostics,
} = vector_outcome;
let mut report = fused.report;
report.rounds = rounds;
if budget_exhausted {
report.budget_exhausted = true;
report.termination = crate::fusion::FusionTermination::BudgetFullMaterialization;
}
let mut diagnostics =
crate::diag::QueryDiagnostics::hybrid(crate::diag::HybridDiagnostics {
snapshot_generation: vector_diagnostics.snapshot_generation,
indexed_through_seq: vector_diagnostics.indexed_through_seq,
plan: work.plans,
approximate: vector_diagnostics.approximate
|| report.termination
== crate::fusion::FusionTermination::ApproximateCandidates,
exact_rescore: vector_diagnostics.exact_rescore,
requested_k: hybrid_query.k,
returned: fused.hits.len(),
scan: work.scan,
graph: work.graph,
lexical: work.lexical,
report,
hybrid: report_window,
tier_resolution,
epoch,
elapsed: self.clock.now().saturating_duration_since(started),
});
diagnostics.counters.lexical_cache_hits = work.lexical_cache_hits;
diagnostics.scan_rescore = options.scan_rescore().map(|_| rescore_counters);
diagnostics.counters.lexical_cache_builds = work.lexical_cache_builds;
diagnostics.timings = crate::diag::QUERY_TIMING_ENABLED.then_some(timings);
let outcome = crate::ingest::StoreHybridSearchOutcome {
hits: fused.hits,
lexical_expansions: lexical_preparation.take_expansions()?,
generation,
diagnostics,
};
finish(
outcome,
&admitted.snapshot,
&admitted.active_segment,
&fusion_cancellation,
addresses.as_ref(),
)
.map_err(HybridPreparationError::Search)
}
#[cfg(any(test, feature = "test-seams"))]
#[must_use]
pub fn lexical_index_cache_counters(&self) -> (u64, u64) {
(
self.lexical_index_cache.hits.load(Ordering::Relaxed),
self.lexical_index_cache.builds.load(Ordering::Relaxed),
)
}
#[cfg(any(test, feature = "test-seams"))]
#[must_use]
pub fn lexical_contribution_cache_counters(&self) -> (u64, u64) {
(
self.lexical_index_cache
.contribution_hits
.load(Ordering::Relaxed),
self.lexical_index_cache
.contribution_builds
.load(Ordering::Relaxed),
)
}
fn consume_hybrid_test_fault(&self, leg: crate::fusion::FusionLeg) -> bool {
#[cfg(any(test, feature = "test-seams"))]
{
let Ok(mut fault) = self.hybrid_leg_fault.lock() else {
return false;
};
if matches!(*fault, Some(HybridLegTestFault::Panic(candidate)) if candidate == leg) {
*fault = None;
return true;
}
}
#[cfg(not(any(test, feature = "test-seams")))]
let _ = leg;
false
}
fn search_with_graph_bound_mode(
&self,
request: crate::ingest::SearchRequest<'_>,
k: usize,
options: SearchOptions,
control: QueryControl,
graph_bound_mode: GraphBoundMode,
) -> Result<crate::ingest::SearchOutcome, QueryError> {
self.search_with_graph_bound_mode_then(
request,
k,
options,
control,
graph_bound_mode,
|outcome, _, _| Ok(outcome),
)
}
fn search_with_graph_bound_mode_then<R>(
&self,
request: crate::ingest::SearchRequest<'_>,
k: usize,
options: SearchOptions,
control: QueryControl,
graph_bound_mode: GraphBoundMode,
finish: impl FnOnce(
crate::ingest::SearchOutcome,
&AdmittedVectorSearch<'_>,
&QueryControl,
) -> Result<R, QueryError>,
) -> Result<R, QueryError> {
let control = control.with_clock(Arc::clone(&self.clock));
let started = self.clock.now();
let admission_started = timing_start(self.clock.as_ref());
let admitted = self.admit_vector_search(options)?;
let admission_time = timing_elapsed(self.clock.as_ref(), admission_started);
let mut vector_preparation = prepared::PreparedVectorQuery::new(
request.vector(),
self.epoch_identity(),
&self.accounting,
);
let score = || {
let mut outcome = search_pinned(
admitted.pool.as_deref(),
&admitted.snapshot,
&admitted.active_segment,
&self.accounting,
admitted.generation,
self.epoch_identity(),
request,
k,
options,
control.clone(),
graph_bound_mode,
options.explicit_tier(),
None,
started,
self.clock.as_ref(),
&mut vector_preparation,
#[cfg(any(test, feature = "test-seams"))]
self.vector_fault_controller.as_ref(),
)?;
if let Some(timings) = &mut outcome.diagnostics.timings {
timings.admission = admission_time;
}
finish(outcome, &admitted, &control)
};
#[cfg(any(test, feature = "test-seams"))]
let result = crate::kernels::vector_fault::run_store_scoring(
self.kernel_fault_controller.as_ref(),
score,
);
#[cfg(not(any(test, feature = "test-seams")))]
let result = score();
#[cfg(any(test, feature = "test-seams"))]
{
if let Some(controller) = self.vector_fault_controller.as_ref() {
controller.finalize_search(result.is_ok());
}
}
result
}
fn ensure_query_pool(&self) -> Result<Arc<pool::QueryPool>, QueryError> {
let mut pool_slot = self.query_pool.lock().map_err(|_| {
QueryError::Store(StoreError::Synchronization {
component: "query pool",
})
})?;
if pool_slot.is_none() {
let capacity = crate::scan::physical_thread_capacity().map_err(|error| {
QueryError::Store(StoreError::QueryPoolCapacity {
source: error.to_string(),
})
})?;
let pool =
pool::QueryPool::start(capacity, &self.accounting).map_err(QueryError::Store)?;
*pool_slot = Some(Arc::new(pool));
}
let pool = pool_slot.as_ref().cloned().ok_or({
QueryError::Store(StoreError::Synchronization {
component: "query pool initialization",
})
})?;
drop(pool_slot);
Ok(pool)
}
fn ensure_lexical_worker(
&self,
) -> Result<Arc<pool::LexicalWorker>, crate::fusion::FusionError> {
let mut worker_slot =
self.lexical_worker
.lock()
.map_err(|_| crate::fusion::FusionError::LegThreadStart {
leg: crate::fusion::FusionLeg::Lexical,
detail: "pooled lexical worker synchronization was poisoned".to_owned(),
})?;
if worker_slot.is_none() {
*worker_slot = Some(Arc::new(pool::LexicalWorker::start()?));
}
let worker = worker_slot.as_ref().cloned().ok_or_else(|| {
crate::fusion::FusionError::LegThreadStart {
leg: crate::fusion::FusionLeg::Lexical,
detail: "pooled lexical worker initialization failed".to_owned(),
}
})?;
drop(worker_slot);
Ok(worker)
}
fn admit_vector_search(
&self,
options: SearchOptions,
) -> Result<AdmittedVectorSearch<'_>, QueryError> {
self.admit_vector_search_for(options, false)
}
fn admit_vector_search_for(
&self,
options: SearchOptions,
hybrid_default: bool,
) -> Result<AdmittedVectorSearch<'_>, QueryError> {
#[cfg(any(test, feature = "test-seams"))]
self.text_materialization_work
.admissions
.fetch_add(1, Ordering::Relaxed);
let state = self
.state
.lock()
.map_err(|_| QueryError::Store(StoreError::Synchronization { component: "state" }))?;
match *state {
StoreState::Open => {}
StoreState::Closing => return Err(QueryError::Store(StoreError::Closing)),
StoreState::Closed => return Err(QueryError::Store(StoreError::Closed)),
}
let active_guard = self.active.lock().map_err(|_| {
QueryError::Store(StoreError::Synchronization {
component: "active segment",
})
})?;
let active_state = active_guard
.as_ref()
.ok_or(QueryError::Store(StoreError::Closed))?;
let generation = active_state.generation;
let active_segment = Arc::clone(&active_state.segment);
let snapshot = self
.snapshot
.read()
.map_err(|_| {
QueryError::Store(StoreError::Synchronization {
component: "published snapshot",
})
})?
.as_ref()
.cloned()
.ok_or(QueryError::Store(StoreError::Closed))?;
let execution_options = if hybrid_default && !snapshot_has_graph(&snapshot) {
options.with_tier(SearchTier::Exact)
} else {
options
};
let needs_query_pool = !active_segment.is_empty()
|| (hybrid_default && !snapshot.segments().is_empty())
|| match execution_options.tier() {
SearchTier::Auto => snapshot.segments().iter().any(|segment| {
!segment.directory().iter().any(|entry| {
entry.kind == crate::segment::layout::RegionKind::GraphNodeBlocks.id()
})
}),
SearchTier::Exact => !snapshot.segments().is_empty(),
SearchTier::Scan => true,
SearchTier::Graph(_) => false,
};
let pool = if needs_query_pool {
Some(self.ensure_query_pool()?)
} else {
None
};
self.active_queries.fetch_add(1, Ordering::Relaxed);
let active_query = ActiveQuery {
count: &self.active_queries,
};
drop(active_guard);
drop(state);
Ok(AdmittedVectorSearch {
pool,
execution_options,
snapshot: snapshot::ReadSnapshot::new(snapshot),
active_segment,
generation,
_active_query: active_query,
})
}
#[cfg(test)]
fn search_independent_for_test(
&self,
request: crate::ingest::SearchRequest<'_>,
k: usize,
options: SearchOptions,
control: QueryControl,
) -> Result<crate::ingest::SearchOutcome, QueryError> {
self.search_with_graph_bound_mode(request, k, options, control, GraphBoundMode::Independent)
}
}
#[derive(Clone, Copy)]
pub(crate) enum StructuredLexicalSource {
Sealed(usize),
Active,
}
#[derive(Clone, Debug)]
pub struct QueryFilter {
predicate: crate::meta::Predicate,
schema: crate::meta::Schema,
}
impl QueryFilter {
pub fn new(
schema: &crate::meta::Schema,
predicate: Option<&crate::meta::Predicate>,
timestamp_range: Option<(i64, i64)>,
) -> Result<Option<Self>, QueryError> {
let Some(predicate) = document_scan_predicate(predicate, timestamp_range)? else {
return Ok(None);
};
crate::planner::validate_predicate(&predicate, schema).map_err(|error| {
QueryError::Store(StoreError::InvalidScan {
detail: error.to_string(),
})
})?;
Ok(Some(Self {
predicate,
schema: schema.clone(),
}))
}
fn active_rows(
&self,
active: &crate::ingest::ActiveSegment,
) -> Result<crate::meta::DocBitmap, QueryError> {
active_scan_rows(active, &self.schema, Some(&self.predicate))
}
}
struct LexicalAssembly {
index: crate::fts::index::LexicalIndex,
alive_sets: Vec<Arc<crate::meta::AliveSet>>,
sources: Vec<StructuredLexicalSource>,
contributions: Vec<Arc<CachedLexicalContribution>>,
vocabulary: Arc<Mutex<Option<Arc<CachedVocabulary>>>>,
_base: Option<Arc<LexicalAssembly>>,
memory: stats::AccountedCounter,
}
struct CachedVocabulary {
view: crate::fts::vocabulary::Vocabulary,
phonetic: Mutex<Option<Arc<expansion::CachedPhonetic>>>,
_memory: stats::AccountedCounter,
}
impl LexicalAssembly {
fn vocabulary(
&self,
accounting: &Arc<stats::Accounting>,
cancellation: &QueryCancellation<'_>,
) -> Result<Arc<CachedVocabulary>, QueryError> {
let mut cached = cancellation.cache_lock(&self.vocabulary, "lexical vocabulary cache")?;
if let Some(value) = cached.as_ref() {
return Ok(Arc::clone(value));
}
let mut memory =
stats::AccountedCounter::new(accounting, stats::AllocationComponent::Cache)
.map_err(QueryError::Store)?;
let mut temporary =
stats::AccountedCounter::new(accounting, stats::AllocationComponent::Temporary)
.map_err(QueryError::Store)?;
let mut work = crate::fts::control::WorkCheck::new(|| {
cancellation.check_graph().map_err(QueryError::Scan)
});
let view = crate::fts::vocabulary::Vocabulary::build_controlled(
self.index.vocabulary_terms(),
|scratch, owned| {
cancellation.check_graph().map_err(QueryError::Scan)?;
prepared_lexical::reserve_weighted_scratch(&mut temporary, scratch)?;
let bytes = owned
.and_then(|bytes| bytes.checked_add(std::mem::size_of::<CachedVocabulary>()))
.ok_or(QueryError::Store(StoreError::BudgetExceeded {
needed: u64::MAX,
budget: u64::MAX,
component: "lexical vocabulary cache",
}))?;
memory.set(bytes).map_err(QueryError::Store)
},
&mut work,
)?;
let value = Arc::new(CachedVocabulary {
view,
phonetic: Mutex::new(None),
_memory: memory,
});
*cached = Some(Arc::clone(&value));
Ok(value)
}
fn owned_bytes(&self) -> Option<usize> {
std::mem::size_of::<Self>()
.checked_add(self.index.resident_spine_bytes()?)?
.checked_add(
self.alive_sets
.capacity()
.checked_mul(std::mem::size_of::<Arc<crate::meta::AliveSet>>())?,
)?
.checked_add(
self.sources
.capacity()
.checked_mul(std::mem::size_of::<StructuredLexicalSource>())?,
)?
.checked_add(
self.contributions
.capacity()
.checked_mul(std::mem::size_of::<Arc<CachedLexicalContribution>>())?,
)
}
}
enum LexicalContributionKey {
Sealed(crate::segment::SegmentMeta, (u64, u64)),
Active(std::sync::Weak<crate::ingest::ActiveSegment>),
}
struct CachedLexicalContribution {
key: LexicalContributionKey,
statistics: Arc<crate::fts::index::LiveSegmentStatistics>,
active_alive: Option<Arc<crate::meta::AliveSet>>,
_memory: stats::AccountedCounter,
}
impl CachedLexicalContribution {
fn build(
key: LexicalContributionKey,
postings: &crate::fts::sealed::SealedSegment,
alive: &Arc<crate::meta::AliveSet>,
ordinal: usize,
accounting: &Arc<stats::Accounting>,
work: &mut crate::fts::control::WorkCheck<impl FnMut() -> Result<(), LexicalAssemblyError>>,
) -> Result<Arc<Self>, LexicalAssemblyError> {
work.check_now()?;
let active = matches!(key, LexicalContributionKey::Active(_));
let alive_bytes = if active {
alive
.resident_bytes()
.and_then(|n| n.checked_add(std::mem::size_of::<crate::meta::AliveSet>()))
.and_then(|n| n.checked_add(2 * std::mem::size_of::<usize>()))
} else {
Some(0)
};
let bytes =
crate::fts::index::LiveSegmentStatistics::allocation_bytes(alive.alive_bitmap())
.and_then(|n| n.checked_add(std::mem::size_of::<Self>()))
.and_then(|n| n.checked_add(2 * std::mem::size_of::<usize>()))
.and_then(|n| n.checked_add(alive_bytes?))
.ok_or(LexicalAssemblyError::Store(StoreError::BudgetExceeded {
needed: u64::MAX,
budget: u64::MAX,
component: "cache",
}))?;
let mut memory =
stats::AccountedCounter::new(accounting, stats::AllocationComponent::Cache)
.map_err(LexicalAssemblyError::Store)?;
memory.set(bytes).map_err(LexicalAssemblyError::Store)?;
let mut statistics = crate::fts::index::LiveSegmentStatistics::build_controlled(
ordinal,
postings,
alive.alive_bitmap(),
work,
)?;
if statistics.rows.cardinality() != 0
&& statistics.rows.cardinality() != u64::from(postings.row_count())
{
statistics.frequency_cache = Some(
crate::fts::live_df::LiveFrequencyCache::new(accounting)
.map_err(LexicalAssemblyError::Store)?,
);
}
let statistics = Arc::new(statistics);
work.check_now()?;
Ok(Arc::new(Self {
key,
statistics,
active_alive: active.then(|| Arc::clone(alive)),
_memory: memory,
}))
}
}
enum LexicalAssemblyError {
Store(StoreError),
Lexical(crate::fts::index::IndexError),
MissingDocumentIdentity(crate::segment::SegmentId),
Cancelled(crate::scan::ScanError),
}
impl From<crate::fts::index::IndexError> for LexicalAssemblyError {
fn from(error: crate::fts::index::IndexError) -> Self {
Self::Lexical(error)
}
}
impl From<StoreError> for LexicalAssemblyError {
fn from(error: StoreError) -> Self {
Self::Store(error)
}
}
struct CachedLexicalAssembly {
generation: u64,
snapshot: std::sync::Weak<PublishedSnapshot>,
active: std::sync::Weak<crate::ingest::ActiveSegment>,
document_identity_verified: bool,
assembly: Arc<LexicalAssembly>,
}
impl CachedLexicalAssembly {
fn matches(
&self,
snapshot: &Arc<PublishedSnapshot>,
active: &Arc<crate::ingest::ActiveSegment>,
) -> bool {
self.snapshot
.upgrade()
.is_some_and(|cached| Arc::ptr_eq(&cached, snapshot))
&& self
.active
.upgrade()
.is_some_and(|cached| Arc::ptr_eq(&cached, active))
}
}
pub(crate) struct LexicalIndexCache {
entry: Mutex<Option<CachedLexicalAssembly>>,
#[cfg(any(test, feature = "test-seams"))]
hits: AtomicU64,
#[cfg(any(test, feature = "test-seams"))]
builds: AtomicU64,
#[cfg(any(test, feature = "test-seams"))]
contribution_hits: AtomicU64,
#[cfg(any(test, feature = "test-seams"))]
contribution_builds: AtomicU64,
}
impl LexicalIndexCache {
pub(crate) const fn new() -> Self {
Self {
entry: Mutex::new(None),
#[cfg(any(test, feature = "test-seams"))]
hits: AtomicU64::new(0),
#[cfg(any(test, feature = "test-seams"))]
builds: AtomicU64::new(0),
#[cfg(any(test, feature = "test-seams"))]
contribution_hits: AtomicU64::new(0),
#[cfg(any(test, feature = "test-seams"))]
contribution_builds: AtomicU64::new(0),
}
}
#[cfg(any(test, feature = "test-seams"))]
fn record_hit(&self) {
self.hits.fetch_add(1, Ordering::Relaxed);
}
#[cfg(not(any(test, feature = "test-seams")))]
const fn record_hit(&self) {}
#[cfg(any(test, feature = "test-seams"))]
fn record_build(&self) {
self.builds.fetch_add(1, Ordering::Relaxed);
}
#[cfg(not(any(test, feature = "test-seams")))]
const fn record_build(&self) {}
fn record_contribution(&self, hit: bool) {
#[cfg(any(test, feature = "test-seams"))]
if hit {
self.contribution_hits.fetch_add(1, Ordering::Relaxed);
} else {
self.contribution_builds.fetch_add(1, Ordering::Relaxed);
}
#[cfg(not(any(test, feature = "test-seams")))]
let _ = hit;
}
fn lock(
&self,
) -> Result<std::sync::MutexGuard<'_, Option<CachedLexicalAssembly>>, LexicalAssemblyError>
{
self.entry.lock().map_err(|_| {
LexicalAssemblyError::Store(StoreError::Synchronization {
component: "lexical index cache",
})
})
}
fn lock_controlled(
&self,
cancellation: Option<&QueryCancellation<'_>>,
) -> Result<std::sync::MutexGuard<'_, Option<CachedLexicalAssembly>>, LexicalAssemblyError>
{
let Some(cancellation) = cancellation else {
return self.lock();
};
cancellation
.check_graph()
.map_err(LexicalAssemblyError::Cancelled)?;
let guard = loop {
match self.entry.try_lock() {
Ok(guard) => break guard,
Err(std::sync::TryLockError::Poisoned(_)) => {
return Err(LexicalAssemblyError::Store(StoreError::Synchronization {
component: "lexical index cache",
}));
}
Err(std::sync::TryLockError::WouldBlock) => {}
}
cancellation
.check_graph()
.map_err(LexicalAssemblyError::Cancelled)?;
std::thread::park_timeout(Duration::from_millis(1));
};
cancellation
.check_graph()
.map_err(LexicalAssemblyError::Cancelled)?;
Ok(guard)
}
}
#[derive(Clone, Copy)]
struct LexicalInputs<'a> {
filter: Option<&'a QueryFilter>,
generation: u64,
cache: &'a LexicalIndexCache,
snapshot: &'a Arc<PublishedSnapshot>,
active: &'a Arc<crate::ingest::ActiveSegment>,
accounting: &'a Arc<stats::Accounting>,
}
struct LexicalAssemblyReceipt {
assembly: Arc<LexicalAssembly>,
cache_hit: bool,
}
impl std::ops::Deref for LexicalAssemblyReceipt {
type Target = LexicalAssembly;
fn deref(&self) -> &Self::Target {
&self.assembly
}
}
fn map_eligibility_error(error: QueryError) -> LexicalAssemblyError {
match error {
QueryError::Store(error) => LexicalAssemblyError::Store(error),
QueryError::Scan(error) => LexicalAssemblyError::Cancelled(error),
error => LexicalAssemblyError::Store(StoreError::InvalidScan {
detail: error.to_string(),
}),
}
}
fn assemble_lexical_index(
inputs: LexicalInputs<'_>,
require_document_identity: bool,
cancellation: Option<&QueryCancellation<'_>>,
) -> Result<LexicalAssemblyReceipt, LexicalAssemblyError> {
let receipt =
assemble_unfiltered_lexical_index(inputs, require_document_identity, cancellation)?;
let Some(filter) = inputs.filter else {
return Ok(receipt);
};
let mut alive_sets = Vec::with_capacity(receipt.alive_sets.len());
for (source, alive) in receipt.sources.iter().zip(&receipt.alive_sets) {
if let Some(control) = cancellation {
control
.check_graph()
.map_err(LexicalAssemblyError::Cancelled)?;
}
let rows = match source {
StructuredLexicalSource::Active => filter.active_rows(inputs.active),
StructuredLexicalSource::Sealed(ordinal) => inputs
.snapshot
.segments()
.get(*ordinal)
.ok_or(QueryError::Store(StoreError::ActiveRowOverflow))
.and_then(|segment| sealed_scan_rows(segment, Some(&filter.predicate))),
}
.map_err(map_eligibility_error)?;
let mut eligible = (**alive).clone();
#[cfg(any(test, feature = "test-seams"))]
crate::fts::preparation_observer::record_filter_work(0, 0, 1);
eligible.retain(&rows);
alive_sets.push(Arc::new(eligible));
}
if let Some(control) = cancellation {
control
.check_graph()
.map_err(LexicalAssemblyError::Cancelled)?;
}
let memory =
stats::AccountedCounter::new(inputs.accounting, stats::AllocationComponent::Temporary)
.map_err(LexicalAssemblyError::Store)?;
let mut assembly = LexicalAssembly {
index: receipt.index.clone(),
alive_sets,
sources: receipt.sources.clone(),
contributions: receipt.contributions.clone(),
vocabulary: Arc::clone(&receipt.vocabulary),
_base: Some(Arc::clone(&receipt.assembly)),
memory,
};
let bytes = assembly
.alive_sets
.iter()
.fold(assembly.owned_bytes(), |bytes, alive| {
bytes?
.checked_add(alive.resident_bytes()?)?
.checked_add(std::mem::size_of::<crate::meta::AliveSet>())?
.checked_add(2 * std::mem::size_of::<usize>())
})
.ok_or(LexicalAssemblyError::Store(StoreError::ActiveRowOverflow))?;
assembly
.memory
.set(bytes)
.map_err(LexicalAssemblyError::Store)?;
Ok(LexicalAssemblyReceipt {
cache_hit: receipt.cache_hit,
assembly: Arc::new(assembly),
})
}
fn assemble_unfiltered_lexical_index(
inputs: LexicalInputs<'_>,
require_document_identity: bool,
cancellation: Option<&QueryCancellation<'_>>,
) -> Result<LexicalAssemblyReceipt, LexicalAssemblyError> {
let LexicalInputs {
filter: _,
generation,
cache,
snapshot,
active,
accounting,
} = inputs;
let mut entry = cache.lock_controlled(cancellation)?;
if let Some(cached) = entry.as_mut()
&& cached.matches(snapshot, active)
{
if require_document_identity && !cached.document_identity_verified {
verify_document_identity(snapshot, cancellation)?;
cached.document_identity_verified = true;
}
let assembly = Arc::clone(&cached.assembly);
cache.record_hit();
return Ok(LexicalAssemblyReceipt {
assembly,
cache_hit: true,
});
}
let previous = entry.as_ref().map(|cached| Arc::clone(&cached.assembly));
let contributions = previous
.as_ref()
.map_or(&[][..], |a| a.contributions.as_slice());
let (mut assembly, owned_bytes) = build_lexical_assembly(
snapshot,
active,
accounting,
require_document_identity,
cancellation,
cache,
contributions,
)?;
cache.record_build();
assembly
.memory
.set(owned_bytes)
.map_err(LexicalAssemblyError::Store)?;
if let Some(cancellation) = cancellation {
cancellation
.check_graph()
.map_err(LexicalAssemblyError::Cancelled)?;
}
let assembly = Arc::new(assembly);
if entry
.as_ref()
.is_none_or(|cached| generation >= cached.generation)
{
*entry = Some(CachedLexicalAssembly {
generation,
snapshot: Arc::downgrade(snapshot),
active: Arc::downgrade(active),
document_identity_verified: require_document_identity,
assembly: Arc::clone(&assembly),
});
}
drop(entry);
Ok(LexicalAssemblyReceipt {
assembly,
cache_hit: false,
})
}
fn check_segment_document_identity(
segment: &crate::segment::reader::SegmentReader,
postings: &crate::fts::sealed::SealedSegment,
) -> Result<(), LexicalAssemblyError> {
if postings.row_count() != 0
&& segment
.document_version(0)
.map_err(StoreError::Segment)
.map_err(LexicalAssemblyError::Store)?
.is_none()
{
return Err(LexicalAssemblyError::MissingDocumentIdentity(
segment.meta().id,
));
}
Ok(())
}
fn verify_document_identity(
snapshot: &PublishedSnapshot,
cancellation: Option<&QueryCancellation<'_>>,
) -> Result<(), LexicalAssemblyError> {
for segment in snapshot.segments() {
if let Some(cancellation) = cancellation {
cancellation
.check_graph()
.map_err(LexicalAssemblyError::Cancelled)?;
}
if let Some(postings) = segment
.query_postings()
.map_err(LexicalAssemblyError::Store)?
{
check_segment_document_identity(segment, &postings)?;
}
}
Ok(())
}
fn build_lexical_assembly(
snapshot: &PublishedSnapshot,
active: &Arc<crate::ingest::ActiveSegment>,
accounting: &Arc<stats::Accounting>,
require_document_identity: bool,
cancellation: Option<&QueryCancellation<'_>>,
cache: &LexicalIndexCache,
previous: &[Arc<CachedLexicalContribution>],
) -> Result<(LexicalAssembly, usize), LexicalAssemblyError> {
let mut work = crate::fts::control::WorkCheck::new(|| {
cancellation
.map_or(Ok(()), |control| control.check_graph())
.map_err(LexicalAssemblyError::Cancelled)
});
work.check_now()?;
let mut index = crate::fts::index::LexicalIndex::new();
let mut alive_sets = Vec::new();
let mut sources = Vec::new();
let mut contributions = Vec::new();
for (ordinal, segment) in snapshot.segments().iter().enumerate() {
if let Some(cancellation) = cancellation {
cancellation
.check_graph()
.map_err(LexicalAssemblyError::Cancelled)?;
}
if let Some(postings) = segment
.query_postings()
.map_err(LexicalAssemblyError::Store)?
{
if require_document_identity {
check_segment_document_identity(segment, &postings)?;
}
let alive = segment.query_alive().map_err(LexicalAssemblyError::Store)?;
let identity = segment
.lexical_postings_identity()
.map_err(StoreError::Segment)
.map_err(LexicalAssemblyError::Store)?;
let reused = previous
.binary_search_by(|value| match &value.key {
LexicalContributionKey::Sealed(meta, _) => meta.id.cmp(&segment.meta().id),
LexicalContributionKey::Active(_) => std::cmp::Ordering::Greater,
})
.ok()
.and_then(|slot| previous.get(slot))
.filter(|value| {
matches!(&value.key, LexicalContributionKey::Sealed(meta, stored)
if meta.same_segment_file(segment.meta()) && *stored == identity)
&& value.statistics.rows == *alive.alive_bitmap()
});
cache.record_contribution(reused.is_some());
let contribution = match reused {
Some(value) => Arc::clone(value),
None => CachedLexicalContribution::build(
LexicalContributionKey::Sealed(segment.meta().clone(), identity),
&postings,
&alive,
index.segments().len(),
accounting,
&mut work,
)?,
};
index.push_shared_with_statistics(postings, Arc::clone(&contribution.statistics));
contributions.push(contribution);
alive_sets.push(alive);
sources.push(StructuredLexicalSource::Sealed(ordinal));
}
}
if active.has_text_controlled(&mut work)? {
let sealed = active.sealed_lexical_controlled(accounting, &mut work)?;
let reused = previous.last().filter(|value| {
matches!(&value.key, LexicalContributionKey::Active(cached)
if cached.upgrade().is_some_and(|cached| Arc::ptr_eq(&cached, active)))
});
cache.record_contribution(reused.is_some());
let contribution = match reused {
Some(value) => Arc::clone(value),
None => {
let alive = Arc::new(active.alive_controlled(&mut work)?);
CachedLexicalContribution::build(
LexicalContributionKey::Active(Arc::downgrade(active)),
&sealed,
&alive,
index.segments().len(),
accounting,
&mut work,
)?
}
};
let alive = contribution
.active_alive
.as_ref()
.ok_or(LexicalAssemblyError::Store(StoreError::Synchronization {
component: "active lexical contribution",
}))?;
index.push_shared_with_statistics(sealed, Arc::clone(&contribution.statistics));
alive_sets.push(Arc::clone(alive));
contributions.push(contribution);
sources.push(StructuredLexicalSource::Active);
}
crate::fts::control::sort_by(
&mut contributions,
|left, right| match (&left.key, &right.key) {
(LexicalContributionKey::Sealed(a, _), LexicalContributionKey::Sealed(b, _)) => {
a.id.cmp(&b.id)
}
(LexicalContributionKey::Active(_), LexicalContributionKey::Active(_)) => {
std::cmp::Ordering::Equal
}
(LexicalContributionKey::Active(_), _) => std::cmp::Ordering::Greater,
(_, LexicalContributionKey::Active(_)) => std::cmp::Ordering::Less,
},
&mut work,
)?;
let assembly = LexicalAssembly {
index,
alive_sets,
sources,
contributions,
vocabulary: Arc::new(Mutex::new(None)),
_base: None,
memory: stats::AccountedCounter::new(accounting, stats::AllocationComponent::Cache)
.map_err(LexicalAssemblyError::Store)?,
};
let owned_bytes = assembly
.owned_bytes()
.and_then(|n| {
n.checked_add(
std::mem::size_of::<Mutex<Option<Arc<CachedVocabulary>>>>()
+ 2 * std::mem::size_of::<usize>(),
)
})
.ok_or(LexicalAssemblyError::Store(StoreError::BudgetExceeded {
needed: u64::MAX,
budget: u64::MAX,
component: "cache",
}))?;
work.check_now()?;
Ok((assembly, owned_bytes))
}
fn map_store_lexical_assembly_error(
error: LexicalAssemblyError,
) -> crate::ingest::StoreLexicalError {
match error {
LexicalAssemblyError::Store(error) => QueryError::Store(error).into(),
LexicalAssemblyError::Lexical(error) => {
crate::planner::LexicalFilterError::from(error).into()
}
LexicalAssemblyError::MissingDocumentIdentity(segment_id) => {
crate::ingest::StoreLexicalError::MissingDocumentIdentity { segment_id }
}
LexicalAssemblyError::Cancelled(error) => QueryError::Scan(error).into(),
}
}
fn map_fusion_lexical_assembly_error(error: LexicalAssemblyError) -> crate::fusion::FusionError {
match error {
LexicalAssemblyError::Store(error) => crate::fusion::FusionError::Leg {
leg: crate::fusion::FusionLeg::Lexical,
kind: crate::fusion::LegFailureKind::Store(error.kind()),
detail: error.to_string(),
},
LexicalAssemblyError::Lexical(error) => crate::fusion::FusionError::Leg {
leg: crate::fusion::FusionLeg::Lexical,
kind: crate::fusion::LegFailureKind::Lexical,
detail: error.to_string(),
},
LexicalAssemblyError::MissingDocumentIdentity(segment_id) => {
crate::fusion::FusionError::Leg {
leg: crate::fusion::FusionLeg::Lexical,
kind: crate::fusion::LegFailureKind::Invariant,
detail: format!("sealed lexical segment {segment_id} has no document identity"),
}
}
LexicalAssemblyError::Cancelled(error) => {
crate::fusion::FusionError::from(QueryError::Scan(error))
}
}
}
fn accumulate_search_counters(
total: &mut crate::fts::search::SearchCounters,
delta: &crate::fts::search::SearchCounters,
) {
total.docs_evaluated = total.docs_evaluated.saturating_add(delta.docs_evaluated);
total.postings_decoded = total
.postings_decoded
.saturating_add(delta.postings_decoded);
total.blocks_decoded = total.blocks_decoded.saturating_add(delta.blocks_decoded);
total.blocks_skipped = total.blocks_skipped.saturating_add(delta.blocks_skipped);
}
fn map_fusion_controlled_lexical_error(
error: crate::fts::search::ControlledSearchError<crate::scan::ScanError>,
) -> crate::fusion::FusionError {
match error {
crate::fts::search::ControlledSearchError::Index(error) => {
crate::fusion::FusionError::Leg {
leg: crate::fusion::FusionLeg::Lexical,
kind: crate::fusion::LegFailureKind::Lexical,
detail: error.to_string(),
}
}
crate::fts::search::ControlledSearchError::Control(error) => {
crate::fusion::FusionError::from(QueryError::Scan(error))
}
}
}
#[derive(Clone, Copy)]
enum PinnedLexicalQuery<'a> {
Term(&'a crate::fts::search::TermQuery),
Structured(&'a crate::fts::query::LexicalQuery),
}
enum LexicalDocumentError {
SourceOrdinalOutOfRange,
RowOverflow,
SealedSourceAbsent,
Segment(crate::segment::SegmentError),
MissingDocumentIdentity(crate::segment::SegmentId),
ActiveRowAbsent,
}
fn structured_lexical_document(
snapshot: &PublishedSnapshot,
active: &crate::ingest::ActiveSegment,
sources: &[StructuredLexicalSource],
doc: crate::fts::search::GlobalDocId,
require_document_identity: bool,
) -> Result<Option<crate::ingest::DocumentVersion>, LexicalDocumentError> {
let source = usize::try_from(doc.segment)
.ok()
.and_then(|slot| sources.get(slot))
.ok_or(LexicalDocumentError::SourceOrdinalOutOfRange)?;
let row = usize::try_from(doc.row).map_err(|_| LexicalDocumentError::RowOverflow)?;
match source {
StructuredLexicalSource::Sealed(ordinal) => {
let segment = snapshot
.segments()
.get(*ordinal)
.ok_or(LexicalDocumentError::SealedSourceAbsent)?;
let document = segment
.document_version(row)
.map_err(LexicalDocumentError::Segment)?;
if require_document_identity && document.is_none() {
return Err(LexicalDocumentError::MissingDocumentIdentity(
segment.meta().id,
));
}
Ok(document)
}
StructuredLexicalSource::Active => {
let document = active.document(row);
if require_document_identity && document.is_none() {
return Err(LexicalDocumentError::ActiveRowAbsent);
}
Ok(document)
}
}
}
fn map_store_lexical_document_error(
error: LexicalDocumentError,
) -> crate::ingest::StoreLexicalError {
match error {
LexicalDocumentError::SourceOrdinalOutOfRange
| LexicalDocumentError::RowOverflow
| LexicalDocumentError::SealedSourceAbsent
| LexicalDocumentError::ActiveRowAbsent => {
QueryError::Store(StoreError::ActiveRowOverflow).into()
}
LexicalDocumentError::Segment(error) => {
QueryError::Store(StoreError::Segment(error)).into()
}
LexicalDocumentError::MissingDocumentIdentity(segment_id) => {
crate::ingest::StoreLexicalError::MissingDocumentIdentity { segment_id }
}
}
}
fn map_fusion_lexical_document_error(
error: LexicalDocumentError,
invariant_kind: crate::fusion::LegFailureKind,
segment_kind: crate::fusion::LegFailureKind,
) -> crate::fusion::FusionError {
let (kind, detail) = match error {
LexicalDocumentError::SourceOrdinalOutOfRange => (
invariant_kind,
"lexical source ordinal is out of range".to_owned(),
),
LexicalDocumentError::RowOverflow => {
(invariant_kind, "lexical row exceeds usize".to_owned())
}
LexicalDocumentError::SealedSourceAbsent => {
(invariant_kind, "lexical sealed source is absent".to_owned())
}
LexicalDocumentError::Segment(error) => (segment_kind, error.to_string()),
LexicalDocumentError::MissingDocumentIdentity(segment_id) => (
invariant_kind,
format!("sealed lexical segment {segment_id} has no document identity"),
),
LexicalDocumentError::ActiveRowAbsent => (
invariant_kind,
"lexical active source row is absent".to_owned(),
),
};
crate::fusion::FusionError::Leg {
leg: crate::fusion::FusionLeg::Lexical,
kind,
detail,
}
}
fn map_structured_leg_document_error(error: LexicalDocumentError) -> crate::fusion::FusionError {
map_fusion_lexical_document_error(
error,
crate::fusion::LegFailureKind::Lexical,
crate::fusion::LegFailureKind::Lexical,
)
}
fn map_term_leg_document_error(error: LexicalDocumentError) -> crate::fusion::FusionError {
map_fusion_lexical_document_error(
error,
crate::fusion::LegFailureKind::Invariant,
crate::fusion::LegFailureKind::Segment,
)
}
fn structured_lexical_row<'a>(
snapshot: &'a PublishedSnapshot,
active: &'a crate::ingest::ActiveSegment,
sources: &[StructuredLexicalSource],
doc: crate::fts::search::GlobalDocId,
) -> Result<(&'a str, crate::ingest::DocumentVersion), crate::ingest::StoreLexicalError> {
let source = usize::try_from(doc.segment)
.ok()
.and_then(|slot| sources.get(slot))
.ok_or(QueryError::Store(StoreError::ActiveRowOverflow))?;
let row =
usize::try_from(doc.row).map_err(|_| QueryError::Store(StoreError::ActiveRowOverflow))?;
match source {
StructuredLexicalSource::Sealed(ordinal) => {
let segment = snapshot
.segments()
.get(*ordinal)
.ok_or(QueryError::Store(StoreError::ActiveRowOverflow))?;
let text = segment
.query_stored_text()
.map_err(StoreError::Segment)
.map_err(QueryError::Store)?
.and_then(|rows| rows.row(row).flatten())
.ok_or(crate::ingest::StoreLexicalError::MissingStoredText {
segment_id: segment.meta().id,
row: doc.row,
})?;
let document = segment
.document_version(row)
.map_err(StoreError::Segment)
.map_err(QueryError::Store)?
.ok_or(crate::ingest::StoreLexicalError::MissingDocumentIdentity {
segment_id: segment.meta().id,
})?;
Ok((text, document))
}
StructuredLexicalSource::Active => {
let text = active
.text(row)
.map_err(QueryError::Store)?
.ok_or(QueryError::Store(StoreError::ActiveRowOverflow))?;
let document = active
.document(row)
.ok_or(QueryError::Store(StoreError::ActiveRowOverflow))?;
Ok((text, document))
}
}
}
type ExactLexicalLeg = (
Vec<hybrid::LexicalHit>,
Arc<LexicalAssembly>,
crate::fts::search::SearchCounters,
bool,
);
fn resolve_hybrid_leg_results<Vector, Lexical>(
vector: Result<Vector, crate::fusion::FusionError>,
lexical: Result<Lexical, crate::fusion::FusionError>,
) -> Result<(Vector, Lexical), crate::fusion::FusionError> {
fn control_priority(error: &crate::fusion::FusionError) -> Option<u8> {
match error {
crate::fusion::FusionError::ReadCancelled { .. } => Some(0),
crate::fusion::FusionError::Timeout { .. } => Some(1),
crate::fusion::FusionError::Cancelled { .. } => Some(2),
_ => None,
}
}
match (vector, lexical) {
(Ok(vector), Ok(lexical)) => Ok((vector, lexical)),
(Err(vector), Ok(_)) => Err(vector),
(Ok(_), Err(lexical)) => Err(lexical),
(Err(vector), Err(lexical)) => {
match (control_priority(&vector), control_priority(&lexical)) {
(Some(vector_priority), Some(lexical_priority)) => {
if vector_priority <= lexical_priority {
Err(vector)
} else {
Err(lexical)
}
}
(Some(_), None) => Err(vector),
(None, Some(_)) => Err(lexical),
(None, None) => Err(vector),
}
}
}
}
fn exact_structured_lexical_leg(
inputs: LexicalInputs<'_>,
_analyzer: &crate::fts::tokenizer::Analyzer,
query: &crate::fts::query::LexicalQuery,
bound: usize,
cancellation: &QueryCancellation<'_>,
preparation: &mut prepared_lexical::PreparedLexicalQuery<'_>,
) -> Result<ExactLexicalLeg, crate::fusion::FusionError> {
let (snapshot, active) = (inputs.snapshot, inputs.active);
let lexical_error = |detail: String| crate::fusion::FusionError::Leg {
leg: crate::fusion::FusionLeg::Lexical,
kind: crate::fusion::LegFailureKind::Lexical,
detail,
};
let (prepared, cache_hit) = preparation.prepare_structured(inputs, cancellation)?;
let assembly = &prepared.assembly;
let index = &assembly.index;
let sources = &assembly.sources;
let Some(query_scoring) = prepared.query.as_ref() else {
return Ok((
Vec::new(),
Arc::clone(assembly),
crate::fts::search::SearchCounters::default(),
cache_hit,
));
};
let allow_lists = assembly
.alive_sets
.iter()
.map(|alive| alive.alive_bitmap())
.collect::<Vec<_>>();
let mut scratch_memory =
stats::AccountedCounter::new(inputs.accounting, stats::AllocationComponent::Temporary)
.map_err(QueryError::Store)?;
let mut phrase = prepared_lexical::PhraseEligibility::new(
index,
Some(query),
inputs.accounting,
cancellation,
)?;
let result = crate::fts::prune::weighted::search(
index,
query_scoring,
bound,
&allow_lists,
|| {
cancellation
.check_graph()
.map_err(QueryError::Scan)
.map_err(crate::fusion::FusionError::from)
},
|doc| {
phrase
.matches(doc, cancellation)
.map_err(crate::fusion::FusionError::from)
},
|bytes| {
prepared_lexical::reserve_weighted_scratch(&mut scratch_memory, bytes)
.map_err(crate::fusion::FusionError::from)
},
)
.map_err(|error| match error {
crate::fts::search::ControlledSearchError::Index(error) => lexical_error(error.to_string()),
crate::fts::search::ControlledSearchError::Control(error) => error,
})?;
let counters = result.counters;
let scored = result.hits;
let mut joined = Vec::with_capacity(scored.len());
for (doc, score) in scored {
let document = structured_lexical_document(snapshot, active, sources, doc, false)
.map_err(map_structured_leg_document_error)?;
joined.push(hybrid::LexicalHit {
doc,
document,
bm25: score,
});
}
Ok((joined, Arc::clone(assembly), counters, cache_hit))
}
fn exact_lexical_leg(
inputs: LexicalInputs<'_>,
bound: usize,
cancellation: &QueryCancellation<'_>,
preparation: &mut prepared_lexical::PreparedLexicalQuery<'_>,
) -> Result<ExactLexicalLeg, crate::fusion::FusionError> {
let (snapshot, active) = (inputs.snapshot, inputs.active);
let (prepared, cache_hit) = preparation.prepare_term(inputs, cancellation)?;
let assembly = &prepared.assembly;
let index = &assembly.index;
let sources = &assembly.sources;
if index.segments().is_empty() {
return Ok((
Vec::new(),
Arc::clone(assembly),
crate::fts::search::SearchCounters::default(),
cache_hit,
));
}
let k = bound.min(usize::try_from(index.document_count()).unwrap_or(usize::MAX));
let allow_lists = assembly
.alive_sets
.iter()
.map(|alive| alive.alive_bitmap())
.collect::<Vec<_>>();
let result =
exact_hybrid_lexical_search(index, prepared.scoring()?, k, &allow_lists, cancellation)?;
let mut joined = Vec::with_capacity(result.hits.len());
for hit in result.hits {
let document = structured_lexical_document(snapshot, active, sources, hit.doc, false)
.map_err(map_term_leg_document_error)?;
joined.push(hybrid::LexicalHit {
doc: hit.doc,
document,
bm25: hit.score,
});
}
Ok((joined, Arc::clone(assembly), result.counters, cache_hit))
}
fn exact_hybrid_lexical_search(
index: &crate::fts::index::LexicalIndex,
prepared: &crate::fts::search::PreparedTermQuery,
k: usize,
allow_lists: &[&crate::meta::DocBitmap],
cancellation: &QueryCancellation<'_>,
) -> Result<crate::fts::search::SearchResult, crate::fusion::FusionError> {
let allowed = allow_lists
.iter()
.map(|allow_list| allow_list.cardinality())
.fold(0_u64, u64::saturating_add);
if allowed.saturating_mul(crate::planner::LEXICAL_ALLOW_LIST_DIVISOR) <= index.document_count()
{
return crate::fts::search::search_allow_list_prepared_controlled(
index,
prepared,
k,
allow_lists,
|| cancellation.check_graph(),
)
.map_err(map_fusion_controlled_lexical_error);
}
crate::fts::prune::search_pruned_prepared_controlled(index, prepared, k, allow_lists, || {
cancellation.check_graph()
})
.map_err(map_fusion_controlled_lexical_error)
}
#[allow(clippy::too_many_arguments)]
fn search_pinned(
pool: Option<&pool::QueryPool>,
snapshot: &Arc<PublishedSnapshot>,
active: &crate::ingest::ActiveSegment,
accounting: &Arc<stats::Accounting>,
generation: u64,
epoch: Option<crate::epoch::EpochIdentity>,
request: crate::ingest::SearchRequest<'_>,
k: usize,
options: SearchOptions,
control: QueryControl,
graph_bound_mode: GraphBoundMode,
requested_tier: Option<SearchTier>,
hybrid_scan_reason: Option<crate::planner::ScanReason>,
started: std::time::Instant,
clock: &dyn MonotonicClock,
vector_preparation: &mut prepared::PreparedVectorQuery<'_>,
#[cfg(any(test, feature = "test-seams"))] vector_fault_controller: Option<
&crate::scan::vector_fault::VectorFaultController,
>,
) -> Result<crate::ingest::SearchOutcome, QueryError> {
use crate::ingest::{GraphSearchStats, RowSource, SearchOutcome};
use crate::quant::prepare_bit4_query;
use crate::scan::{ScanQuery, ScanRequest, ScanRows, ScanStats};
let vector_started = timing_start(clock);
let scan_options = options.scan();
let rescore_options = options.scan_rescore();
let mut rescore_counters = crate::diag::ScanRescoreCounters::default();
if let Some(options) = rescore_options {
options.frontier(k, 0)?;
}
let query = request.vector();
prepared::validate_vector_coordinates(query, crate::kernels::MAX_DOT_I8_DIMENSION, |_| {
Ok::<(), std::convert::Infallible>(())
})
.map_err(|error| match error {
prepared::VectorValidationError::Data(error) => {
QueryError::Scan(crate::scan::ScanError::Quant(error))
}
prepared::VectorValidationError::Control(never) => match never {},
})?;
vector_preparation.validate_binding(query, epoch)?;
let bit4_query = std::cell::OnceCell::new();
let int8_query = std::cell::OnceCell::new();
let mut candidates = Vec::new();
let mut dims_touched = 0_u64;
let mut bytes_read = 0_u64;
let mut worst_squared_l2: Option<f64> = None;
let mut worst_exhaustive = true;
let mut worker_thread_ids = Vec::new();
let mut graph_stats = GraphSearchStats::default();
let mut plans = Vec::new();
let auto_graph_options = if matches!(options.tier(), SearchTier::Auto)
&& snapshot.segments().iter().any(|segment| {
segment
.directory()
.iter()
.any(|entry| entry.kind == crate::segment::layout::RegionKind::GraphNodeBlocks.id())
}) {
Some(auto_graph_search_options(snapshot)?)
} else {
None
};
if matches!(options.tier(), SearchTier::Graph(_)) {
for segment in snapshot.segments() {
if !segment
.directory()
.iter()
.any(|entry| entry.kind == crate::segment::layout::RegionKind::GraphNodeBlocks.id())
{
return Err(QueryError::Store(StoreError::GraphUnavailable {
segment_id: segment.meta().id,
}));
}
}
}
let full_precision = matches!(options.tier(), SearchTier::Exact | SearchTier::Graph(_))
|| (matches!(options.tier(), SearchTier::Auto)
&& auto_uses_full_precision(snapshot, active));
let exact_lane = if full_precision {
pool.map_or(0, |pool| {
pool.reserve_exact_lane(scan_options.thread_budget)
})
} else {
0
};
if !active.is_empty() {
let mut alive = active.alive().map_err(QueryError::Store)?;
if let Some(filter) = request.filter {
alive.retain(&filter.active_rows(active)?);
}
let mut exact_score = full_precision;
let mut approximate = false;
let frontier = rescore_options
.map(|options| options.frontier(k, alive.live_count()))
.transpose()?
.unwrap_or(k);
let mut exact_memory = ExactScanMemory::new(accounting)?;
let mut outcome = match options.tier() {
SearchTier::Auto if full_precision => {
let lease = SnapshotLease::new_at(Arc::clone(snapshot), generation);
scan_active_squared_l2(
active,
&alive,
request.vector(),
frontier,
ExactScanDispatch {
pool,
accounting,
control: &control,
lease: &lease,
thread_budget: scan_options.thread_budget,
lane: exact_lane,
},
&mut exact_memory,
#[cfg(any(test, feature = "test-seams"))]
vector_fault_controller,
#[cfg(any(test, feature = "test-seams"))]
vector_fault_tier(options.tier()),
)?
}
SearchTier::Auto | SearchTier::Scan => {
let query_pool = pool.ok_or(QueryError::Store(StoreError::Synchronization {
component: "active scan-tier query pool",
}))?;
let bit4_query = bit4_query
.get_or_init(|| prepare_bit4_query(request.vector(), 0))
.as_ref()
.map_err(|error| crate::scan::ScanError::Quant(error.clone()))
.map_err(QueryError::Scan)?;
execute_store_scan(
query_pool,
ScanRequest {
query: ScanQuery::Bit4(bit4_query),
rows: ScanRows::Bit4RowMajor {
codes: active.codes(),
factors: active.factors(),
},
row_mask: Some(alive.scan_mask()),
},
frontier,
scan_options,
control.clone(),
SnapshotLease::new_at(Arc::clone(snapshot), generation),
#[cfg(any(test, feature = "test-seams"))]
vector_fault_controller,
#[cfg(any(test, feature = "test-seams"))]
crate::scan::vector_fault::VectorRowSource::Active,
#[cfg(any(test, feature = "test-seams"))]
vector_fault_tier(options.tier()),
)?
}
SearchTier::Exact | SearchTier::Graph(_) => {
let lease = SnapshotLease::new_at(Arc::clone(snapshot), generation);
scan_active_squared_l2(
active,
&alive,
request.vector(),
frontier,
ExactScanDispatch {
pool,
accounting,
control: &control,
lease: &lease,
thread_budget: scan_options.thread_budget,
lane: exact_lane,
},
&mut exact_memory,
#[cfg(any(test, feature = "test-seams"))]
vector_fault_controller,
#[cfg(any(test, feature = "test-seams"))]
vector_fault_tier(options.tier()),
)?
}
};
if let Some(rescore) = rescore_options {
let lease = SnapshotLease::new_at(Arc::clone(snapshot), generation);
let cancellation = QueryCancellation::new(&control, &lease);
let (counters, exhaustive, worst) = rescored_scan::rescore(
&mut outcome.candidates,
&mut outcome.stats,
active.vectors(),
query,
k,
alive.live_count(),
rescore,
&cancellation,
accounting,
&mut exact_memory,
)?;
outcome.worst_score = worst;
exact_score = true;
approximate = !exhaustive;
rescored_scan::accumulate(&mut rescore_counters, counters)?;
}
fold_worst_squared_l2(&mut worst_squared_l2, &mut worst_exhaustive, &outcome);
merge_store_outcome(
outcome,
RowSource::Active,
|row| Ok(active.document(row)),
&mut candidates,
&mut dims_touched,
&mut bytes_read,
&mut worker_thread_ids,
exact_score,
#[cfg(any(test, feature = "test-seams"))]
vector_fault_controller,
#[cfg(any(test, feature = "test-seams"))]
vector_fault_tier(options.tier()),
)?;
let mut plan = crate::planner::SegmentPlan::unfiltered_scan(
RowSource::Active,
crate::planner::SegmentTier::ActiveScan,
alive.live_count(),
crate::planner::ScanReason::ActiveSegment,
);
plan.approximate = approximate;
plans.push(plan);
if matches!(graph_bound_mode, GraphBoundMode::Shared) {
retain_global_top_k(&mut candidates, k);
}
}
let mut ordered_segments = snapshot.segments().iter().collect::<Vec<_>>();
if matches!(graph_bound_mode, GraphBoundMode::Shared) {
ordered_segments.sort_unstable_by(|left, right| {
right
.meta()
.row_count
.cmp(&left.meta().row_count)
.then_with(|| left.meta().id.cmp(&right.meta().id))
});
}
for segment in ordered_segments {
let original_alive = segment.query_alive().map_err(QueryError::Store)?;
let mut filtered_alive;
let alive = if let Some(filter) = request.filter {
filtered_alive = (*original_alive).clone();
filtered_alive.retain(&sealed_scan_rows(segment, Some(&filter.predicate))?);
&filtered_alive
} else {
original_alive.as_ref()
};
let source = RowSource::Sealed(segment.meta().id);
let graph_options = match options.tier() {
SearchTier::Graph(graph_options) => Some(graph_options),
SearchTier::Auto
if segment.directory().iter().any(|entry| {
entry.kind == crate::segment::layout::RegionKind::GraphNodeBlocks.id()
}) =>
{
auto_graph_options
}
SearchTier::Auto | SearchTier::Exact | SearchTier::Scan => None,
};
if let Some(graph_options) = graph_options {
worst_exhaustive = false;
let competitive_distance = global_competitive_distance(&candidates, k);
let MergedSegmentGraph {
plan,
traversed,
_result,
_scratch,
_lease,
} = traverse_segment_graph(
segment,
alive,
source,
graph_options,
graph_bound_mode,
request,
k,
vector_preparation,
accounting,
snapshot,
generation,
&control,
competitive_distance,
#[cfg(any(test, feature = "test-seams"))]
vector_fault_controller,
options.tier(),
)?
.merge_into(
segment,
alive,
source,
&mut candidates,
&mut dims_touched,
&mut bytes_read,
&mut worker_thread_ids,
&mut graph_stats,
#[cfg(any(test, feature = "test-seams"))]
vector_fault_controller,
options.tier(),
)?;
plans.push(plan);
if traversed && matches!(graph_bound_mode, GraphBoundMode::Shared) {
retain_global_top_k(&mut candidates, k);
}
continue;
}
let mut exact_memory = ExactScanMemory::new(accounting)?;
let frontier = rescore_options
.map(|options| options.frontier(k, alive.live_count()))
.transpose()?
.unwrap_or(k);
let mut outcome = scan_sealed_segment(
pool,
snapshot,
segment,
alive,
generation,
request,
frontier,
scan_options,
&control,
&bit4_query,
&int8_query,
full_precision,
exact_lane,
accounting,
&mut exact_memory,
source,
options.tier(),
#[cfg(any(test, feature = "test-seams"))]
vector_fault_controller,
)?;
let mut exact_score = full_precision || segment.meta().scheme == 0;
let mut approximate = false;
if let Some(rescore) = rescore_options {
let lease = SnapshotLease::new_at(Arc::clone(snapshot), generation);
let cancellation = QueryCancellation::new(&control, &lease);
let vectors = exact_rescore_rows_for_search(
segment,
#[cfg(any(test, feature = "test-seams"))]
vector_fault_controller,
#[cfg(any(test, feature = "test-seams"))]
vector_fault_tier(options.tier()),
)?;
let (counters, exhaustive, worst) = rescored_scan::rescore(
&mut outcome.candidates,
&mut outcome.stats,
vectors,
query,
k,
alive.live_count(),
rescore,
&cancellation,
accounting,
&mut exact_memory,
)?;
outcome.worst_score = worst;
exact_score = true;
approximate = !exhaustive;
rescored_scan::accumulate(&mut rescore_counters, counters)?;
}
fold_worst_squared_l2(&mut worst_squared_l2, &mut worst_exhaustive, &outcome);
merge_store_outcome(
outcome,
source,
|row| {
segment
.document_version(row)
.map_err(StoreError::Segment)
.map_err(QueryError::Store)
},
&mut candidates,
&mut dims_touched,
&mut bytes_read,
&mut worker_thread_ids,
exact_score,
#[cfg(any(test, feature = "test-seams"))]
vector_fault_controller,
#[cfg(any(test, feature = "test-seams"))]
vector_fault_tier(options.tier()),
)?;
let tier = if segment
.directory()
.iter()
.any(|entry| entry.kind == crate::segment::layout::RegionKind::GraphNodeBlocks.id())
{
crate::planner::SegmentTier::SealedGraph
} else {
crate::planner::SegmentTier::SealedScan
};
let mut plan = crate::planner::SegmentPlan::unfiltered_scan(
source,
tier,
alive.live_count(),
sealed_scan_reason(segment, requested_tier, hybrid_scan_reason),
);
plan.approximate = approximate;
plans.push(plan);
if matches!(graph_bound_mode, GraphBoundMode::Shared) {
retain_global_top_k(&mut candidates, k);
}
}
retain_global_top_k(&mut candidates, k);
let stats = ScanStats {
dims_touched,
bytes_read,
threads_used: worker_thread_ids.len(),
worker_thread_ids,
};
accounting.record_plans(&plans).map_err(QueryError::Store)?;
let vector_ceiling = Some(exact_vector_ceiling(snapshot, active, query)?);
let mut diagnostics = crate::diag::QueryDiagnostics::vector(crate::diag::VectorDiagnostics {
snapshot_generation: generation,
indexed_through_seq: active
.indexed_through_seq()
.max(crate::wal::LogSeq::new(snapshot.absorbed_through())),
approximate: plans.iter().any(|plan| plan.approximate),
exact_rescore: candidates.iter().all(|candidate| candidate.exact_score()),
requested_k: k,
returned: candidates.len(),
budget_exhausted: false,
plan: plans,
scan: stats.clone(),
graph: graph_stats,
epoch,
elapsed: clock.now().saturating_duration_since(started),
});
diagnostics.scan_rescore = rescore_options.map(|_| rescore_counters);
diagnostics.timings = crate::diag::QUERY_TIMING_ENABLED.then(|| crate::diag::QueryTimings {
vector: timing_elapsed(clock, vector_started),
..Default::default()
});
Ok(SearchOutcome {
candidates,
vector_ceiling,
stats,
graph_stats,
generation,
epoch,
diagnostics,
})
}
const GRAPH_WIDENING_CAP_NUMERATOR: usize = 1;
const GRAPH_WIDENING_CAP_DENOMINATOR: usize = 4;
fn graph_round_is_worth_it(k: usize, corpus_rows: usize) -> bool {
if corpus_rows == 0 {
return true;
}
let ef = graph_round_ef(k);
let cap =
corpus_rows.saturating_mul(GRAPH_WIDENING_CAP_NUMERATOR) / GRAPH_WIDENING_CAP_DENOMINATOR;
ef <= cap
}
fn graph_round_ef(k: usize) -> usize {
k.saturating_mul(4).max(140)
}
fn sealed_scan_reason(
segment: &crate::segment::reader::SegmentReader,
requested_tier: Option<SearchTier>,
hybrid_scan_reason: Option<crate::planner::ScanReason>,
) -> crate::planner::ScanReason {
use crate::planner::{ExplicitScanTier, ScanReason};
let has_graph = segment
.directory()
.iter()
.any(|entry| entry.kind == crate::segment::layout::RegionKind::GraphNodeBlocks.id());
if has_graph && let Some(reason) = hybrid_scan_reason {
return reason;
}
match requested_tier {
Some(SearchTier::Exact) => return ScanReason::ExplicitTier(ExplicitScanTier::Exact),
Some(SearchTier::Scan) => return ScanReason::ExplicitTier(ExplicitScanTier::Scan),
Some(SearchTier::Auto | SearchTier::Graph(_)) | None => {}
}
let rows = segment.meta().row_count;
let min_rows = crate::tier::thresholds::for_bucket(segment.meta().dims, segment.meta().scheme)
.graph_min_rows;
if rows < min_rows {
ScanReason::BelowGraphThreshold { rows, min_rows }
} else {
ScanReason::GraphPending
}
}
fn snapshot_has_graph(snapshot: &PublishedSnapshot) -> bool {
snapshot.segments().iter().any(|segment| {
segment
.directory()
.iter()
.any(|entry| entry.kind == crate::segment::layout::RegionKind::GraphNodeBlocks.id())
})
}
fn exact_vector_ceiling(
snapshot: &PublishedSnapshot,
active: &crate::ingest::ActiveSegment,
query: &[f32],
) -> Result<f64, QueryError> {
fold_exact_vector_ceiling(
snapshot,
active,
query,
remembered_segment_vector_ceiling_norm_range,
)
}
fn remembered_segment_vector_ceiling_norm_range(
segment: &crate::segment::reader::SegmentReader,
) -> Result<crate::graph::search::GraphSegmentNormRange, QueryError> {
if let Some(range) = segment.cached_vector_ceiling_norm_range() {
return Ok(range);
}
let range = compute_segment_vector_ceiling_norm_range(segment)?;
segment.remember_vector_ceiling_norm_range(range);
Ok(range)
}
fn fold_exact_vector_ceiling(
snapshot: &PublishedSnapshot,
active: &crate::ingest::ActiveSegment,
query: &[f32],
mut segment_range: impl FnMut(
&crate::segment::reader::SegmentReader,
)
-> Result<crate::graph::search::GraphSegmentNormRange, QueryError>,
) -> Result<f64, QueryError> {
let mut ceiling = 0.0_f64;
if !active.is_empty() {
ceiling = fold_vector_ceiling(
ceiling,
crate::graph::search::GraphSegmentNormRange::from_factors(active.factors()),
query,
)?;
}
for segment in snapshot.segments() {
let range = segment_range(segment)?;
ceiling = fold_vector_ceiling(ceiling, range, query)?;
}
Ok(ceiling)
}
fn compute_segment_vector_ceiling_norm_range(
segment: &crate::segment::reader::SegmentReader,
) -> Result<crate::graph::search::GraphSegmentNormRange, QueryError> {
match segment.meta().scheme {
4 => Ok(crate::graph::search::GraphSegmentNormRange::from_factors(
segment
.query_bit4_factors()
.map_err(StoreError::Segment)
.map_err(QueryError::Store)?,
)),
0 | 2 => Ok(
crate::graph::search::GraphSegmentNormRange::from_exact_rows(
segment
.query_rescore_f32()
.map_err(StoreError::Segment)
.map_err(QueryError::Store)?,
segment.meta().dims as usize,
),
),
scheme => Err(QueryError::Store(StoreError::Segment(
crate::segment::SegmentError::Geometry(format!(
"vector norm enclosure does not support scheme {scheme}"
)),
))),
}
}
fn fold_vector_ceiling(
current: f64,
range: crate::graph::search::GraphSegmentNormRange,
query: &[f32],
) -> Result<f64, QueryError> {
let ceiling = range.squared_l2_upper_bound(query);
if !ceiling.is_finite() {
return Err(QueryError::Store(StoreError::Segment(
crate::segment::SegmentError::Geometry(
"vector norm enclosure did not produce a finite squared-L2 ceiling".to_owned(),
),
)));
}
Ok(current.max(ceiling))
}
fn fold_worst_squared_l2(
worst: &mut Option<f64>,
exhaustive: &mut bool,
outcome: &crate::scan::ScanOutcome,
) {
match outcome.worst_score {
Some(score) => {
let squared_l2 = -f64::from(score);
*worst = Some(worst.map_or(squared_l2, |current: f64| current.max(squared_l2)));
}
None => {
if !outcome.candidates.is_empty() {
*exhaustive = false;
}
}
}
}
#[allow(clippy::large_enum_variant)]
enum SegmentGraphResult<'a> {
Pruned {
plan: crate::planner::SegmentPlan,
graph_stats: crate::ingest::GraphSearchStats,
_lease: SnapshotLease,
},
Traversed {
result: crate::graph::search::GraphSearchResult,
plan: crate::planner::SegmentPlan,
dims_touched: u64,
bytes_read: u64,
caller_thread: std::thread::ThreadId,
graph_stats: crate::ingest::GraphSearchStats,
_scratch: graph_cache::GraphScratchLease<'a>,
_lease: SnapshotLease,
},
}
struct MergedSegmentGraph<'a> {
plan: crate::planner::SegmentPlan,
traversed: bool,
_result: Option<crate::graph::search::GraphSearchResult>,
_scratch: Option<graph_cache::GraphScratchLease<'a>>,
_lease: SnapshotLease,
}
impl<'a> SegmentGraphResult<'a> {
#[allow(clippy::too_many_arguments)]
fn merge_into(
self,
segment: &crate::segment::reader::SegmentReader,
alive: &crate::meta::AliveSet,
source: crate::ingest::RowSource,
candidates: &mut Vec<crate::ingest::SearchCandidate>,
dims_touched: &mut u64,
bytes_read: &mut u64,
worker_thread_ids: &mut Vec<std::thread::ThreadId>,
graph_stats: &mut crate::ingest::GraphSearchStats,
#[cfg(any(test, feature = "test-seams"))] vector_fault_controller: Option<
&crate::scan::vector_fault::VectorFaultController,
>,
_tier: SearchTier,
) -> Result<MergedSegmentGraph<'a>, QueryError> {
match self {
Self::Pruned {
plan,
graph_stats: delta,
_lease,
} => {
add_traversal_stats(graph_stats, delta)?;
Ok(MergedSegmentGraph {
plan,
traversed: false,
_result: None,
_scratch: None,
_lease,
})
}
Self::Traversed {
result,
plan,
dims_touched: traversal_dims_touched,
bytes_read: traversal_bytes_read,
caller_thread,
graph_stats: delta,
_scratch,
_lease,
} => {
*dims_touched = dims_touched
.checked_add(traversal_dims_touched)
.ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
*bytes_read = bytes_read
.checked_add(traversal_bytes_read)
.ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
if !worker_thread_ids.contains(&caller_thread) {
worker_thread_ids.push(caller_thread);
}
add_traversal_stats(graph_stats, delta)?;
reserve_global_candidates(
candidates,
result
.candidates()
.iter()
.filter(|candidate| alive.is_alive(candidate.row_id()))
.count(),
#[cfg(any(test, feature = "test-seams"))]
vector_fault_controller,
#[cfg(any(test, feature = "test-seams"))]
vector_fault_tier(_tier),
)?;
for candidate in result.candidates() {
if !alive.is_alive(candidate.row_id()) {
continue;
}
let score = -(candidate.distance() as f32);
if !score.is_finite() {
return Err(QueryError::Graph(
crate::graph::search::GraphSearchError::Geometry(format!(
"graph result distance for row {} is not representable as f32",
candidate.row_id()
)),
));
}
candidates.push(crate::ingest::SearchCandidate::new(
crate::ingest::GlobalRowId::new(source, candidate.row_id()),
segment
.document_version(candidate.row_id() as usize)
.map_err(StoreError::Segment)
.map_err(QueryError::Store)?,
score,
true,
));
}
Ok(MergedSegmentGraph {
plan,
traversed: true,
_result: Some(result),
_scratch: Some(_scratch),
_lease,
})
}
}
}
}
#[allow(clippy::too_many_arguments)]
fn traverse_segment_graph<'a>(
segment: &'a crate::segment::reader::SegmentReader,
alive: &crate::meta::AliveSet,
source: crate::ingest::RowSource,
graph_options: GraphSearchOptions,
graph_bound_mode: GraphBoundMode,
request: crate::ingest::SearchRequest<'_>,
k: usize,
vector_preparation: &mut prepared::PreparedVectorQuery<'_>,
accounting: &Arc<stats::Accounting>,
snapshot: &Arc<PublishedSnapshot>,
generation: u64,
control: &QueryControl,
competitive_distance: Option<f32>,
#[cfg(any(test, feature = "test-seams"))] vector_fault_controller: Option<
&crate::scan::vector_fault::VectorFaultController,
>,
_tier: SearchTier,
) -> Result<SegmentGraphResult<'a>, QueryError> {
let lease = SnapshotLease::new_at(Arc::clone(snapshot), generation);
let cancellation = QueryCancellation::new(control, &lease);
cancellation.check_graph().map_err(map_scan_error)?;
let (graph, graph_validated, prepared_entry_seed_discovered, norm_range) =
match graph_bound_mode {
GraphBoundMode::Shared => {
let prepared = segment
.graph_search_cache
.prepare_shared(segment, &cancellation)
.map_err(map_graph_cache_error)?;
(
prepared.graph,
prepared.graph_validated,
prepared.entry_seed_discovered,
Some(prepared.norm_range),
)
}
#[cfg(test)]
GraphBoundMode::Independent => {
let (graph, graph_validated) = segment
.graph_search_cache
.bind_graph(segment)
.map_err(map_graph_cache_error)?;
(graph, graph_validated, false, None)
}
};
let rescore = query_rescore_rows_for_search(
segment,
#[cfg(any(test, feature = "test-seams"))]
vector_fault_controller,
#[cfg(any(test, feature = "test-seams"))]
vector_fault_tier(_tier),
)?;
let node_count = graph.node_count() as usize;
let segment_k = k.min(node_count);
let has_tombstones = alive.tombstone_count() != 0;
let target_live_k = if has_tombstones {
usize::try_from(alive.live_count())
.map_err(|_| QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?
.min(k)
} else {
segment_k
};
let maximum_candidate_k = if has_tombstones {
let tombstone_count = usize::try_from(alive.tombstone_count())
.map_err(|_| QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
target_live_k
.checked_add(tombstone_count)
.ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?
.min(node_count)
.max(segment_k)
} else {
segment_k
};
if let Some(ef) = graph_options.ef()
&& ef < segment_k
{
return Err(QueryError::Graph(
crate::graph::search::GraphSearchError::AdaptiveEf(
crate::graph::search::AdaptiveEfError::ExplicitBelowK { k: segment_k, ef },
),
));
}
let request_for = |candidate_k| {
let graph_request = crate::graph::search::GraphSearchRequest::new(
request.vector(),
candidate_k,
graph_options.seed(),
)
.with_profile(graph_options.profile());
graph_options.ef().map_or(graph_request, |ef| {
graph_request.with_ef(ef.max(candidate_k))
})
};
let mut candidate_k = segment_k;
let mut graph_request = request_for(candidate_k);
let effective_ef = graph_request
.effective_ef(graph.node_count() as usize)
.map_err(crate::graph::search::GraphSearchError::AdaptiveEf)
.map_err(QueryError::Graph)?;
let requestable_max_k = maximum_candidate_k;
let scratch_ef = if has_tombstones && requestable_max_k > candidate_k {
request_for(requestable_max_k)
.effective_ef(node_count)
.map_err(crate::graph::search::GraphSearchError::AdaptiveEf)
.map_err(QueryError::Graph)?
} else {
effective_ef
};
if let (Some(norm_range), Some(competitive_distance)) = (norm_range, competitive_distance)
&& norm_range.squared_l2_upper_bound(request.vector()) <= f64::from(f32::MAX)
&& (norm_range.squared_l2_lower_bound(request.vector()) as f32) > competitive_distance
{
cancellation.check_graph().map_err(map_scan_error)?;
let graph_stats = crate::ingest::GraphSearchStats {
graph_validations: usize::from(graph_validated),
entry_seed_discoveries: usize::from(prepared_entry_seed_discovered),
segments_pruned_by_bound: 1,
..crate::ingest::GraphSearchStats::default()
};
let plan = crate::planner::SegmentPlan::unfiltered_pruned(source, alive.live_count());
return Ok(SegmentGraphResult::Pruned {
plan,
graph_stats,
_lease: lease,
});
}
let mut scratch = segment
.graph_search_cache
.checkout(graph, scratch_ef, accounting, &cancellation)
.map_err(map_graph_cache_error)?;
let entry_seed_discovered = prepared_entry_seed_discovered || scratch.entry_seed_discovered();
let entries = scratch.entries();
let mut searcher = crate::graph::search::GraphSearcher::with_entry_row_ids(
graph,
rescore,
entries,
scratch.scratch_mut().map_err(map_graph_error)?,
)
.map_err(map_graph_error)?
.with_rescore_validator(segment);
#[cfg(any(test, feature = "test-seams"))]
if let Some(after_hops) = graph_options.cancel_after_hops() {
let token = cancellation.cancel_token_for_test().ok_or_else(|| {
QueryError::Graph(crate::graph::search::GraphSearchError::Geometry(
"test hop cancellation requires a Cancel query control".to_owned(),
))
})?;
let _observed_hops = searcher.cancel_after_hops(after_hops, token);
}
#[cfg(any(test, feature = "test-seams"))]
if let Some(controller) = vector_fault_controller {
searcher = searcher.with_vector_fault_controller(
controller,
vector_fault_source(source),
crate::scan::vector_fault::VectorSearchTier::Graph,
);
}
let mut traversal_dims_touched = 0_u64;
let mut traversal_bytes_read = 0_u64;
let mut traversal_epoch_clears = 0_usize;
let mut traversal_candidates_scored = 0_usize;
let mut traversal_candidates_rescored = 0_usize;
let result = loop {
let result = searcher.search_reusing(
graph_request,
Some(&cancellation),
vector_preparation,
segment,
source,
)?;
let counters = result.counters();
traversal_dims_touched = traversal_dims_touched
.checked_add(counters.dims_touched())
.ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
traversal_bytes_read = traversal_bytes_read
.checked_add(counters.bytes_read())
.ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
traversal_epoch_clears = traversal_epoch_clears
.checked_add(usize::from(counters.visited_epoch_cleared()))
.ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
traversal_candidates_scored = traversal_candidates_scored
.checked_add(counters.candidates_scored())
.ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
traversal_candidates_rescored = traversal_candidates_rescored
.checked_add(counters.candidates_rescored())
.ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
if !has_tombstones {
break result;
}
let retained_live = result
.candidates()
.iter()
.filter(|candidate| alive.is_alive(candidate.row_id()))
.count();
if retained_live >= target_live_k {
break result;
}
if candidate_k >= requestable_max_k {
return Err(QueryError::Graph(
crate::graph::search::GraphSearchError::Geometry(format!(
"graph traversal exhausted {candidate_k} candidates but retained {retained_live} live rows, expected {target_live_k}"
)),
));
}
let widened_k = candidate_k
.saturating_mul(2)
.max(candidate_k.saturating_add(1))
.min(requestable_max_k);
if widened_k <= candidate_k {
return Err(QueryError::Graph(
crate::graph::search::GraphSearchError::Geometry(format!(
"graph traversal could not widen beyond {candidate_k} candidates"
)),
));
}
candidate_k = widened_k;
graph_request = request_for(candidate_k);
};
let graph_stats = crate::ingest::GraphSearchStats {
segments_traversed: 1,
graph_validations: usize::from(graph_validated),
entry_seed_discoveries: usize::from(entry_seed_discovered),
visited_epoch_clears: traversal_epoch_clears,
candidates_scored: traversal_candidates_scored,
candidates_rescored: traversal_candidates_rescored,
segments_pruned_by_bound: 0,
};
let plan = crate::planner::SegmentPlan::unfiltered_graph(
source,
alive.live_count(),
graph_options.ef(),
effective_ef,
graph_options.profile(),
);
let caller_thread = std::thread::current().id();
#[cfg(any(test, feature = "test-seams"))]
drop(searcher);
Ok(SegmentGraphResult::Traversed {
result,
plan,
dims_touched: traversal_dims_touched,
bytes_read: traversal_bytes_read,
caller_thread,
graph_stats,
_scratch: scratch,
_lease: lease,
})
}
fn add_traversal_stats(
stats: &mut crate::ingest::GraphSearchStats,
delta: crate::ingest::GraphSearchStats,
) -> Result<(), QueryError> {
stats.segments_traversed = stats
.segments_traversed
.checked_add(delta.segments_traversed)
.ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
stats.graph_validations = stats
.graph_validations
.checked_add(delta.graph_validations)
.ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
stats.entry_seed_discoveries = stats
.entry_seed_discoveries
.checked_add(delta.entry_seed_discoveries)
.ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
stats.visited_epoch_clears = stats
.visited_epoch_clears
.checked_add(delta.visited_epoch_clears)
.ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
stats.candidates_scored = stats
.candidates_scored
.checked_add(delta.candidates_scored)
.ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
stats.candidates_rescored = stats
.candidates_rescored
.checked_add(delta.candidates_rescored)
.ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
stats.segments_pruned_by_bound = stats
.segments_pruned_by_bound
.checked_add(delta.segments_pruned_by_bound)
.ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
Ok(())
}
#[allow(clippy::too_many_arguments)]
fn scan_sealed_segment(
pool: Option<&pool::QueryPool>,
snapshot: &Arc<PublishedSnapshot>,
segment: &crate::segment::reader::SegmentReader,
alive: &crate::meta::AliveSet,
generation: u64,
request: crate::ingest::SearchRequest<'_>,
k: usize,
scan_options: crate::scan::ScanOptions,
control: &QueryControl,
bit4_query: &std::cell::OnceCell<Result<crate::quant::Bit4Query, crate::quant::QuantError>>,
int8_query: &std::cell::OnceCell<Result<crate::quant::Int8Query, crate::quant::QuantError>>,
full_precision: bool,
exact_lane: usize,
accounting: &Arc<stats::Accounting>,
exact_memory: &mut ExactScanMemory,
_source: crate::ingest::RowSource,
_tier: SearchTier,
#[cfg(any(test, feature = "test-seams"))] vector_fault_controller: Option<
&crate::scan::vector_fault::VectorFaultController,
>,
) -> Result<crate::scan::ScanOutcome, QueryError> {
use crate::scan::{ScanQuery, ScanRequest, ScanRows};
if full_precision {
let lease = SnapshotLease::new_at(Arc::clone(snapshot), generation);
let vectors = exact_rescore_rows_for_search(
segment,
#[cfg(any(test, feature = "test-seams"))]
vector_fault_controller,
#[cfg(any(test, feature = "test-seams"))]
vector_fault_tier(_tier),
)?;
return scan_squared_l2(
vectors,
segment.meta().row_count as usize,
alive,
request.vector(),
k,
ExactScanDispatch {
pool,
accounting,
control,
lease: &lease,
thread_budget: scan_options.thread_budget,
lane: exact_lane,
},
exact_memory,
#[cfg(any(test, feature = "test-seams"))]
vector_fault_controller,
#[cfg(any(test, feature = "test-seams"))]
vector_fault_source(_source),
#[cfg(any(test, feature = "test-seams"))]
vector_fault_tier(_tier),
);
}
let query_pool = pool.ok_or(QueryError::Store(StoreError::Synchronization {
component: "sealed scan-tier query pool",
}))?;
match segment.meta().scheme {
0 => execute_store_scan(
query_pool,
ScanRequest {
query: ScanQuery::F32(request.vector()),
rows: ScanRows::F32BorrowedRowMajor(
segment
.query_f32_codes()
.map_err(StoreError::Segment)
.map_err(QueryError::Store)?,
),
row_mask: Some(alive.scan_mask()),
},
k,
scan_options,
control.clone(),
SnapshotLease::new_at(Arc::clone(snapshot), generation),
#[cfg(any(test, feature = "test-seams"))]
vector_fault_controller,
#[cfg(any(test, feature = "test-seams"))]
vector_fault_source(_source),
#[cfg(any(test, feature = "test-seams"))]
vector_fault_tier(_tier),
),
4 => execute_store_scan(
query_pool,
ScanRequest {
query: ScanQuery::Bit4(
bit4_query
.get_or_init(|| crate::quant::prepare_bit4_query(request.vector(), 0))
.as_ref()
.map_err(|error| crate::scan::ScanError::Quant(error.clone()))
.map_err(QueryError::Scan)?,
),
rows: ScanRows::Bit4RowMajor {
codes: segment
.query_bit4_codes()
.map_err(StoreError::Segment)
.map_err(QueryError::Store)?,
factors: segment
.query_bit4_factors()
.map_err(StoreError::Segment)
.map_err(QueryError::Store)?,
},
row_mask: Some(alive.scan_mask()),
},
k,
scan_options,
control.clone(),
SnapshotLease::new_at(Arc::clone(snapshot), generation),
#[cfg(any(test, feature = "test-seams"))]
vector_fault_controller,
#[cfg(any(test, feature = "test-seams"))]
vector_fault_source(_source),
#[cfg(any(test, feature = "test-seams"))]
vector_fault_tier(_tier),
),
2 => {
let factors = segment.query_int8_factors().map_err(QueryError::Store)?;
execute_store_scan(
query_pool,
ScanRequest {
query: ScanQuery::Int8(
int8_query
.get_or_init(|| crate::quant::prepare_int8_query(request.vector()))
.as_ref()
.map_err(|error| crate::scan::ScanError::Quant(error.clone()))
.map_err(QueryError::Scan)?,
),
rows: ScanRows::Int8RowMajor {
codes: segment
.query_int8_codes()
.map_err(StoreError::Segment)
.map_err(QueryError::Store)?,
factors: factors.as_slice(),
},
row_mask: Some(alive.scan_mask()),
},
k,
scan_options,
control.clone(),
SnapshotLease::new_at(Arc::clone(snapshot), generation),
#[cfg(any(test, feature = "test-seams"))]
vector_fault_controller,
#[cfg(any(test, feature = "test-seams"))]
vector_fault_source(_source),
#[cfg(any(test, feature = "test-seams"))]
vector_fault_tier(_tier),
)
}
scheme => Err(QueryError::Store(StoreError::Segment(
crate::segment::SegmentError::Geometry(format!(
"store search does not support sealed scheme {scheme}"
)),
))),
}
}
#[allow(clippy::too_many_arguments)]
fn execute_store_scan(
pool: &pool::QueryPool,
request: crate::scan::ScanRequest<'_>,
k: usize,
options: crate::scan::ScanOptions,
control: QueryControl,
lease: SnapshotLease,
#[cfg(any(test, feature = "test-seams"))] controller: Option<
&crate::scan::vector_fault::VectorFaultController,
>,
#[cfg(any(test, feature = "test-seams"))] source: crate::scan::vector_fault::VectorRowSource,
#[cfg(any(test, feature = "test-seams"))] tier: crate::scan::vector_fault::VectorSearchTier,
) -> Result<crate::scan::ScanOutcome, QueryError> {
#[cfg(any(test, feature = "test-seams"))]
if let Some(controller) = controller {
let geometry = crate::scan::scan_geometry(request).map_err(QueryError::Scan)?;
let cancellation = QueryCancellation::new(&control, &lease);
let partition = crate::scan::scan_partition_with_vector_faults(
request,
k,
0..geometry.row_count,
Some(&cancellation),
controller,
source,
tier,
)
.map_err(map_scan_error)?;
return Ok(crate::scan::ScanOutcome {
candidates: partition.candidates,
worst_score: None,
stats: crate::scan::ScanStats {
dims_touched: partition.dims_touched,
bytes_read: partition.bytes_read,
threads_used: 1,
worker_thread_ids: vec![partition.worker_thread_id],
},
});
}
pool.execute(request, k, options, control, lease)
}
#[cfg(any(test, feature = "test-seams"))]
fn vector_fault_tier(tier: SearchTier) -> crate::scan::vector_fault::VectorSearchTier {
match tier {
SearchTier::Exact => crate::scan::vector_fault::VectorSearchTier::Exact,
SearchTier::Scan => crate::scan::vector_fault::VectorSearchTier::Scan,
SearchTier::Graph(_) => crate::scan::vector_fault::VectorSearchTier::Graph,
SearchTier::Auto => crate::scan::vector_fault::VectorSearchTier::Auto,
}
}
#[cfg(any(test, feature = "test-seams"))]
fn vector_fault_source(
source: crate::ingest::RowSource,
) -> crate::scan::vector_fault::VectorRowSource {
match source {
crate::ingest::RowSource::Active => crate::scan::vector_fault::VectorRowSource::Active,
crate::ingest::RowSource::Sealed(segment) => {
crate::scan::vector_fault::VectorRowSource::Sealed(*segment.as_bytes())
}
}
}
pub(crate) fn auto_uses_full_precision(
snapshot: &PublishedSnapshot,
active: &crate::ingest::ActiveSegment,
) -> bool {
let mut has_exact = false;
let mut has_estimated = !active.is_empty();
for segment in snapshot.segments() {
if segment
.directory()
.iter()
.any(|entry| entry.kind == crate::segment::layout::RegionKind::GraphNodeBlocks.id())
{
return true;
}
if segment.meta().row_count == 0 {
continue;
}
match segment.meta().scheme {
0 => has_exact = true,
2 | 4 => has_estimated = true,
_ => {}
}
}
has_exact && has_estimated
}
pub(crate) fn exact_rescore_rows(
segment: &crate::segment::reader::SegmentReader,
) -> Result<&[f32], QueryError> {
segment
.query_rescore_f32()
.map_err(|source| QueryError::Store(StoreError::Segment(source)))
}
#[cfg(not(any(test, feature = "test-seams")))]
use exact_rescore_rows as exact_rescore_rows_for_search;
#[cfg(any(test, feature = "test-seams"))]
fn exact_rescore_rows_for_search<'a>(
segment: &'a crate::segment::reader::SegmentReader,
#[cfg(any(test, feature = "test-seams"))] controller: Option<
&crate::scan::vector_fault::VectorFaultController,
>,
#[cfg(any(test, feature = "test-seams"))] tier: crate::scan::vector_fault::VectorSearchTier,
) -> Result<&'a [f32], QueryError> {
#[cfg(any(test, feature = "test-seams"))]
if let Some(available_rows) = controller.and_then(|controller| {
controller.missing_rescore_rows(
crate::scan::vector_fault::VectorRowSource::Sealed(*segment.meta().id.as_bytes()),
tier,
crate::scan::vector_fault::MissingRescoreSite::ExactRescoreRows,
segment.meta().row_count,
)
}) {
return Err(missing_rescore_rows_error(segment, available_rows));
}
exact_rescore_rows(segment)
}
pub(crate) fn query_rescore_rows(
segment: &crate::segment::reader::SegmentReader,
) -> Result<&[f32], QueryError> {
segment
.query_rescore_f32()
.map_err(|source| QueryError::Store(StoreError::Segment(source)))
}
#[cfg(not(any(test, feature = "test-seams")))]
use query_rescore_rows as query_rescore_rows_for_search;
#[cfg(any(test, feature = "test-seams"))]
fn query_rescore_rows_for_search<'a>(
segment: &'a crate::segment::reader::SegmentReader,
#[cfg(any(test, feature = "test-seams"))] controller: Option<
&crate::scan::vector_fault::VectorFaultController,
>,
#[cfg(any(test, feature = "test-seams"))] tier: crate::scan::vector_fault::VectorSearchTier,
) -> Result<&'a [f32], QueryError> {
#[cfg(any(test, feature = "test-seams"))]
if let Some(available_rows) = controller.and_then(|controller| {
controller.missing_rescore_rows(
crate::scan::vector_fault::VectorRowSource::Sealed(*segment.meta().id.as_bytes()),
tier,
crate::scan::vector_fault::MissingRescoreSite::QueryRescoreRows,
segment.meta().row_count,
)
}) {
return Err(missing_rescore_rows_error(segment, available_rows));
}
query_rescore_rows(segment)
}
#[cfg(any(test, feature = "test-seams"))]
fn missing_rescore_rows_error(
segment: &crate::segment::reader::SegmentReader,
available_rows: u32,
) -> QueryError {
QueryError::Store(StoreError::Segment(crate::segment::SegmentError::Geometry(
format!(
"exact scores unavailable for segment {}: expected {} rows, got {available_rows}",
segment.meta().id,
segment.meta().row_count,
),
)))
}
fn retain_global_top_k(candidates: &mut Vec<crate::ingest::SearchCandidate>, k: usize) {
candidates.sort_unstable_by(crate::ingest::compare_search_candidates);
candidates.truncate(k);
}
fn global_competitive_distance(
candidates: &[crate::ingest::SearchCandidate],
k: usize,
) -> Option<f32> {
if k == 0 || candidates.len() < k {
return None;
}
candidates
.get(k.saturating_sub(1))
.map(|candidate| -candidate.score())
.filter(|distance| distance.is_finite() && *distance >= 0.0)
}
fn map_graph_cache_error(error: graph_cache::GraphCacheError) -> QueryError {
match error {
graph_cache::GraphCacheError::Store(error) => QueryError::Store(error),
graph_cache::GraphCacheError::Search(error) => map_graph_error(error),
}
}
pub(crate) fn map_graph_error(error: crate::graph::search::GraphSearchError) -> QueryError {
match error {
crate::graph::search::GraphSearchError::ExactRescoreUnavailable(detail) => {
QueryError::Store(StoreError::Segment(crate::segment::SegmentError::Geometry(
detail,
)))
}
crate::graph::search::GraphSearchError::Cancelled { partial } => {
QueryError::Cancelled { partial }
}
crate::graph::search::GraphSearchError::Timeout { partial } => {
QueryError::Timeout { partial }
}
crate::graph::search::GraphSearchError::ReadCancelled { partial } => {
QueryError::ReadCancelled { partial }
}
crate::graph::search::GraphSearchError::Scan(error) => QueryError::Scan(error),
error => QueryError::Graph(error),
}
}
#[allow(
clippy::too_many_arguments,
reason = "scoped dispatch and test source identity stay explicit"
)]
fn scan_active_squared_l2(
active: &crate::ingest::ActiveSegment,
alive: &crate::meta::AliveSet,
query: &[f32],
k: usize,
dispatch: ExactScanDispatch<'_>,
memory: &mut ExactScanMemory,
#[cfg(any(test, feature = "test-seams"))] controller: Option<
&crate::scan::vector_fault::VectorFaultController,
>,
#[cfg(any(test, feature = "test-seams"))] tier: crate::scan::vector_fault::VectorSearchTier,
) -> Result<crate::scan::ScanOutcome, QueryError> {
scan_squared_l2(
active.vectors(),
active.row_count(),
alive,
query,
k,
dispatch,
memory,
#[cfg(any(test, feature = "test-seams"))]
controller,
#[cfg(any(test, feature = "test-seams"))]
crate::scan::vector_fault::VectorRowSource::Active,
#[cfg(any(test, feature = "test-seams"))]
tier,
)
}
#[derive(Debug)]
enum ExactScanError {
Query(QueryError),
Narrowing { row_id: usize, score: f64 },
}
impl From<QueryError> for ExactScanError {
fn from(error: QueryError) -> Self {
Self::Query(error)
}
}
impl ExactScanError {
fn into_query(self) -> QueryError {
match self {
Self::Query(error) => error,
Self::Narrowing { row_id, .. } => {
QueryError::Scan(crate::scan::ScanError::NonFiniteScore { row_id })
}
}
}
}
pub(crate) struct ExactScanMemory {
candidates: stats::AccountedCounter,
worker_ids: stats::AccountedCounter,
}
impl ExactScanMemory {
pub(crate) fn new(accounting: &Arc<stats::Accounting>) -> Result<Self, QueryError> {
Ok(Self {
candidates: stats::AccountedCounter::new(
accounting,
stats::AllocationComponent::Temporary,
)
.map_err(QueryError::Store)?,
worker_ids: stats::AccountedCounter::new(
accounting,
stats::AllocationComponent::Temporary,
)
.map_err(QueryError::Store)?,
})
}
}
struct ExactScanDispatch<'a> {
pool: Option<&'a pool::QueryPool>,
accounting: &'a Arc<stats::Accounting>,
control: &'a QueryControl,
lease: &'a SnapshotLease,
thread_budget: usize,
lane: usize,
}
#[allow(
clippy::too_many_arguments,
reason = "scoped dispatch and test source identity stay explicit"
)]
fn scan_squared_l2(
vectors: &[f32],
row_count: usize,
alive: &crate::meta::AliveSet,
query: &[f32],
k: usize,
dispatch: ExactScanDispatch<'_>,
memory: &mut ExactScanMemory,
#[cfg(any(test, feature = "test-seams"))] controller: Option<
&crate::scan::vector_fault::VectorFaultController,
>,
#[cfg(any(test, feature = "test-seams"))] source: crate::scan::vector_fault::VectorRowSource,
#[cfg(any(test, feature = "test-seams"))] tier: crate::scan::vector_fault::VectorSearchTier,
) -> Result<crate::scan::ScanOutcome, QueryError> {
if let Some(pool) = dispatch.pool
&& pool.exact_workers(row_count, query.len(), dispatch.thread_budget) > 1
{
return pool.execute_exact(
vectors,
row_count,
alive,
query,
k,
dispatch.thread_budget,
dispatch.lane,
dispatch.control,
dispatch.lease,
dispatch.accounting,
memory,
#[cfg(any(test, feature = "test-seams"))]
controller,
#[cfg(any(test, feature = "test-seams"))]
source,
#[cfg(any(test, feature = "test-seams"))]
tier,
);
}
let cancellation = QueryCancellation::new(dispatch.control, dispatch.lease);
#[cfg(any(test, feature = "test-seams"))]
if controller.is_some_and(|controller| controller.has_exact_row_hooks()) {
return scan_exact_partition::<false>(
vectors,
row_count,
0..row_count,
alive,
query,
k,
&cancellation,
memory,
#[cfg(any(test, feature = "test-seams"))]
controller,
#[cfg(any(test, feature = "test-seams"))]
source,
#[cfg(any(test, feature = "test-seams"))]
tier,
)
.map_err(ExactScanError::into_query);
}
scan_serial_squared_l2(
vectors,
row_count,
alive,
query,
k,
&cancellation,
&mut memory.candidates,
&mut memory.worker_ids,
#[cfg(any(test, feature = "test-seams"))]
controller,
#[cfg(any(test, feature = "test-seams"))]
source,
#[cfg(any(test, feature = "test-seams"))]
tier,
)
}
#[allow(clippy::too_many_arguments)]
fn scan_exact_partition<const OFFSET_ROWS: bool>(
vectors: &[f32],
row_count: usize,
range: std::ops::Range<usize>,
alive: &crate::meta::AliveSet,
query: &[f32],
k: usize,
cancellation: &QueryCancellation<'_>,
memory: &mut ExactScanMemory,
#[cfg(any(test, feature = "test-seams"))] controller: Option<
&crate::scan::vector_fault::VectorFaultController,
>,
#[cfg(any(test, feature = "test-seams"))] source: crate::scan::vector_fault::VectorRowSource,
#[cfg(any(test, feature = "test-seams"))] tier: crate::scan::vector_fault::VectorSearchTier,
) -> Result<crate::scan::ScanOutcome, ExactScanError> {
if query.is_empty() {
return Err(QueryError::Scan(crate::scan::ScanError::ZeroDimension).into());
}
let expected = row_count
.checked_mul(query.len())
.ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
if vectors.len() != expected {
return Err(QueryError::Scan(crate::scan::ScanError::RowDataLength {
dimension: query.len(),
actual: vectors.len(),
})
.into());
}
if range.start > range.end || range.end > row_count || (!OFFSET_ROWS && range.start != 0) {
return Err(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow).into());
}
let mut live_rows = 0;
for row in range.clone() {
if row.is_multiple_of(64) {
cancellation.check_graph().map_err(map_scan_error)?;
}
let local_row = u32::try_from(row)
.map_err(|_| QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
live_rows += usize::from(alive.is_alive(local_row));
}
cancellation.check_graph().map_err(map_scan_error)?;
let mut collector = crate::scan::topk::ExactTopK::try_new(k, live_rows, &mut memory.candidates)
.map_err(QueryError::Store)?;
let mut worst_score: Option<f32> = None;
let mut narrowing_error: Option<(usize, f64)> = None;
let start = range
.start
.checked_mul(query.len())
.ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
let end = range
.end
.checked_mul(query.len())
.ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
let partition_vectors = vectors
.get(start..end)
.ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
let rows = partition_vectors
.chunks_exact(query.len())
.enumerate()
.map(|(local, vector)| {
(
if OFFSET_ROWS {
range.start + local
} else {
local
},
vector,
)
})
.filter(|(row, _)| u32::try_from(*row).is_ok_and(|row| alive.is_alive(row)));
#[cfg(any(test, feature = "test-seams"))]
let mut trace =
controller.and_then(|controller| controller.trace_rows(source, tier, range.clone()));
let scored_rows = crate::quant::exact_squared_l2_with_sink(
query,
rows,
|_row, is_checkpoint| {
#[cfg(any(test, feature = "test-seams"))]
if let Some(controller) = controller {
if let Some(trace) = &mut trace {
trace.checked(_row);
}
if is_checkpoint {
cancellation.check_graph().map_err(map_scan_error)?;
}
if controller.after_eligible_row(source, tier, _row) {
return Err(QueryError::Cancelled { partial: false });
}
return Ok(());
}
if is_checkpoint {
cancellation.check_graph().map_err(map_scan_error)?;
}
Ok(())
},
|row_id, exact| {
let score = exact as f32;
if !score.is_finite() {
if narrowing_error.is_none_or(|(row, previous)| {
exact.total_cmp(&previous).is_gt()
|| (exact.total_cmp(&previous).is_eq() && row_id < row)
}) {
narrowing_error = Some((row_id, exact));
}
return Ok(());
}
if worst_score.is_none_or(|worst| score.total_cmp(&worst).is_lt()) {
worst_score = Some(score);
}
collector
.try_push(
crate::scan::ScanCandidate { row_id, score },
&mut memory.candidates,
)
.map_err(QueryError::Store)
},
)
.map_err(|error| match error {
crate::quant::RescoreCheckError::Rescore(error) => map_l2_rescore_error(error),
crate::quant::RescoreCheckError::Check(error) => error,
})?;
cancellation.check_graph().map_err(map_scan_error)?;
if let Some((row_id, score)) = narrowing_error {
return Err(ExactScanError::Narrowing { row_id, score });
}
let (candidates, _collector_capacity) = collector
.try_into_sorted_with_ties(&mut memory.candidates)
.map_err(QueryError::Store)?;
#[cfg(any(test, feature = "test-seams"))]
if let Some(controller) = controller {
controller.record_exact_scan(crate::scan::vector_fault::ExactScanWork {
sorted_items: candidates.len(),
collector_capacity: _collector_capacity,
scored_rows,
worst_score,
..Default::default()
});
}
let dims = u64::try_from(query.len())
.map_err(|_| QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
let scored_rows = u64::try_from(scored_rows)
.map_err(|_| QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
let dims_touched = scored_rows
.checked_mul(dims)
.ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
let bytes_read = dims_touched
.checked_mul(std::mem::size_of::<f32>() as u64)
.ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
memory
.worker_ids
.set(std::mem::size_of::<std::thread::ThreadId>())
.map_err(QueryError::Store)?;
let mut worker_thread_ids = Vec::new();
worker_thread_ids.try_reserve_exact(1).map_err(|_| {
QueryError::Store(StoreError::AllocationFailed {
needed: memory.worker_ids.bytes(),
component: "exact worker IDs",
})
})?;
worker_thread_ids.push(std::thread::current().id());
Ok(crate::scan::ScanOutcome {
candidates,
worst_score,
stats: crate::scan::ScanStats {
dims_touched,
bytes_read,
threads_used: 1,
worker_thread_ids,
},
})
}
#[allow(clippy::too_many_arguments)]
fn scan_serial_squared_l2(
vectors: &[f32],
row_count: usize,
alive: &crate::meta::AliveSet,
query: &[f32],
k: usize,
cancellation: &QueryCancellation<'_>,
memory: &mut stats::AccountedCounter,
worker_memory: &mut stats::AccountedCounter,
#[cfg(any(test, feature = "test-seams"))] controller: Option<
&crate::scan::vector_fault::VectorFaultController,
>,
#[cfg(any(test, feature = "test-seams"))] source: crate::scan::vector_fault::VectorRowSource,
#[cfg(any(test, feature = "test-seams"))] tier: crate::scan::vector_fault::VectorSearchTier,
) -> Result<crate::scan::ScanOutcome, QueryError> {
if query.is_empty() {
return Err(QueryError::Scan(crate::scan::ScanError::ZeroDimension));
}
let expected = row_count
.checked_mul(query.len())
.ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
if vectors.len() != expected {
return Err(QueryError::Scan(crate::scan::ScanError::RowDataLength {
dimension: query.len(),
actual: vectors.len(),
}));
}
let mut live_rows = 0;
for row in 0..row_count {
if row.is_multiple_of(64) {
cancellation.check_graph().map_err(map_scan_error)?;
}
let local_row = u32::try_from(row)
.map_err(|_| QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
live_rows += usize::from(alive.is_alive(local_row));
}
cancellation.check_graph().map_err(map_scan_error)?;
let mut collector =
crate::scan::topk::ExactTopK::try_new(k, live_rows, memory).map_err(QueryError::Store)?;
let mut worst_score: Option<f32> = None;
let mut narrowing_error: Option<(usize, f64)> = None;
let rows = vectors
.chunks_exact(query.len())
.enumerate()
.filter(|(row, _)| u32::try_from(*row).is_ok_and(|row| alive.is_alive(row)));
let scored_rows = crate::quant::exact_squared_l2_with_sink(
query,
rows,
|_row, is_checkpoint| {
if is_checkpoint {
cancellation.check_graph().map_err(map_scan_error)?;
}
#[cfg(any(test, feature = "test-seams"))]
if controller
.is_some_and(|controller| controller.after_eligible_row(source, tier, _row))
{
return Err(QueryError::Cancelled { partial: false });
}
Ok(())
},
|row_id, exact| {
let score = exact as f32;
if !score.is_finite() {
if narrowing_error.is_none_or(|(row, previous)| {
exact.total_cmp(&previous).is_gt()
|| (exact.total_cmp(&previous).is_eq() && row_id < row)
}) {
narrowing_error = Some((row_id, exact));
}
return Ok(());
}
if worst_score.is_none_or(|worst| score.total_cmp(&worst).is_lt()) {
worst_score = Some(score);
}
collector
.try_push(crate::scan::ScanCandidate { row_id, score }, memory)
.map_err(QueryError::Store)
},
)
.map_err(|error| match error {
crate::quant::RescoreCheckError::Rescore(error) => map_l2_rescore_error(error),
crate::quant::RescoreCheckError::Check(error) => error,
})?;
cancellation.check_graph().map_err(map_scan_error)?;
if let Some((row_id, _)) = narrowing_error {
return Err(QueryError::Scan(crate::scan::ScanError::NonFiniteScore {
row_id,
}));
}
let (candidates, _collector_capacity) = collector
.try_into_sorted_with_ties(memory)
.map_err(QueryError::Store)?;
#[cfg(any(test, feature = "test-seams"))]
if let Some(controller) = controller {
controller.record_exact_scan(crate::scan::vector_fault::ExactScanWork {
sorted_items: candidates.len(),
collector_capacity: _collector_capacity,
scored_rows,
worst_score,
..Default::default()
});
}
let dims = u64::try_from(query.len())
.map_err(|_| QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
let scored_rows = u64::try_from(scored_rows)
.map_err(|_| QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
let dims_touched = scored_rows
.checked_mul(dims)
.ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
let bytes_read = dims_touched
.checked_mul(std::mem::size_of::<f32>() as u64)
.ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
worker_memory
.set(std::mem::size_of::<std::thread::ThreadId>())
.map_err(QueryError::Store)?;
let mut worker_thread_ids = Vec::new();
worker_thread_ids.try_reserve_exact(1).map_err(|_| {
QueryError::Store(StoreError::AllocationFailed {
needed: worker_memory.bytes(),
component: "exact worker IDs",
})
})?;
worker_thread_ids.push(std::thread::current().id());
Ok(crate::scan::ScanOutcome {
candidates,
worst_score,
stats: crate::scan::ScanStats {
dims_touched,
bytes_read,
threads_used: 1,
worker_thread_ids,
},
})
}
fn map_l2_rescore_error(error: crate::quant::RescoreError) -> QueryError {
let error = match error {
crate::quant::RescoreError::ZeroDimension => crate::scan::ScanError::ZeroDimension,
crate::quant::RescoreError::QueryDimension { expected, actual }
| crate::quant::RescoreError::CoarseScoreCount { expected, actual }
| crate::quant::RescoreError::CandidateRowCount { expected, actual } => {
crate::scan::ScanError::RowDataLength {
dimension: expected,
actual,
}
}
crate::quant::RescoreError::RowDataLength { dimension, actual } => {
crate::scan::ScanError::RowDataLength { dimension, actual }
}
crate::quant::RescoreError::CandidateRowOutOfRange {
row_index,
row_count,
..
} => crate::scan::ScanError::RowDataLength {
dimension: row_count,
actual: row_index,
},
crate::quant::RescoreError::NonFiniteCoarseScore { index }
| crate::quant::RescoreError::NonFiniteExactScore { row_index: index } => {
crate::scan::ScanError::NonFiniteScore { row_id: index }
}
crate::quant::RescoreError::ZeroK
| crate::quant::RescoreError::ZeroOversample
| crate::quant::RescoreError::InsufficientCandidates { .. }
| crate::quant::RescoreError::ArithmeticOverflow => {
crate::scan::ScanError::ArithmeticOverflow
}
error @ (crate::quant::RescoreError::ZeroCandidateLimit
| crate::quant::RescoreError::CandidateLimitExceeded { .. }) => {
crate::scan::ScanError::Rescore(error)
}
};
QueryError::Scan(error)
}
fn map_scan_error(error: crate::scan::ScanError) -> QueryError {
match error {
crate::scan::ScanError::Cancelled { partial } => QueryError::Cancelled { partial },
crate::scan::ScanError::Timeout { partial } => QueryError::Timeout { partial },
crate::scan::ScanError::ReadCancelled { partial } => QueryError::ReadCancelled { partial },
error => QueryError::Scan(error),
}
}
#[allow(clippy::too_many_arguments)]
fn merge_store_outcome(
outcome: crate::scan::ScanOutcome,
source: crate::ingest::RowSource,
document: impl Fn(usize) -> Result<Option<crate::ingest::DocumentVersion>, QueryError>,
candidates: &mut Vec<crate::ingest::SearchCandidate>,
dims_touched: &mut u64,
bytes_read: &mut u64,
worker_thread_ids: &mut Vec<std::thread::ThreadId>,
exact_score: bool,
#[cfg(any(test, feature = "test-seams"))] vector_fault_controller: Option<
&crate::scan::vector_fault::VectorFaultController,
>,
#[cfg(any(test, feature = "test-seams"))]
vector_fault_tier: crate::scan::vector_fault::VectorSearchTier,
) -> Result<(), QueryError> {
*dims_touched = dims_touched
.checked_add(outcome.stats.dims_touched)
.ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
*bytes_read = bytes_read
.checked_add(outcome.stats.bytes_read)
.ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
for worker in outcome.stats.worker_thread_ids {
if !worker_thread_ids.contains(&worker) {
worker_thread_ids.push(worker);
}
}
reserve_global_candidates(
candidates,
outcome.candidates.len(),
#[cfg(any(test, feature = "test-seams"))]
vector_fault_controller,
#[cfg(any(test, feature = "test-seams"))]
vector_fault_tier,
)?;
for candidate in outcome.candidates {
let local_row = u32::try_from(candidate.row_id)
.map_err(|_| QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
candidates.push(crate::ingest::SearchCandidate::new(
crate::ingest::GlobalRowId::new(source, local_row),
document(candidate.row_id)?,
candidate.score,
exact_score,
));
}
Ok(())
}
pub(crate) fn reserve_global_candidates(
candidates: &mut Vec<crate::ingest::SearchCandidate>,
additional: usize,
#[cfg(any(test, feature = "test-seams"))] vector_fault_controller: Option<
&crate::scan::vector_fault::VectorFaultController,
>,
#[cfg(any(test, feature = "test-seams"))]
vector_fault_tier: crate::scan::vector_fault::VectorSearchTier,
) -> Result<(), QueryError> {
if additional == 0 {
return Ok(());
}
let items = u64::try_from(additional)
.map_err(|_| QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
let candidate_bytes = u64::try_from(std::mem::size_of::<crate::ingest::SearchCandidate>())
.map_err(|_| QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
let needed = items
.checked_mul(candidate_bytes)
.ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
#[cfg(any(test, feature = "test-seams"))]
if vector_fault_controller.is_some_and(|controller| {
controller.deny_global_candidate_allocation(vector_fault_tier, items, needed)
}) {
return Err(QueryError::Store(StoreError::AllocationFailed {
needed,
component: "vector search global candidates",
}));
}
candidates.try_reserve_exact(additional).map_err(|_| {
QueryError::Store(StoreError::AllocationFailed {
needed,
component: "vector search global candidates",
})
})
}
struct AdmittedLexicalQuery<'a> {
generation: u64,
active: Arc<crate::ingest::ActiveSegment>,
snapshot: snapshot::ReadSnapshot,
active_query: ActiveQuery<'a>,
}
struct AdmittedDocumentRead<'a> {
generation: u64,
active: Arc<crate::ingest::ActiveSegment>,
snapshot: snapshot::ReadSnapshot,
active_query: ActiveQuery<'a>,
}
struct AdmittedVectorSearch<'a> {
execution_options: SearchOptions,
pool: Option<Arc<pool::QueryPool>>,
snapshot: snapshot::ReadSnapshot,
active_segment: Arc<crate::ingest::ActiveSegment>,
generation: u64,
_active_query: ActiveQuery<'a>,
}
struct ActiveQuery<'a> {
count: &'a AtomicU64,
}
impl Drop for ActiveQuery<'_> {
fn drop(&mut self) {
self.count.fetch_sub(1, Ordering::Relaxed);
}
}
impl Drop for Store {
fn drop(&mut self) {
self.close_best_effort();
}
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
struct DocumentScanRow {
source: DocumentScanSource,
row: u32,
timestamp: i64,
doc_id: crate::ingest::DocId,
attribute_key: Option<u64>,
}
fn validate_scan_order_attribute(
schema: &crate::meta::Schema,
column: crate::meta::ColumnId,
) -> Result<(), QueryError> {
use crate::meta::ColumnType;
let id = column.get();
let detail = if column == crate::meta::TIMESTAMP_COLUMN {
Some(format!(
"scan order attribute {id} is the document timestamp; use a timestamp order"
))
} else {
match schema
.column(column)
.map(crate::meta::ColumnDefinition::column_type)
{
None => Some(format!(
"scan order attribute {id} is not a declared attribute"
)),
Some(ColumnType::U64 | ColumnType::I64 | ColumnType::F64) => None,
Some(other) => Some(format!(
"scan order attribute {id} has type {other:?}; only u64, i64 and f64 attributes are orderable"
)),
}
};
detail.map_or(Ok(()), |detail| {
Err(QueryError::Store(StoreError::InvalidScan { detail }))
})
}
fn attribute_order_key(column: &crate::meta::Column, row: u32) -> Result<Option<u64>, QueryError> {
use crate::meta::Column;
const SIGN: u64 = 1 << 63;
match column {
Column::U64(values) => Ok(values.get(row)),
Column::I64(values) => Ok(values.get(row).map(|value| value.cast_unsigned() ^ SIGN)),
Column::F64(values) => Ok(values
.get(row)
.filter(|value| !value.is_nan())
.map(|value| {
let bits = (value + 0.0).to_bits();
if bits & SIGN == 0 { bits | SIGN } else { !bits }
})),
Column::Id128(_) | Column::Bool(_) | Column::DictionaryString(_) | Column::RawString(_) => {
Err(QueryError::Store(StoreError::InvalidScan {
detail: "scan order attribute column is not numeric".to_owned(),
}))
}
}
}
fn keyed_scan_row(
mut row: DocumentScanRow,
columns: Option<(&crate::meta::ColumnStore, crate::meta::ColumnId)>,
) -> Result<DocumentScanRow, QueryError> {
if let Some((columns, column)) = columns {
let values = columns.column(column).ok_or_else(|| {
QueryError::Store(StoreError::InvalidScan {
detail: format!("scan order attribute {} has no stored column", column.get()),
})
})?;
row.attribute_key = attribute_order_key(values, row.row)?;
}
Ok(row)
}
fn document_scan_predicate(
predicate: Option<&crate::meta::Predicate>,
timestamp_range: Option<(i64, i64)>,
) -> Result<Option<crate::meta::Predicate>, QueryError> {
use crate::meta::{Predicate, PredicateValue, RangeBound, RangePredicate, TIMESTAMP_COLUMN};
let timestamp = timestamp_range
.map(|(start, end)| {
if start > end {
return Err(QueryError::Store(StoreError::InvalidScan {
detail: format!("timestamp range start {start} exceeds end {end}"),
}));
}
Ok(Predicate::Range(RangePredicate {
column: TIMESTAMP_COLUMN,
lower: Some(RangeBound::inclusive(PredicateValue::I64(start))),
upper: Some(RangeBound::exclusive(PredicateValue::I64(end))),
}))
})
.transpose()?;
Ok(match (predicate, timestamp) {
(Some(predicate), Some(timestamp)) => {
Some(Predicate::And(vec![predicate.clone(), timestamp]))
}
(Some(predicate), None) => Some(predicate.clone()),
(None, Some(timestamp)) => Some(timestamp),
(None, None) => None,
})
}
fn scan_control_checkpoint(
control: &QueryControl,
lease: &SnapshotLease,
work: &mut usize,
) -> Result<(), QueryError> {
*work = work.saturating_add(1);
if (*work).is_multiple_of(64) {
control.checkpoint()?;
lease.check_active().map_err(QueryError::Store)?;
}
Ok(())
}
fn active_scan_columns(
active: &crate::ingest::ActiveSegment,
schema: &crate::meta::Schema,
) -> Result<crate::meta::ColumnStore, QueryError> {
let mut builder = crate::meta::ColumnStoreBuilder::new(schema.clone());
for (row, timestamp) in active.timestamps().iter().copied().enumerate() {
let values = active.column_values(row).map_err(QueryError::Store)?;
builder
.push_row(timestamp, &crate::ingest::column_inputs(&values))
.map_err(|error| {
QueryError::Store(StoreError::Segment(crate::segment::SegmentError::Columns(
error.to_string(),
)))
})?;
}
builder.finish().map_err(|error| {
QueryError::Store(StoreError::Segment(crate::segment::SegmentError::Columns(
error.to_string(),
)))
})
}
fn sealed_scan_rows(
segment: &crate::segment::reader::SegmentReader,
predicate: Option<&crate::meta::Predicate>,
) -> Result<crate::meta::DocBitmap, QueryError> {
if predicate.is_some_and(|predicate| {
!crate::planner::segment_may_match(segment.meta().clustering_key_range, predicate)
}) {
return Ok(crate::meta::DocBitmap::new());
}
let alive = segment.query_alive().map_err(QueryError::Store)?;
match predicate {
Some(predicate) => {
let columns = segment.query_columns().map_err(QueryError::Store)?;
crate::meta::evaluate(predicate, &columns, &alive).map_err(|error| {
QueryError::Store(StoreError::Segment(crate::segment::SegmentError::Columns(
error.to_string(),
)))
})
}
None => Ok(alive.alive_bitmap().clone()),
}
}
fn active_scan_rows(
active: &crate::ingest::ActiveSegment,
schema: &crate::meta::Schema,
predicate: Option<&crate::meta::Predicate>,
) -> Result<crate::meta::DocBitmap, QueryError> {
let alive = active.alive().map_err(QueryError::Store)?;
match predicate {
Some(predicate) => {
let columns = active_scan_columns(active, schema)?;
crate::meta::evaluate(predicate, &columns, &alive).map_err(|error| {
QueryError::Store(StoreError::Segment(crate::segment::SegmentError::Columns(
error.to_string(),
)))
})
}
None => Ok(alive.alive_bitmap().clone()),
}
}
fn sealed_document_scan_row(
segment: &crate::segment::reader::SegmentReader,
row: u32,
) -> Result<DocumentScanRow, QueryError> {
let version = segment
.document_version(row as usize)
.map_err(StoreError::Segment)
.map_err(QueryError::Store)?
.ok_or_else(|| {
QueryError::Store(StoreError::Segment(crate::segment::SegmentError::Geometry(
format!("sealed document row {row} has no identity"),
)))
})?;
let columns = segment.query_columns().map_err(QueryError::Store)?;
let timestamp = columns.timestamp(row).ok_or_else(|| {
QueryError::Store(StoreError::Segment(crate::segment::SegmentError::Geometry(
format!("sealed document row {row} has no timestamp"),
)))
})?;
Ok(DocumentScanRow {
source: DocumentScanSource::Sealed(segment.meta().id),
row,
timestamp,
doc_id: version.doc_id(),
attribute_key: None,
})
}
fn active_document_scan_row(
active: &crate::ingest::ActiveSegment,
row: u32,
) -> Result<DocumentScanRow, QueryError> {
let index = row as usize;
let version = active
.document(index)
.ok_or(QueryError::Store(StoreError::ActiveRowOverflow))?;
let timestamp = active
.timestamps()
.get(index)
.copied()
.ok_or(QueryError::Store(StoreError::ActiveRowOverflow))?;
Ok(DocumentScanRow {
source: DocumentScanSource::Active,
row,
timestamp,
doc_id: version.doc_id(),
attribute_key: None,
})
}
#[allow(clippy::too_many_arguments)]
fn collect_storage_scan_rows(
snapshot: &PublishedSnapshot,
active: &crate::ingest::ActiveSegment,
schema: &crate::meta::Schema,
predicate: Option<&crate::meta::Predicate>,
cursor: Option<DocumentScanCursor>,
limit: usize,
control: &QueryControl,
lease: &SnapshotLease,
work: &mut usize,
) -> Result<Vec<DocumentScanRow>, QueryError> {
let capacity = limit.saturating_add(1);
let mut result = Vec::new();
result.try_reserve_exact(capacity).map_err(|_| {
QueryError::Store(StoreError::AllocationFailed {
needed: allocation_bytes::<DocumentScanRow>(capacity),
component: "scan positions",
})
})?;
let mut started = cursor.is_none();
'segments: for segment in snapshot.segments() {
let source = DocumentScanSource::Sealed(segment.meta().id);
let rows = sealed_scan_rows(segment, predicate)?;
for row in rows.iter() {
scan_control_checkpoint(control, lease, work)?;
if !started {
started =
cursor.is_some_and(|cursor| cursor.source == source && cursor.next_row == row);
if !started {
continue;
}
}
result.push(sealed_document_scan_row(segment, row)?);
if result.len() == capacity {
break 'segments;
}
}
}
if result.len() < capacity {
let source = DocumentScanSource::Active;
let rows = active_scan_rows(active, schema, predicate)?;
for row in rows.iter() {
scan_control_checkpoint(control, lease, work)?;
if !started {
started =
cursor.is_some_and(|cursor| cursor.source == source && cursor.next_row == row);
if !started {
continue;
}
}
result.push(active_document_scan_row(active, row)?);
if result.len() == capacity {
break;
}
}
}
if !started {
return Err(QueryError::Store(StoreError::InvalidScan {
detail: "scan cursor does not name a matching row".to_owned(),
}));
}
Ok(result)
}
fn document_scan_row_cmp(
left: &DocumentScanRow,
right: &DocumentScanRow,
order: ScanOrder,
) -> std::cmp::Ordering {
let primary = match order {
ScanOrder::Storage | ScanOrder::TimestampAscending => left.timestamp.cmp(&right.timestamp),
ScanOrder::TimestampDescending => right.timestamp.cmp(&left.timestamp),
ScanOrder::Attribute { direction, .. } => match (left.attribute_key, right.attribute_key) {
(Some(left), Some(right)) => match direction {
ScanDirection::Ascending => left.cmp(&right),
ScanDirection::Descending => right.cmp(&left),
},
(Some(_), None) => std::cmp::Ordering::Less,
(None, Some(_)) => std::cmp::Ordering::Greater,
(None, None) => std::cmp::Ordering::Equal,
},
};
primary
.then_with(|| left.doc_id.cmp(&right.doc_id))
.then_with(|| document_scan_source_cmp(left.source, right.source))
.then_with(|| left.row.cmp(&right.row))
}
fn document_scan_source_cmp(
left: DocumentScanSource,
right: DocumentScanSource,
) -> std::cmp::Ordering {
match (left, right) {
(DocumentScanSource::Sealed(left), DocumentScanSource::Sealed(right)) => left.cmp(&right),
(DocumentScanSource::Sealed(_), DocumentScanSource::Active) => std::cmp::Ordering::Less,
(DocumentScanSource::Active, DocumentScanSource::Sealed(_)) => std::cmp::Ordering::Greater,
(DocumentScanSource::Active, DocumentScanSource::Active) => std::cmp::Ordering::Equal,
}
}
fn insert_timestamp_scan_row(
rows: &mut Vec<DocumentScanRow>,
candidate: DocumentScanRow,
capacity: usize,
order: ScanOrder,
) {
let position = rows.partition_point(|row| {
document_scan_row_cmp(row, &candidate, order) != std::cmp::Ordering::Greater
});
if position < capacity {
rows.insert(position, candidate);
if rows.len() > capacity {
let _ = rows.pop();
}
}
}
fn clustering_range_cmp(
left: crate::segment::ClusteringKeyRange,
right: crate::segment::ClusteringKeyRange,
order: ScanOrder,
) -> std::cmp::Ordering {
use crate::segment::ClusteringKeyRange::{Bounded, Empty, Unstamped};
match (left, right) {
(Bounded { min_ts: left, .. }, Bounded { min_ts: right, .. })
if order == ScanOrder::TimestampAscending =>
{
left.cmp(&right)
}
(Bounded { max_ts: left, .. }, Bounded { max_ts: right, .. }) => right.cmp(&left),
(Unstamped, Unstamped) | (Empty, Empty) => std::cmp::Ordering::Equal,
(Unstamped, _) => std::cmp::Ordering::Less,
(_, Unstamped) => std::cmp::Ordering::Greater,
(Empty, _) => std::cmp::Ordering::Greater,
(_, Empty) => std::cmp::Ordering::Less,
}
}
#[allow(clippy::too_many_arguments)]
fn collect_sorted_scan_rows(
snapshot: &PublishedSnapshot,
active: &crate::ingest::ActiveSegment,
schema: &crate::meta::Schema,
predicate: Option<&crate::meta::Predicate>,
cursor: Option<DocumentScanCursor>,
limit: usize,
order: ScanOrder,
control: &QueryControl,
lease: &SnapshotLease,
work: &mut usize,
) -> Result<Vec<DocumentScanRow>, QueryError> {
let capacity = limit.saturating_add(1);
let mut result = Vec::new();
result.try_reserve_exact(capacity).map_err(|_| {
QueryError::Store(StoreError::AllocationFailed {
needed: allocation_bytes::<DocumentScanRow>(capacity),
component: "scan positions",
})
})?;
let attribute = match order {
ScanOrder::Attribute { column, .. } => Some(column),
ScanOrder::Storage | ScanOrder::TimestampAscending | ScanOrder::TimestampDescending => None,
};
let active_columns = attribute
.map(|_| active_scan_columns(active, schema))
.transpose()?;
let active_keys = active_columns.as_ref().zip(attribute);
let sealed_columns = |segment: &crate::segment::reader::SegmentReader| {
attribute
.map(|_| segment.query_columns().map_err(QueryError::Store))
.transpose()
};
let cursor_row = cursor
.map(|cursor| match cursor.source {
DocumentScanSource::Active => keyed_scan_row(
active_document_scan_row(active, cursor.next_row)?,
active_keys,
),
DocumentScanSource::Sealed(segment_id) => {
let segment = snapshot
.segments()
.iter()
.find(|segment| segment.meta().id == segment_id)
.ok_or_else(|| {
QueryError::Store(StoreError::InvalidScan {
detail: format!("scan cursor names absent segment {segment_id}"),
})
})?;
let columns = sealed_columns(segment)?;
keyed_scan_row(
sealed_document_scan_row(segment, cursor.next_row)?,
columns.as_deref().zip(attribute),
)
}
})
.transpose()?;
let mut segments = snapshot.segments().iter().collect::<Vec<_>>();
segments.sort_by(|left, right| {
clustering_range_cmp(
left.meta().clustering_key_range,
right.meta().clustering_key_range,
order,
)
.then_with(|| left.meta().id.cmp(&right.meta().id))
});
for segment in segments {
let rows = sealed_scan_rows(segment, predicate)?;
let columns = sealed_columns(segment)?;
for row in rows.iter() {
scan_control_checkpoint(control, lease, work)?;
let candidate = keyed_scan_row(
sealed_document_scan_row(segment, row)?,
columns.as_deref().zip(attribute),
)?;
if cursor_row.as_ref().is_none_or(|cursor_row| {
document_scan_row_cmp(&candidate, cursor_row, order) != std::cmp::Ordering::Less
}) {
insert_timestamp_scan_row(&mut result, candidate, capacity, order);
}
}
}
for row in active_scan_rows(active, schema, predicate)?.iter() {
scan_control_checkpoint(control, lease, work)?;
let candidate = keyed_scan_row(active_document_scan_row(active, row)?, active_keys)?;
if cursor_row.as_ref().is_none_or(|cursor_row| {
document_scan_row_cmp(&candidate, cursor_row, order) != std::cmp::Ordering::Less
}) {
insert_timestamp_scan_row(&mut result, candidate, capacity, order);
}
}
if let Some(cursor) = cursor {
let starts_at_cursor = result
.first()
.is_some_and(|row| row.source == cursor.source && row.row == cursor.next_row);
if !starts_at_cursor {
return Err(QueryError::Store(StoreError::InvalidScan {
detail: "scan cursor does not name a matching row".to_owned(),
}));
}
}
Ok(result)
}
fn materialize_active_document(
active: &crate::ingest::ActiveSegment,
row: usize,
fields: DocumentFields,
) -> Result<StoredDocument, StoreError> {
let version = active.document(row).ok_or(StoreError::ActiveRowOverflow)?;
let timestamp = active
.timestamps()
.get(row)
.copied()
.ok_or(StoreError::ActiveRowOverflow)?;
let vector = if fields.contains(DocumentFields::VECTOR) {
let dimensions = active.dims().ok_or(StoreError::ActiveRowOverflow)?;
let start = row
.checked_mul(dimensions)
.ok_or(StoreError::ActiveRowOverflow)?;
let end = start
.checked_add(dimensions)
.ok_or(StoreError::ActiveRowOverflow)?;
let values = active
.vectors()
.get(start..end)
.ok_or(StoreError::ActiveRowOverflow)?;
Some(copy_slice(values, "get vector")?)
} else {
None
};
let text = if fields.contains(DocumentFields::TEXT) {
active
.text(row)?
.map(|value| copy_string(value, "get text"))
.transpose()?
} else {
None
};
let metadata = if fields.contains(DocumentFields::METADATA) {
active
.metadata(row)
.map(|value| copy_slice(value, "get metadata"))
.transpose()?
} else {
None
};
let attributes = if fields.contains(DocumentFields::ATTRIBUTES) {
Some(active.column_values(row)?)
} else {
None
};
Ok(StoredDocument {
doc_id: version.doc_id(),
revision: version.revision(),
timestamp,
vector,
text,
metadata,
attributes,
})
}
fn materialize_sealed_document(
segment: &crate::segment::reader::SegmentReader,
row: usize,
fields: DocumentFields,
) -> Result<StoredDocument, StoreError> {
let row_u32 = u32::try_from(row).map_err(|_| {
StoreError::Segment(crate::segment::SegmentError::Geometry(
"sealed document row exceeds u32".to_owned(),
))
})?;
let version = segment
.document_version(row)
.map_err(StoreError::Segment)?
.ok_or_else(|| {
StoreError::Segment(crate::segment::SegmentError::Geometry(format!(
"sealed document row {row} has no identity"
)))
})?;
let columns = segment.query_columns()?;
let timestamp = columns.timestamp(row_u32).ok_or_else(|| {
StoreError::Segment(crate::segment::SegmentError::Geometry(format!(
"sealed document row {row} has no timestamp"
)))
})?;
let vector = if fields.contains(DocumentFields::VECTOR) {
segment
.validate_rescore_rows(&[row_u32])
.map_err(StoreError::Segment)?;
let dimensions = segment.meta().dims as usize;
let start = row.checked_mul(dimensions).ok_or_else(|| {
StoreError::Segment(crate::segment::SegmentError::Geometry(
"sealed vector row offset overflow".to_owned(),
))
})?;
let end = start.checked_add(dimensions).ok_or_else(|| {
StoreError::Segment(crate::segment::SegmentError::Geometry(
"sealed vector row end overflow".to_owned(),
))
})?;
let rescore = segment.query_rescore_f32().map_err(StoreError::Segment)?;
let values = rescore.get(start..end).ok_or_else(|| {
StoreError::Segment(crate::segment::SegmentError::Geometry(format!(
"sealed vector row {row} is missing"
)))
})?;
Some(copy_slice(values, "get vector")?)
} else {
None
};
let text = if fields.contains(DocumentFields::TEXT) {
match segment.query_stored_text().map_err(StoreError::Segment)? {
Some(rows) => rows
.row(row)
.ok_or_else(|| {
StoreError::Segment(crate::segment::SegmentError::Geometry(format!(
"stored-text row {row} is invalid"
)))
})?
.map(|value| copy_string(value, "get text"))
.transpose()?,
None => None,
}
} else {
None
};
let metadata = if fields.contains(DocumentFields::METADATA) {
match segment
.query_stored_metadata()
.map_err(StoreError::Segment)?
{
Some(rows) => Some(copy_slice(
rows.row(row).ok_or_else(|| {
StoreError::Segment(crate::segment::SegmentError::Geometry(format!(
"stored-metadata row {row} is invalid"
)))
})?,
"get metadata",
)?),
None => None,
}
} else {
None
};
let attributes = fields
.contains(DocumentFields::ATTRIBUTES)
.then(|| materialize_sealed_attributes(&columns, row_u32))
.transpose()?;
Ok(StoredDocument {
doc_id: version.doc_id(),
revision: version.revision(),
timestamp,
vector,
text,
metadata,
attributes,
})
}
fn materialize_sealed_attributes(
columns: &crate::meta::ColumnStore,
row: u32,
) -> Result<Vec<(crate::meta::ColumnId, crate::meta::PredicateValue)>, StoreError> {
let mut attributes = Vec::new();
attributes
.try_reserve_exact(columns.schema().user_column_count())
.map_err(|_| StoreError::AllocationFailed {
needed: allocation_bytes::<(crate::meta::ColumnId, crate::meta::PredicateValue)>(
columns.schema().user_column_count(),
),
component: "get attributes",
})?;
for definition in columns
.schema()
.columns()
.iter()
.filter(|definition| definition.id() != crate::meta::TIMESTAMP_COLUMN)
{
let column = columns.column(definition.id()).ok_or_else(|| {
StoreError::Segment(crate::segment::SegmentError::Geometry(format!(
"sealed attribute column {} is missing",
definition.id().get()
)))
})?;
let value = predicate_value_for_row(column, row)?;
if value.is_none() && !definition.is_nullable() {
return Err(StoreError::Segment(crate::segment::SegmentError::Geometry(
format!(
"sealed required attribute {} is null at row {row}",
definition.id().get()
),
)));
}
if let Some(value) = value {
attributes.push((definition.id(), value));
}
}
Ok(attributes)
}
fn predicate_value_for_row(
column: &crate::meta::Column,
row: u32,
) -> Result<Option<crate::meta::PredicateValue>, StoreError> {
use crate::meta::{Column, PredicateValue};
match column {
Column::Id128(values) => Ok(values.get(row).map(PredicateValue::Id128)),
Column::U64(values) => Ok(values.get(row).map(PredicateValue::U64)),
Column::I64(values) => Ok(values.get(row).map(PredicateValue::I64)),
Column::F64(values) => Ok(values.get(row).map(PredicateValue::F64)),
Column::Bool(values) => Ok(values.get(row).map(PredicateValue::Bool)),
Column::DictionaryString(values) => values
.get(row)
.map(|value| copy_string(value, "get attribute string").map(PredicateValue::String))
.transpose(),
Column::RawString(values) => values
.get(row)
.map(|value| copy_string(value, "get attribute string").map(PredicateValue::String))
.transpose(),
}
}
fn copy_slice<T: Copy>(values: &[T], component: &'static str) -> Result<Vec<T>, StoreError> {
let mut copied = Vec::new();
copied
.try_reserve_exact(values.len())
.map_err(|_| StoreError::AllocationFailed {
needed: allocation_bytes::<T>(values.len()),
component,
})?;
copied.extend_from_slice(values);
Ok(copied)
}
fn copy_string(value: &str, component: &'static str) -> Result<String, StoreError> {
let mut copied = String::new();
copied
.try_reserve_exact(value.len())
.map_err(|_| StoreError::AllocationFailed {
needed: u64::try_from(value.len()).unwrap_or(u64::MAX),
component,
})?;
copied.push_str(value);
Ok(copied)
}
fn allocation_bytes<T>(count: usize) -> u64 {
count
.checked_mul(std::mem::size_of::<T>())
.and_then(|bytes| u64::try_from(bytes).ok())
.unwrap_or(u64::MAX)
}
#[cfg(test)]
mod lexical_assembly_tests;
#[cfg(test)]
mod get_tests;
#[cfg(test)]
mod group_count_tests;
#[cfg(test)]
mod scan_attribute_order_tests;
#[cfg(test)]
mod scan_document_tests;
#[cfg(test)]
#[allow(clippy::expect_used, clippy::panic)]
mod tests {
use std::error::Error;
use std::sync::Arc;
use tempfile::tempdir;
use crate::lifecycle::durability::DurabilityPolicy;
use crate::manifest::io::commit_manifest;
use crate::manifest::{Manifest, ManifestError};
use crate::meta::Schema;
use crate::vfs::StdVfs;
use crate::vfs::crash::MemoryVfs;
use super::durability::{CommitTier, DurabilityMode, DurabilityPolicyError};
use super::{
CancelToken, GraphSearchOptions, ManualMonotonicClock, OpenOptions, QueryControl,
SearchOptions, SearchTier, Store, StoreError, StoreTestDependencies,
resolve_hybrid_leg_results,
};
#[test]
fn astra_12_vocabulary_charge_survives_eviction_and_close() {
use crate::fts::{index::DEFAULT_FIELD, query::LexicalQuery};
use crate::ingest::{DocId, DocumentVersion, IngestBatch, IngestDocument, Revision};
let directory = tempdir().expect("directory");
let store = Store::open(directory.path(), OpenOptions::default()).expect("open");
store
.ingest(IngestBatch::new(vec![
IngestDocument::new(
DocumentVersion::new(DocId::new(1), Revision::new(1)),
vec![1.0, 0.0],
)
.with_text("alpha alphabet beta"),
]))
.expect("ingest");
store
.search_lexical_structured(
&LexicalQuery::prefix(b"alph".to_vec(), DEFAULT_FIELD),
1,
64,
QueryControl::Cancel(CancelToken::new()),
)
.expect("build cache");
let dictionary = {
let entry = store.lexical_index_cache.entry.lock().expect("assembly");
let vocabulary = entry
.as_ref()
.expect("assembly")
.assembly
.vocabulary
.lock()
.expect("vocabulary");
Arc::clone(vocabulary.as_ref().expect("cached"))
};
let bytes = dictionary._memory.bytes();
assert!(bytes > 17);
let second_owner = Arc::clone(&dictionary);
drop(
store
.lexical_index_cache
.entry
.lock()
.expect("evict")
.take(),
);
assert_eq!(store.stats().expect("held bytes").cache_bytes, bytes);
store
.search_lexical_structured(
&LexicalQuery::prefix(b"alph".to_vec(), DEFAULT_FIELD),
1,
64,
QueryControl::Cancel(CancelToken::new()),
)
.expect("replacement cache");
assert!(store.stats().expect("both dictionaries").cache_bytes > bytes);
store.close().expect("close");
assert_eq!(
store
.accounting
.audit()
.expect("closed accounting")
.cache_bytes,
bytes
);
drop(dictionary);
assert_eq!(
store.accounting.audit().expect("second owner").cache_bytes,
bytes
);
assert_eq!(
second_owner.view.iter().collect::<Vec<_>>(),
vec![b"alpha".as_slice(), b"alphabet", b"beta"]
);
drop(second_owner);
assert_eq!(store.accounting.audit().expect("last owner").cache_bytes, 0);
assert_eq!(
store.accounting.audit().expect("scratch").temporary_bytes,
0
);
}
#[test]
fn astra_12_refused_vocabulary_is_not_published() {
use crate::fts::{index::DEFAULT_FIELD, query::LexicalQuery};
use crate::ingest::{
DocId, DocumentVersion, IngestBatch, IngestDocument, Revision, StoreLexicalError,
};
let directory = tempdir().expect("directory");
let writer = Store::open(directory.path(), OpenOptions::default()).expect("open");
writer
.ingest(IngestBatch::new(vec![
IngestDocument::new(
DocumentVersion::new(DocId::new(1), Revision::new(1)),
vec![1.0, 0.0],
)
.with_text("alpha alphabet beta"),
]))
.expect("ingest");
writer.seal().expect("seal");
writer.close().expect("close writer");
let store = Store::open(
directory.path(),
OpenOptions::default().with_max_temp_bytes(1),
)
.expect("small budget");
for _ in 0..2 {
let result = store.search_lexical_structured(
&LexicalQuery::prefix(b"alph".to_vec(), DEFAULT_FIELD),
1,
64,
QueryControl::Cancel(CancelToken::new()),
);
assert!(
matches!(
result,
Err(StoreLexicalError::Query(super::QueryError::Store(
StoreError::BudgetExceeded {
component: "temporary",
..
}
)))
),
"{result:?}"
);
assert_eq!(store.stats().expect("released").temporary_bytes, 0);
let entry = store.lexical_index_cache.entry.lock().expect("assembly");
assert!(
entry
.as_ref()
.expect("assembled")
.assembly
.vocabulary
.lock()
.expect("dictionary")
.is_none()
);
}
store.close().expect("close");
assert_eq!(
store
.accounting
.audit()
.expect("released cache")
.cache_bytes,
0
);
}
#[test]
fn astra_03_lexical_assembly_charge_survives_cache_eviction() {
use crate::fts::index::DEFAULT_FIELD;
use crate::fts::search::TermQuery;
use crate::ingest::{DocId, DocumentVersion, IngestBatch, IngestDocument, Revision};
let directory = tempdir().expect("directory");
let store = Store::open(directory.path(), OpenOptions::default()).expect("open");
store
.ingest(IngestBatch::new(vec![
IngestDocument::new(
DocumentVersion::new(DocId::new(1), Revision::new(1)),
vec![1.0, 0.0],
)
.with_text("alpha beta"),
]))
.expect("ingest");
store
.search_lexical(
&TermQuery::flat(vec![b"alpha".to_vec()], &[DEFAULT_FIELD]),
1,
QueryControl::Cancel(CancelToken::new()),
)
.expect("build lexical cache");
let (first, second) = {
let entry = store.lexical_index_cache.entry.lock().expect("cache lock");
let assembly = &entry.as_ref().expect("cached assembly").assembly;
(Arc::clone(assembly), Arc::clone(assembly))
};
let charged = store.stats().expect("charged stats").cache_bytes;
let evicted = store
.lexical_index_cache
.entry
.lock()
.expect("cache lock")
.take();
drop(evicted);
assert_eq!(
first.index.document_count(),
1,
"eviction left the query data alive"
);
assert_eq!(
store.stats().expect("held stats").cache_bytes,
charged,
"cache eviction must not release a query-owned assembly's charge"
);
drop(first);
assert_eq!(
store.stats().expect("second owner stats").cache_bytes,
charged
);
drop(second);
let released = store.stats().expect("released stats").cache_bytes;
assert_eq!(
released, 0,
"the final owner must release the full assembly charge"
);
println!("assembly cache bytes: retained={charged}, after final owner={released}");
store.close().expect("close");
}
const CEILING_ROWS: [[f32; 4]; 3] = [
[0.5, -0.25, 0.75, 0.125],
[3.0, 1.5, -2.5, 0.25],
[0.0, 0.0, 0.0, 0.0],
];
fn publish_f32_segment(directory: &std::path::Path) {
let created = Store::open(directory, OpenOptions::default()).expect("create f32 store");
created.close().expect("close f32 store");
let schema = Schema::timestamp_only();
let mut columns = crate::meta::ColumnStoreBuilder::new(schema.clone());
for row in 0..CEILING_ROWS.len() {
columns
.push_row(i64::try_from(row).expect("row timestamp"), &[])
.expect("timestamp row");
}
let columns = columns.finish().expect("columns");
let alive =
crate::meta::AliveSet::new(u32::try_from(CEILING_ROWS.len()).expect("row count"));
let rescore = CEILING_ROWS.iter().flatten().copied().collect::<Vec<f32>>();
let codes = rescore
.iter()
.flat_map(|value| value.to_bits().to_le_bytes())
.collect::<Vec<u8>>();
let policy = DurabilityPolicy::new(DurabilityMode::Derived, CommitTier::None)
.expect("f32 fixture durability");
let meta = crate::segment::writer::write_segment(
&StdVfs,
directory,
crate::segment::writer::SegmentBuild {
id: crate::segment::SegmentId::new(11, [7; 10]),
scheme: 0,
dims: 4,
codes: &codes,
factors: crate::segment::writer::SegmentFactors::F32,
rescore: &rescore,
columns: &columns,
alive: &alive,
},
policy,
)
.expect("write scheme-0 segment");
commit_manifest(
&StdVfs,
directory,
&Manifest {
generation: 1,
log_seq: 0,
segments: vec![meta],
epochs: Vec::new(),
epoch_alias: None,
schema,
},
policy,
)
.expect("publish scheme-0 segment");
}
fn ingest_ceiling_rows(store: &Store, keep_active: usize, one_segment_per_row: bool) {
use crate::ingest::{DocId, DocumentVersion, IngestBatch, IngestDocument, Revision};
let sealed = CEILING_ROWS.len() - keep_active;
for (index, row) in CEILING_ROWS.iter().enumerate() {
let version = DocumentVersion::new(
DocId::new(u128::try_from(index).expect("document id")),
Revision::new(1),
);
store
.ingest(IngestBatch::new(vec![IngestDocument::new(
version,
row.to_vec(),
)]))
.expect("ingest ceiling row");
let seal_here = if one_segment_per_row {
index < sealed
} else {
index + 1 == sealed
};
if seal_here {
store.seal().expect("seal ceiling rows");
}
}
}
fn assert_remembered_ceiling_is_exact(store: &Store, expected_segments: usize) {
let queries: [Vec<f32>; 3] = [
vec![0.0, 0.0, 0.0, 0.0],
vec![1.0, -1.0, 0.5, 2.0],
vec![7.5, 0.25, -3.5, 0.125],
];
let admitted = store
.admit_vector_search(SearchOptions::default())
.expect("admit ceiling query");
let snapshot = &admitted.snapshot;
let active = &admitted.active_segment;
assert_eq!(snapshot.segments().len(), expected_segments);
for segment in snapshot.segments() {
assert!(
segment.cached_vector_ceiling_norm_range().is_none(),
"a segment remembered an enclosure before any query walked it"
);
}
for query in &queries {
let cold = super::exact_vector_ceiling(snapshot, active, query).expect("cold ceiling");
for segment in snapshot.segments() {
let remembered = segment
.cached_vector_ceiling_norm_range()
.expect("a walked segment must remember its enclosure");
let recomputed = super::compute_segment_vector_ceiling_norm_range(segment)
.expect("recomputed enclosure");
assert_eq!(
remembered.endpoint_bits(),
recomputed.endpoint_bits(),
"remembered enclosure differs from the recomputed one"
);
}
let warm = super::exact_vector_ceiling(snapshot, active, query).expect("warm ceiling");
let recomputed = super::fold_exact_vector_ceiling(
snapshot,
active,
query,
super::compute_segment_vector_ceiling_norm_range,
)
.expect("recomputed ceiling");
assert_eq!(
cold.to_bits(),
warm.to_bits(),
"the remembered ceiling changed between queries"
);
assert_eq!(
cold.to_bits(),
recomputed.to_bits(),
"the remembered ceiling is not the recomputed ceiling"
);
}
}
#[test]
fn astra_00_elapsed_includes_ceiling_work() {
use std::time::{Duration, Instant};
struct CeilingClock {
base: Instant,
snapshot: Arc<super::PublishedSnapshot>,
}
impl super::MonotonicClock for CeilingClock {
fn now(&self) -> Instant {
if self
.snapshot
.segments()
.iter()
.all(|segment| segment.cached_vector_ceiling_norm_range().is_some())
{
self.base + Duration::from_secs(7)
} else {
self.base
}
}
}
let directory = tempdir().expect("ceiling timer store");
let mut store = Store::open(directory.path(), OpenOptions::default()).expect("open");
ingest_ceiling_rows(&store, 0, false);
let admitted = store
.admit_vector_search(SearchOptions::default())
.expect("snapshot");
assert_eq!(admitted.snapshot.segments().len(), 1);
assert!(
admitted
.snapshot
.segments()
.iter()
.all(|segment| { segment.cached_vector_ceiling_norm_range().is_none() })
);
let snapshot = Arc::clone(&admitted.snapshot);
drop(admitted);
store.clock = Arc::new(CeilingClock {
base: Instant::now(),
snapshot,
});
let outcome = store
.search(
crate::ingest::SearchRequest::new(&[1.0, 0.0, 0.0, 0.0]),
1,
SearchOptions::default().with_tier(SearchTier::Exact),
QueryControl::Cancel(CancelToken::new()),
)
.expect("exact query");
assert_eq!(outcome.diagnostics.elapsed, Duration::from_secs(7));
store.clock = Arc::new(super::SystemMonotonicClock);
store.close().expect("close");
}
#[test]
fn the_remembered_vector_ceiling_is_the_recomputed_ceiling_for_every_scheme() {
let bit4_directory = tempdir().expect("Bit4 ceiling store directory");
let bit4 = Store::open(bit4_directory.path(), OpenOptions::default())
.expect("open Bit4 ceiling store");
ingest_ceiling_rows(&bit4, 0, false);
assert_remembered_ceiling_is_exact(&bit4, 1);
bit4.close().expect("close Bit4 ceiling store");
let split_directory = tempdir().expect("multi-segment ceiling store directory");
let split = Store::open(split_directory.path(), OpenOptions::default())
.expect("open multi-segment ceiling store");
ingest_ceiling_rows(&split, 0, true);
assert_remembered_ceiling_is_exact(&split, CEILING_ROWS.len());
split.close().expect("close multi-segment ceiling store");
let int8_directory = tempdir().expect("Int8 ceiling store directory");
let int8 = Store::open_with_test_dependencies(
int8_directory.path(),
OpenOptions::default(),
StoreTestDependencies::new(Arc::new(StdVfs), Arc::new(super::SystemMonotonicClock))
.with_vector_seal_scheme(crate::quant::QuantScheme::Int8),
)
.expect("open Int8 ceiling store");
ingest_ceiling_rows(&int8, 0, false);
assert_remembered_ceiling_is_exact(&int8, 1);
int8.close().expect("close Int8 ceiling store");
let f32_directory = tempdir().expect("F32 ceiling store directory");
publish_f32_segment(f32_directory.path());
let f32_store = Store::open(f32_directory.path(), OpenOptions::default())
.expect("open F32 ceiling store");
assert_remembered_ceiling_is_exact(&f32_store, 1);
f32_store.close().expect("close F32 ceiling store");
let active_directory = tempdir().expect("active ceiling store directory");
let active = Store::open(active_directory.path(), OpenOptions::default())
.expect("open active ceiling store");
ingest_ceiling_rows(&active, 2, false);
assert!(
!active
.admit_vector_search(SearchOptions::default())
.expect("admit active ceiling query")
.active_segment
.is_empty(),
"the active-segment case sealed every row"
);
assert_remembered_ceiling_is_exact(&active, 1);
active.close().expect("close active ceiling store");
}
#[test]
fn exact_and_graph_tiers_preserve_quantizer_validation_errors() {
let directory = tempdir().expect("store directory");
let store = Store::open(directory.path(), OpenOptions::default()).expect("open");
let query = vec![0.0_f32; crate::kernels::MAX_DOT_I8_DIMENSION + 1];
for tier in [
SearchTier::Exact,
SearchTier::Graph(GraphSearchOptions::default()),
] {
let error = store
.search(
crate::ingest::SearchRequest::new(&query),
1,
SearchOptions::default().with_tier(tier),
QueryControl::Cancel(CancelToken::new()),
)
.expect_err("oversized quantized query must retain its error");
assert!(matches!(
error,
super::QueryError::Scan(crate::scan::ScanError::Quant(
crate::quant::QuantError::DimensionTooLarge { actual, maximum }
)) if actual == query.len()
&& maximum == crate::kernels::MAX_DOT_I8_DIMENSION
));
}
}
#[cfg(feature = "allocation-audit")]
#[test]
fn exact_tier_query_does_not_prepare_quantized_forms() {
let directory = tempdir().expect("store directory");
let store = Store::open(directory.path(), OpenOptions::default()).expect("open");
let query = [0.25_f32; 128];
let options = SearchOptions::default().with_tier(SearchTier::Exact);
let warm = store
.search(
crate::ingest::SearchRequest::new(&query),
1,
options,
QueryControl::Cancel(CancelToken::new()),
)
.expect("warm exact-tier search");
assert!(warm.candidates.is_empty());
let control = QueryControl::Cancel(CancelToken::new());
let (outcome, report) = crate::allocation_audit::audit_engine_path(|| {
store.search(
crate::ingest::SearchRequest::new(&query),
1,
options,
control,
)
});
let outcome = outcome.expect("audited exact-tier search");
assert!(outcome.candidates.is_empty());
assert_eq!(report.allocations, 0, "Exact tier prepared quantized forms");
assert_eq!(report.attributed_bytes, 0);
assert_eq!(report.unattributed_bytes, 0);
}
#[test]
fn astra_18_scoped_work_is_joined_before_cancel_returns() {
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::mpsc;
let worker = super::pool::LexicalWorker::start().expect("worker");
let token = CancelToken::new();
let work_token = token.clone();
let finished = AtomicBool::new(false);
let mut borrowed_rows = Vec::new();
let (started_tx, started_rx) = mpsc::sync_channel(0);
let (release_tx, release_rx) = mpsc::sync_channel(0);
let worker_rows = &mut borrowed_rows;
let worker_finished = &finished;
let (caller, lexical) = worker
.run_scoped(
move || {
started_tx.send(()).expect("entered borrowed job");
release_rx.recv().expect("caller has cancelled");
assert!(work_token.is_cancelled());
worker_rows.push(42_u32);
worker_finished.store(true, Ordering::Release);
Err::<(), _>(crate::fusion::FusionError::Cancelled { partial: false })
},
|| {
started_rx.recv().expect("worker entered");
token.cancel();
assert!(!finished.load(Ordering::Acquire));
release_tx.send(()).expect("release borrowed job");
Err::<(), _>(crate::fusion::FusionError::Cancelled { partial: false })
},
)
.expect("scoped dispatch");
assert!(finished.load(Ordering::Acquire));
assert_eq!(borrowed_rows, [42]);
assert_eq!(
resolve_hybrid_leg_results(caller, lexical.expect("worker returned")),
Err(crate::fusion::FusionError::Cancelled { partial: false }),
);
worker.stop_and_join().expect("join worker");
}
#[test]
fn astra_18_control_error_precedence_is_completion_order_independent() {
use crate::fusion::{FusionError, FusionLeg, LegFailureKind};
use std::sync::mpsc;
let failures = [
FusionError::Leg {
leg: FusionLeg::Vector,
kind: LegFailureKind::Scan,
detail: "scan".to_owned(),
},
FusionError::Cancelled { partial: false },
FusionError::Timeout { partial: false },
FusionError::ReadCancelled { partial: false },
];
let worker = super::pool::LexicalWorker::start().expect("worker");
for (a, vector) in failures.iter().enumerate() {
for (b, lexical) in failures.iter().enumerate() {
for lexical_first in [false, true] {
let (lexical_tx, lexical_rx) = mpsc::sync_channel(1);
let (vector_tx, vector_rx) = mpsc::sync_channel(1);
let (vector_result, lexical_result) = worker
.run_scoped(
move || {
if !lexical_first {
vector_rx.recv().expect("vector completed");
}
if lexical_first {
lexical_tx.send(()).expect("lexical completed");
}
Err::<(), _>(lexical.clone())
},
|| {
if lexical_first {
lexical_rx.recv().expect("lexical completed");
}
if !lexical_first {
vector_tx.send(()).expect("vector completed");
}
Err::<(), _>(vector.clone())
},
)
.expect("dispatch");
assert_eq!(
resolve_hybrid_leg_results(vector_result, lexical_result.expect("joined")),
Err(failures.get(a.max(b)).expect("known failure index").clone()),
"vector={a} lexical={b} lexical_first={lexical_first}",
);
}
}
}
worker.stop_and_join().expect("join worker");
}
#[test]
fn hybrid_dual_failure_precedence_is_independent_of_leg_completion_order() {
use crate::fusion::{FusionError, FusionLeg, LegFailureKind};
let vector_failure = FusionError::Leg {
leg: FusionLeg::Vector,
kind: LegFailureKind::Scan,
detail: "vector failure".to_owned(),
};
let lexical_failure = FusionError::Leg {
leg: FusionLeg::Lexical,
kind: LegFailureKind::Lexical,
detail: "lexical failure".to_owned(),
};
assert_eq!(
resolve_hybrid_leg_results::<(), ()>(
Err(vector_failure.clone()),
Err(lexical_failure.clone()),
),
Err(vector_failure.clone()),
);
assert_eq!(
resolve_hybrid_leg_results::<(), ()>(
Err(vector_failure.clone()),
Err(FusionError::Cancelled { partial: false }),
),
Err(FusionError::Cancelled { partial: false }),
);
assert_eq!(
resolve_hybrid_leg_results::<(), ()>(
Err(FusionError::Timeout { partial: false }),
Err(FusionError::ReadCancelled { partial: false }),
),
Err(FusionError::ReadCancelled { partial: false }),
);
assert_eq!(
resolve_hybrid_leg_results::<(), ()>(
Err(FusionError::Cancelled { partial: false }),
Err(FusionError::Timeout { partial: false }),
),
Err(FusionError::Timeout { partial: false }),
);
assert_eq!(
resolve_hybrid_leg_results::<(), ()>(
Err(FusionError::ReadCancelled { partial: false }),
Err(lexical_failure),
),
Err(FusionError::ReadCancelled { partial: false }),
);
}
#[test]
fn accounting_errors_report_the_typed_context() {
let statistics = StoreError::Statistics {
component: "mapped resident bytes",
source: std::io::Error::other("mincore failed"),
};
assert_eq!(
statistics.to_string(),
"store statistics mapped resident bytes: mincore failed"
);
assert_eq!(
statistics.source().map(ToString::to_string),
Some("mincore failed".to_owned())
);
let budget = StoreError::BudgetExceeded {
needed: 65,
budget: 64,
component: "wal",
};
assert_eq!(
budget.to_string(),
"store wal allocation needs 65 bytes, budget is 64 bytes"
);
assert!(budget.source().is_none());
let allocation = StoreError::AllocationFailed {
needed: 23,
component: "snapshot",
};
assert_eq!(
allocation.to_string(),
"store snapshot allocator rejected 23 bytes"
);
assert!(allocation.source().is_none());
}
#[test]
fn open_options_resolve_durability_before_filesystem_mutation() {
let parent = tempdir().expect("parent directory");
let store_path = parent.path().join("not-created");
let options =
OpenOptions::new().with_durability(DurabilityMode::Attached, CommitTier::Durable);
let error = match Store::open(&store_path, options) {
Ok(store) => {
drop(store);
panic!("attached durability unexpectedly opened")
}
Err(error) => error,
};
assert!(matches!(
error,
StoreError::Durability(DurabilityPolicyError::AttachedNotYetSupported)
));
assert!(!store_path.exists(), "rejected open created store files");
}
#[test]
fn memory_vfs_writable_open_requires_host_lease_directory() {
let parent = tempdir().expect("parent directory");
let store_path = parent.path().join("memory-only-store");
let dependencies = StoreTestDependencies::new(
Arc::new(MemoryVfs::new()),
Arc::new(ManualMonotonicClock::new()),
);
let result =
Store::open_with_test_dependencies(&store_path, OpenOptions::default(), dependencies);
assert!(
matches!(result, Err(StoreError::Lock(super::lock::StoreLockError::Io { ref source, .. })) if source.kind() == std::io::ErrorKind::NotFound),
"writable open must refuse the missing host lease directory"
);
assert!(!store_path.exists());
}
#[test]
fn memory_vfs_store_retains_real_reader_lease() {
let directory = tempdir().expect("host lease directory");
let dependencies = StoreTestDependencies::new(
Arc::new(MemoryVfs::new()),
Arc::new(ManualMonotonicClock::new()),
);
let store = Store::open_with_test_dependencies(
directory.path(),
OpenOptions::default(),
dependencies,
)
.expect("MemoryVfs store with host lease");
let snapshot = store.open_snapshot().expect("snapshot view");
assert!(matches!(
super::lock::StoreLock::reclaim_exclusive(directory.path()),
Err(super::lock::StoreLockError::Io { source, .. })
if source.kind() == std::io::ErrorKind::WouldBlock
));
drop(store);
assert!(matches!(
super::lock::StoreLock::reclaim_exclusive(directory.path()),
Err(super::lock::StoreLockError::Io { source, .. })
if source.kind() == std::io::ErrorKind::WouldBlock
));
drop(snapshot);
drop(super::lock::StoreLock::reclaim_exclusive(directory.path()).expect("released lease"));
}
#[test]
fn open_refuses_a_manifest_ahead_of_the_checked_wal() {
let directory = tempdir().expect("store directory");
let policy = DurabilityPolicy::new(DurabilityMode::Derived, CommitTier::Ordered)
.expect("derived policy");
commit_manifest(
&StdVfs,
directory.path(),
&Manifest {
generation: 1,
log_seq: 1,
segments: Vec::new(),
epochs: Vec::new(),
epoch_alias: None,
schema: Schema::new(Vec::new()).expect("schema"),
},
policy,
)
.expect("manifest");
let error = match Store::open(directory.path(), OpenOptions::default()) {
Ok(store) => {
drop(store);
panic!("manifest ahead of missing WAL unexpectedly opened")
}
Err(error) => error,
};
assert!(matches!(
error,
StoreError::Manifest(ManifestError::AheadOfLog {
snapshot: 1,
durable: 0
})
));
}
}