1mod budget;
4mod cancel;
5mod clock;
6mod close;
7pub mod durability;
8mod expansion;
9pub(crate) mod graph_cache;
10mod hybrid;
11pub mod lock;
12mod materialize;
13mod pool;
14pub(crate) mod prepared;
15mod prepared_lexical;
16pub(crate) mod rescored_scan;
17#[cfg(test)]
18mod shared_bound_tests;
19mod snapshot;
20pub(crate) mod stats;
21
22use crate::diag::{timing_elapsed, timing_start};
23pub use cancel::{
24 CancelToken, Deadline, DeadlineError, QueryCancellation, QueryControl, QueryError,
25};
26#[cfg(any(test, feature = "test-support"))]
27pub use clock::ManualMonotonicClock;
28pub use clock::{MonotonicClock, SystemMonotonicClock};
29#[cfg(any(test, feature = "test-support"))]
30pub use hybrid::{
31 HybridScoreTestObservations, begin_hybrid_fresh_round_test_observations,
32 begin_hybrid_score_test_observations, take_hybrid_score_test_observations,
33};
34#[cfg(any(test, feature = "test-support"))]
35pub use materialize::MaterializationTestCounters;
36pub use materialize::{
37 HybridPreparationError, MaterializationError, MaterializedRow, QueryMaterializer,
38};
39
40#[cfg(test)]
41mod hybrid_overlap_tests;
42pub use rescored_scan::ScanRescoreOptions;
43pub use snapshot::{
44 InMemorySegment, InMemorySegmentFactors, PreparedSegment, PublishedSnapshot, SnapshotLease,
45};
46pub use stats::Stats;
47
48use std::collections::HashSet;
49#[cfg(any(test, feature = "test-support"))]
50use std::io::IoSlice;
51use std::path::{Path, PathBuf};
52use std::sync::atomic::{AtomicU64, Ordering};
53use std::sync::{Arc, Condvar, Mutex, RwLock};
54use std::time::Duration;
55
56use durability::{
57 CommitTier, DurabilityMode, DurabilityPolicy, DurabilityPolicyError, SyncRequirement,
58};
59use lock::{StoreLock, StoreLockError};
60
61use self::close::BackgroundThread;
62
63#[cfg(any(test, feature = "test-support"))]
65#[derive(Clone)]
66pub struct StoreTestDependencies {
67 vfs: Arc<dyn crate::vfs::Vfs>,
68 clock: Arc<dyn MonotonicClock>,
69 hybrid_leg_fault: Option<HybridLegTestFault>,
70 storage_fault_controller: Option<StorageFaultController>,
71 ingest_retention_fault_controller: Option<crate::ingest::IngestRetentionFaultController>,
72 metadata_test_controller: Option<Arc<crate::planner::MetadataTestController>>,
73 pub(crate) vector_fault_controller: Option<crate::scan::vector_fault::VectorFaultController>,
74 pub(crate) kernel_fault_controller: Option<crate::kernels::vector_fault::KernelFaultController>,
75 vector_seal_scheme: Option<crate::quant::QuantScheme>,
76}
77
78#[cfg(any(test, feature = "test-support"))]
79impl StoreTestDependencies {
80 #[must_use]
83 pub fn new(vfs: Arc<dyn crate::vfs::Vfs>, clock: Arc<dyn MonotonicClock>) -> Self {
84 Self {
85 vfs,
86 clock,
87 hybrid_leg_fault: None,
88 storage_fault_controller: None,
89 ingest_retention_fault_controller: None,
90 metadata_test_controller: None,
91 vector_fault_controller: None,
92 kernel_fault_controller: None,
93 vector_seal_scheme: None,
94 }
95 }
96
97 #[must_use]
99 pub const fn with_hybrid_leg_fault(mut self, fault: HybridLegTestFault) -> Self {
100 self.hybrid_leg_fault = Some(fault);
101 self
102 }
103
104 #[must_use]
106 pub fn with_storage_fault_controller(mut self, controller: StorageFaultController) -> Self {
107 self.storage_fault_controller = Some(controller);
108 self
109 }
110
111 #[must_use]
113 pub fn with_ingest_retention_fault_controller(
114 mut self,
115 controller: crate::ingest::IngestRetentionFaultController,
116 ) -> Self {
117 self.ingest_retention_fault_controller = Some(controller);
118 self
119 }
120
121 #[must_use]
124 pub fn with_metadata_test_controller(
125 mut self,
126 controller: Arc<crate::planner::MetadataTestController>,
127 ) -> Self {
128 self.metadata_test_controller = Some(controller);
129 self
130 }
131
132 #[must_use]
134 pub fn with_vector_fault_controller(
135 mut self,
136 controller: crate::scan::vector_fault::VectorFaultController,
137 ) -> Self {
138 self.vector_fault_controller = Some(controller);
139 self
140 }
141
142 #[must_use]
144 pub fn with_kernel_fault_controller(
145 mut self,
146 controller: crate::kernels::vector_fault::KernelFaultController,
147 ) -> Self {
148 self.kernel_fault_controller = Some(controller);
149 self
150 }
151
152 #[must_use]
157 pub const fn with_vector_seal_scheme(mut self, scheme: crate::quant::QuantScheme) -> Self {
158 self.vector_seal_scheme = Some(scheme);
159 self
160 }
161}
162
163#[cfg(any(test, feature = "test-support"))]
165#[derive(Clone, Debug, Eq, PartialEq)]
166#[doc(hidden)]
167pub enum StorageTestFault {
168 TornWalHeader,
169 TornWalBody,
170 TornWalChecksum,
171 PostCommitError,
172 ManifestPreRename,
173 ManifestPostRename,
174 CorruptSegmentRegion,
175 WrongManifestObject,
176 WrongSegmentObject,
177 ListOmission { file_name: String },
178 DeleteOmission { file_name: String },
179}
180
181#[cfg(any(test, feature = "test-support"))]
182impl StorageTestFault {
183 const fn operation(&self) -> &'static str {
184 match self {
185 Self::TornWalHeader | Self::TornWalBody | Self::TornWalChecksum => "wal-prefix",
186 Self::PostCommitError => "retry",
187 Self::ManifestPreRename | Self::ManifestPostRename => "publication",
188 Self::CorruptSegmentRegion | Self::WrongManifestObject | Self::WrongSegmentObject => {
189 "format-check"
190 }
191 Self::ListOmission { .. } | Self::DeleteOmission { .. } => "orphan-cleanup",
192 }
193 }
194
195 const fn key(&self) -> &'static str {
196 match self {
197 Self::TornWalHeader => "torn-wal-header",
198 Self::TornWalBody => "torn-wal-body",
199 Self::TornWalChecksum => "torn-wal-checksum",
200 Self::PostCommitError => "post-commit-error",
201 Self::ManifestPreRename => "manifest-pre-rename-crash",
202 Self::ManifestPostRename => "manifest-post-rename-crash",
203 Self::CorruptSegmentRegion => "corrupt-segment-region",
204 Self::WrongManifestObject => "wrong-manifest-object",
205 Self::WrongSegmentObject => "wrong-segment-object",
206 Self::ListOmission { .. } | Self::DeleteOmission { .. } => "list-delete-omission",
207 }
208 }
209}
210
211#[cfg(any(test, feature = "test-support"))]
213#[derive(Clone, Debug, Eq, PartialEq)]
214#[doc(hidden)]
215pub struct StorageFaultPlan {
216 op_index: u32,
217 artifact: String,
218 offset: Option<u64>,
219 segment: Option<crate::segment::SegmentId>,
220 region_kind: Option<u16>,
221 chunk: Option<u32>,
222}
223
224#[cfg(any(test, feature = "test-support"))]
225impl StorageFaultPlan {
226 #[must_use]
228 pub fn new(op_index: u32, artifact: impl Into<String>) -> Self {
229 Self {
230 op_index,
231 artifact: artifact.into(),
232 offset: None,
233 segment: None,
234 region_kind: None,
235 chunk: None,
236 }
237 }
238
239 #[must_use]
241 pub const fn with_offset(mut self, offset: u64) -> Self {
242 self.offset = Some(offset);
243 self
244 }
245
246 #[must_use]
248 pub const fn with_segment_region(
249 mut self,
250 segment: crate::segment::SegmentId,
251 region_kind: u16,
252 chunk: u32,
253 ) -> Self {
254 self.segment = Some(segment);
255 self.region_kind = Some(region_kind);
256 self.chunk = Some(chunk);
257 self
258 }
259
260 #[must_use]
262 pub const fn with_segment(mut self, segment: crate::segment::SegmentId) -> Self {
263 self.segment = Some(segment);
264 self
265 }
266
267 #[must_use]
269 pub const fn op_index(&self) -> u32 {
270 self.op_index
271 }
272
273 #[must_use]
275 pub fn artifact(&self) -> &str {
276 &self.artifact
277 }
278
279 #[must_use]
281 pub const fn offset(&self) -> Option<u64> {
282 self.offset
283 }
284
285 #[must_use]
287 pub const fn segment(&self) -> Option<crate::segment::SegmentId> {
288 self.segment
289 }
290
291 #[must_use]
293 pub const fn region_kind(&self) -> Option<u16> {
294 self.region_kind
295 }
296
297 #[must_use]
299 pub const fn chunk(&self) -> Option<u32> {
300 self.chunk
301 }
302}
303
304#[cfg(any(test, feature = "test-support"))]
306#[derive(Clone, Copy, Debug, Eq, PartialEq)]
307#[doc(hidden)]
308pub enum StorageReceiptSite {
309 WalOpenHeaderValidation,
310 WalOpenRecordValidation,
311 WalOpenRecordChecksum,
312 WalCommitAppendAfterInnerSuccess,
313 ManifestCommitBeforeRename,
314 ManifestCommitAfterRename,
315 SegmentReadRegionChecksum,
316 ManifestOpenFamilyValidation,
317 SegmentOpenFamilyValidation,
318 SegmentOpenObjectIdentity,
319 OrphanCleanupList,
320 OrphanCleanupDelete,
321}
322
323#[cfg(any(test, feature = "test-support"))]
324impl StorageReceiptSite {
325 const fn key(self) -> &'static str {
326 match self {
327 Self::WalOpenHeaderValidation => "WalOpen.HeaderValidation",
328 Self::WalOpenRecordValidation => "WalOpen.RecordValidation",
329 Self::WalOpenRecordChecksum => "WalOpen.RecordChecksum",
330 Self::WalCommitAppendAfterInnerSuccess => "WalCommit.AppendAfterInnerSuccess",
331 Self::ManifestCommitBeforeRename => "ManifestCommit.BeforeRename",
332 Self::ManifestCommitAfterRename => "ManifestCommit.AfterRename",
333 Self::SegmentReadRegionChecksum => "SegmentRead.RegionChecksum",
334 Self::ManifestOpenFamilyValidation => "ManifestOpen.FamilyValidation",
335 Self::SegmentOpenFamilyValidation => "SegmentOpen.FamilyValidation",
336 Self::SegmentOpenObjectIdentity => "SegmentOpen.ObjectIdentity",
337 Self::OrphanCleanupList => "OrphanCleanup.List",
338 Self::OrphanCleanupDelete => "OrphanCleanup.Delete",
339 }
340 }
341}
342
343#[cfg(any(test, feature = "test-support"))]
345#[derive(Clone, Debug, Eq, PartialEq)]
346#[doc(hidden)]
347pub enum StorageReceiptObserved {
348 WalHeader {
349 artifact: String,
350 reason: crate::wal::header::WalHeaderError,
351 },
352 WalRecord {
353 artifact: String,
354 offset: u64,
355 reason: crate::wal::replay::CorruptionReason,
356 },
357 WalAppend {
358 artifact: String,
359 encoded_len: u64,
360 first_seq: u64,
361 last_seq: u64,
362 inner_append_completed: bool,
363 caller_saw_error: bool,
364 },
365 ManifestRename {
366 temporary: String,
367 committed: String,
368 rename_performed: bool,
369 new_segment_final: bool,
370 directory_sync_returned: bool,
371 },
372 SegmentChecksum {
373 artifact: String,
374 segment: crate::segment::SegmentId,
375 region_kind: u16,
376 chunk: u32,
377 expected_checksum: u64,
378 actual_checksum: u64,
379 },
380 Format {
381 artifact: String,
382 check: crate::format::frame::FormatCheck,
383 expected_family: Option<u16>,
384 actual_family: Option<u16>,
385 expected_id: Option<crate::segment::SegmentId>,
386 actual_id: Option<crate::segment::SegmentId>,
387 },
388 Omission {
389 artifact: String,
390 deletion_observed: bool,
391 },
392}
393
394#[cfg(any(test, feature = "test-support"))]
396#[derive(Clone, Debug, Eq, PartialEq)]
397#[doc(hidden)]
398pub struct StorageFaultReceipt {
399 operation: &'static str,
400 fault: &'static str,
401 site: StorageReceiptSite,
402 plan: StorageFaultPlan,
403 observed: StorageReceiptObserved,
404 cardinality: u32,
405}
406
407#[cfg(any(test, feature = "test-support"))]
408impl StorageFaultReceipt {
409 #[must_use]
410 pub const fn campaign(&self) -> &'static str {
411 "storage-durability"
412 }
413
414 #[must_use]
415 pub const fn operation(&self) -> &'static str {
416 self.operation
417 }
418
419 #[must_use]
420 pub const fn fault(&self) -> &'static str {
421 self.fault
422 }
423
424 #[must_use]
425 pub const fn site(&self) -> &'static str {
426 self.site.key()
427 }
428
429 #[must_use]
430 pub const fn typed_site(&self) -> StorageReceiptSite {
431 self.site
432 }
433
434 #[must_use]
435 pub const fn cardinality(&self) -> u32 {
436 self.cardinality
437 }
438
439 #[must_use]
440 pub const fn plan(&self) -> &StorageFaultPlan {
441 &self.plan
442 }
443
444 #[must_use]
445 pub const fn observed(&self) -> &StorageReceiptObserved {
446 &self.observed
447 }
448
449 pub fn from_manifest_abort_acknowledgment(acknowledgment: &str) -> Result<Self, String> {
455 let body = acknowledgment
456 .strip_suffix('\n')
457 .filter(|body| !body.contains('\n'))
458 .ok_or_else(|| {
459 "manifest abort acknowledgment must end in exactly one newline".to_owned()
460 })?;
461 let fields = body
462 .split('|')
463 .map(|field| {
464 field.split_once('=').ok_or_else(|| {
465 format!("manifest abort acknowledgment field lacks '=': {field}")
466 })
467 })
468 .collect::<Result<Vec<_>, _>>()?;
469 let expected_keys = [
470 "campaign",
471 "operation",
472 "fault",
473 "site",
474 "op_index",
475 "cardinality",
476 "artifact",
477 "temporary",
478 "committed",
479 "rename_performed",
480 "new_segment_final",
481 "directory_sync_returned",
482 ];
483 if fields.len() != expected_keys.len()
484 || fields
485 .iter()
486 .map(|(key, _)| *key)
487 .ne(expected_keys.iter().copied())
488 {
489 return Err("manifest abort acknowledgment field inventory differs".to_owned());
490 }
491 let mut values = fields.into_iter().map(|(_, value)| value);
492 let campaign = values
493 .next()
494 .ok_or_else(|| "manifest abort acknowledgment lacks campaign".to_owned())?;
495 let operation = values
496 .next()
497 .ok_or_else(|| "manifest abort acknowledgment lacks operation".to_owned())?;
498 let fault = values
499 .next()
500 .ok_or_else(|| "manifest abort acknowledgment lacks fault".to_owned())?;
501 let site = values
502 .next()
503 .ok_or_else(|| "manifest abort acknowledgment lacks site".to_owned())?;
504 let op_index = values
505 .next()
506 .ok_or_else(|| "manifest abort acknowledgment lacks op index".to_owned())?
507 .parse::<u32>()
508 .map_err(|error| format!("manifest abort op index is invalid: {error}"))?;
509 let cardinality = values
510 .next()
511 .ok_or_else(|| "manifest abort acknowledgment lacks cardinality".to_owned())?
512 .parse::<u32>()
513 .map_err(|error| format!("manifest abort cardinality is invalid: {error}"))?;
514 let artifact = values
515 .next()
516 .ok_or_else(|| "manifest abort acknowledgment lacks artifact".to_owned())?;
517 let temporary = values
518 .next()
519 .ok_or_else(|| "manifest abort acknowledgment lacks temporary path".to_owned())?;
520 let committed = values
521 .next()
522 .ok_or_else(|| "manifest abort acknowledgment lacks committed path".to_owned())?;
523 let parse_bool = |field: &str, value: &str| match value {
524 "true" => Ok(true),
525 "false" => Ok(false),
526 _ => Err(format!("manifest abort {field} is not a bool: {value}")),
527 };
528 let rename_performed = parse_bool(
529 "rename_performed",
530 values
531 .next()
532 .ok_or_else(|| "manifest abort acknowledgment lacks rename result".to_owned())?,
533 )?;
534 let new_segment_final = parse_bool(
535 "new_segment_final",
536 values.next().ok_or_else(|| {
537 "manifest abort acknowledgment lacks final-segment result".to_owned()
538 })?,
539 )?;
540 let directory_sync_returned = parse_bool(
541 "directory_sync_returned",
542 values.next().ok_or_else(|| {
543 "manifest abort acknowledgment lacks directory-sync result".to_owned()
544 })?,
545 )?;
546 if campaign != "storage-durability"
547 || operation != "publication"
548 || cardinality != 1
549 || artifact != crate::manifest::io::MANIFEST_FILE
550 || temporary != crate::manifest::io::MANIFEST_TEMP_FILE
551 || committed != crate::manifest::io::MANIFEST_FILE
552 || !new_segment_final
553 || directory_sync_returned
554 {
555 return Err("manifest abort acknowledgment fixed facts differ".to_owned());
556 }
557 let (fault, site, expected_rename) = match (fault, site) {
558 ("manifest-pre-rename-crash", "ManifestCommit.BeforeRename") => (
559 "manifest-pre-rename-crash",
560 StorageReceiptSite::ManifestCommitBeforeRename,
561 false,
562 ),
563 ("manifest-post-rename-crash", "ManifestCommit.AfterRename") => (
564 "manifest-post-rename-crash",
565 StorageReceiptSite::ManifestCommitAfterRename,
566 true,
567 ),
568 _ => return Err("manifest abort fault and production site differ".to_owned()),
569 };
570 if rename_performed != expected_rename {
571 return Err("manifest abort rename result differs from production site".to_owned());
572 }
573 Ok(Self {
574 operation: "publication",
575 fault,
576 site,
577 plan: StorageFaultPlan::new(op_index, artifact),
578 observed: StorageReceiptObserved::ManifestRename {
579 temporary: temporary.to_owned(),
580 committed: committed.to_owned(),
581 rename_performed,
582 new_segment_final,
583 directory_sync_returned,
584 },
585 cardinality,
586 })
587 }
588}
589
590#[cfg(any(test, feature = "test-support"))]
591#[derive(Debug)]
592struct StorageFaultState {
593 fault: StorageTestFault,
594 plan: StorageFaultPlan,
595 consumed: bool,
596 receipt: Option<StorageFaultReceipt>,
597 cleanup_report: Option<StorageCleanupReport>,
598 #[cfg(unix)]
599 child_abort_ack: Option<std::os::unix::net::UnixStream>,
600}
601
602#[cfg(any(test, feature = "test-support"))]
604#[derive(Clone, Debug, Eq, PartialEq)]
605#[doc(hidden)]
606pub struct StorageCleanupReport {
607 reclaimed_bytes: u64,
608 deleted_paths: Vec<String>,
609 retained_eligible_paths: Vec<String>,
610 directory_synced: bool,
611}
612
613#[cfg(any(test, feature = "test-support"))]
614impl StorageCleanupReport {
615 #[must_use]
617 pub const fn reclaimed_bytes(&self) -> u64 {
618 self.reclaimed_bytes
619 }
620
621 #[must_use]
623 pub fn deleted_paths(&self) -> &[String] {
624 &self.deleted_paths
625 }
626
627 #[must_use]
629 pub fn retained_eligible_paths(&self) -> &[String] {
630 &self.retained_eligible_paths
631 }
632
633 #[must_use]
635 pub const fn directory_synced(&self) -> bool {
636 self.directory_synced
637 }
638}
639
640#[cfg(any(test, feature = "test-support"))]
642#[derive(Clone, Debug)]
643#[doc(hidden)]
644pub struct StorageFaultController {
645 state: Arc<Mutex<StorageFaultState>>,
646}
647
648#[cfg(any(test, feature = "test-support"))]
649impl StorageFaultController {
650 #[must_use]
652 pub fn new(fault: StorageTestFault, plan: StorageFaultPlan) -> Self {
653 Self {
654 state: Arc::new(Mutex::new(StorageFaultState {
655 fault,
656 plan,
657 consumed: false,
658 receipt: None,
659 cleanup_report: None,
660 #[cfg(unix)]
661 child_abort_ack: None,
662 })),
663 }
664 }
665
666 #[cfg(unix)]
668 #[must_use]
669 pub fn with_child_abort_ack(self, acknowledgment: std::os::unix::net::UnixStream) -> Self {
670 if let Ok(mut state) = self.state.lock() {
671 state.child_abort_ack = Some(acknowledgment);
672 }
673 self
674 }
675
676 pub fn take_receipt(&self) -> Option<StorageFaultReceipt> {
678 self.state
679 .lock()
680 .ok()
681 .and_then(|mut state| state.receipt.take())
682 }
683
684 pub fn take_cleanup_report(&self) -> Option<StorageCleanupReport> {
686 self.state
687 .lock()
688 .ok()
689 .and_then(|mut state| state.cleanup_report.take())
690 }
691
692 fn record_cleanup_report(&self, report: crate::manifest::io::OrphanCleanupReport) {
693 let normalize = |paths: Vec<PathBuf>| {
694 paths
695 .into_iter()
696 .filter_map(|path| {
697 path.file_name()
698 .and_then(|name| name.to_str())
699 .map(str::to_owned)
700 })
701 .collect::<Vec<_>>()
702 };
703 if let Ok(mut state) = self.state.lock() {
704 state.cleanup_report = Some(StorageCleanupReport {
705 reclaimed_bytes: report.reclaimed_bytes,
706 deleted_paths: normalize(report.deleted_paths),
707 retained_eligible_paths: normalize(report.retained_eligible_paths),
708 directory_synced: report.directory_synced,
709 });
710 }
711 }
712
713 fn armed_fault(&self) -> Option<StorageTestFault> {
714 self.state
715 .lock()
716 .ok()
717 .and_then(|state| (!state.consumed).then(|| state.fault.clone()))
718 }
719
720 fn planned_segment(&self) -> Option<crate::segment::SegmentId> {
721 self.state
722 .lock()
723 .ok()
724 .and_then(|state| state.plan.segment())
725 }
726
727 fn emit(&self, site: StorageReceiptSite, observed: StorageReceiptObserved) -> bool {
728 let Ok(mut state) = self.state.lock() else {
729 return false;
730 };
731 if state.consumed {
732 return false;
733 }
734 state.consumed = true;
735 state.receipt = Some(StorageFaultReceipt {
736 operation: state.fault.operation(),
737 fault: state.fault.key(),
738 site,
739 plan: state.plan.clone(),
740 observed,
741 cardinality: 1,
742 });
743 true
744 }
745
746 fn abort_after_manifest_receipt(&self) -> ! {
747 #[cfg(unix)]
748 if let Ok(mut state) = self.state.lock() {
749 let wire_receipt = state
750 .receipt
751 .as_ref()
752 .and_then(storage_manifest_abort_wire_receipt);
753 if let (Some(wire_receipt), Some(acknowledgment)) =
754 (wire_receipt, state.child_abort_ack.as_mut())
755 {
756 use std::io::Write as _;
757 let _ = acknowledgment.write_all(wire_receipt.as_bytes());
758 let _ = acknowledgment.flush();
759 }
760 }
761 std::process::abort()
762 }
763}
764
765#[cfg(any(test, feature = "test-support"))]
766thread_local! {
767 static STORAGE_OPEN_CONTROLLER: std::cell::RefCell<
768 Option<std::sync::Weak<Mutex<StorageFaultState>>>
769 > = const { std::cell::RefCell::new(None) };
770}
771
772#[cfg(any(test, feature = "test-support"))]
773fn with_storage_open_controller<R>(
774 controller: Option<&StorageFaultController>,
775 operation: impl FnOnce() -> R,
776) -> R {
777 STORAGE_OPEN_CONTROLLER.with(|slot| {
778 let previous = slot.replace(controller.map(|controller| Arc::downgrade(&controller.state)));
779 let result = operation();
780 let _ = slot.replace(previous);
781 result
782 })
783}
784
785#[cfg(any(test, feature = "test-support"))]
786fn storage_open_controller() -> Option<StorageFaultController> {
787 STORAGE_OPEN_CONTROLLER.with(|slot| {
788 slot.borrow()
789 .as_ref()
790 .and_then(std::sync::Weak::upgrade)
791 .map(|state| StorageFaultController { state })
792 })
793}
794
795#[cfg(any(test, feature = "test-support"))]
796pub(crate) fn record_storage_wal_recovery_fault(error: &crate::wal::WalRecoveryError) {
797 let Some(controller) = storage_open_controller() else {
798 return;
799 };
800 match (controller.armed_fault(), error) {
801 (
802 Some(StorageTestFault::TornWalHeader),
803 crate::wal::WalRecoveryError::InvalidHeader(reason),
804 ) => {
805 let _ = controller.emit(
806 StorageReceiptSite::WalOpenHeaderValidation,
807 StorageReceiptObserved::WalHeader {
808 artifact: "wal.ze".to_owned(),
809 reason: *reason,
810 },
811 );
812 }
813 (
814 Some(StorageTestFault::TornWalBody),
815 crate::wal::WalRecoveryError::CorruptAt {
816 offset,
817 reason:
818 reason @ crate::wal::replay::CorruptionReason::Record {
819 error:
820 crate::wal::record::RecordError::HeaderTruncated { .. }
821 | crate::wal::record::RecordError::BodyTruncated { .. }
822 | crate::wal::record::RecordError::LengthOverflow { .. },
823 ..
824 },
825 },
826 ) => record_storage_wal_record_fault(
827 &controller,
828 StorageReceiptSite::WalOpenRecordValidation,
829 *offset,
830 *reason,
831 ),
832 (
833 Some(StorageTestFault::TornWalChecksum),
834 crate::wal::WalRecoveryError::CorruptAt {
835 offset,
836 reason:
837 reason @ crate::wal::replay::CorruptionReason::Record {
838 error: crate::wal::record::RecordError::ChecksumMismatch { .. },
839 ..
840 },
841 },
842 ) => record_storage_wal_record_fault(
843 &controller,
844 StorageReceiptSite::WalOpenRecordChecksum,
845 *offset,
846 *reason,
847 ),
848 _ => {}
849 }
850}
851
852#[cfg(any(test, feature = "test-support"))]
853fn record_storage_wal_record_fault(
854 controller: &StorageFaultController,
855 site: StorageReceiptSite,
856 offset: usize,
857 reason: crate::wal::replay::CorruptionReason,
858) {
859 let Ok(offset) = u64::try_from(offset) else {
860 return;
861 };
862 let _ = controller.emit(
863 site,
864 StorageReceiptObserved::WalRecord {
865 artifact: "wal.ze".to_owned(),
866 offset,
867 reason,
868 },
869 );
870}
871
872#[cfg(any(test, feature = "test-support"))]
873pub(crate) fn record_storage_manifest_format_fault(
874 error: &crate::format::frame::FormatError,
875 actual_family: Option<u16>,
876) {
877 let Some(controller) = storage_open_controller() else {
878 return;
879 };
880 if matches!(
881 controller.armed_fault(),
882 Some(StorageTestFault::WrongManifestObject)
883 ) && error.check() == crate::format::frame::FormatCheck::Family
884 {
885 let _ = controller.emit(
886 StorageReceiptSite::ManifestOpenFamilyValidation,
887 StorageReceiptObserved::Format {
888 artifact: storage_artifact_name(error.artifact()).to_owned(),
889 check: error.check(),
890 expected_family: Some(crate::format::FormatFamily::Manifest.id()),
891 actual_family,
892 expected_id: None,
893 actual_id: None,
894 },
895 );
896 }
897}
898
899#[cfg(any(test, feature = "test-support"))]
900pub(crate) fn record_storage_segment_format_fault(
901 error: &crate::format::frame::FormatError,
902 actual_family: Option<u16>,
903 expected_id: crate::segment::SegmentId,
904) {
905 let Some(controller) = storage_open_controller() else {
906 return;
907 };
908 if matches!(
909 controller.armed_fault(),
910 Some(StorageTestFault::WrongSegmentObject)
911 ) && error.check() == crate::format::frame::FormatCheck::Family
912 {
913 let _ = controller.emit(
914 StorageReceiptSite::SegmentOpenFamilyValidation,
915 StorageReceiptObserved::Format {
916 artifact: storage_artifact_name(error.artifact()).to_owned(),
917 check: error.check(),
918 expected_family: Some(crate::format::FormatFamily::Segment.id()),
919 actual_family,
920 expected_id: Some(expected_id),
921 actual_id: None,
922 },
923 );
924 }
925}
926
927#[cfg(any(test, feature = "test-support"))]
928pub(crate) fn record_storage_segment_identity_fault(
929 artifact: &str,
930 expected: crate::segment::SegmentId,
931 actual: crate::segment::SegmentId,
932) {
933 let Some(controller) = storage_open_controller() else {
934 return;
935 };
936 if matches!(
937 controller.armed_fault(),
938 Some(StorageTestFault::WrongSegmentObject)
939 ) {
940 let _ = controller.emit(
941 StorageReceiptSite::SegmentOpenObjectIdentity,
942 StorageReceiptObserved::Format {
943 artifact: storage_artifact_name(artifact).to_owned(),
944 check: crate::format::frame::FormatCheck::ObjectIdentity,
945 expected_family: None,
946 actual_family: None,
947 expected_id: Some(expected),
948 actual_id: Some(actual),
949 },
950 );
951 }
952}
953
954#[cfg(any(test, feature = "test-support"))]
955fn storage_manifest_abort_wire_receipt(receipt: &StorageFaultReceipt) -> Option<String> {
956 let StorageReceiptObserved::ManifestRename {
957 temporary,
958 committed,
959 rename_performed,
960 new_segment_final,
961 directory_sync_returned,
962 } = &receipt.observed
963 else {
964 return None;
965 };
966 Some(format!(
967 "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",
968 receipt.campaign(),
969 receipt.operation,
970 receipt.fault,
971 receipt.site.key(),
972 receipt.plan.op_index,
973 receipt.cardinality,
974 receipt.plan.artifact,
975 ))
976}
977
978#[cfg(any(test, feature = "test-support"))]
979fn storage_artifact_name(artifact: &str) -> &str {
980 Path::new(artifact)
981 .file_name()
982 .and_then(|name| name.to_str())
983 .unwrap_or(artifact)
984}
985
986#[cfg(any(test, feature = "test-support"))]
987type StorageSegmentControllerRegistry = Mutex<
988 Vec<(
989 PathBuf,
990 crate::segment::SegmentId,
991 std::sync::Weak<Mutex<StorageFaultState>>,
992 )>,
993>;
994
995#[cfg(any(test, feature = "test-support"))]
996static STORAGE_SEGMENT_CONTROLLERS: std::sync::OnceLock<StorageSegmentControllerRegistry> =
997 std::sync::OnceLock::new();
998
999#[cfg(any(test, feature = "test-support"))]
1000fn register_storage_segment_controller(
1001 store_directory: &Path,
1002 segment: crate::segment::SegmentId,
1003 controller: &StorageFaultController,
1004) {
1005 let registry = STORAGE_SEGMENT_CONTROLLERS.get_or_init(|| Mutex::new(Vec::new()));
1006 if let Ok(mut entries) = registry.lock() {
1007 entries.retain(|(_, _, state)| state.strong_count() != 0);
1008 entries.push((
1009 store_directory.to_path_buf(),
1010 segment,
1011 Arc::downgrade(&controller.state),
1012 ));
1013 }
1014}
1015
1016#[cfg(any(test, feature = "test-support"))]
1017pub(crate) fn record_storage_segment_checksum_fault(
1018 store_directory: &Path,
1019 segment: crate::segment::SegmentId,
1020 kind: crate::segment::layout::RegionKind,
1021 chunk: u32,
1022 expected: u64,
1023 actual: u64,
1024) {
1025 let Some(registry) = STORAGE_SEGMENT_CONTROLLERS.get() else {
1026 return;
1027 };
1028 let states = match registry.lock() {
1029 Ok(mut entries) => {
1030 entries.retain(|(_, _, state)| state.strong_count() != 0);
1031 entries
1032 .iter()
1033 .filter(|(directory, registered, _)| {
1034 directory == store_directory && *registered == segment
1035 })
1036 .filter_map(|(_, _, state)| state.upgrade())
1037 .collect::<Vec<_>>()
1038 }
1039 Err(_) => return,
1040 };
1041 for state in states {
1042 let controller = StorageFaultController { state };
1043 if matches!(
1044 controller.armed_fault(),
1045 Some(StorageTestFault::CorruptSegmentRegion)
1046 ) {
1047 let _ = controller.emit(
1048 StorageReceiptSite::SegmentReadRegionChecksum,
1049 StorageReceiptObserved::SegmentChecksum {
1050 artifact: segment.file_name(),
1051 segment,
1052 region_kind: kind.id(),
1053 chunk,
1054 expected_checksum: expected,
1055 actual_checksum: actual,
1056 },
1057 );
1058 }
1059 }
1060}
1061
1062#[cfg(any(test, feature = "test-support"))]
1063struct StorageFaultVfs {
1064 inner: Arc<dyn crate::vfs::Vfs>,
1065 controller: StorageFaultController,
1066}
1067
1068#[cfg(any(test, feature = "test-support"))]
1069struct StorageFaultFile {
1070 inner: Box<dyn crate::vfs::VfsFile>,
1071 path: PathBuf,
1072 controller: StorageFaultController,
1073}
1074
1075#[cfg(any(test, feature = "test-support"))]
1076impl StorageFaultFile {
1077 fn after_append(&self, encoded_len: u64, first_seq: u64, last_seq: u64) -> std::io::Result<()> {
1078 if matches!(
1079 self.controller.armed_fault(),
1080 Some(StorageTestFault::PostCommitError)
1081 ) && self.path.file_name().is_some_and(|name| name == "wal.ze")
1082 && self.controller.emit(
1083 StorageReceiptSite::WalCommitAppendAfterInnerSuccess,
1084 StorageReceiptObserved::WalAppend {
1085 artifact: "wal.ze".to_owned(),
1086 encoded_len,
1087 first_seq,
1088 last_seq,
1089 inner_append_completed: true,
1090 caller_saw_error: true,
1091 },
1092 )
1093 {
1094 return Err(std::io::Error::other(
1095 "scheduled post-commit error after inner append",
1096 ));
1097 }
1098 Ok(())
1099 }
1100}
1101
1102#[cfg(any(test, feature = "test-support"))]
1103impl crate::vfs::VfsFile for StorageFaultFile {
1104 fn append(&mut self, bytes: &[u8]) -> std::io::Result<()> {
1105 let record_facts = storage_wal_append_facts(bytes)?;
1106 self.inner.append(bytes)?;
1107 record_facts.map_or(Ok(()), |(encoded_len, first_seq, last_seq)| {
1108 self.after_append(encoded_len, first_seq, last_seq)
1109 })
1110 }
1111
1112 fn append_vectored(&mut self, buffers: &mut [IoSlice<'_>]) -> std::io::Result<()> {
1113 let mut aggregate = None::<(u64, u64, u64)>;
1114 for bytes in buffers.iter() {
1115 if let Some((encoded_len, first_seq, last_seq)) = storage_wal_append_facts(bytes)? {
1116 aggregate = Some(match aggregate {
1117 None => (encoded_len, first_seq, last_seq),
1118 Some((total, first, _)) => (
1119 total.checked_add(encoded_len).ok_or_else(|| {
1120 std::io::Error::other("WAL record append length overflow")
1121 })?,
1122 first,
1123 last_seq,
1124 ),
1125 });
1126 }
1127 }
1128 self.inner.append_vectored(buffers)?;
1129 aggregate.map_or(Ok(()), |(encoded_len, first_seq, last_seq)| {
1130 self.after_append(encoded_len, first_seq, last_seq)
1131 })
1132 }
1133
1134 fn sync(&self, kind: crate::vfs::SyncKind) -> std::io::Result<()> {
1135 self.inner.sync(kind)
1136 }
1137}
1138
1139#[cfg(any(test, feature = "test-support"))]
1140impl crate::vfs::Vfs for StorageFaultVfs {
1141 fn segment_data_read_counter(&self) -> Option<Arc<AtomicU64>> {
1142 self.inner.segment_data_read_counter()
1143 }
1144
1145 fn ensure_directory(&self, path: &Path, create: bool) -> std::io::Result<bool> {
1146 self.inner.ensure_directory(path, create)
1147 }
1148
1149 fn open(&self, path: &Path) -> std::io::Result<u64> {
1150 self.inner.open(path)
1151 }
1152
1153 fn open_for_map(&self, path: &Path) -> std::io::Result<std::fs::File> {
1154 self.inner.open_for_map(path)
1155 }
1156
1157 fn read(&self, path: &Path) -> std::io::Result<Vec<u8>> {
1158 self.inner.read(path)
1159 }
1160
1161 fn read_range(&self, path: &Path, offset: u64, length: usize) -> std::io::Result<Vec<u8>> {
1162 self.inner.read_range(path, offset, length)
1163 }
1164
1165 fn write(&self, path: &Path, bytes: &[u8]) -> std::io::Result<()> {
1166 self.inner.write(path, bytes)
1167 }
1168
1169 fn open_append(&self, path: &Path) -> std::io::Result<Box<dyn crate::vfs::VfsFile>> {
1170 Ok(Box::new(StorageFaultFile {
1171 inner: self.inner.open_append(path)?,
1172 path: path.to_path_buf(),
1173 controller: self.controller.clone(),
1174 }))
1175 }
1176
1177 fn rename(&self, from: &Path, to: &Path) -> std::io::Result<()> {
1178 let is_manifest = from
1179 .file_name()
1180 .is_some_and(|name| name == crate::manifest::io::MANIFEST_TEMP_FILE)
1181 && to
1182 .file_name()
1183 .is_some_and(|name| name == crate::manifest::io::MANIFEST_FILE);
1184 if is_manifest
1185 && matches!(
1186 self.controller.armed_fault(),
1187 Some(StorageTestFault::ManifestPreRename)
1188 )
1189 && self.controller.emit(
1190 StorageReceiptSite::ManifestCommitBeforeRename,
1191 StorageReceiptObserved::ManifestRename {
1192 temporary: crate::manifest::io::MANIFEST_TEMP_FILE.to_owned(),
1193 committed: crate::manifest::io::MANIFEST_FILE.to_owned(),
1194 rename_performed: false,
1195 new_segment_final: storage_planned_segment_is_final(
1196 &*self.inner,
1197 from,
1198 self.controller.planned_segment(),
1199 ),
1200 directory_sync_returned: false,
1201 },
1202 )
1203 {
1204 self.controller.abort_after_manifest_receipt();
1205 }
1206 self.inner.rename(from, to)?;
1207 if is_manifest
1208 && matches!(
1209 self.controller.armed_fault(),
1210 Some(StorageTestFault::ManifestPostRename)
1211 )
1212 && self.controller.emit(
1213 StorageReceiptSite::ManifestCommitAfterRename,
1214 StorageReceiptObserved::ManifestRename {
1215 temporary: crate::manifest::io::MANIFEST_TEMP_FILE.to_owned(),
1216 committed: crate::manifest::io::MANIFEST_FILE.to_owned(),
1217 rename_performed: true,
1218 new_segment_final: storage_planned_segment_is_final(
1219 &*self.inner,
1220 to,
1221 self.controller.planned_segment(),
1222 ),
1223 directory_sync_returned: false,
1224 },
1225 )
1226 {
1227 self.controller.abort_after_manifest_receipt();
1228 }
1229 Ok(())
1230 }
1231
1232 fn sync(&self, path: &Path, kind: crate::vfs::SyncKind) -> std::io::Result<()> {
1233 self.inner.sync(path, kind)
1234 }
1235
1236 fn list(&self, directory: &Path) -> std::io::Result<Vec<PathBuf>> {
1237 let mut paths = self.inner.list(directory)?;
1238 if let Some(StorageTestFault::ListOmission { file_name }) = self.controller.armed_fault()
1239 && let Some(position) = paths.iter().position(|path| {
1240 path.file_name()
1241 .is_some_and(|name| name == file_name.as_str())
1242 })
1243 {
1244 let _ = paths.remove(position);
1245 let _ = self.controller.emit(
1246 StorageReceiptSite::OrphanCleanupList,
1247 StorageReceiptObserved::Omission {
1248 artifact: file_name,
1249 deletion_observed: false,
1250 },
1251 );
1252 }
1253 Ok(paths)
1254 }
1255
1256 fn delete(&self, path: &Path) -> std::io::Result<()> {
1257 if let Some(StorageTestFault::DeleteOmission { file_name }) = self.controller.armed_fault()
1258 && path
1259 .file_name()
1260 .is_some_and(|name| name == file_name.as_str())
1261 && self.controller.emit(
1262 StorageReceiptSite::OrphanCleanupDelete,
1263 StorageReceiptObserved::Omission {
1264 artifact: file_name,
1265 deletion_observed: false,
1266 },
1267 )
1268 {
1269 return Ok(());
1270 }
1271 self.inner.delete(path)
1272 }
1273}
1274
1275#[cfg(any(test, feature = "test-support"))]
1276fn storage_wal_append_facts(bytes: &[u8]) -> std::io::Result<Option<(u64, u64, u64)>> {
1277 const RECORD_HEADER_LEN: usize = 14;
1278 const RECORD_CHECKSUM_LEN: usize = 8;
1279 let has_wal_header = bytes
1280 .get(..8)
1281 .is_some_and(|magic| magic == crate::format::frame::FILE_MAGIC)
1282 && bytes
1283 .get(8..10)
1284 .is_some_and(|family| family == crate::format::FormatFamily::Wal.id().to_le_bytes());
1285 let mut offset = if has_wal_header {
1286 crate::wal::header::WAL_HEADER_LEN
1287 } else {
1288 0
1289 };
1290 if offset == bytes.len() {
1291 return Ok(None);
1292 }
1293 let mut encoded_len = 0_u64;
1294 let mut first_seq = None;
1295 let mut last_seq = None;
1296 while offset < bytes.len() {
1297 let header_end = offset
1298 .checked_add(RECORD_HEADER_LEN)
1299 .ok_or_else(|| std::io::Error::other("WAL record header offset overflow"))?;
1300 let header = bytes
1301 .get(offset..header_end)
1302 .ok_or_else(|| std::io::Error::other("WAL append has a torn record header"))?;
1303 let payload_len = usize::try_from(u32::from_le_bytes(
1304 header
1305 .get(..4)
1306 .ok_or_else(|| std::io::Error::other("WAL payload length is absent"))?
1307 .try_into()
1308 .map_err(|_| std::io::Error::other("WAL payload length width changed"))?,
1309 ))
1310 .map_err(|_| std::io::Error::other("WAL payload length does not fit usize"))?;
1311 let seq = u64::from_le_bytes(
1312 header
1313 .get(4..12)
1314 .ok_or_else(|| std::io::Error::other("WAL sequence is absent"))?
1315 .try_into()
1316 .map_err(|_| std::io::Error::other("WAL sequence width changed"))?,
1317 );
1318 let record_len = RECORD_HEADER_LEN
1319 .checked_add(payload_len)
1320 .and_then(|length| length.checked_add(RECORD_CHECKSUM_LEN))
1321 .ok_or_else(|| std::io::Error::other("WAL record length overflow"))?;
1322 offset = offset
1323 .checked_add(record_len)
1324 .ok_or_else(|| std::io::Error::other("WAL record end overflow"))?;
1325 if offset > bytes.len() {
1326 return Err(std::io::Error::other("WAL append has a torn record body"));
1327 }
1328 encoded_len = encoded_len
1329 .checked_add(
1330 u64::try_from(record_len)
1331 .map_err(|_| std::io::Error::other("WAL record length does not fit u64"))?,
1332 )
1333 .ok_or_else(|| std::io::Error::other("WAL record append length overflow"))?;
1334 first_seq.get_or_insert(seq);
1335 last_seq = Some(seq);
1336 }
1337 match (first_seq, last_seq) {
1338 (Some(first), Some(last)) => Ok(Some((encoded_len, first, last))),
1339 _ => Ok(None),
1340 }
1341}
1342
1343#[cfg(any(test, feature = "test-support"))]
1344fn storage_planned_segment_is_final(
1345 vfs: &dyn crate::vfs::Vfs,
1346 artifact: &Path,
1347 planned: Option<crate::segment::SegmentId>,
1348) -> bool {
1349 let Some(file_name) = planned.map(crate::segment::SegmentId::file_name) else {
1350 return false;
1351 };
1352 artifact.parent().is_some_and(|directory| {
1353 vfs.list(directory).is_ok_and(|paths| {
1354 paths.iter().any(|path| {
1355 path.file_name()
1356 .and_then(|name| name.to_str())
1357 .is_some_and(|name| name == file_name)
1358 })
1359 })
1360 })
1361}
1362
1363#[cfg(any(test, feature = "test-support"))]
1365#[derive(Clone, Copy, Debug, Eq, PartialEq)]
1366#[doc(hidden)]
1367pub enum HybridLegTestFault {
1368 Panic(crate::fusion::FusionLeg),
1370}
1371
1372#[allow(clippy::panic)]
1373fn maybe_trigger_hybrid_leg_panic(armed: bool, detail: &'static str) {
1374 #[cfg(any(test, feature = "test-support"))]
1375 if armed {
1376 std::panic::panic_any(detail);
1377 }
1378 #[cfg(not(any(test, feature = "test-support")))]
1379 let _ = (armed, detail);
1380}
1381
1382#[cfg(any(test, feature = "test-support"))]
1384#[derive(Clone, Debug, Eq, PartialEq)]
1385#[doc(hidden)]
1386pub struct HybridExecutionReceipt {
1387 pub vector_thread: std::thread::ThreadId,
1389 pub lexical_thread_name: Option<String>,
1391 pub vector_completed: bool,
1393 pub lexical_completed: bool,
1395}
1396
1397pub const DEFAULT_READER_DRAIN_TIMEOUT: Duration = Duration::from_millis(250);
1399
1400#[derive(Clone, Copy, Debug, Eq, PartialEq)]
1402pub enum StoreState {
1403 Open,
1405 Closing,
1407 Closed,
1409}
1410
1411#[derive(Clone, Debug, Eq, PartialEq)]
1413pub struct OpenOptions {
1414 access_mode: AccessMode,
1415 durability_mode: DurabilityMode,
1416 commit_tier: CommitTier,
1417 reader_drain_timeout: Duration,
1418 max_resident_bytes: u64,
1419 max_temp_bytes: u64,
1420 epoch: Option<crate::epoch::StoreEpoch>,
1421 tokenizer: Option<crate::fts::tokenizer::TokenizerConfig>,
1422 schema: Option<crate::meta::Schema>,
1423}
1424
1425#[derive(Clone, Copy, Debug, Eq, PartialEq)]
1427pub enum AccessMode {
1428 ReadWrite,
1430 ReadOnly,
1432}
1433
1434impl OpenOptions {
1435 #[must_use]
1437 pub const fn new() -> Self {
1438 Self {
1439 access_mode: AccessMode::ReadWrite,
1440 durability_mode: DurabilityMode::Derived,
1441 commit_tier: CommitTier::Ordered,
1442 reader_drain_timeout: DEFAULT_READER_DRAIN_TIMEOUT,
1443 max_resident_bytes: u64::MAX,
1444 max_temp_bytes: u64::MAX,
1445 epoch: None,
1446 tokenizer: None,
1447 schema: None,
1448 }
1449 }
1450
1451 #[must_use]
1453 pub const fn read_only() -> Self {
1454 Self {
1455 access_mode: AccessMode::ReadOnly,
1456 durability_mode: DurabilityMode::Derived,
1457 commit_tier: CommitTier::Ordered,
1458 reader_drain_timeout: DEFAULT_READER_DRAIN_TIMEOUT,
1459 max_resident_bytes: u64::MAX,
1460 max_temp_bytes: u64::MAX,
1461 epoch: None,
1462 tokenizer: None,
1463 schema: None,
1464 }
1465 }
1466
1467 #[must_use]
1469 pub const fn with_durability(mut self, mode: DurabilityMode, tier: CommitTier) -> Self {
1470 self.durability_mode = mode;
1471 self.commit_tier = tier;
1472 self
1473 }
1474
1475 #[must_use]
1477 pub const fn with_reader_drain_timeout(mut self, timeout: Duration) -> Self {
1478 self.reader_drain_timeout = timeout;
1479 self
1480 }
1481
1482 #[must_use]
1484 pub const fn with_max_resident_bytes(mut self, bytes: u64) -> Self {
1485 self.max_resident_bytes = bytes;
1486 self
1487 }
1488
1489 #[must_use]
1493 pub fn with_schema(mut self, schema: crate::meta::Schema) -> Self {
1494 self.schema = Some(schema);
1495 self
1496 }
1497
1498 #[must_use]
1500 pub const fn with_max_temp_bytes(mut self, bytes: u64) -> Self {
1501 self.max_temp_bytes = bytes;
1502 self
1503 }
1504
1505 #[must_use]
1511 pub fn with_epoch(mut self, epoch: crate::epoch::StoreEpoch) -> Self {
1512 self.epoch = Some(epoch);
1513 self
1514 }
1515
1516 #[must_use]
1521 pub fn with_tokenizer(mut self, tokenizer: crate::fts::tokenizer::TokenizerConfig) -> Self {
1522 self.tokenizer = Some(tokenizer);
1523 self
1524 }
1525}
1526
1527impl Default for OpenOptions {
1528 fn default() -> Self {
1529 Self::new()
1530 }
1531}
1532
1533#[derive(Clone, Copy, Debug, Eq, PartialEq)]
1535pub struct GraphSearchOptions {
1536 profile: crate::graph::search::GraphSearchProfile,
1537 ef: Option<usize>,
1538 seed: u64,
1539 #[cfg(any(test, feature = "test-support"))]
1540 cancel_after_hops: Option<usize>,
1541}
1542
1543impl GraphSearchOptions {
1544 #[must_use]
1546 pub const fn new(profile: crate::graph::search::GraphSearchProfile) -> Self {
1547 Self {
1548 profile,
1549 ef: None,
1550 seed: 0,
1551 #[cfg(any(test, feature = "test-support"))]
1552 cancel_after_hops: None,
1553 }
1554 }
1555
1556 #[must_use]
1558 pub const fn with_ef(mut self, ef: usize) -> Self {
1559 self.ef = Some(ef);
1560 self
1561 }
1562
1563 #[must_use]
1565 pub const fn with_seed(mut self, seed: u64) -> Self {
1566 self.seed = seed;
1567 self
1568 }
1569
1570 #[cfg(any(test, feature = "test-support"))]
1572 #[must_use]
1573 pub const fn with_cancel_after_hops(mut self, hops: usize) -> Self {
1574 self.cancel_after_hops = Some(hops);
1575 self
1576 }
1577
1578 pub(crate) const fn profile(self) -> crate::graph::search::GraphSearchProfile {
1579 self.profile
1580 }
1581
1582 pub(crate) const fn ef(self) -> Option<usize> {
1583 self.ef
1584 }
1585
1586 pub(crate) const fn seed(self) -> u64 {
1587 self.seed
1588 }
1589
1590 #[cfg(any(test, feature = "test-support"))]
1591 pub(crate) const fn cancel_after_hops(self) -> Option<usize> {
1592 self.cancel_after_hops
1593 }
1594}
1595
1596impl Default for GraphSearchOptions {
1597 fn default() -> Self {
1598 Self::new(crate::graph::search::GraphSearchProfile::SiftClass)
1599 }
1600}
1601
1602pub(crate) fn auto_graph_search_options(
1603 snapshot: &PublishedSnapshot,
1604) -> Result<GraphSearchOptions, QueryError> {
1605 snapshot
1606 .graph_profile()
1607 .map(|profile| GraphSearchOptions::new(profile.search_profile()))
1608 .map_err(crate::graph::search::GraphSearchError::Profile)
1609 .map_err(QueryError::Graph)
1610}
1611
1612#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
1614pub enum SearchTier {
1615 #[default]
1617 Auto,
1618 Exact,
1622 Scan,
1624 Graph(GraphSearchOptions),
1626}
1627
1628#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
1630pub struct SearchOptions {
1631 scan: crate::scan::ScanOptions,
1632 tier: Option<SearchTier>,
1633 scan_rescore: Option<ScanRescoreOptions>,
1634}
1635
1636#[derive(Clone, Copy, Debug, Eq, PartialEq)]
1637enum GraphBoundMode {
1638 Shared,
1639 #[cfg(test)]
1640 Independent,
1641}
1642
1643impl SearchOptions {
1644 #[must_use]
1646 pub const fn new(scan: crate::scan::ScanOptions) -> Self {
1647 Self {
1648 scan,
1649 tier: None,
1650 scan_rescore: None,
1651 }
1652 }
1653
1654 #[must_use]
1656 pub const fn with_tier(mut self, tier: SearchTier) -> Self {
1657 self.tier = Some(tier);
1658 self.scan_rescore = None;
1659 self
1660 }
1661
1662 #[must_use]
1666 pub const fn with_scan_rescore(mut self, options: ScanRescoreOptions) -> Self {
1667 self.tier = Some(SearchTier::Scan);
1668 self.scan_rescore = Some(options);
1669 self
1670 }
1671
1672 #[must_use]
1674 pub const fn scan_rescore(self) -> Option<ScanRescoreOptions> {
1675 self.scan_rescore
1676 }
1677
1678 #[must_use]
1680 pub const fn scan(self) -> crate::scan::ScanOptions {
1681 self.scan
1682 }
1683
1684 #[must_use]
1687 pub const fn tier(self) -> SearchTier {
1688 match self.tier {
1689 Some(tier) => tier,
1690 None => SearchTier::Auto,
1691 }
1692 }
1693
1694 pub(crate) const fn explicit_tier(self) -> Option<SearchTier> {
1695 self.tier
1696 }
1697}
1698
1699impl From<crate::scan::ScanOptions> for SearchOptions {
1700 fn from(scan: crate::scan::ScanOptions) -> Self {
1701 Self::new(scan)
1702 }
1703}
1704
1705#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
1707pub struct DocumentFields(u8);
1708
1709impl DocumentFields {
1710 pub const NONE: Self = Self(0);
1712 pub const VECTOR: Self = Self(1 << 0);
1714 pub const TEXT: Self = Self(1 << 1);
1716 pub const METADATA: Self = Self(1 << 2);
1718 pub const ATTRIBUTES: Self = Self(1 << 3);
1720 pub const ALL: Self =
1722 Self(Self::VECTOR.0 | Self::TEXT.0 | Self::METADATA.0 | Self::ATTRIBUTES.0);
1723
1724 #[must_use]
1726 pub const fn contains(self, fields: Self) -> bool {
1727 self.0 & fields.0 == fields.0
1728 }
1729}
1730
1731impl std::ops::BitOr for DocumentFields {
1732 type Output = Self;
1733
1734 fn bitor(self, rhs: Self) -> Self::Output {
1735 Self(self.0 | rhs.0)
1736 }
1737}
1738
1739#[derive(Clone, Debug, PartialEq)]
1741pub struct StoredDocument {
1742 pub doc_id: crate::ingest::DocId,
1744 pub revision: crate::ingest::Revision,
1746 pub timestamp: i64,
1748 pub vector: Option<Vec<f32>>,
1750 pub text: Option<String>,
1752 pub metadata: Option<Vec<u8>>,
1754 pub attributes: Option<Vec<(crate::meta::ColumnId, crate::meta::PredicateValue)>>,
1756}
1757
1758#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
1760pub enum ScanOrder {
1761 #[default]
1764 Storage,
1765 TimestampAscending,
1767 TimestampDescending,
1769}
1770
1771#[derive(Clone, Copy, Debug, Eq, PartialEq)]
1773pub enum DocumentScanSource {
1774 Sealed(crate::segment::SegmentId),
1776 Active,
1778}
1779
1780#[derive(Clone, Copy, Debug, Eq, PartialEq)]
1782pub struct DocumentScanCursor {
1783 pub generation: u64,
1785 pub source: DocumentScanSource,
1787 pub next_row: u32,
1789}
1790
1791#[derive(Clone, Debug)]
1793pub struct DocumentScanRequest<'a> {
1794 limit: usize,
1795 fields: DocumentFields,
1796 order: ScanOrder,
1797 cursor: Option<DocumentScanCursor>,
1798 timestamp_range: Option<(i64, i64)>,
1799 predicate: Option<&'a crate::meta::Predicate>,
1800 control: QueryControl,
1801}
1802
1803impl<'a> DocumentScanRequest<'a> {
1804 #[must_use]
1806 pub fn new(limit: usize, fields: DocumentFields, control: QueryControl) -> Self {
1807 Self {
1808 limit,
1809 fields,
1810 order: ScanOrder::Storage,
1811 cursor: None,
1812 timestamp_range: None,
1813 predicate: None,
1814 control,
1815 }
1816 }
1817
1818 #[must_use]
1820 pub const fn with_order(mut self, order: ScanOrder) -> Self {
1821 self.order = order;
1822 self
1823 }
1824
1825 #[must_use]
1827 pub const fn with_cursor(mut self, cursor: DocumentScanCursor) -> Self {
1828 self.cursor = Some(cursor);
1829 self
1830 }
1831
1832 #[must_use]
1834 pub const fn with_timestamp_range(mut self, start: i64, end: i64) -> Self {
1835 self.timestamp_range = Some((start, end));
1836 self
1837 }
1838
1839 #[must_use]
1841 pub const fn with_predicate(mut self, predicate: &'a crate::meta::Predicate) -> Self {
1842 self.predicate = Some(predicate);
1843 self
1844 }
1845}
1846
1847#[derive(Clone, Debug, PartialEq)]
1849pub struct DocumentScanPage {
1850 pub documents: Vec<StoredDocument>,
1852 pub generation: u64,
1854 pub continuation: Option<DocumentScanCursor>,
1856}
1857
1858#[derive(Clone, Copy, Debug, Eq, PartialEq)]
1860pub struct DocumentCount {
1861 pub count: u64,
1863 pub generation: u64,
1865}
1866
1867pub const MAX_DOCUMENT_SCAN_LIMIT: usize = 1 << 20;
1869
1870#[derive(Debug)]
1872pub enum StoreError {
1873 Io {
1875 path: PathBuf,
1877 source: std::io::Error,
1879 },
1880 NotDirectory {
1882 path: PathBuf,
1884 },
1885 StoreBusy {
1887 path: PathBuf,
1889 },
1890 Lock(StoreLockError),
1892 Durability(DurabilityPolicyError),
1894 Kernel(crate::kernels::KernelInitError),
1896 Manifest(crate::manifest::ManifestError),
1898 Tokenizer(crate::fts::tokenizer::TokenizerError),
1900 EpochMismatch(crate::epoch::EpochMismatch),
1902 EpochUndeclared,
1904 EpochUnstamped,
1906 SchemaMismatch {
1908 persisted: crate::meta::Schema,
1910 declared: crate::meta::Schema,
1912 },
1913 ScanStale {
1915 cursor_generation: u64,
1917 current_generation: u64,
1919 },
1920 InvalidScan {
1922 detail: String,
1924 },
1925 Segment(crate::segment::SegmentError),
1927 DuplicateLiveDocument {
1929 doc_id: crate::ingest::DocId,
1931 },
1932 GraphUnavailable {
1934 segment_id: crate::segment::SegmentId,
1936 },
1937 Wal(crate::wal::WalReadError),
1939 WalRecovery(crate::wal::WalRecoveryError),
1941 WalRecord {
1943 seq: crate::wal::LogSeq,
1945 source: crate::wal::VisibleRecordError,
1947 },
1948 WalMutation {
1950 seq: crate::wal::LogSeq,
1952 op: u16,
1954 source: crate::ingest::wal_payload::PayloadError,
1956 },
1957 WalRevisionOrder {
1959 seq: crate::wal::LogSeq,
1961 doc_id: crate::ingest::DocId,
1963 current: crate::ingest::Revision,
1965 attempted: crate::ingest::Revision,
1967 },
1968 UnsupportedWalMutation {
1970 seq: crate::wal::LogSeq,
1972 op: u16,
1974 },
1975 WalVector {
1977 seq: crate::wal::LogSeq,
1979 source: crate::quant::QuantError,
1981 },
1982 WalWrite(crate::wal::WalWriteError),
1984 WalRetire(crate::wal::WalRetireError),
1986 Statistics {
1988 component: &'static str,
1990 source: std::io::Error,
1992 },
1993 BudgetExceeded {
1995 needed: u64,
1997 budget: u64,
1999 component: &'static str,
2001 },
2002 AllocationFailed {
2004 needed: u64,
2006 component: &'static str,
2008 },
2009 DimensionMismatch {
2011 expected: usize,
2013 actual: usize,
2015 },
2016 ActiveRowOverflow,
2018 EmptyActiveSegment,
2020 SealCancelled,
2022 ReadOnly,
2024 SealedTombstoneRecoveryRequired,
2026 ForeignPreparedSegment,
2028 GenerationOverflow,
2030 ConcurrentActiveMutation {
2032 expected_generation: u64,
2034 actual_generation: u64,
2036 },
2037 PartitionBytesOverflow,
2039 PurgeRecovery {
2041 detail: String,
2043 },
2044 Closing,
2046 Closed,
2048 ReadCancelled,
2050 BackgroundStart {
2052 source: std::io::Error,
2054 },
2055 BackgroundHandshake,
2057 BackgroundThreadPanicked,
2059 QueryPoolStart {
2061 source: std::io::Error,
2063 },
2064 QueryPoolHandshake,
2066 QueryPoolThreadPanicked,
2068 QueryPoolCapacity {
2070 source: String,
2072 },
2073 Synchronization {
2075 component: &'static str,
2077 },
2078}
2079
2080impl std::fmt::Display for StoreError {
2081 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
2082 match self {
2083 Self::Io { path, source } => {
2084 write!(formatter, "store I/O {}: {source}", path.display())
2085 }
2086 Self::NotDirectory { path } => {
2087 write!(
2088 formatter,
2089 "store path is not a directory: {}",
2090 path.display()
2091 )
2092 }
2093 Self::StoreBusy { path } => {
2094 write!(formatter, "store already has a writer: {}", path.display())
2095 }
2096 Self::Lock(error) => error.fmt(formatter),
2097 Self::Durability(error) => error.fmt(formatter),
2098 Self::Kernel(error) => error.fmt(formatter),
2099 Self::Manifest(error) => error.fmt(formatter),
2100 Self::Tokenizer(error) => write!(formatter, "store tokenizer: {error}"),
2101 Self::EpochMismatch(error) => error.fmt(formatter),
2102 Self::EpochUndeclared => {
2103 formatter.write_str("store epoch is persisted but the caller declared none")
2104 }
2105 Self::EpochUnstamped => formatter
2106 .write_str("store manifest predates epoch identity and cannot adopt a declaration"),
2107 Self::SchemaMismatch {
2108 persisted,
2109 declared,
2110 } => write!(
2111 formatter,
2112 "declared store schema {declared:?} does not match persisted schema {persisted:?}"
2113 ),
2114 Self::ScanStale {
2115 cursor_generation,
2116 current_generation,
2117 } => write!(
2118 formatter,
2119 "scan cursor generation {cursor_generation} differs from current generation {current_generation}"
2120 ),
2121 Self::InvalidScan { detail } => write!(formatter, "invalid document scan: {detail}"),
2122 Self::Segment(error) => error.fmt(formatter),
2123 Self::DuplicateLiveDocument { doc_id } => {
2124 write!(
2125 formatter,
2126 "document {} has multiple live rows",
2127 doc_id.get()
2128 )
2129 }
2130 Self::GraphUnavailable { segment_id } => {
2131 write!(formatter, "sealed segment {segment_id} has no graph region")
2132 }
2133 Self::Wal(error) => error.fmt(formatter),
2134 Self::WalRecovery(error) => error.fmt(formatter),
2135 Self::WalRecord { seq, source } => {
2136 write!(
2137 formatter,
2138 "WAL sequence {} retained record: {source}",
2139 seq.get()
2140 )
2141 }
2142 Self::WalMutation { seq, op, source } => write!(
2143 formatter,
2144 "WAL sequence {} operation {op} payload: {source}",
2145 seq.get()
2146 ),
2147 Self::WalRevisionOrder {
2148 seq,
2149 doc_id,
2150 current,
2151 attempted,
2152 } => write!(
2153 formatter,
2154 "WAL sequence {} document {} revision {} does not follow {}",
2155 seq.get(),
2156 doc_id.get(),
2157 attempted.get(),
2158 current.get()
2159 ),
2160 Self::UnsupportedWalMutation { seq, op } => write!(
2161 formatter,
2162 "WAL sequence {} operation {op} has no active-segment recovery path",
2163 seq.get()
2164 ),
2165 Self::WalVector { seq, source } => {
2166 write!(formatter, "WAL sequence {} vector: {source}", seq.get())
2167 }
2168 Self::WalWrite(error) => error.fmt(formatter),
2169 Self::WalRetire(error) => error.fmt(formatter),
2170 Self::Statistics { component, source } => {
2171 write!(formatter, "store statistics {component}: {source}")
2172 }
2173 Self::BudgetExceeded {
2174 needed,
2175 budget,
2176 component,
2177 } => write!(
2178 formatter,
2179 "store {component} allocation needs {needed} bytes, budget is {budget} bytes"
2180 ),
2181 Self::AllocationFailed { needed, component } => write!(
2182 formatter,
2183 "store {component} allocator rejected {needed} bytes"
2184 ),
2185 Self::DimensionMismatch { expected, actual } => write!(
2186 formatter,
2187 "active vector dimension {actual} does not match {expected}"
2188 ),
2189 Self::ActiveRowOverflow => {
2190 formatter.write_str("active segment row or byte geometry overflow")
2191 }
2192 Self::EmptyActiveSegment => formatter.write_str("active segment is empty"),
2193 Self::SealCancelled => formatter.write_str("seal was cancelled before commit"),
2194 Self::ReadOnly => formatter.write_str("store handle is read-only"),
2195 Self::SealedTombstoneRecoveryRequired => formatter.write_str(
2196 "store requires writable sealed-tombstone recovery before reads are safe",
2197 ),
2198 Self::ForeignPreparedSegment => {
2199 formatter.write_str("prepared segment belongs to another store")
2200 }
2201 Self::GenerationOverflow => formatter.write_str("store snapshot generation overflow"),
2202 Self::ConcurrentActiveMutation {
2203 expected_generation,
2204 actual_generation,
2205 } => write!(
2206 formatter,
2207 "active generation changed from {expected_generation} to {actual_generation} while preparing a commit"
2208 ),
2209 Self::PartitionBytesOverflow => {
2210 formatter.write_str("partition reclaimed-byte count overflow")
2211 }
2212 Self::PurgeRecovery { detail } => {
2213 write!(formatter, "physical purge recovery failed: {detail}")
2214 }
2215 Self::Closing => formatter.write_str("store is closing"),
2216 Self::Closed => formatter.write_str("store is closed"),
2217 Self::ReadCancelled => formatter.write_str("store close cancelled the admitted read"),
2218 Self::BackgroundStart { source } => {
2219 write!(
2220 formatter,
2221 "store lifecycle thread could not start: {source}"
2222 )
2223 }
2224 Self::BackgroundHandshake => {
2225 formatter.write_str("store lifecycle thread startup handshake failed")
2226 }
2227 Self::BackgroundThreadPanicked => {
2228 formatter.write_str("store lifecycle thread panicked")
2229 }
2230 Self::QueryPoolStart { source } => {
2231 write!(formatter, "store query worker could not start: {source}")
2232 }
2233 Self::QueryPoolHandshake => {
2234 formatter.write_str("store query worker startup handshake failed")
2235 }
2236 Self::QueryPoolThreadPanicked => formatter.write_str("store query worker panicked"),
2237 Self::QueryPoolCapacity { source } => {
2238 write!(formatter, "store query worker capacity failed: {source}")
2239 }
2240 Self::Synchronization { component } => {
2241 write!(
2242 formatter,
2243 "store lifecycle synchronization poisoned: {component}"
2244 )
2245 }
2246 }
2247 }
2248}
2249
2250#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
2255pub enum StoreErrorKind {
2256 Io,
2258 InvalidArgument,
2260 StoreBusy,
2262 Unsupported,
2264 Corrupt,
2266 BudgetExceeded,
2268 OutOfMemory,
2270 DimensionMismatch,
2272 EpochMismatch,
2274 EpochUndeclared,
2276 EpochUnstamped,
2278 Internal,
2280 EmptyBatch,
2282 Cancelled,
2284 ReadOnly,
2286 Closing,
2288 Closed,
2290 Panic,
2292 Synchronization,
2294}
2295
2296impl StoreError {
2297 #[must_use]
2299 pub fn kind(&self) -> StoreErrorKind {
2300 match self {
2301 Self::Io { .. }
2302 | Self::Lock(_)
2303 | Self::Statistics { .. }
2304 | Self::BackgroundStart { .. }
2305 | Self::QueryPoolStart { .. }
2306 | Self::WalWrite(_)
2307 | Self::WalRetire(_) => StoreErrorKind::Io,
2308 Self::NotDirectory { .. }
2309 | Self::Tokenizer(_)
2310 | Self::SchemaMismatch { .. }
2311 | Self::ScanStale { .. }
2312 | Self::InvalidScan { .. } => StoreErrorKind::InvalidArgument,
2313 Self::StoreBusy { .. } => StoreErrorKind::StoreBusy,
2314 Self::Durability(_)
2315 | Self::Kernel(_)
2316 | Self::GraphUnavailable { .. }
2317 | Self::UnsupportedWalMutation { .. } => StoreErrorKind::Unsupported,
2318 Self::Manifest(_)
2319 | Self::Segment(_)
2320 | Self::DuplicateLiveDocument { .. }
2321 | Self::Wal(_)
2322 | Self::WalRecovery(_)
2323 | Self::WalRecord { .. }
2324 | Self::WalMutation { .. }
2325 | Self::WalRevisionOrder { .. }
2326 | Self::WalVector { .. }
2327 | Self::PurgeRecovery { .. } => StoreErrorKind::Corrupt,
2328 Self::BudgetExceeded { .. } => StoreErrorKind::BudgetExceeded,
2329 Self::AllocationFailed { .. } => StoreErrorKind::OutOfMemory,
2330 Self::DimensionMismatch { .. } => StoreErrorKind::DimensionMismatch,
2331 Self::EpochMismatch(_) => StoreErrorKind::EpochMismatch,
2332 Self::EpochUndeclared => StoreErrorKind::EpochUndeclared,
2333 Self::EpochUnstamped => StoreErrorKind::EpochUnstamped,
2334 Self::ActiveRowOverflow
2335 | Self::GenerationOverflow
2336 | Self::ConcurrentActiveMutation { .. }
2337 | Self::PartitionBytesOverflow
2338 | Self::ForeignPreparedSegment
2339 | Self::BackgroundHandshake
2340 | Self::QueryPoolHandshake
2341 | Self::QueryPoolCapacity { .. } => StoreErrorKind::Internal,
2342 Self::EmptyActiveSegment => StoreErrorKind::EmptyBatch,
2343 Self::SealCancelled | Self::ReadCancelled => StoreErrorKind::Cancelled,
2344 Self::ReadOnly | Self::SealedTombstoneRecoveryRequired => StoreErrorKind::ReadOnly,
2345 Self::Closing => StoreErrorKind::Closing,
2346 Self::Closed => StoreErrorKind::Closed,
2347 Self::BackgroundThreadPanicked | Self::QueryPoolThreadPanicked => StoreErrorKind::Panic,
2348 Self::Synchronization { .. } => StoreErrorKind::Synchronization,
2349 }
2350 }
2351}
2352
2353impl std::error::Error for StoreError {
2354 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
2355 match self {
2356 Self::Io { source, .. } => Some(source),
2357 Self::Lock(error) => Some(error),
2358 Self::Durability(error) => Some(error),
2359 Self::Kernel(error) => Some(error),
2360 Self::Manifest(error) => Some(error),
2361 Self::Tokenizer(error) => Some(error),
2362 Self::EpochMismatch(error) => Some(error),
2363 Self::Segment(error) => Some(error),
2364 Self::Wal(error) => Some(error),
2365 Self::WalRecovery(error) => Some(error),
2366 Self::WalRecord { source, .. } => Some(source),
2367 Self::WalMutation { source, .. } => Some(source),
2368 Self::WalVector { source, .. } => Some(source),
2369 Self::WalWrite(error) => Some(error),
2370 Self::WalRetire(error) => Some(error),
2371 Self::Statistics { source, .. } => Some(source),
2372 Self::BackgroundStart { source } => Some(source),
2373 Self::QueryPoolStart { source } => Some(source),
2374 Self::NotDirectory { .. }
2375 | Self::StoreBusy { .. }
2376 | Self::EpochUndeclared
2377 | Self::EpochUnstamped
2378 | Self::SchemaMismatch { .. }
2379 | Self::ScanStale { .. }
2380 | Self::InvalidScan { .. }
2381 | Self::GraphUnavailable { .. }
2382 | Self::DuplicateLiveDocument { .. }
2383 | Self::WalRevisionOrder { .. }
2384 | Self::UnsupportedWalMutation { .. }
2385 | Self::BudgetExceeded { .. }
2386 | Self::AllocationFailed { .. }
2387 | Self::DimensionMismatch { .. }
2388 | Self::ActiveRowOverflow
2389 | Self::EmptyActiveSegment
2390 | Self::SealCancelled
2391 | Self::ReadOnly
2392 | Self::SealedTombstoneRecoveryRequired
2393 | Self::ForeignPreparedSegment
2394 | Self::GenerationOverflow
2395 | Self::ConcurrentActiveMutation { .. }
2396 | Self::PartitionBytesOverflow
2397 | Self::PurgeRecovery { .. }
2398 | Self::Closing
2399 | Self::Closed
2400 | Self::ReadCancelled
2401 | Self::BackgroundHandshake
2402 | Self::BackgroundThreadPanicked
2403 | Self::QueryPoolHandshake
2404 | Self::QueryPoolThreadPanicked
2405 | Self::QueryPoolCapacity { .. }
2406 | Self::Synchronization { .. } => None,
2407 }
2408 }
2409}
2410
2411pub struct Store {
2413 pub(crate) directory: PathBuf,
2414 pub(crate) vfs: Arc<dyn crate::vfs::Vfs>,
2415 pub(crate) clock: Arc<dyn MonotonicClock>,
2416 pub(crate) state: Mutex<StoreState>,
2417 pub(crate) state_changed: Condvar,
2418 pub(crate) background: Mutex<Option<BackgroundThread>>,
2419 pub(crate) query_pool: Mutex<Option<Arc<pool::QueryPool>>>,
2420 pub(crate) lexical_worker: Mutex<Option<Arc<pool::LexicalWorker>>>,
2421 pub(crate) snapshot: RwLock<Option<Arc<PublishedSnapshot>>>,
2422 pub(crate) active: Mutex<Option<crate::ingest::ActiveState>>,
2425 pub(crate) wal_writer: Mutex<Option<crate::ingest::StoreWal>>,
2426 pub(crate) writer_lock: Mutex<Option<StoreLock>>,
2427 pub(crate) maintenance: Mutex<()>,
2428 pub(crate) health_state: Mutex<crate::diag::HealthState>,
2429 pub(crate) durability_policy: DurabilityPolicy,
2430 pub(crate) reader_drain_timeout: Duration,
2431 pub(crate) accounting: Arc<stats::Accounting>,
2432 pub(crate) active_queries: AtomicU64,
2433 #[cfg(any(test, feature = "test-support"))]
2434 text_materialization_work: materialize::TestMaterializationWork,
2435 pub(crate) epoch: Option<crate::epoch::StoreEpoch>,
2436 pub(crate) epoch_alias: crate::epoch::EpochAliasCell,
2437 pub(crate) tokenizer: crate::fts::tokenizer::Analyzer,
2438 pub(crate) schema: crate::meta::Schema,
2439 lexical_index_cache: LexicalIndexCache,
2443 #[cfg(any(test, feature = "test-support"))]
2444 pub(crate) ingest_retention_fault_controller:
2445 Option<crate::ingest::IngestRetentionFaultController>,
2446 #[cfg(any(test, feature = "test-support"))]
2447 hybrid_leg_fault: Mutex<Option<HybridLegTestFault>>,
2448 #[cfg(any(test, feature = "test-support"))]
2449 hybrid_execution_receipt: Mutex<Option<HybridExecutionReceipt>>,
2450 #[cfg(any(test, feature = "test-support"))]
2451 pub(crate) metadata_test_controller: Option<Arc<crate::planner::MetadataTestController>>,
2452 #[cfg(any(test, feature = "test-support"))]
2453 pub(crate) vector_fault_controller: Option<crate::scan::vector_fault::VectorFaultController>,
2454 #[cfg(any(test, feature = "test-support"))]
2455 pub(crate) kernel_fault_controller: Option<crate::kernels::vector_fault::KernelFaultController>,
2456 #[cfg(any(test, feature = "test-support"))]
2457 pub(crate) vector_seal_scheme: Option<crate::quant::QuantScheme>,
2458 #[cfg(test)]
2459 pub(crate) teardown_probe: Arc<close::TeardownProbe>,
2460}
2461
2462fn acquire_writer_lock(
2463 path: &Path,
2464 access_mode: AccessMode,
2465) -> Result<Option<StoreLock>, StoreError> {
2466 Ok(match access_mode {
2467 AccessMode::ReadWrite => Some(StoreLock::acquire(path).map_err(|error| match error {
2468 StoreLockError::Io { path, source }
2469 if source.kind() == std::io::ErrorKind::WouldBlock =>
2470 {
2471 StoreError::StoreBusy { path }
2472 }
2473 StoreLockError::Io { path, source } => {
2474 StoreError::Lock(StoreLockError::Io { path, source })
2475 }
2476 })?),
2477 AccessMode::ReadOnly => None,
2478 })
2479}
2480
2481fn resolve_open_schema(
2482 manifest_exists: bool,
2483 snapshot: &PublishedSnapshot,
2484 declared: Option<&crate::meta::Schema>,
2485) -> Result<crate::meta::Schema, StoreError> {
2486 if manifest_exists {
2487 let persisted = snapshot.schema().clone();
2488 if let Some(declared) = declared
2489 && declared != &persisted
2490 {
2491 return Err(StoreError::SchemaMismatch {
2492 persisted,
2493 declared: declared.clone(),
2494 });
2495 }
2496 Ok(snapshot.schema().clone())
2497 } else {
2498 Ok(declared
2499 .cloned()
2500 .unwrap_or_else(crate::meta::Schema::timestamp_only))
2501 }
2502}
2503
2504fn validate_epoch_identity(
2505 persisted: Option<crate::epoch::EpochIdentity>,
2506 declared: Option<crate::epoch::EpochIdentity>,
2507 manifest_exists: bool,
2508 access_mode: AccessMode,
2509) -> Result<(), StoreError> {
2510 match (persisted, declared) {
2511 (Some(expected), Some(declared)) if expected != declared => {
2512 Err(StoreError::EpochMismatch(crate::epoch::EpochMismatch {
2513 expected,
2514 declared,
2515 }))
2516 }
2517 (Some(_), None) => Err(StoreError::EpochUndeclared),
2518 (None, Some(_)) if manifest_exists || access_mode == AccessMode::ReadOnly => {
2519 Err(StoreError::EpochUnstamped)
2520 }
2521 (Some(_), Some(_)) | (None, Some(_)) | (None, None) => Ok(()),
2522 }
2523}
2524
2525fn validate_tokenizer_epoch(
2526 expected: Option<crate::epoch::EpochIdentity>,
2527 tokenizer: &crate::fts::tokenizer::Analyzer,
2528) -> Result<(), StoreError> {
2529 if let Some(expected) = expected
2530 && tokenizer.epoch() != expected.tokenizer
2531 {
2532 return Err(StoreError::EpochMismatch(crate::epoch::EpochMismatch {
2533 expected,
2534 declared: crate::epoch::EpochIdentity {
2535 embedding: expected.embedding,
2536 tokenizer: tokenizer.epoch(),
2537 },
2538 }));
2539 }
2540 Ok(())
2541}
2542
2543fn cleanup_open_orphans(
2544 store: &Store,
2545) -> Result<crate::manifest::io::OrphanCleanupReport, StoreError> {
2546 let reachable_segments = store
2547 .snapshot
2548 .read()
2549 .map_err(|_| StoreError::Synchronization {
2550 component: "published snapshot",
2551 })?
2552 .as_ref()
2553 .map(|snapshot| {
2554 snapshot
2555 .all_segments()
2556 .iter()
2557 .map(|segment| store.directory.join(segment.meta().id.file_name()))
2558 .collect::<HashSet<_>>()
2559 })
2560 .ok_or(StoreError::Closed)?;
2561 crate::manifest::io::cleanup_store_orphans(
2562 store.vfs.as_ref(),
2563 &store.directory,
2564 &reachable_segments,
2565 store.durability_policy,
2566 )
2567 .map_err(StoreError::Manifest)
2568}
2569
2570impl Store {
2571 pub fn open(path: impl AsRef<Path>, options: OpenOptions) -> Result<Self, StoreError> {
2573 Self::open_with_infrastructure(
2574 path,
2575 options,
2576 Arc::new(crate::vfs::StdVfs),
2577 Arc::new(SystemMonotonicClock),
2578 #[cfg(any(test, feature = "test-support"))]
2579 None,
2580 #[cfg(any(test, feature = "test-support"))]
2581 None,
2582 #[cfg(any(test, feature = "test-support"))]
2583 None,
2584 #[cfg(any(test, feature = "test-support"))]
2585 None,
2586 #[cfg(any(test, feature = "test-support"))]
2587 None,
2588 #[cfg(any(test, feature = "test-support"))]
2589 None,
2590 #[cfg(any(test, feature = "test-support"))]
2591 None,
2592 )
2593 }
2594
2595 #[allow(
2596 clippy::too_many_arguments,
2597 reason = "test-support controllers are explicit optional infrastructure dependencies"
2598 )]
2599 fn open_with_infrastructure(
2600 path: impl AsRef<Path>,
2601 options: OpenOptions,
2602 vfs: Arc<dyn crate::vfs::Vfs>,
2603 clock: Arc<dyn MonotonicClock>,
2604 #[cfg(any(test, feature = "test-support"))] hybrid_leg_fault: Option<HybridLegTestFault>,
2605 #[cfg(any(test, feature = "test-support"))] storage_fault_controller: Option<
2606 StorageFaultController,
2607 >,
2608 #[cfg(any(test, feature = "test-support"))] ingest_retention_fault_controller: Option<
2609 crate::ingest::IngestRetentionFaultController,
2610 >,
2611 #[cfg(any(test, feature = "test-support"))] metadata_test_controller: Option<
2612 Arc<crate::planner::MetadataTestController>,
2613 >,
2614 #[cfg(any(test, feature = "test-support"))] vector_fault_controller: Option<
2615 crate::scan::vector_fault::VectorFaultController,
2616 >,
2617 #[cfg(any(test, feature = "test-support"))] kernel_fault_controller: Option<
2618 crate::kernels::vector_fault::KernelFaultController,
2619 >,
2620 #[cfg(any(test, feature = "test-support"))] vector_seal_scheme: Option<
2621 crate::quant::QuantScheme,
2622 >,
2623 ) -> Result<Self, StoreError> {
2624 let path = path.as_ref();
2625 #[cfg(any(test, feature = "test-support"))]
2626 match kernel_fault_controller.as_ref() {
2627 Some(controller) => {
2628 controller
2629 .initialize_for_store()
2630 .map_err(StoreError::Kernel)?;
2631 }
2632 None => {
2633 crate::kernels::initialize().map_err(StoreError::Kernel)?;
2634 }
2635 }
2636 #[cfg(not(any(test, feature = "test-support")))]
2637 crate::kernels::initialize().map_err(StoreError::Kernel)?;
2638 #[cfg(any(test, feature = "test-support"))]
2639 let vfs = match storage_fault_controller.as_ref() {
2640 Some(controller) => Arc::new(StorageFaultVfs {
2641 inner: vfs,
2642 controller: controller.clone(),
2643 }) as Arc<dyn crate::vfs::Vfs>,
2644 None => vfs,
2645 };
2646 let durability_policy = DurabilityPolicy::new(options.durability_mode, options.commit_tier)
2647 .map_err(StoreError::Durability)?;
2648 let accounting = Arc::new(stats::Accounting::new(
2649 options.max_resident_bytes,
2650 options.max_temp_bytes,
2651 ));
2652 let is_directory = vfs
2653 .ensure_directory(path, options.access_mode == AccessMode::ReadWrite)
2654 .map_err(|source| StoreError::Io {
2655 path: path.to_path_buf(),
2656 source,
2657 })?;
2658 if !is_directory {
2659 return Err(StoreError::NotDirectory {
2660 path: path.to_path_buf(),
2661 });
2662 }
2663 let writer_lock = acquire_writer_lock(path, options.access_mode)?;
2664 let manifest_path = path.join(crate::manifest::io::MANIFEST_FILE);
2665 let manifest_exists = match vfs.open(&manifest_path) {
2666 Ok(_) => true,
2667 Err(source) if source.kind() == std::io::ErrorKind::NotFound => false,
2668 Err(source) => {
2669 return Err(StoreError::Io {
2670 path: manifest_path,
2671 source,
2672 });
2673 }
2674 };
2675 let snapshot = PublishedSnapshot::load_for_open_on_vfs(path, &accounting, vfs.as_ref())?;
2680 let schema = resolve_open_schema(manifest_exists, &snapshot, options.schema.as_ref())?;
2681 let persisted_epoch = snapshot.epoch_alias();
2682 let declared_epoch = options
2683 .epoch
2684 .as_ref()
2685 .map(crate::epoch::StoreEpoch::identity);
2686 validate_epoch_identity(
2687 persisted_epoch,
2688 declared_epoch,
2689 manifest_exists,
2690 options.access_mode,
2691 )?;
2692 let tokenizer = crate::fts::tokenizer::Analyzer::new(
2693 options
2694 .tokenizer
2695 .clone()
2696 .unwrap_or_else(crate::fts::tokenizer::TokenizerConfig::text_default),
2697 )
2698 .map_err(StoreError::Tokenizer)?;
2699 validate_tokenizer_epoch(persisted_epoch.or(declared_epoch), &tokenizer)?;
2700 let absorbed_through = snapshot.absorbed_through();
2701 let wal_path = path.join("wal.ze");
2702 let (active, recovered_wal, sealed_tombstones) = crate::ingest::ActiveState::recover(
2703 vfs.as_ref(),
2704 &wal_path,
2705 snapshot.generation(),
2706 absorbed_through,
2707 &accounting,
2708 &schema,
2709 &tokenizer,
2710 )?;
2711 if options.access_mode == AccessMode::ReadWrite && manifest_exists {
2712 match durability_policy.directory_sync() {
2716 SyncRequirement::Skip => {}
2717 SyncRequirement::Sync(kind) => {
2718 vfs.sync(path, kind).map_err(|source| StoreError::Io {
2719 path: path.to_path_buf(),
2720 source,
2721 })?;
2722 }
2723 }
2724 }
2725 if options.access_mode == AccessMode::ReadWrite
2726 && !manifest_exists
2727 && (options.epoch.is_some() || options.schema.is_some())
2728 {
2729 crate::manifest::io::commit_manifest(
2730 vfs.as_ref(),
2731 path,
2732 &crate::manifest::Manifest {
2733 generation: active.generation,
2734 log_seq: 0,
2735 segments: Vec::new(),
2736 epochs: options
2737 .epoch
2738 .as_ref()
2739 .cloned()
2740 .map(crate::manifest::EpochMeta::from)
2741 .into_iter()
2742 .collect(),
2743 epoch_alias: options
2744 .epoch
2745 .as_ref()
2746 .map(crate::epoch::StoreEpoch::identity),
2747 schema: schema.clone(),
2748 },
2749 durability_policy,
2750 )
2751 .map_err(StoreError::Manifest)?;
2752 }
2753 let wal_writer = match options.access_mode {
2754 AccessMode::ReadWrite => Some(match recovered_wal {
2755 Some(recovered) => crate::ingest::StoreWal::resume(
2756 vfs.as_ref(),
2757 &wal_path,
2758 recovered,
2759 durability_policy,
2760 absorbed_through,
2761 &accounting,
2762 )?,
2763 None => crate::ingest::StoreWal::create(
2764 Arc::clone(&vfs),
2765 path,
2766 &wal_path,
2767 durability_policy,
2768 &accounting,
2769 )?,
2770 }),
2771 AccessMode::ReadOnly => {
2772 drop(recovered_wal);
2773 None
2774 }
2775 };
2776 let background = match options.access_mode {
2777 AccessMode::ReadWrite => Some(BackgroundThread::start()?),
2778 AccessMode::ReadOnly => None,
2779 };
2780 #[cfg(test)]
2781 let (snapshot, background, teardown_probe) = {
2782 let mut snapshot = snapshot;
2783 let mut background = background;
2784 let teardown_probe = Arc::new(close::TeardownProbe::new());
2785 snapshot.set_teardown_probe(Arc::clone(&teardown_probe));
2786 if let Some(background) = background.as_mut() {
2787 background.set_teardown_probe(Arc::clone(&teardown_probe));
2788 }
2789 (snapshot, background, teardown_probe)
2790 };
2791 let store = Self {
2792 directory: path.to_path_buf(),
2793 vfs,
2794 clock,
2795 state: Mutex::new(StoreState::Open),
2796 state_changed: Condvar::new(),
2797 background: Mutex::new(background),
2798 query_pool: Mutex::new(None),
2799 lexical_worker: Mutex::new(None),
2800 snapshot: RwLock::new(None),
2801 active: Mutex::new(Some(active)),
2802 wal_writer: Mutex::new(wal_writer),
2803 writer_lock: Mutex::new(writer_lock),
2804 maintenance: Mutex::new(()),
2805 health_state: Mutex::new(crate::diag::HealthState::default()),
2806 durability_policy,
2807 reader_drain_timeout: options.reader_drain_timeout,
2808 accounting,
2809 active_queries: AtomicU64::new(0),
2810 #[cfg(any(test, feature = "test-support"))]
2811 text_materialization_work: materialize::TestMaterializationWork::default(),
2812 epoch_alias: crate::epoch::EpochAliasCell::new(persisted_epoch.or(declared_epoch)),
2813 epoch: options.epoch,
2814 tokenizer,
2815 schema,
2816 lexical_index_cache: LexicalIndexCache::new(),
2817 #[cfg(any(test, feature = "test-support"))]
2818 ingest_retention_fault_controller,
2819 #[cfg(any(test, feature = "test-support"))]
2820 hybrid_leg_fault: Mutex::new(hybrid_leg_fault),
2821 #[cfg(any(test, feature = "test-support"))]
2822 hybrid_execution_receipt: Mutex::new(None),
2823 #[cfg(any(test, feature = "test-support"))]
2824 metadata_test_controller,
2825 #[cfg(any(test, feature = "test-support"))]
2826 vector_fault_controller,
2827 #[cfg(any(test, feature = "test-support"))]
2828 kernel_fault_controller,
2829 #[cfg(any(test, feature = "test-support"))]
2830 vector_seal_scheme,
2831 #[cfg(test)]
2832 teardown_probe,
2833 };
2834 store.publish_snapshot(snapshot)?;
2835 store
2836 .recover_pending_physical_purge()
2837 .map_err(|error| StoreError::PurgeRecovery {
2838 detail: error.to_string(),
2839 })?;
2840 store.recover_sealed_tombstones(&sealed_tombstones)?;
2841 if options.access_mode == AccessMode::ReadWrite {
2842 let cleanup_report = cleanup_open_orphans(&store)?;
2843 #[cfg(any(test, feature = "test-support"))]
2844 if let Some(controller) = storage_fault_controller.as_ref() {
2845 controller.record_cleanup_report(cleanup_report);
2846 }
2847 #[cfg(not(any(test, feature = "test-support")))]
2848 let _ = cleanup_report;
2849 }
2850 #[cfg(any(test, feature = "test-support"))]
2851 if let Some(controller) = storage_fault_controller.as_ref() {
2852 let published = store
2853 .snapshot
2854 .read()
2855 .map_err(|_| StoreError::Synchronization {
2856 component: "published snapshot",
2857 })?;
2858 let snapshot = published.as_ref().ok_or(StoreError::Closed)?;
2859 for segment in snapshot.all_segments() {
2860 register_storage_segment_controller(
2861 &store.directory,
2862 segment.meta().id,
2863 controller,
2864 );
2865 }
2866 }
2867 Ok(store)
2868 }
2869
2870 #[cfg(any(test, feature = "test-support"))]
2872 #[doc(hidden)]
2873 pub fn open_with_test_dependencies(
2874 path: impl AsRef<Path>,
2875 options: OpenOptions,
2876 dependencies: StoreTestDependencies,
2877 ) -> Result<Self, StoreError> {
2878 let controller = dependencies.storage_fault_controller.clone();
2879 with_storage_open_controller(controller.as_ref(), || {
2880 Self::open_with_infrastructure(
2881 path,
2882 options,
2883 dependencies.vfs,
2884 dependencies.clock,
2885 dependencies.hybrid_leg_fault,
2886 dependencies.storage_fault_controller,
2887 dependencies.ingest_retention_fault_controller,
2888 dependencies.metadata_test_controller,
2889 dependencies.vector_fault_controller,
2890 dependencies.kernel_fault_controller,
2891 dependencies.vector_seal_scheme,
2892 )
2893 })
2894 }
2895
2896 #[cfg(any(test, feature = "test-support"))]
2898 #[doc(hidden)]
2899 pub fn take_hybrid_execution_receipt(&self) -> Option<HybridExecutionReceipt> {
2900 self.hybrid_execution_receipt
2901 .lock()
2902 .ok()
2903 .and_then(|mut receipt| receipt.take())
2904 }
2905
2906 pub(crate) fn epoch_registry(
2914 &self,
2915 prior: &[crate::manifest::EpochMeta],
2916 ) -> Vec<crate::manifest::EpochMeta> {
2917 if prior.is_empty() {
2918 self.epoch
2919 .as_ref()
2920 .map(crate::manifest::EpochMeta::from)
2921 .into_iter()
2922 .collect()
2923 } else {
2924 prior.to_vec()
2925 }
2926 }
2927
2928 #[must_use]
2931 pub fn epoch_identity(&self) -> Option<crate::epoch::EpochIdentity> {
2932 self.epoch_alias.load()
2933 }
2934
2935 #[must_use]
2937 pub const fn schema(&self) -> &crate::meta::Schema {
2938 &self.schema
2939 }
2940
2941 pub fn state(&self) -> Result<StoreState, StoreError> {
2943 self.state
2944 .lock()
2945 .map(|state| *state)
2946 .map_err(|_| StoreError::Synchronization { component: "state" })
2947 }
2948
2949 pub fn snapshot(&self) -> Result<SnapshotLease, StoreError> {
2951 let state = self
2952 .state
2953 .lock()
2954 .map_err(|_| StoreError::Synchronization { component: "state" })?;
2955 match *state {
2956 StoreState::Open => {}
2957 StoreState::Closing => return Err(StoreError::Closing),
2958 StoreState::Closed => return Err(StoreError::Closed),
2959 }
2960 let active = self
2961 .active
2962 .lock()
2963 .map_err(|_| StoreError::Synchronization {
2964 component: "active segment",
2965 })?;
2966 let generation = active.as_ref().ok_or(StoreError::Closed)?.generation;
2967 let snapshot = self
2968 .snapshot
2969 .read()
2970 .map_err(|_| StoreError::Synchronization {
2971 component: "published snapshot",
2972 })?
2973 .as_ref()
2974 .cloned()
2975 .ok_or(StoreError::Closed)?;
2976 drop(active);
2977 drop(state);
2978 Ok(SnapshotLease::new_at(snapshot, generation))
2979 }
2980
2981 #[cfg(any(test, feature = "test-support"))]
2987 #[doc(hidden)]
2988 pub fn test_metadata_row_values(
2989 &self,
2990 row_id: crate::ingest::GlobalRowId,
2991 ) -> Result<Vec<(crate::meta::ColumnId, Option<crate::meta::PredicateValue>)>, StoreError> {
2992 match row_id.source() {
2993 crate::ingest::RowSource::Active => {
2994 let row = usize::try_from(row_id.local_row())
2995 .map_err(|_| StoreError::ActiveRowOverflow)?;
2996 let active = self
2997 .active
2998 .lock()
2999 .map_err(|_| StoreError::Synchronization {
3000 component: "active segment",
3001 })?;
3002 let segment = &active.as_ref().ok_or(StoreError::Closed)?.segment;
3003 let timestamp = segment
3004 .timestamps()
3005 .get(row)
3006 .copied()
3007 .ok_or(StoreError::ActiveRowOverflow)?;
3008 let mut values = segment.column_values(row)?;
3009 values.push((
3010 crate::meta::TIMESTAMP_COLUMN,
3011 crate::meta::PredicateValue::I64(timestamp),
3012 ));
3013 self.schema
3014 .columns()
3015 .iter()
3016 .map(|definition| {
3017 let value = values
3018 .iter()
3019 .find(|(column, _)| *column == definition.id())
3020 .map(|(_, value)| value.clone());
3021 Ok((definition.id(), value))
3022 })
3023 .collect()
3024 }
3025 crate::ingest::RowSource::Sealed(segment_id) => {
3026 let snapshot = self.snapshot()?;
3027 let segment = snapshot
3028 .segments()
3029 .iter()
3030 .find(|segment| segment.meta().id == segment_id)
3031 .ok_or_else(|| {
3032 StoreError::Segment(crate::segment::SegmentError::Geometry(format!(
3033 "public metadata row names absent segment {segment_id}"
3034 )))
3035 })?;
3036 let columns = segment.columns().map_err(StoreError::Segment)?;
3037 if row_id.local_row() >= columns.row_count() {
3038 return Err(StoreError::Segment(crate::segment::SegmentError::Geometry(
3039 format!(
3040 "public metadata row {} is outside {} rows",
3041 row_id.local_row(),
3042 columns.row_count()
3043 ),
3044 )));
3045 }
3046 columns
3047 .schema()
3048 .columns()
3049 .iter()
3050 .map(|definition| {
3051 let column = columns.column(definition.id()).ok_or_else(|| {
3052 StoreError::Segment(crate::segment::SegmentError::Geometry(format!(
3053 "public metadata row is missing column {}",
3054 definition.id().get()
3055 )))
3056 })?;
3057 let value = match column {
3058 crate::meta::Column::U64(column) => column
3059 .get(row_id.local_row())
3060 .map(crate::meta::PredicateValue::U64),
3061 crate::meta::Column::I64(column) => column
3062 .get(row_id.local_row())
3063 .map(crate::meta::PredicateValue::I64),
3064 crate::meta::Column::F64(column) => column
3065 .get(row_id.local_row())
3066 .map(crate::meta::PredicateValue::F64),
3067 crate::meta::Column::Bool(column) => column
3068 .get(row_id.local_row())
3069 .map(crate::meta::PredicateValue::Bool),
3070 crate::meta::Column::DictionaryString(column) => column
3071 .get(row_id.local_row())
3072 .map(|value| crate::meta::PredicateValue::String(value.to_owned())),
3073 crate::meta::Column::RawString(column) => column
3074 .get(row_id.local_row())
3075 .map(|value| crate::meta::PredicateValue::String(value.to_owned())),
3076 };
3077 Ok((definition.id(), value))
3078 })
3079 .collect()
3080 }
3081 }
3082 }
3083
3084 #[cfg(any(test, feature = "test-support"))]
3086 #[doc(hidden)]
3087 pub fn test_metadata_evaluate_source(
3088 &self,
3089 source: crate::ingest::RowSource,
3090 predicate: &crate::meta::Predicate,
3091 ) -> Result<(u32, Vec<u32>, Vec<u32>), StoreError> {
3092 let evaluate = |columns: &crate::meta::ColumnStore,
3093 alive: &crate::meta::AliveSet|
3094 -> Result<(u32, Vec<u32>, Vec<u32>), StoreError> {
3095 let matches = crate::meta::evaluate(predicate, columns, alive).map_err(|error| {
3096 StoreError::Segment(crate::segment::SegmentError::Columns(format!(
3097 "test metadata source evaluation: {error}"
3098 )))
3099 })?;
3100 Ok((
3101 columns.row_count(),
3102 alive.iter_alive().collect(),
3103 matches.iter().collect(),
3104 ))
3105 };
3106 match source {
3107 crate::ingest::RowSource::Active => {
3108 let active = self
3109 .active
3110 .lock()
3111 .map_err(|_| StoreError::Synchronization {
3112 component: "active segment",
3113 })?;
3114 let segment = &active.as_ref().ok_or(StoreError::Closed)?.segment;
3115 let mut builder = crate::meta::ColumnStoreBuilder::new(self.schema.clone());
3116 for (row, timestamp) in segment.timestamps().iter().copied().enumerate() {
3117 let values = segment.column_values(row)?;
3118 builder
3119 .push_row(timestamp, &crate::ingest::column_inputs(&values))
3120 .map_err(|error| {
3121 StoreError::Segment(crate::segment::SegmentError::Columns(
3122 error.to_string(),
3123 ))
3124 })?;
3125 }
3126 let columns = builder.finish().map_err(|error| {
3127 StoreError::Segment(crate::segment::SegmentError::Columns(error.to_string()))
3128 })?;
3129 let alive = segment.alive()?;
3130 evaluate(&columns, &alive)
3131 }
3132 crate::ingest::RowSource::Sealed(segment_id) => {
3133 let snapshot = self.snapshot()?;
3134 let segment = snapshot
3135 .segments()
3136 .iter()
3137 .find(|segment| segment.meta().id == segment_id)
3138 .ok_or_else(|| {
3139 StoreError::Segment(crate::segment::SegmentError::Geometry(format!(
3140 "metadata source evaluation names absent segment {segment_id}"
3141 )))
3142 })?;
3143 let columns = segment.columns().map_err(StoreError::Segment)?;
3144 let alive = segment.alive().map_err(StoreError::Segment)?;
3145 evaluate(&columns, &alive)
3146 }
3147 }
3148 }
3149
3150 pub(crate) fn publish_snapshot(&self, snapshot: PublishedSnapshot) -> Result<(), StoreError> {
3151 let state = self
3152 .state
3153 .lock()
3154 .map_err(|_| StoreError::Synchronization { component: "state" })?;
3155 match *state {
3156 StoreState::Open => {}
3157 StoreState::Closing => return Err(StoreError::Closing),
3158 StoreState::Closed => return Err(StoreError::Closed),
3159 }
3160 let generation = snapshot.generation();
3161 let mut active = self
3162 .active
3163 .lock()
3164 .map_err(|_| StoreError::Synchronization {
3165 component: "active segment",
3166 })?;
3167 let active_generation = &mut active.as_mut().ok_or(StoreError::Closed)?.generation;
3168 *active_generation = (*active_generation).max(generation);
3169 let mut published = self
3170 .snapshot
3171 .write()
3172 .map_err(|_| StoreError::Synchronization {
3173 component: "published snapshot",
3174 })?;
3175 *published = Some(Arc::new(snapshot));
3176 drop(active);
3177 drop(state);
3178 Ok(())
3179 }
3180
3181 pub fn top_k_with_options(
3186 &self,
3187 request: crate::scan::ScanRequest<'_>,
3188 k: usize,
3189 options: crate::scan::ScanOptions,
3190 control: QueryControl,
3191 ) -> Result<crate::scan::ScanOutcome, QueryError> {
3192 let control = control.with_clock(Arc::clone(&self.clock));
3193 let state = self
3194 .state
3195 .lock()
3196 .map_err(|_| QueryError::Store(StoreError::Synchronization { component: "state" }))?;
3197 match *state {
3198 StoreState::Open => {}
3199 StoreState::Closing => return Err(QueryError::Store(StoreError::Closing)),
3200 StoreState::Closed => return Err(QueryError::Store(StoreError::Closed)),
3201 }
3202 let snapshot = self
3203 .snapshot
3204 .read()
3205 .map_err(|_| {
3206 QueryError::Store(StoreError::Synchronization {
3207 component: "published snapshot",
3208 })
3209 })?
3210 .as_ref()
3211 .cloned()
3212 .ok_or(QueryError::Store(StoreError::Closed))?;
3213 let mut pool_slot = self.query_pool.lock().map_err(|_| {
3214 QueryError::Store(StoreError::Synchronization {
3215 component: "query pool",
3216 })
3217 })?;
3218 if pool_slot.is_none() {
3219 let capacity = crate::scan::physical_thread_capacity().map_err(|error| {
3220 QueryError::Store(StoreError::QueryPoolCapacity {
3221 source: error.to_string(),
3222 })
3223 })?;
3224 let pool =
3225 pool::QueryPool::start(capacity, &self.accounting).map_err(QueryError::Store)?;
3226 *pool_slot = Some(Arc::new(pool));
3227 }
3228 let pool = pool_slot.as_ref().cloned().ok_or({
3229 QueryError::Store(StoreError::Synchronization {
3230 component: "query pool initialization",
3231 })
3232 })?;
3233 self.active_queries.fetch_add(1, Ordering::Relaxed);
3234 let active = ActiveQuery {
3235 count: &self.active_queries,
3236 };
3237 drop(pool_slot);
3238 drop(state);
3239
3240 let result = pool
3241 .execute(request, k, options, control, SnapshotLease::new(snapshot))
3242 .map(|mut outcome| {
3243 outcome.candidates.truncate(k);
3244 outcome
3245 });
3246 drop(active);
3247 result
3248 }
3249
3250 #[cfg(any(test, feature = "test-support"))]
3255 pub fn panic_query_worker_for_test(&self) -> Result<(), StoreError> {
3256 let pool = self
3257 .query_pool
3258 .lock()
3259 .map_err(|_| StoreError::Synchronization {
3260 component: "query pool",
3261 })?
3262 .as_ref()
3263 .cloned()
3264 .ok_or(StoreError::Synchronization {
3265 component: "query pool initialization",
3266 })?;
3267 pool.panic_one_and_join()
3268 }
3269
3270 pub fn search(
3279 &self,
3280 request: crate::ingest::SearchRequest<'_>,
3281 k: usize,
3282 options: impl Into<SearchOptions>,
3283 control: QueryControl,
3284 ) -> Result<crate::ingest::SearchOutcome, QueryError> {
3285 self.search_with_graph_bound_mode(
3286 request,
3287 k,
3288 options.into(),
3289 control,
3290 GraphBoundMode::Shared,
3291 )
3292 }
3293
3294 fn admit_lexical_query(&self) -> Result<AdmittedLexicalQuery<'_>, QueryError> {
3295 #[cfg(any(test, feature = "test-support"))]
3296 self.text_materialization_work
3297 .admissions
3298 .fetch_add(1, Ordering::Relaxed);
3299 #[cfg(all(feature = "query-timing", any(test, feature = "test-support")))]
3300 let lock_started = std::time::Instant::now();
3301 let state = self
3302 .state
3303 .lock()
3304 .map_err(|_| QueryError::Store(StoreError::Synchronization { component: "state" }))?;
3305 #[cfg(all(feature = "query-timing", any(test, feature = "test-support")))]
3306 self.text_materialization_work
3307 .observe_lexical_lock(lock_started);
3308 match *state {
3309 StoreState::Open => {}
3310 StoreState::Closing => return Err(QueryError::Store(StoreError::Closing)),
3311 StoreState::Closed => return Err(QueryError::Store(StoreError::Closed)),
3312 }
3313 #[cfg(all(feature = "query-timing", any(test, feature = "test-support")))]
3314 let lock_started = std::time::Instant::now();
3315 let active_guard = self.active.lock().map_err(|_| {
3316 QueryError::Store(StoreError::Synchronization {
3317 component: "active segment",
3318 })
3319 })?;
3320 #[cfg(all(feature = "query-timing", any(test, feature = "test-support")))]
3321 self.text_materialization_work
3322 .observe_lexical_lock(lock_started);
3323 let active_state = active_guard
3324 .as_ref()
3325 .ok_or(QueryError::Store(StoreError::Closed))?;
3326 let generation = active_state.generation;
3327 let active = Arc::clone(&active_state.segment);
3328 #[cfg(all(feature = "query-timing", any(test, feature = "test-support")))]
3329 let lock_started = std::time::Instant::now();
3330 let snapshot = self
3331 .snapshot
3332 .read()
3333 .map_err(|_| {
3334 QueryError::Store(StoreError::Synchronization {
3335 component: "published snapshot",
3336 })
3337 })
3338 .inspect(|_| {
3339 #[cfg(all(feature = "query-timing", any(test, feature = "test-support")))]
3340 self.text_materialization_work
3341 .observe_lexical_lock(lock_started);
3342 })?
3343 .as_ref()
3344 .cloned()
3345 .ok_or(QueryError::Store(StoreError::Closed))?;
3346 self.active_queries.fetch_add(1, Ordering::Relaxed);
3347 let active_query = ActiveQuery {
3348 count: &self.active_queries,
3349 };
3350 drop(active_guard);
3351 drop(state);
3352 Ok(AdmittedLexicalQuery {
3353 generation,
3354 active,
3355 snapshot: snapshot::ReadSnapshot::new(snapshot),
3356 active_query,
3357 })
3358 }
3359
3360 fn admit_document_read(&self) -> Result<AdmittedDocumentRead<'_>, StoreError> {
3361 let state = self
3362 .state
3363 .lock()
3364 .map_err(|_| StoreError::Synchronization { component: "state" })?;
3365 match *state {
3366 StoreState::Open => {}
3367 StoreState::Closing => return Err(StoreError::Closing),
3368 StoreState::Closed => return Err(StoreError::Closed),
3369 }
3370 let active_guard = self
3371 .active
3372 .lock()
3373 .map_err(|_| StoreError::Synchronization {
3374 component: "active segment",
3375 })?;
3376 let active_state = active_guard.as_ref().ok_or(StoreError::Closed)?;
3377 let generation = active_state.generation;
3378 let active = Arc::clone(&active_state.segment);
3379 let snapshot = self
3380 .snapshot
3381 .read()
3382 .map_err(|_| StoreError::Synchronization {
3383 component: "published snapshot",
3384 })?
3385 .as_ref()
3386 .cloned()
3387 .ok_or(StoreError::Closed)?;
3388 self.active_queries.fetch_add(1, Ordering::Relaxed);
3389 let active_query = ActiveQuery {
3390 count: &self.active_queries,
3391 };
3392 drop(active_guard);
3393 drop(state);
3394 Ok(AdmittedDocumentRead {
3395 generation,
3396 active,
3397 snapshot: snapshot::ReadSnapshot::new(snapshot),
3398 active_query,
3399 })
3400 }
3401
3402 pub fn search_lexical(
3404 &self,
3405 query: &crate::fts::search::TermQuery,
3406 k: usize,
3407 control: QueryControl,
3408 ) -> Result<crate::ingest::StoreLexicalSearchOutcome, crate::ingest::StoreLexicalError> {
3409 self.search_lexical_then(query, k, control, |outcome, _, _, _, _, _| Ok(outcome))
3410 }
3411
3412 fn search_lexical_then<R>(
3413 &self,
3414 query: &crate::fts::search::TermQuery,
3415 k: usize,
3416 control: QueryControl,
3417 finish: impl FnOnce(
3418 crate::ingest::StoreLexicalSearchOutcome,
3419 &PublishedSnapshot,
3420 &crate::ingest::ActiveSegment,
3421 &LexicalAssembly,
3422 &[crate::fts::search::ScoredDoc],
3423 &QueryCancellation<'_>,
3424 ) -> Result<R, crate::ingest::StoreLexicalError>,
3425 ) -> Result<R, crate::ingest::StoreLexicalError> {
3426 let control = control.with_clock(Arc::clone(&self.clock));
3427 let started = self.clock.now();
3428 let admission_started = timing_start(self.clock.as_ref());
3429 let AdmittedLexicalQuery {
3430 generation,
3431 active,
3432 snapshot,
3433 active_query,
3434 } = self.admit_lexical_query()?;
3435 let admission_time = timing_elapsed(self.clock.as_ref(), admission_started);
3436 let lexical_started = timing_start(self.clock.as_ref());
3437
3438 let lease = SnapshotLease::new_at(Arc::clone(&snapshot), generation);
3439 let cancellation = QueryCancellation::new(&control, &lease);
3440 cancellation.check_graph().map_err(QueryError::Scan)?;
3441 let assembly = assemble_lexical_index(
3442 LexicalInputs {
3443 generation,
3444 cache: &self.lexical_index_cache,
3445 snapshot: &snapshot,
3446 active: &active,
3447 accounting: &self.accounting,
3448 },
3449 true,
3450 Some(&cancellation),
3451 )
3452 .map_err(map_store_lexical_assembly_error)?;
3453 let mut work = crate::fts::control::WorkCheck::new(|| {
3454 cancellation
3455 .check_graph()
3456 .map_err(QueryError::Scan)
3457 .map_err(crate::ingest::StoreLexicalError::from)
3458 });
3459 work.check_now()?;
3460 let mut allow_lists = Vec::with_capacity(assembly.alive_sets.len());
3461 for alive in &assembly.alive_sets {
3462 work.step()?;
3463 allow_lists.push(alive.alive_bitmap());
3464 }
3465 let lexical = crate::planner::search_lexical_filtered_refs_controlled(
3466 &assembly.index,
3467 query,
3468 k,
3469 crate::fts::bm25::Bm25Params::beir(),
3470 &allow_lists,
3471 None,
3472 || {
3473 cancellation
3474 .check_graph()
3475 .map_err(QueryError::Scan)
3476 .map_err(crate::ingest::StoreLexicalError::from)
3477 },
3478 )?;
3479 cancellation.check_graph().map_err(QueryError::Scan)?;
3480 let mut candidates = Vec::with_capacity(lexical.result.hits.len());
3481 for hit in &lexical.result.hits {
3482 work.step()?;
3483 let document =
3484 structured_lexical_document(&snapshot, &active, &assembly.sources, hit.doc, true)
3485 .map_err(map_store_lexical_document_error)?
3486 .ok_or(QueryError::Store(StoreError::ActiveRowOverflow))?;
3487 candidates.push(crate::ingest::LexicalCandidate {
3488 document,
3489 score: hit.score,
3490 });
3491 }
3492 let mut diagnostics =
3493 crate::diag::QueryDiagnostics::lexical(crate::diag::LexicalDiagnostics {
3494 snapshot_generation: generation,
3495 indexed_through_seq: active
3496 .indexed_through_seq()
3497 .max(crate::wal::LogSeq::new(snapshot.absorbed_through())),
3498 requested_k: k,
3499 returned: candidates.len(),
3500 counters: lexical.result.counters,
3501 elapsed: self.clock.now().saturating_duration_since(started),
3502 });
3503 diagnostics.counters.lexical_cache_hits = usize::from(assembly.cache_hit);
3504 diagnostics.counters.lexical_cache_builds = usize::from(!assembly.cache_hit);
3505 diagnostics.timings =
3506 crate::diag::QUERY_TIMING_ENABLED.then(|| crate::diag::QueryTimings {
3507 admission: admission_time,
3508 lexical: timing_elapsed(self.clock.as_ref(), lexical_started),
3509 ..Default::default()
3510 });
3511 let outcome = crate::ingest::StoreLexicalSearchOutcome {
3512 candidates,
3513 generation,
3514 diagnostics,
3515 };
3516 work.check_now()?;
3517 let result = finish(
3518 outcome,
3519 &snapshot,
3520 &active,
3521 &assembly,
3522 &lexical.result.hits,
3523 &cancellation,
3524 );
3525 drop(active_query);
3526 result
3527 }
3528
3529 pub fn stored_text(
3535 &self,
3536 version: crate::ingest::DocumentVersion,
3537 ) -> Result<Option<String>, crate::ingest::StoreLexicalError> {
3538 let AdmittedLexicalQuery {
3539 active,
3540 snapshot,
3541 active_query,
3542 ..
3543 } = self.admit_lexical_query()?;
3544 #[cfg(any(test, feature = "test-support"))]
3545 self.text_materialization_work
3546 .reverse_version_lookups
3547 .fetch_add(1, Ordering::Relaxed);
3548 if let Some((row, active_version, _)) = active.existing(version.doc_id())
3549 && active_version == version
3550 {
3551 let text = active
3552 .text(row)
3553 .map_err(crate::ingest::StoreLexicalError::from)?;
3554 drop(active_query);
3555 return Ok(text.map(str::to_owned));
3556 }
3557 for segment in snapshot.segments() {
3558 #[cfg(any(test, feature = "test-support"))]
3559 self.text_materialization_work
3560 .reverse_version_lookups
3561 .fetch_add(1, Ordering::Relaxed);
3562 let Some(row) = segment
3563 .query_row_for_document_version(version)
3564 .map_err(crate::ingest::StoreLexicalError::from)?
3565 else {
3566 continue;
3567 };
3568 let text = segment
3569 .query_stored_text()
3570 .map_err(StoreError::Segment)
3571 .map_err(crate::ingest::StoreLexicalError::from)?
3572 .and_then(|rows| rows.row(row).flatten())
3573 .map(str::to_owned);
3574 drop(active_query);
3575 return Ok(text);
3576 }
3577 drop(active_query);
3578 Ok(None)
3579 }
3580
3581 pub fn get_documents(
3586 &self,
3587 ids: &[crate::ingest::DocId],
3588 fields: DocumentFields,
3589 ) -> Result<Vec<Option<StoredDocument>>, StoreError> {
3590 self.get_documents_with_generation(ids, fields)
3591 .map(|(_, documents)| documents)
3592 }
3593
3594 #[doc(hidden)]
3598 pub fn get_documents_with_generation(
3599 &self,
3600 ids: &[crate::ingest::DocId],
3601 fields: DocumentFields,
3602 ) -> Result<(u64, Vec<Option<StoredDocument>>), StoreError> {
3603 let AdmittedDocumentRead {
3604 generation,
3605 active,
3606 snapshot,
3607 active_query,
3608 } = self.admit_document_read()?;
3609 let mut documents = Vec::new();
3610 documents
3611 .try_reserve_exact(ids.len())
3612 .map_err(|_| StoreError::AllocationFailed {
3613 needed: allocation_bytes::<Option<StoredDocument>>(ids.len()),
3614 component: "get documents",
3615 })?;
3616 for id in ids {
3617 let mut document = active
3618 .existing(*id)
3619 .filter(|(row, _, _)| !active.is_tombstoned(*row))
3620 .map(|(row, _, _)| materialize_active_document(&active, row, fields))
3621 .transpose()?;
3622 for segment in snapshot.segments() {
3623 for row in segment.query_rows_for_doc_id(*id)? {
3624 if document.is_some() {
3625 return Err(StoreError::DuplicateLiveDocument { doc_id: *id });
3626 }
3627 document = Some(materialize_sealed_document(segment, row, fields)?);
3628 }
3629 }
3630 documents.push(document);
3631 }
3632 drop(active_query);
3633 Ok((generation, documents))
3634 }
3635
3636 pub fn scan_documents(
3641 &self,
3642 request: DocumentScanRequest<'_>,
3643 ) -> Result<DocumentScanPage, QueryError> {
3644 let DocumentScanRequest {
3645 limit,
3646 fields,
3647 order,
3648 cursor,
3649 timestamp_range,
3650 predicate,
3651 control,
3652 } = request;
3653 if limit == 0 || limit > MAX_DOCUMENT_SCAN_LIMIT {
3654 return Err(QueryError::Store(StoreError::InvalidScan {
3655 detail: format!("limit must be in 1..={MAX_DOCUMENT_SCAN_LIMIT}, received {limit}"),
3656 }));
3657 }
3658 let predicate = document_scan_predicate(predicate, timestamp_range)?;
3659 if let Some(predicate) = &predicate {
3660 crate::planner::validate_predicate(predicate, &self.schema).map_err(|error| {
3661 QueryError::Store(StoreError::InvalidScan {
3662 detail: error.to_string(),
3663 })
3664 })?;
3665 }
3666 let control = control.with_clock(Arc::clone(&self.clock));
3667 control.checkpoint()?;
3668 let AdmittedDocumentRead {
3669 generation,
3670 active,
3671 snapshot,
3672 active_query,
3673 } = self.admit_document_read().map_err(QueryError::Store)?;
3674 if let Some(cursor) = cursor.filter(|cursor| cursor.generation != generation) {
3675 return Err(QueryError::Store(StoreError::ScanStale {
3676 cursor_generation: cursor.generation,
3677 current_generation: generation,
3678 }));
3679 }
3680 let lease = SnapshotLease::new_at(Arc::clone(&snapshot), generation);
3681 let mut work = 0_usize;
3682 let rows = match order {
3683 ScanOrder::Storage => collect_storage_scan_rows(
3684 &snapshot,
3685 &active,
3686 &self.schema,
3687 predicate.as_ref(),
3688 cursor,
3689 limit,
3690 &control,
3691 &lease,
3692 &mut work,
3693 )?,
3694 ScanOrder::TimestampAscending | ScanOrder::TimestampDescending => {
3695 collect_timestamp_scan_rows(
3696 &snapshot,
3697 &active,
3698 &self.schema,
3699 predicate.as_ref(),
3700 cursor,
3701 limit,
3702 order,
3703 &control,
3704 &lease,
3705 &mut work,
3706 )?
3707 }
3708 };
3709 let continuation = rows.get(limit).map(|row| DocumentScanCursor {
3710 generation,
3711 source: row.source,
3712 next_row: row.row,
3713 });
3714 let result_len = rows.len().min(limit);
3715 let mut documents = Vec::new();
3716 documents.try_reserve_exact(result_len).map_err(|_| {
3717 QueryError::Store(StoreError::AllocationFailed {
3718 needed: allocation_bytes::<StoredDocument>(result_len),
3719 component: "scan documents",
3720 })
3721 })?;
3722 for row in rows.into_iter().take(result_len) {
3723 scan_control_checkpoint(&control, &lease, &mut work)?;
3724 let document = match row.source {
3725 DocumentScanSource::Active => {
3726 materialize_active_document(&active, row.row as usize, fields)
3727 .map_err(QueryError::Store)?
3728 }
3729 DocumentScanSource::Sealed(segment_id) => {
3730 let segment = snapshot
3731 .segments()
3732 .iter()
3733 .find(|segment| segment.meta().id == segment_id)
3734 .ok_or_else(|| {
3735 QueryError::Store(StoreError::InvalidScan {
3736 detail: format!("scan row names absent segment {segment_id}"),
3737 })
3738 })?;
3739 materialize_sealed_document(segment, row.row as usize, fields)
3740 .map_err(QueryError::Store)?
3741 }
3742 };
3743 documents.push(document);
3744 }
3745 drop(active_query);
3746 Ok(DocumentScanPage {
3747 documents,
3748 generation,
3749 continuation,
3750 })
3751 }
3752
3753 pub fn count_documents(
3755 &self,
3756 predicate: Option<&crate::meta::Predicate>,
3757 timestamp_range: Option<(i64, i64)>,
3758 ) -> Result<DocumentCount, QueryError> {
3759 let predicate = document_scan_predicate(predicate, timestamp_range)?;
3760 if let Some(predicate) = &predicate {
3761 crate::planner::validate_predicate(predicate, &self.schema).map_err(|error| {
3762 QueryError::Store(StoreError::InvalidScan {
3763 detail: error.to_string(),
3764 })
3765 })?;
3766 }
3767 let AdmittedDocumentRead {
3768 generation,
3769 active,
3770 snapshot,
3771 active_query,
3772 } = self.admit_document_read().map_err(QueryError::Store)?;
3773 let mut count = 0_u64;
3774 for segment in snapshot.segments() {
3775 count = count
3776 .checked_add(sealed_scan_rows(segment, predicate.as_ref())?.cardinality())
3777 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
3778 }
3779 count = count
3780 .checked_add(active_scan_rows(&active, &self.schema, predicate.as_ref())?.cardinality())
3781 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
3782 drop(active_query);
3783 Ok(DocumentCount { count, generation })
3784 }
3785
3786 pub fn search_lexical_structured(
3789 &self,
3790 query: &crate::fts::query::LexicalQuery,
3791 k: usize,
3792 snippet_bytes: usize,
3793 control: QueryControl,
3794 ) -> Result<crate::ingest::StoreStructuredLexicalSearchOutcome, crate::ingest::StoreLexicalError>
3795 {
3796 let control = control.with_clock(Arc::clone(&self.clock));
3797 let started = self.clock.now();
3798 let admission_started = timing_start(self.clock.as_ref());
3799 let AdmittedLexicalQuery {
3800 generation,
3801 active,
3802 snapshot,
3803 active_query,
3804 } = self.admit_lexical_query()?;
3805 let admission_time = timing_elapsed(self.clock.as_ref(), admission_started);
3806 let lexical_started = timing_start(self.clock.as_ref());
3807
3808 let lease = SnapshotLease::new_at(Arc::clone(&snapshot), generation);
3809 let cancellation = QueryCancellation::new(&control, &lease);
3810 cancellation.check_graph().map_err(QueryError::Scan)?;
3811 let assembly = assemble_lexical_index(
3812 LexicalInputs {
3813 generation,
3814 cache: &self.lexical_index_cache,
3815 snapshot: &snapshot,
3816 active: &active,
3817 accounting: &self.accounting,
3818 },
3819 true,
3820 Some(&cancellation),
3821 )
3822 .map_err(map_store_lexical_assembly_error)?;
3823 let index = &assembly.index;
3824 let sources = &assembly.sources;
3825 let mut expansions = assembly
3826 .expand(query, &self.accounting, &cancellation)
3827 .map_err(|error| match error {
3828 expansion::ExpansionError::Query(error) => {
3829 crate::ingest::StoreLexicalError::Query(error)
3830 }
3831 expansion::ExpansionError::Shape(error) => {
3832 crate::ingest::StoreLexicalError::Structured(error)
3833 }
3834 })?;
3835 let allow_lists = assembly
3836 .alive_sets
3837 .iter()
3838 .map(|alive| alive.alive_bitmap())
3839 .collect::<Vec<_>>();
3840 let mut prepared_memory =
3841 stats::AccountedCounter::new(&self.accounting, stats::AllocationComponent::Temporary)
3842 .map_err(QueryError::Store)?;
3843 let mut work = crate::fts::control::WorkCheck::new(|| cancellation.check_graph());
3844 let fields = query
3845 .fields_controlled(&mut work)
3846 .map_err(QueryError::Scan)?;
3847 prepared_memory
3848 .set(
3849 crate::fts::query::PreparedWeightedQuery::allocation_bytes_controlled(
3850 &expansions,
3851 &fields,
3852 &mut work,
3853 )
3854 .map_err(QueryError::Scan)?
3855 .ok_or(QueryError::Store(StoreError::ActiveRowOverflow))?,
3856 )
3857 .map_err(QueryError::Store)?;
3858 let mut prepared = if index.segments().is_empty() || expansions.is_empty() {
3859 None
3860 } else {
3861 Some(
3862 crate::fts::query::PreparedWeightedQuery::new_controlled(
3863 index,
3864 std::mem::take(&mut expansions),
3865 fields,
3866 || cancellation.check_graph(),
3867 )
3868 .map_err(|error| match error {
3869 crate::fts::search::ControlledSearchError::Index(error) => {
3870 crate::ingest::StoreLexicalError::from(
3871 crate::planner::LexicalFilterError::from(error),
3872 )
3873 }
3874 crate::fts::search::ControlledSearchError::Control(error) => {
3875 crate::ingest::StoreLexicalError::from(QueryError::Scan(error))
3876 }
3877 })?,
3878 )
3879 };
3880 let mut scratch_memory =
3881 stats::AccountedCounter::new(&self.accounting, stats::AllocationComponent::Temporary)
3882 .map_err(QueryError::Store)?;
3883 let mut phrase = prepared_lexical::PhraseEligibility::new(
3884 index,
3885 Some(query),
3886 &self.accounting,
3887 &cancellation,
3888 )?;
3889 let result = match &prepared {
3890 None => crate::fts::prune::weighted::WeightedResult::default(),
3891 Some(prepared) => crate::fts::prune::weighted::search(
3892 index,
3893 prepared,
3894 k,
3895 &allow_lists,
3896 || {
3897 cancellation
3898 .check_graph()
3899 .map_err(QueryError::Scan)
3900 .map_err(crate::ingest::StoreLexicalError::from)
3901 },
3902 |doc| {
3903 phrase
3904 .matches(doc, &cancellation)
3905 .map_err(crate::ingest::StoreLexicalError::from)
3906 },
3907 |bytes| {
3908 prepared_lexical::reserve_weighted_scratch(&mut scratch_memory, bytes)
3909 .map_err(crate::ingest::StoreLexicalError::from)
3910 },
3911 )
3912 .map_err(|error| match error {
3913 crate::fts::search::ControlledSearchError::Index(error) => {
3914 crate::ingest::StoreLexicalError::from(
3915 crate::planner::LexicalFilterError::from(error),
3916 )
3917 }
3918 crate::fts::search::ControlledSearchError::Control(error) => error,
3919 })?,
3920 };
3921 let mut counters = result.counters;
3922 let scored = result.hits;
3923 let mut provenance_memory =
3924 stats::AccountedCounter::new(&self.accounting, stats::AllocationComponent::Temporary)
3925 .map_err(QueryError::Store)?;
3926 let mut candidates = Vec::with_capacity(scored.len());
3927 let mut provenance_work = crate::fts::control::WorkCheck::new(|| {
3928 cancellation.check_graph().map_err(QueryError::Scan)
3929 });
3930 for (doc, score) in scored {
3931 let mut provenance = Vec::new();
3932 let segment = index
3933 .segments()
3934 .get(doc.segment as usize)
3935 .ok_or(QueryError::Store(StoreError::ActiveRowOverflow))?;
3936 let scoring = prepared
3937 .as_ref()
3938 .ok_or(QueryError::Store(StoreError::ActiveRowOverflow))?;
3939 for expansion in scoring.expansions() {
3940 cancellation.check_graph().map_err(QueryError::Scan)?;
3941 let fields = &scoring.scoring().query().fields;
3942 prepared_lexical::reserve_weighted_scratch(
3943 &mut provenance_memory,
3944 crate::fts::sealed::TermStream::allocation_bytes(
3945 segment,
3946 &expansion.term,
3947 fields,
3948 ),
3949 )?;
3950 if let Some(mut stream) = crate::fts::sealed::TermStream::open_sized_controlled(
3951 segment,
3952 &expansion.term,
3953 fields,
3954 &mut provenance_work,
3955 )? {
3956 stream.seek_controlled(doc.row, &mut provenance_work)?;
3957 if stream.current_row() == Some(doc.row) {
3958 provenance.push(expansion.clone());
3959 counters.postings_decoded = counters.postings_decoded.saturating_add(1);
3960 }
3961 counters.blocks_decoded = counters
3962 .blocks_decoded
3963 .saturating_add(stream.blocks_decoded());
3964 counters.blocks_skipped = counters
3965 .blocks_skipped
3966 .saturating_add(stream.blocks_skipped());
3967 }
3968 provenance_memory.set(0).map_err(QueryError::Store)?;
3969 }
3970 cancellation.check_graph().map_err(QueryError::Scan)?;
3971 let (text, document) = structured_lexical_row(&snapshot, &active, sources, doc)?;
3972 let terms = provenance
3973 .iter()
3974 .map(|entry| entry.term.clone())
3975 .collect::<Vec<_>>();
3976 let snippet = crate::fts::snippet::best_window(
3977 &self.tokenizer,
3978 text,
3979 &terms,
3980 snippet_bytes,
3981 true,
3982 )?
3983 .ok_or(crate::ingest::StoreLexicalError::MissingSnippetMatch {
3984 segment: doc.segment,
3985 row: doc.row,
3986 })?;
3987 let snippet_text = snippet
3988 .text(text)
3989 .ok_or(QueryError::Store(StoreError::ActiveRowOverflow))?;
3990 candidates.push(crate::ingest::ExplainedLexicalCandidate {
3991 document,
3992 score,
3993 provenance,
3994 snippet: crate::fts::query::OwnedLexicalSnippet {
3995 text: snippet_text.to_owned(),
3996 source: snippet.window,
3997 highlights: snippet.highlights,
3998 },
3999 });
4000 }
4001 cancellation.check_graph().map_err(QueryError::Scan)?;
4002 let mut diagnostics =
4003 crate::diag::QueryDiagnostics::lexical(crate::diag::LexicalDiagnostics {
4004 snapshot_generation: generation,
4005 indexed_through_seq: active
4006 .indexed_through_seq()
4007 .max(crate::wal::LogSeq::new(snapshot.absorbed_through())),
4008 requested_k: k,
4009 returned: candidates.len(),
4010 counters,
4011 elapsed: self.clock.now().saturating_duration_since(started),
4012 });
4013 diagnostics.counters.lexical_cache_hits = usize::from(assembly.cache_hit);
4014 diagnostics.counters.lexical_cache_builds = usize::from(!assembly.cache_hit);
4015 diagnostics.timings =
4016 crate::diag::QUERY_TIMING_ENABLED.then(|| crate::diag::QueryTimings {
4017 admission: admission_time,
4018 lexical: timing_elapsed(self.clock.as_ref(), lexical_started),
4019 ..Default::default()
4020 });
4021 drop(active_query);
4022 Ok(crate::ingest::StoreStructuredLexicalSearchOutcome {
4023 candidates,
4024 expansions: prepared
4025 .as_mut()
4026 .map_or(expansions, |query| query.take_expansions()),
4027 generation,
4028 diagnostics,
4029 })
4030 }
4031
4032 pub fn search_hybrid(
4035 &self,
4036 vector_query: crate::ingest::SearchRequest<'_>,
4037 lexical_query: &crate::fts::search::TermQuery,
4038 hybrid_query: &crate::fusion::HybridQuery,
4039 options: impl Into<SearchOptions>,
4040 control: QueryControl,
4041 ) -> Result<crate::ingest::StoreHybridSearchOutcome, crate::fusion::FusionError> {
4042 self.search_hybrid_inner(
4043 vector_query,
4044 PinnedLexicalQuery::Term(lexical_query),
4045 hybrid_query,
4046 options.into(),
4047 control,
4048 )
4049 }
4050
4051 pub fn search_hybrid_structured(
4054 &self,
4055 vector_query: crate::ingest::SearchRequest<'_>,
4056 lexical_query: &crate::fts::query::LexicalQuery,
4057 hybrid_query: &crate::fusion::HybridQuery,
4058 options: impl Into<SearchOptions>,
4059 control: QueryControl,
4060 ) -> Result<crate::ingest::StoreHybridSearchOutcome, crate::fusion::FusionError> {
4061 self.search_hybrid_inner(
4062 vector_query,
4063 PinnedLexicalQuery::Structured(lexical_query),
4064 hybrid_query,
4065 options.into(),
4066 control,
4067 )
4068 }
4069
4070 fn search_hybrid_inner(
4071 &self,
4072 vector_query: crate::ingest::SearchRequest<'_>,
4073 lexical_query: PinnedLexicalQuery<'_>,
4074 hybrid_query: &crate::fusion::HybridQuery,
4075 options: SearchOptions,
4076 control: QueryControl,
4077 ) -> Result<crate::ingest::StoreHybridSearchOutcome, crate::fusion::FusionError> {
4078 self.search_hybrid_then(
4079 vector_query,
4080 lexical_query,
4081 hybrid_query,
4082 options,
4083 control,
4084 false,
4085 |outcome, _, _, _, _| Ok(outcome),
4086 )
4087 }
4088
4089 #[allow(clippy::too_many_arguments)] fn search_hybrid_then<R>(
4091 &self,
4092 vector_query: crate::ingest::SearchRequest<'_>,
4093 lexical_query: PinnedLexicalQuery<'_>,
4094 hybrid_query: &crate::fusion::HybridQuery,
4095 options: SearchOptions,
4096 control: QueryControl,
4097 capture_rows: bool,
4098 finish: impl FnOnce(
4099 crate::ingest::StoreHybridSearchOutcome,
4100 &PublishedSnapshot,
4101 &crate::ingest::ActiveSegment,
4102 &QueryCancellation<'_>,
4103 Option<&materialize::HybridAddresses>,
4104 ) -> Result<R, crate::fusion::FusionError>,
4105 ) -> Result<R, crate::fusion::FusionError> {
4106 self.search_hybrid_prepared_then(
4107 || Ok::<_, std::convert::Infallible>(vector_query),
4108 lexical_query,
4109 hybrid_query,
4110 options,
4111 control,
4112 capture_rows,
4113 finish,
4114 )
4115 .map_err(|error| match error {
4116 HybridPreparationError::Preparation(never) => match never {},
4117 HybridPreparationError::Search(error) => error,
4118 })
4119 }
4120
4121 #[allow(clippy::too_many_arguments)] fn search_hybrid_prepared_then<'vector, E, R>(
4123 &self,
4124 prepare_vector: impl FnOnce() -> Result<crate::ingest::SearchRequest<'vector>, E>,
4125 lexical_query: PinnedLexicalQuery<'_>,
4126 hybrid_query: &crate::fusion::HybridQuery,
4127 options: SearchOptions,
4128 control: QueryControl,
4129 capture_rows: bool,
4130 finish: impl FnOnce(
4131 crate::ingest::StoreHybridSearchOutcome,
4132 &PublishedSnapshot,
4133 &crate::ingest::ActiveSegment,
4134 &QueryCancellation<'_>,
4135 Option<&materialize::HybridAddresses>,
4136 ) -> Result<R, crate::fusion::FusionError>,
4137 ) -> Result<R, HybridPreparationError<E>> {
4138 let control = control.with_clock(Arc::clone(&self.clock));
4139 let started = self.clock.now();
4140 let admission_started = timing_start(self.clock.as_ref());
4141 let mut options = options;
4142 let requested_tier = options.explicit_tier();
4143 let admitted = self
4159 .admit_vector_search_for(options, requested_tier.is_none())
4160 .map_err(crate::fusion::FusionError::from)?;
4161 options = admitted.execution_options;
4162 let tier_resolution = if requested_tier.is_none() {
4163 if snapshot_has_graph(&admitted.snapshot) {
4164 Some(crate::diag::HybridTierResolution::Auto)
4165 } else {
4166 options = options.with_tier(SearchTier::Exact);
4167 Some(crate::diag::HybridTierResolution::Exact)
4168 }
4169 } else {
4170 None
4171 };
4172 let corpus_rows = hybrid::corpus_rows(&admitted.snapshot, &admitted.active_segment)?;
4173 let mut width = if hybrid_query.max_rounds == 0 {
4178 corpus_rows
4179 } else {
4180 hybrid::hybrid_window(hybrid_query.k, corpus_rows)?.width
4181 };
4182 let panic_vector = self.consume_hybrid_test_fault(crate::fusion::FusionLeg::Vector);
4183 let panic_lexical = self.consume_hybrid_test_fault(crate::fusion::FusionLeg::Lexical);
4184 let lexical_worker = self.ensure_lexical_worker()?;
4185 let mut timings = crate::diag::QueryTimings {
4186 admission: timing_elapsed(self.clock.as_ref(), admission_started),
4187 ..Default::default()
4188 };
4189 let lexical_inputs = LexicalInputs {
4190 generation: admitted.generation,
4191 cache: &self.lexical_index_cache,
4192 snapshot: &admitted.snapshot,
4193 active: &admitted.active_segment,
4194 accounting: &self.accounting,
4195 };
4196 #[cfg(any(test, feature = "test-support"))]
4197 let preparation_observer = crate::fts::preparation_observer::current();
4198 let mut lexical_preparation = prepared_lexical::PreparedLexicalQuery::new(lexical_query);
4199 #[cfg(any(test, feature = "test-support"))]
4200 let fresh_preparation = hybrid::fresh_round_test_control();
4201 let run_lexical_leg = |preparation: &mut prepared_lexical::PreparedLexicalQuery<'_>,
4202 bound,
4203 queued| {
4204 #[cfg(any(test, feature = "test-support"))]
4205 if fresh_preparation {
4206 preparation.reset_for_test();
4207 }
4208 #[cfg(any(test, feature = "test-support"))]
4209 let _preparation_scope =
4210 crate::fts::preparation_observer::enter(preparation_observer.clone());
4211 let queue_time = timing_elapsed(self.clock.as_ref(), queued);
4212 let lexical_started = timing_start(self.clock.as_ref());
4213 let name = std::thread::current().name().map(str::to_owned);
4214 let result = (|| {
4215 maybe_trigger_hybrid_leg_panic(panic_lexical, "injected lexical hybrid leg panic");
4216 let lease =
4217 SnapshotLease::new_at(Arc::clone(&admitted.snapshot), admitted.generation);
4218 let cancellation = QueryCancellation::new(&control, &lease);
4219 cancellation
4220 .check_graph()
4221 .map_err(QueryError::Scan)
4222 .map_err(crate::fusion::FusionError::from)?;
4223 match lexical_query {
4224 PinnedLexicalQuery::Term(_) => {
4225 exact_lexical_leg(lexical_inputs, bound, &cancellation, preparation)
4226 }
4227 PinnedLexicalQuery::Structured(query) => exact_structured_lexical_leg(
4228 lexical_inputs,
4229 &self.tokenizer,
4230 query,
4231 bound,
4232 &cancellation,
4233 preparation,
4234 ),
4235 }
4236 })();
4237 (
4238 name,
4239 result,
4240 queue_time,
4241 timing_elapsed(self.clock.as_ref(), lexical_started),
4242 )
4243 };
4244 let caller_chose_tier = requested_tier.is_some();
4245 let run_vector_leg = |vector_query: crate::ingest::SearchRequest<'_>,
4246 vector_preparation: &mut prepared::PreparedVectorQuery<'_>,
4247 k: usize| {
4248 let graph_available = snapshot_has_graph(&admitted.snapshot);
4249 let scan_reason_override = if caller_chose_tier || !graph_available {
4250 None
4251 } else if k > corpus_rows {
4252 Some(crate::planner::ScanReason::FullMaterialization)
4253 } else if graph_round_is_worth_it(k, corpus_rows) {
4254 None
4255 } else {
4256 Some(crate::planner::ScanReason::WideningCap {
4257 ef: graph_round_ef(k),
4258 rows: corpus_rows,
4259 })
4260 };
4261 search_pinned(
4262 admitted.pool.as_deref(),
4263 &admitted.snapshot,
4264 &admitted.active_segment,
4265 &self.accounting,
4266 admitted.generation,
4267 self.epoch_identity(),
4268 vector_query,
4269 k,
4270 if caller_chose_tier || graph_round_is_worth_it(k, corpus_rows) {
4282 options
4283 } else {
4284 options.with_tier(SearchTier::Exact)
4285 },
4286 control.clone(),
4287 GraphBoundMode::Shared,
4288 requested_tier,
4289 scan_reason_override,
4290 started,
4291 self.clock.as_ref(),
4292 vector_preparation,
4293 #[cfg(any(test, feature = "test-support"))]
4294 self.vector_fault_controller.as_ref(),
4295 )
4296 .map_err(crate::fusion::FusionError::from)
4297 };
4298 let lexical_queued = timing_start(self.clock.as_ref());
4299 let (prepared_vector_result, lexical_result) = lexical_worker.run_scoped(
4300 || {
4301 run_lexical_leg(
4302 &mut lexical_preparation,
4303 width.saturating_add(1),
4304 lexical_queued,
4305 )
4306 },
4307 || {
4308 std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
4309 let vector_query =
4310 prepare_vector().map_err(HybridPreparationError::Preparation)?;
4311 let mut vector_preparation = prepared::PreparedVectorQuery::new(
4312 vector_query.vector(),
4313 self.epoch_identity(),
4314 &self.accounting,
4315 );
4316 maybe_trigger_hybrid_leg_panic(
4317 panic_vector,
4318 "injected vector hybrid leg panic",
4319 );
4320 let result = run_vector_leg(
4321 vector_query,
4322 &mut vector_preparation,
4323 width.saturating_add(1),
4324 );
4325 Ok((vector_query, vector_preparation, result))
4326 }))
4327 .unwrap_or(Err(HybridPreparationError::Search(
4328 crate::fusion::FusionError::LegPanic {
4329 leg: crate::fusion::FusionLeg::Vector,
4330 detail: "vector hybrid leg panicked",
4331 },
4332 )))
4333 },
4334 )?;
4335 let (lexical_thread_name, lexical_result, queue_time, lexical_time) = lexical_result
4336 .unwrap_or_else(|error| {
4337 (
4338 Some("zeppelin-fts".to_owned()),
4339 Err(error),
4340 std::time::Duration::ZERO,
4341 std::time::Duration::ZERO,
4342 )
4343 });
4344 #[cfg(any(test, feature = "test-support"))]
4345 if let Ok(mut receipt) = self.hybrid_execution_receipt.lock() {
4346 *receipt = Some(HybridExecutionReceipt {
4347 vector_thread: std::thread::current().id(),
4348 lexical_thread_name,
4349 vector_completed: true,
4350 lexical_completed: true,
4351 });
4352 }
4353 #[cfg(not(any(test, feature = "test-support")))]
4354 let _ = lexical_thread_name;
4355 let (vector_query, mut vector_preparation, vector_result) = match prepared_vector_result {
4359 Ok(prepared) => prepared,
4360 Err(error) => {
4361 let lease =
4362 SnapshotLease::new_at(Arc::clone(&admitted.snapshot), admitted.generation);
4363 QueryCancellation::new(&control, &lease)
4364 .check_graph()
4365 .map_err(map_scan_error)?;
4366 if let Err(lexical_error) = lexical_result
4367 && matches!(
4368 lexical_error,
4369 crate::fusion::FusionError::ReadCancelled { .. }
4370 | crate::fusion::FusionError::Timeout { .. }
4371 | crate::fusion::FusionError::Cancelled { .. }
4372 )
4373 {
4374 return Err(HybridPreparationError::Search(lexical_error));
4375 }
4376 return Err(error);
4377 }
4378 };
4379 let (
4380 mut vector_outcome,
4381 (mut lexical_hits, mut lexical_assembly, mut lexical_counters, mut lexical_cache_hit),
4382 ) = resolve_hybrid_leg_results(vector_result, lexical_result)?;
4383
4384 timings.lexical_queue = queue_time;
4385 timings.lexical = lexical_time;
4386 let anchors = hybrid::HybridAnchors {
4387 vector_ceiling: vector_outcome.vector_ceiling,
4388 lexical_maximum: lexical_hits.first().map_or(0.0, |hit| hit.bm25),
4389 };
4390 let fusion_lease =
4391 SnapshotLease::new_at(Arc::clone(&admitted.snapshot), admitted.generation);
4392 let fusion_cancellation = QueryCancellation::new(&control, &fusion_lease);
4393 let mut work = hybrid::HybridWork::new();
4394 let mut cross_scores = hybrid::CrossScoreCache::new(vector_outcome.epoch);
4395 let mut round_buffers = hybrid::HybridRoundBuffers::new();
4396 let mut fusion_scratch =
4397 crate::fusion::StoreFusionScratch::new(Arc::clone(&self.accounting));
4398 let mut rescore_counters = crate::diag::ScanRescoreCounters::default();
4399 let mut rounds = 1_usize;
4400 let mut budget_exhausted = false;
4401 let (fused, report_window, addresses) = loop {
4402 if let Some(vector_timings) = vector_outcome.diagnostics.timings {
4403 timings.vector += vector_timings.vector;
4404 }
4405 let fusion_started = timing_start(self.clock.as_ref());
4406 let round = hybrid::build_round(
4407 &admitted.snapshot,
4408 &admitted.active_segment,
4409 &lexical_assembly,
4410 &lexical_preparation,
4411 &self.tokenizer,
4412 lexical_query,
4413 vector_query.vector(),
4414 &vector_outcome.candidates,
4415 anchors,
4416 &lexical_hits,
4417 width,
4418 &fusion_cancellation,
4419 capture_rows,
4420 &self.accounting,
4421 &mut cross_scores,
4422 &mut round_buffers,
4423 )?;
4424 let cross_fill_dims = u64::try_from(round.vector_scores_computed)
4428 .ok()
4429 .and_then(|rows| rows.checked_mul(vector_query.vector().len() as u64))
4430 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
4431 let cross_fill_bytes = cross_fill_dims
4432 .checked_mul(4)
4433 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
4434 vector_outcome.stats.dims_touched = vector_outcome
4435 .stats
4436 .dims_touched
4437 .checked_add(cross_fill_dims)
4438 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
4439 vector_outcome.stats.bytes_read = vector_outcome
4440 .stats
4441 .bytes_read
4442 .checked_add(cross_fill_bytes)
4443 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
4444 if cross_fill_dims != 0 {
4445 let caller = std::thread::current().id();
4446 if !vector_outcome.stats.worker_thread_ids.contains(&caller) {
4447 vector_outcome.stats.worker_thread_ids.push(caller);
4448 }
4449 vector_outcome.stats.threads_used = vector_outcome.stats.worker_thread_ids.len();
4450 }
4451 accumulate_search_counters(&mut lexical_counters, &round.lexical_counters);
4452 if let Some(counters) = vector_outcome.diagnostics.scan_rescore {
4453 rescored_scan::accumulate(&mut rescore_counters, counters)?;
4454 }
4455 work.complete_round(
4456 &vector_outcome.stats,
4457 vector_outcome.graph_stats,
4458 lexical_counters,
4459 )?;
4460 work.lexical_cache_hits += usize::from(lexical_cache_hit);
4461 work.lexical_cache_builds += usize::from(!lexical_cache_hit);
4462 work.plans.append(&mut vector_outcome.diagnostics.plan);
4463 work.cross_filled_vector = work
4464 .cross_filled_vector
4465 .checked_add(round.vector_scores_computed)
4466 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
4467 work.cross_filled_lexical = work
4468 .cross_filled_lexical
4469 .checked_add(round.lexical_scores_computed)
4470 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
4471 work.vector_candidates_produced = work
4472 .vector_candidates_produced
4473 .checked_add(vector_outcome.candidates.len())
4474 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
4475 work.lexical_candidates_produced = work
4476 .lexical_candidates_produced
4477 .checked_add(lexical_hits.len())
4478 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
4479 let provenance = hybrid::round_provenance(&vector_outcome, lexical_hits.len(), width);
4480 let fused = hybrid::fuse_round(
4481 hybrid_query,
4482 &round,
4483 provenance,
4484 requested_tier == Some(SearchTier::Exact),
4485 &mut fusion_scratch,
4486 )?;
4487 timings.fusion_cross_fill += timing_elapsed(self.clock.as_ref(), fusion_started);
4488 if fused.report.termination != crate::fusion::FusionTermination::WindowUnproven
4489 || width >= corpus_rows
4490 {
4491 break (
4492 fused,
4493 crate::diag::HybridReport {
4494 normalization_policy_version:
4495 crate::fusion::HYBRID_NORMALIZATION_POLICY_VERSION,
4496 lexical_full_materializations: 0,
4497 provenance,
4498 window: width,
4499 vector_returned: round.vector.len(),
4500 lexical_returned: round.lexical.len(),
4501 cross_filled_vector: round.cross_filled_vector,
4502 cross_filled_lexical: round.cross_filled_lexical,
4503 total_cross_filled_vector: work.cross_filled_vector,
4504 total_cross_filled_lexical: work.cross_filled_lexical,
4505 vector_candidates_produced: work.vector_candidates_produced,
4506 lexical_candidates_produced: work.lexical_candidates_produced,
4507 },
4508 round.addresses,
4509 );
4510 }
4511 width = if rounds >= hybrid_query.max_rounds {
4512 budget_exhausted = true;
4513 corpus_rows
4514 } else {
4515 width.saturating_mul(2).min(corpus_rows)
4516 };
4517 rounds = rounds.saturating_add(1);
4518 let lexical_queued = timing_start(self.clock.as_ref());
4519 let (vector_result, lexical_result) = lexical_worker.run_scoped(
4520 || {
4521 run_lexical_leg(
4522 &mut lexical_preparation,
4523 width.saturating_add(1),
4524 lexical_queued,
4525 )
4526 },
4527 || {
4528 run_vector_leg(
4529 vector_query,
4530 &mut vector_preparation,
4531 width.saturating_add(1),
4532 )
4533 },
4534 )?;
4535 let (_, lexical_result, queue_time, lexical_time) =
4536 lexical_result.unwrap_or_else(|error| {
4537 (
4538 Some("zeppelin-fts".to_owned()),
4539 Err(error),
4540 std::time::Duration::ZERO,
4541 std::time::Duration::ZERO,
4542 )
4543 });
4544 let (next_vector, (next_hits, next_sources, next_counters, next_cache_hit)) =
4545 resolve_hybrid_leg_results(vector_result, lexical_result)?;
4546 timings.lexical_queue += queue_time;
4547 timings.lexical += lexical_time;
4548 vector_outcome = next_vector;
4549 lexical_hits = next_hits;
4550 lexical_assembly = next_sources;
4551 lexical_counters = next_counters;
4552 lexical_cache_hit = next_cache_hit;
4553 };
4554
4555 let crate::ingest::SearchOutcome {
4556 candidates: _,
4557 vector_ceiling: _,
4558 stats: _,
4559 graph_stats: _,
4560 generation,
4561 epoch,
4562 diagnostics: vector_diagnostics,
4563 } = vector_outcome;
4564 let mut report = fused.report;
4565 report.rounds = rounds;
4566 if budget_exhausted {
4567 report.budget_exhausted = true;
4568 report.termination = crate::fusion::FusionTermination::BudgetFullMaterialization;
4569 }
4570 let mut diagnostics =
4571 crate::diag::QueryDiagnostics::hybrid(crate::diag::HybridDiagnostics {
4572 snapshot_generation: vector_diagnostics.snapshot_generation,
4573 indexed_through_seq: vector_diagnostics.indexed_through_seq,
4574 plan: work.plans,
4575 approximate: vector_diagnostics.approximate
4576 || report.termination
4577 == crate::fusion::FusionTermination::ApproximateCandidates,
4578 exact_rescore: vector_diagnostics.exact_rescore,
4579 requested_k: hybrid_query.k,
4580 returned: fused.hits.len(),
4581 scan: work.scan,
4582 graph: work.graph,
4583 lexical: work.lexical,
4584 report,
4585 hybrid: report_window,
4586 tier_resolution,
4587 epoch,
4588 elapsed: self.clock.now().saturating_duration_since(started),
4589 });
4590 diagnostics.counters.lexical_cache_hits = work.lexical_cache_hits;
4591 diagnostics.scan_rescore = options.scan_rescore().map(|_| rescore_counters);
4592 diagnostics.counters.lexical_cache_builds = work.lexical_cache_builds;
4593 diagnostics.timings = crate::diag::QUERY_TIMING_ENABLED.then_some(timings);
4594 let outcome = crate::ingest::StoreHybridSearchOutcome {
4595 hits: fused.hits,
4596 lexical_expansions: lexical_preparation.take_expansions()?,
4597 generation,
4598 diagnostics,
4599 };
4600 finish(
4601 outcome,
4602 &admitted.snapshot,
4603 &admitted.active_segment,
4604 &fusion_cancellation,
4605 addresses.as_ref(),
4606 )
4607 .map_err(HybridPreparationError::Search)
4608 }
4609
4610 #[cfg(any(test, feature = "test-support"))]
4616 #[must_use]
4617 pub fn lexical_index_cache_counters(&self) -> (u64, u64) {
4618 (
4619 self.lexical_index_cache.hits.load(Ordering::Relaxed),
4620 self.lexical_index_cache.builds.load(Ordering::Relaxed),
4621 )
4622 }
4623
4624 #[cfg(any(test, feature = "test-support"))]
4626 #[must_use]
4627 pub fn lexical_contribution_cache_counters(&self) -> (u64, u64) {
4628 (
4629 self.lexical_index_cache
4630 .contribution_hits
4631 .load(Ordering::Relaxed),
4632 self.lexical_index_cache
4633 .contribution_builds
4634 .load(Ordering::Relaxed),
4635 )
4636 }
4637
4638 fn consume_hybrid_test_fault(&self, leg: crate::fusion::FusionLeg) -> bool {
4639 #[cfg(any(test, feature = "test-support"))]
4640 {
4641 let Ok(mut fault) = self.hybrid_leg_fault.lock() else {
4642 return false;
4643 };
4644 if matches!(*fault, Some(HybridLegTestFault::Panic(candidate)) if candidate == leg) {
4645 *fault = None;
4646 return true;
4647 }
4648 }
4649 #[cfg(not(any(test, feature = "test-support")))]
4650 let _ = leg;
4651 false
4652 }
4653
4654 fn search_with_graph_bound_mode(
4655 &self,
4656 request: crate::ingest::SearchRequest<'_>,
4657 k: usize,
4658 options: SearchOptions,
4659 control: QueryControl,
4660 graph_bound_mode: GraphBoundMode,
4661 ) -> Result<crate::ingest::SearchOutcome, QueryError> {
4662 self.search_with_graph_bound_mode_then(
4663 request,
4664 k,
4665 options,
4666 control,
4667 graph_bound_mode,
4668 |outcome, _, _| Ok(outcome),
4669 )
4670 }
4671
4672 fn search_with_graph_bound_mode_then<R>(
4673 &self,
4674 request: crate::ingest::SearchRequest<'_>,
4675 k: usize,
4676 options: SearchOptions,
4677 control: QueryControl,
4678 graph_bound_mode: GraphBoundMode,
4679 finish: impl FnOnce(
4680 crate::ingest::SearchOutcome,
4681 &AdmittedVectorSearch<'_>,
4682 &QueryControl,
4683 ) -> Result<R, QueryError>,
4684 ) -> Result<R, QueryError> {
4685 let control = control.with_clock(Arc::clone(&self.clock));
4686 let started = self.clock.now();
4687 let admission_started = timing_start(self.clock.as_ref());
4688 let admitted = self.admit_vector_search(options)?;
4689 let admission_time = timing_elapsed(self.clock.as_ref(), admission_started);
4690 let mut vector_preparation = prepared::PreparedVectorQuery::new(
4691 request.vector(),
4692 self.epoch_identity(),
4693 &self.accounting,
4694 );
4695 let score = || {
4696 let mut outcome = search_pinned(
4697 admitted.pool.as_deref(),
4698 &admitted.snapshot,
4699 &admitted.active_segment,
4700 &self.accounting,
4701 admitted.generation,
4702 self.epoch_identity(),
4703 request,
4704 k,
4705 options,
4706 control.clone(),
4707 graph_bound_mode,
4708 options.explicit_tier(),
4709 None,
4710 started,
4711 self.clock.as_ref(),
4712 &mut vector_preparation,
4713 #[cfg(any(test, feature = "test-support"))]
4714 self.vector_fault_controller.as_ref(),
4715 )?;
4716 if let Some(timings) = &mut outcome.diagnostics.timings {
4717 timings.admission = admission_time;
4718 }
4719 finish(outcome, &admitted, &control)
4720 };
4721 #[cfg(any(test, feature = "test-support"))]
4722 let result = crate::kernels::vector_fault::run_store_scoring(
4723 self.kernel_fault_controller.as_ref(),
4724 score,
4725 );
4726 #[cfg(not(any(test, feature = "test-support")))]
4727 let result = score();
4728 #[cfg(any(test, feature = "test-support"))]
4729 {
4730 if let Some(controller) = self.vector_fault_controller.as_ref() {
4731 controller.finalize_search(result.is_ok());
4732 }
4733 }
4734 result
4735 }
4736
4737 fn ensure_query_pool(&self) -> Result<Arc<pool::QueryPool>, QueryError> {
4738 let mut pool_slot = self.query_pool.lock().map_err(|_| {
4739 QueryError::Store(StoreError::Synchronization {
4740 component: "query pool",
4741 })
4742 })?;
4743 if pool_slot.is_none() {
4744 let capacity = crate::scan::physical_thread_capacity().map_err(|error| {
4745 QueryError::Store(StoreError::QueryPoolCapacity {
4746 source: error.to_string(),
4747 })
4748 })?;
4749 let pool =
4750 pool::QueryPool::start(capacity, &self.accounting).map_err(QueryError::Store)?;
4751 *pool_slot = Some(Arc::new(pool));
4752 }
4753 let pool = pool_slot.as_ref().cloned().ok_or({
4754 QueryError::Store(StoreError::Synchronization {
4755 component: "query pool initialization",
4756 })
4757 })?;
4758 drop(pool_slot);
4759 Ok(pool)
4760 }
4761
4762 fn ensure_lexical_worker(
4763 &self,
4764 ) -> Result<Arc<pool::LexicalWorker>, crate::fusion::FusionError> {
4765 let mut worker_slot =
4766 self.lexical_worker
4767 .lock()
4768 .map_err(|_| crate::fusion::FusionError::LegThreadStart {
4769 leg: crate::fusion::FusionLeg::Lexical,
4770 detail: "pooled lexical worker synchronization was poisoned".to_owned(),
4771 })?;
4772 if worker_slot.is_none() {
4773 *worker_slot = Some(Arc::new(pool::LexicalWorker::start()?));
4774 }
4775 let worker = worker_slot.as_ref().cloned().ok_or_else(|| {
4776 crate::fusion::FusionError::LegThreadStart {
4777 leg: crate::fusion::FusionLeg::Lexical,
4778 detail: "pooled lexical worker initialization failed".to_owned(),
4779 }
4780 })?;
4781 drop(worker_slot);
4782 Ok(worker)
4783 }
4784
4785 fn admit_vector_search(
4786 &self,
4787 options: SearchOptions,
4788 ) -> Result<AdmittedVectorSearch<'_>, QueryError> {
4789 self.admit_vector_search_for(options, false)
4790 }
4791
4792 fn admit_vector_search_for(
4793 &self,
4794 options: SearchOptions,
4795 hybrid_default: bool,
4796 ) -> Result<AdmittedVectorSearch<'_>, QueryError> {
4797 #[cfg(any(test, feature = "test-support"))]
4798 self.text_materialization_work
4799 .admissions
4800 .fetch_add(1, Ordering::Relaxed);
4801 let state = self
4802 .state
4803 .lock()
4804 .map_err(|_| QueryError::Store(StoreError::Synchronization { component: "state" }))?;
4805 match *state {
4806 StoreState::Open => {}
4807 StoreState::Closing => return Err(QueryError::Store(StoreError::Closing)),
4808 StoreState::Closed => return Err(QueryError::Store(StoreError::Closed)),
4809 }
4810 let active_guard = self.active.lock().map_err(|_| {
4811 QueryError::Store(StoreError::Synchronization {
4812 component: "active segment",
4813 })
4814 })?;
4815 let active_state = active_guard
4816 .as_ref()
4817 .ok_or(QueryError::Store(StoreError::Closed))?;
4818 let generation = active_state.generation;
4819 let active_segment = Arc::clone(&active_state.segment);
4820 let snapshot = self
4821 .snapshot
4822 .read()
4823 .map_err(|_| {
4824 QueryError::Store(StoreError::Synchronization {
4825 component: "published snapshot",
4826 })
4827 })?
4828 .as_ref()
4829 .cloned()
4830 .ok_or(QueryError::Store(StoreError::Closed))?;
4831 let execution_options = if hybrid_default && !snapshot_has_graph(&snapshot) {
4832 options.with_tier(SearchTier::Exact)
4833 } else {
4834 options
4835 };
4836 let needs_query_pool = !active_segment.is_empty()
4839 || (hybrid_default && !snapshot.segments().is_empty())
4840 || match execution_options.tier() {
4841 SearchTier::Auto => snapshot.segments().iter().any(|segment| {
4842 !segment.directory().iter().any(|entry| {
4843 entry.kind == crate::segment::layout::RegionKind::GraphNodeBlocks.id()
4844 })
4845 }),
4846 SearchTier::Exact => !snapshot.segments().is_empty(),
4847 SearchTier::Scan => true,
4848 SearchTier::Graph(_) => false,
4849 };
4850 let pool = if needs_query_pool {
4851 Some(self.ensure_query_pool()?)
4852 } else {
4853 None
4854 };
4855 self.active_queries.fetch_add(1, Ordering::Relaxed);
4856 let active_query = ActiveQuery {
4857 count: &self.active_queries,
4858 };
4859 drop(active_guard);
4860 drop(state);
4861 Ok(AdmittedVectorSearch {
4862 pool,
4863 execution_options,
4864 snapshot: snapshot::ReadSnapshot::new(snapshot),
4865 active_segment,
4866 generation,
4867 _active_query: active_query,
4868 })
4869 }
4870
4871 #[cfg(test)]
4872 fn search_independent_for_test(
4873 &self,
4874 request: crate::ingest::SearchRequest<'_>,
4875 k: usize,
4876 options: SearchOptions,
4877 control: QueryControl,
4878 ) -> Result<crate::ingest::SearchOutcome, QueryError> {
4879 self.search_with_graph_bound_mode(request, k, options, control, GraphBoundMode::Independent)
4880 }
4881}
4882
4883#[derive(Clone, Copy)]
4884pub(crate) enum StructuredLexicalSource {
4885 Sealed(usize),
4886 Active,
4887}
4888
4889struct LexicalAssembly {
4890 index: crate::fts::index::LexicalIndex,
4891 alive_sets: Vec<Arc<crate::meta::AliveSet>>,
4892 sources: Vec<StructuredLexicalSource>,
4893 contributions: Vec<Arc<CachedLexicalContribution>>,
4894 vocabulary: Mutex<Option<Arc<CachedVocabulary>>>,
4895 memory: stats::AccountedCounter,
4897}
4898
4899struct CachedVocabulary {
4900 view: crate::fts::vocabulary::Vocabulary,
4901 phonetic: Mutex<Option<Arc<expansion::CachedPhonetic>>>,
4902 _memory: stats::AccountedCounter,
4904}
4905
4906impl LexicalAssembly {
4907 fn vocabulary(
4908 &self,
4909 accounting: &Arc<stats::Accounting>,
4910 cancellation: &QueryCancellation<'_>,
4911 ) -> Result<Arc<CachedVocabulary>, QueryError> {
4912 let mut cached = cancellation.cache_lock(&self.vocabulary, "lexical vocabulary cache")?;
4913 if let Some(value) = cached.as_ref() {
4914 return Ok(Arc::clone(value));
4915 }
4916 let mut memory =
4917 stats::AccountedCounter::new(accounting, stats::AllocationComponent::Cache)
4918 .map_err(QueryError::Store)?;
4919 let mut temporary =
4920 stats::AccountedCounter::new(accounting, stats::AllocationComponent::Temporary)
4921 .map_err(QueryError::Store)?;
4922 let mut work = crate::fts::control::WorkCheck::new(|| {
4923 cancellation.check_graph().map_err(QueryError::Scan)
4924 });
4925 let view = crate::fts::vocabulary::Vocabulary::build_controlled(
4926 self.index.vocabulary_terms(),
4927 |scratch, owned| {
4928 cancellation.check_graph().map_err(QueryError::Scan)?;
4929 prepared_lexical::reserve_weighted_scratch(&mut temporary, scratch)?;
4930 let bytes = owned
4931 .and_then(|bytes| bytes.checked_add(std::mem::size_of::<CachedVocabulary>()))
4932 .ok_or(QueryError::Store(StoreError::BudgetExceeded {
4933 needed: u64::MAX,
4934 budget: u64::MAX,
4935 component: "lexical vocabulary cache",
4936 }))?;
4937 memory.set(bytes).map_err(QueryError::Store)
4938 },
4939 &mut work,
4940 )?;
4941 let value = Arc::new(CachedVocabulary {
4942 view,
4943 phonetic: Mutex::new(None),
4944 _memory: memory,
4945 });
4946 *cached = Some(Arc::clone(&value));
4947 Ok(value)
4948 }
4949
4950 fn owned_bytes(&self) -> Option<usize> {
4953 std::mem::size_of::<Self>()
4954 .checked_add(self.index.resident_spine_bytes()?)?
4955 .checked_add(
4956 self.alive_sets
4957 .capacity()
4958 .checked_mul(std::mem::size_of::<Arc<crate::meta::AliveSet>>())?,
4959 )?
4960 .checked_add(
4961 self.sources
4962 .capacity()
4963 .checked_mul(std::mem::size_of::<StructuredLexicalSource>())?,
4964 )?
4965 .checked_add(
4966 self.contributions
4967 .capacity()
4968 .checked_mul(std::mem::size_of::<Arc<CachedLexicalContribution>>())?,
4969 )
4970 }
4971}
4972
4973enum LexicalContributionKey {
4974 Sealed(crate::segment::SegmentMeta, (u64, u64)),
4977 Active(std::sync::Weak<crate::ingest::ActiveSegment>),
4978}
4979
4980struct CachedLexicalContribution {
4981 key: LexicalContributionKey,
4982 statistics: Arc<crate::fts::index::LiveSegmentStatistics>,
4983 active_alive: Option<Arc<crate::meta::AliveSet>>,
4986 _memory: stats::AccountedCounter,
4987}
4988
4989impl CachedLexicalContribution {
4990 fn build(
4991 key: LexicalContributionKey,
4992 postings: &crate::fts::sealed::SealedSegment,
4993 alive: &Arc<crate::meta::AliveSet>,
4994 ordinal: usize,
4995 accounting: &Arc<stats::Accounting>,
4996 work: &mut crate::fts::control::WorkCheck<impl FnMut() -> Result<(), LexicalAssemblyError>>,
4997 ) -> Result<Arc<Self>, LexicalAssemblyError> {
4998 work.check_now()?;
4999 let active = matches!(key, LexicalContributionKey::Active(_));
5000 let alive_bytes = if active {
5001 alive
5002 .resident_bytes()
5003 .and_then(|n| n.checked_add(std::mem::size_of::<crate::meta::AliveSet>()))
5004 .and_then(|n| n.checked_add(2 * std::mem::size_of::<usize>()))
5005 } else {
5006 Some(0)
5007 };
5008 let bytes =
5009 crate::fts::index::LiveSegmentStatistics::allocation_bytes(alive.alive_bitmap())
5010 .and_then(|n| n.checked_add(std::mem::size_of::<Self>()))
5011 .and_then(|n| n.checked_add(2 * std::mem::size_of::<usize>()))
5012 .and_then(|n| n.checked_add(alive_bytes?))
5013 .ok_or(LexicalAssemblyError::Store(StoreError::BudgetExceeded {
5014 needed: u64::MAX,
5015 budget: u64::MAX,
5016 component: "cache",
5017 }))?;
5018 let mut memory =
5019 stats::AccountedCounter::new(accounting, stats::AllocationComponent::Cache)
5020 .map_err(LexicalAssemblyError::Store)?;
5021 memory.set(bytes).map_err(LexicalAssemblyError::Store)?;
5022 let mut statistics = crate::fts::index::LiveSegmentStatistics::build_controlled(
5023 ordinal,
5024 postings,
5025 alive.alive_bitmap(),
5026 work,
5027 )?;
5028 if statistics.rows.cardinality() != 0
5029 && statistics.rows.cardinality() != u64::from(postings.row_count())
5030 {
5031 statistics.frequency_cache = Some(
5032 crate::fts::live_df::LiveFrequencyCache::new(accounting)
5033 .map_err(LexicalAssemblyError::Store)?,
5034 );
5035 }
5036 let statistics = Arc::new(statistics);
5037 work.check_now()?;
5038 Ok(Arc::new(Self {
5039 key,
5040 statistics,
5041 active_alive: active.then(|| Arc::clone(alive)),
5042 _memory: memory,
5043 }))
5044 }
5045}
5046
5047enum LexicalAssemblyError {
5048 Store(StoreError),
5049 Lexical(crate::fts::index::IndexError),
5050 MissingDocumentIdentity(crate::segment::SegmentId),
5051 Cancelled(crate::scan::ScanError),
5052}
5053
5054impl From<crate::fts::index::IndexError> for LexicalAssemblyError {
5055 fn from(error: crate::fts::index::IndexError) -> Self {
5056 Self::Lexical(error)
5057 }
5058}
5059
5060impl From<StoreError> for LexicalAssemblyError {
5061 fn from(error: StoreError) -> Self {
5062 Self::Store(error)
5063 }
5064}
5065
5066struct CachedLexicalAssembly {
5093 generation: u64,
5094 snapshot: std::sync::Weak<PublishedSnapshot>,
5095 active: std::sync::Weak<crate::ingest::ActiveSegment>,
5096 document_identity_verified: bool,
5101 assembly: Arc<LexicalAssembly>,
5102}
5103
5104impl CachedLexicalAssembly {
5105 fn matches(
5106 &self,
5107 snapshot: &Arc<PublishedSnapshot>,
5108 active: &Arc<crate::ingest::ActiveSegment>,
5109 ) -> bool {
5110 self.snapshot
5111 .upgrade()
5112 .is_some_and(|cached| Arc::ptr_eq(&cached, snapshot))
5113 && self
5114 .active
5115 .upgrade()
5116 .is_some_and(|cached| Arc::ptr_eq(&cached, active))
5117 }
5118}
5119
5120pub(crate) struct LexicalIndexCache {
5126 entry: Mutex<Option<CachedLexicalAssembly>>,
5127 #[cfg(any(test, feature = "test-support"))]
5128 hits: AtomicU64,
5129 #[cfg(any(test, feature = "test-support"))]
5130 builds: AtomicU64,
5131 #[cfg(any(test, feature = "test-support"))]
5132 contribution_hits: AtomicU64,
5133 #[cfg(any(test, feature = "test-support"))]
5134 contribution_builds: AtomicU64,
5135}
5136
5137impl LexicalIndexCache {
5138 pub(crate) const fn new() -> Self {
5139 Self {
5140 entry: Mutex::new(None),
5141 #[cfg(any(test, feature = "test-support"))]
5142 hits: AtomicU64::new(0),
5143 #[cfg(any(test, feature = "test-support"))]
5144 builds: AtomicU64::new(0),
5145 #[cfg(any(test, feature = "test-support"))]
5146 contribution_hits: AtomicU64::new(0),
5147 #[cfg(any(test, feature = "test-support"))]
5148 contribution_builds: AtomicU64::new(0),
5149 }
5150 }
5151
5152 #[cfg(any(test, feature = "test-support"))]
5153 fn record_hit(&self) {
5154 self.hits.fetch_add(1, Ordering::Relaxed);
5155 }
5156
5157 #[cfg(not(any(test, feature = "test-support")))]
5158 const fn record_hit(&self) {}
5159
5160 #[cfg(any(test, feature = "test-support"))]
5161 fn record_build(&self) {
5162 self.builds.fetch_add(1, Ordering::Relaxed);
5163 }
5164
5165 #[cfg(not(any(test, feature = "test-support")))]
5166 const fn record_build(&self) {}
5167
5168 fn record_contribution(&self, hit: bool) {
5169 #[cfg(any(test, feature = "test-support"))]
5170 if hit {
5171 self.contribution_hits.fetch_add(1, Ordering::Relaxed);
5172 } else {
5173 self.contribution_builds.fetch_add(1, Ordering::Relaxed);
5174 }
5175 #[cfg(not(any(test, feature = "test-support")))]
5176 let _ = hit;
5177 }
5178
5179 fn lock(
5180 &self,
5181 ) -> Result<std::sync::MutexGuard<'_, Option<CachedLexicalAssembly>>, LexicalAssemblyError>
5182 {
5183 self.entry.lock().map_err(|_| {
5184 LexicalAssemblyError::Store(StoreError::Synchronization {
5185 component: "lexical index cache",
5186 })
5187 })
5188 }
5189
5190 fn lock_controlled(
5191 &self,
5192 cancellation: Option<&QueryCancellation<'_>>,
5193 ) -> Result<std::sync::MutexGuard<'_, Option<CachedLexicalAssembly>>, LexicalAssemblyError>
5194 {
5195 let Some(cancellation) = cancellation else {
5196 return self.lock();
5197 };
5198 cancellation
5199 .check_graph()
5200 .map_err(LexicalAssemblyError::Cancelled)?;
5201 let guard = loop {
5202 match self.entry.try_lock() {
5203 Ok(guard) => break guard,
5204 Err(std::sync::TryLockError::Poisoned(_)) => {
5205 return Err(LexicalAssemblyError::Store(StoreError::Synchronization {
5206 component: "lexical index cache",
5207 }));
5208 }
5209 Err(std::sync::TryLockError::WouldBlock) => {}
5210 }
5211 cancellation
5212 .check_graph()
5213 .map_err(LexicalAssemblyError::Cancelled)?;
5214 std::thread::park_timeout(Duration::from_millis(1));
5215 };
5216 cancellation
5217 .check_graph()
5218 .map_err(LexicalAssemblyError::Cancelled)?;
5219 Ok(guard)
5220 }
5221}
5222
5223#[derive(Clone, Copy)]
5227struct LexicalInputs<'a> {
5228 generation: u64,
5229 cache: &'a LexicalIndexCache,
5230 snapshot: &'a Arc<PublishedSnapshot>,
5231 active: &'a Arc<crate::ingest::ActiveSegment>,
5232 accounting: &'a Arc<stats::Accounting>,
5233}
5234
5235struct LexicalAssemblyReceipt {
5236 assembly: Arc<LexicalAssembly>,
5237 cache_hit: bool,
5238}
5239
5240impl std::ops::Deref for LexicalAssemblyReceipt {
5241 type Target = LexicalAssembly;
5242 fn deref(&self) -> &Self::Target {
5243 &self.assembly
5244 }
5245}
5246
5247fn assemble_lexical_index(
5248 inputs: LexicalInputs<'_>,
5249 require_document_identity: bool,
5250 cancellation: Option<&QueryCancellation<'_>>,
5251) -> Result<LexicalAssemblyReceipt, LexicalAssemblyError> {
5252 let LexicalInputs {
5253 generation,
5254 cache,
5255 snapshot,
5256 active,
5257 accounting,
5258 } = inputs;
5259 let mut entry = cache.lock_controlled(cancellation)?;
5260 if let Some(cached) = entry.as_mut()
5261 && cached.matches(snapshot, active)
5262 {
5263 if require_document_identity && !cached.document_identity_verified {
5264 verify_document_identity(snapshot, cancellation)?;
5265 cached.document_identity_verified = true;
5266 }
5267 let assembly = Arc::clone(&cached.assembly);
5268 cache.record_hit();
5269 return Ok(LexicalAssemblyReceipt {
5270 assembly,
5271 cache_hit: true,
5272 });
5273 }
5274 let previous = entry.as_ref().map(|cached| Arc::clone(&cached.assembly));
5278 let contributions = previous
5279 .as_ref()
5280 .map_or(&[][..], |a| a.contributions.as_slice());
5281 let (mut assembly, owned_bytes) = build_lexical_assembly(
5282 snapshot,
5283 active,
5284 accounting,
5285 require_document_identity,
5286 cancellation,
5287 cache,
5288 contributions,
5289 )?;
5290 cache.record_build();
5291 assembly
5292 .memory
5293 .set(owned_bytes)
5294 .map_err(LexicalAssemblyError::Store)?;
5295 if let Some(cancellation) = cancellation {
5296 cancellation
5297 .check_graph()
5298 .map_err(LexicalAssemblyError::Cancelled)?;
5299 }
5300 let assembly = Arc::new(assembly);
5301 if entry
5302 .as_ref()
5303 .is_none_or(|cached| generation >= cached.generation)
5304 {
5305 *entry = Some(CachedLexicalAssembly {
5306 generation,
5307 snapshot: Arc::downgrade(snapshot),
5308 active: Arc::downgrade(active),
5309 document_identity_verified: require_document_identity,
5310 assembly: Arc::clone(&assembly),
5311 });
5312 }
5313 drop(entry);
5314 Ok(LexicalAssemblyReceipt {
5315 assembly,
5316 cache_hit: false,
5317 })
5318}
5319
5320fn check_segment_document_identity(
5321 segment: &crate::segment::reader::SegmentReader,
5322 postings: &crate::fts::sealed::SealedSegment,
5323) -> Result<(), LexicalAssemblyError> {
5324 if postings.row_count() != 0
5325 && segment
5326 .document_version(0)
5327 .map_err(StoreError::Segment)
5328 .map_err(LexicalAssemblyError::Store)?
5329 .is_none()
5330 {
5331 return Err(LexicalAssemblyError::MissingDocumentIdentity(
5332 segment.meta().id,
5333 ));
5334 }
5335 Ok(())
5336}
5337
5338fn verify_document_identity(
5344 snapshot: &PublishedSnapshot,
5345 cancellation: Option<&QueryCancellation<'_>>,
5346) -> Result<(), LexicalAssemblyError> {
5347 for segment in snapshot.segments() {
5348 if let Some(cancellation) = cancellation {
5349 cancellation
5350 .check_graph()
5351 .map_err(LexicalAssemblyError::Cancelled)?;
5352 }
5353 if let Some(postings) = segment
5354 .query_postings()
5355 .map_err(LexicalAssemblyError::Store)?
5356 {
5357 check_segment_document_identity(segment, &postings)?;
5358 }
5359 }
5360 Ok(())
5361}
5362
5363fn build_lexical_assembly(
5364 snapshot: &PublishedSnapshot,
5365 active: &Arc<crate::ingest::ActiveSegment>,
5366 accounting: &Arc<stats::Accounting>,
5367 require_document_identity: bool,
5368 cancellation: Option<&QueryCancellation<'_>>,
5369 cache: &LexicalIndexCache,
5370 previous: &[Arc<CachedLexicalContribution>],
5371) -> Result<(LexicalAssembly, usize), LexicalAssemblyError> {
5372 let mut work = crate::fts::control::WorkCheck::new(|| {
5373 cancellation
5374 .map_or(Ok(()), |control| control.check_graph())
5375 .map_err(LexicalAssemblyError::Cancelled)
5376 });
5377 work.check_now()?;
5378 let mut index = crate::fts::index::LexicalIndex::new();
5379 let mut alive_sets = Vec::new();
5380 let mut sources = Vec::new();
5381 let mut contributions = Vec::new();
5382 for (ordinal, segment) in snapshot.segments().iter().enumerate() {
5383 if let Some(cancellation) = cancellation {
5384 cancellation
5385 .check_graph()
5386 .map_err(LexicalAssemblyError::Cancelled)?;
5387 }
5388 if let Some(postings) = segment
5389 .query_postings()
5390 .map_err(LexicalAssemblyError::Store)?
5391 {
5392 if require_document_identity {
5393 check_segment_document_identity(segment, &postings)?;
5394 }
5395 let alive = segment.query_alive().map_err(LexicalAssemblyError::Store)?;
5396 let identity = segment
5397 .lexical_postings_identity()
5398 .map_err(StoreError::Segment)
5399 .map_err(LexicalAssemblyError::Store)?;
5400 let reused = previous
5401 .binary_search_by(|value| match &value.key {
5402 LexicalContributionKey::Sealed(meta, _) => meta.id.cmp(&segment.meta().id),
5403 LexicalContributionKey::Active(_) => std::cmp::Ordering::Greater,
5404 })
5405 .ok()
5406 .and_then(|slot| previous.get(slot))
5407 .filter(|value| {
5408 matches!(&value.key, LexicalContributionKey::Sealed(meta, stored)
5409 if meta.same_segment_file(segment.meta()) && *stored == identity)
5410 && value.statistics.rows == *alive.alive_bitmap()
5411 });
5412 cache.record_contribution(reused.is_some());
5413 let contribution = match reused {
5414 Some(value) => Arc::clone(value),
5415 None => CachedLexicalContribution::build(
5416 LexicalContributionKey::Sealed(segment.meta().clone(), identity),
5417 &postings,
5418 &alive,
5419 index.segments().len(),
5420 accounting,
5421 &mut work,
5422 )?,
5423 };
5424 index.push_shared_with_statistics(postings, Arc::clone(&contribution.statistics));
5425 contributions.push(contribution);
5426 alive_sets.push(alive);
5427 sources.push(StructuredLexicalSource::Sealed(ordinal));
5428 }
5429 }
5430 if active.has_text_controlled(&mut work)? {
5431 let sealed = active.sealed_lexical_controlled(accounting, &mut work)?;
5432 let reused = previous.last().filter(|value| {
5433 matches!(&value.key, LexicalContributionKey::Active(cached)
5434 if cached.upgrade().is_some_and(|cached| Arc::ptr_eq(&cached, active)))
5435 });
5436 cache.record_contribution(reused.is_some());
5437 let contribution = match reused {
5438 Some(value) => Arc::clone(value),
5439 None => {
5440 let alive = Arc::new(active.alive_controlled(&mut work)?);
5441 CachedLexicalContribution::build(
5442 LexicalContributionKey::Active(Arc::downgrade(active)),
5443 &sealed,
5444 &alive,
5445 index.segments().len(),
5446 accounting,
5447 &mut work,
5448 )?
5449 }
5450 };
5451 let alive = contribution
5452 .active_alive
5453 .as_ref()
5454 .ok_or(LexicalAssemblyError::Store(StoreError::Synchronization {
5455 component: "active lexical contribution",
5456 }))?;
5457 index.push_shared_with_statistics(sealed, Arc::clone(&contribution.statistics));
5458 alive_sets.push(Arc::clone(alive));
5459 contributions.push(contribution);
5460 sources.push(StructuredLexicalSource::Active);
5461 }
5462 crate::fts::control::sort_by(
5463 &mut contributions,
5464 |left, right| match (&left.key, &right.key) {
5465 (LexicalContributionKey::Sealed(a, _), LexicalContributionKey::Sealed(b, _)) => {
5466 a.id.cmp(&b.id)
5467 }
5468 (LexicalContributionKey::Active(_), LexicalContributionKey::Active(_)) => {
5469 std::cmp::Ordering::Equal
5470 }
5471 (LexicalContributionKey::Active(_), _) => std::cmp::Ordering::Greater,
5472 (_, LexicalContributionKey::Active(_)) => std::cmp::Ordering::Less,
5473 },
5474 &mut work,
5475 )?;
5476 let assembly = LexicalAssembly {
5477 index,
5478 alive_sets,
5479 sources,
5480 contributions,
5481 vocabulary: Mutex::new(None),
5482 memory: stats::AccountedCounter::new(accounting, stats::AllocationComponent::Cache)
5483 .map_err(LexicalAssemblyError::Store)?,
5484 };
5485 let owned_bytes =
5486 assembly
5487 .owned_bytes()
5488 .ok_or(LexicalAssemblyError::Store(StoreError::BudgetExceeded {
5489 needed: u64::MAX,
5490 budget: u64::MAX,
5491 component: "cache",
5492 }))?;
5493 work.check_now()?;
5494 Ok((assembly, owned_bytes))
5495}
5496
5497fn map_store_lexical_assembly_error(
5498 error: LexicalAssemblyError,
5499) -> crate::ingest::StoreLexicalError {
5500 match error {
5501 LexicalAssemblyError::Store(error) => QueryError::Store(error).into(),
5502 LexicalAssemblyError::Lexical(error) => {
5503 crate::planner::LexicalFilterError::from(error).into()
5504 }
5505 LexicalAssemblyError::MissingDocumentIdentity(segment_id) => {
5506 crate::ingest::StoreLexicalError::MissingDocumentIdentity { segment_id }
5507 }
5508 LexicalAssemblyError::Cancelled(error) => QueryError::Scan(error).into(),
5509 }
5510}
5511
5512fn map_fusion_lexical_assembly_error(error: LexicalAssemblyError) -> crate::fusion::FusionError {
5513 match error {
5514 LexicalAssemblyError::Store(error) => crate::fusion::FusionError::Leg {
5515 leg: crate::fusion::FusionLeg::Lexical,
5516 kind: crate::fusion::LegFailureKind::Store(error.kind()),
5517 detail: error.to_string(),
5518 },
5519 LexicalAssemblyError::Lexical(error) => crate::fusion::FusionError::Leg {
5520 leg: crate::fusion::FusionLeg::Lexical,
5521 kind: crate::fusion::LegFailureKind::Lexical,
5522 detail: error.to_string(),
5523 },
5524 LexicalAssemblyError::MissingDocumentIdentity(segment_id) => {
5525 crate::fusion::FusionError::Leg {
5526 leg: crate::fusion::FusionLeg::Lexical,
5527 kind: crate::fusion::LegFailureKind::Invariant,
5528 detail: format!("sealed lexical segment {segment_id} has no document identity"),
5529 }
5530 }
5531 LexicalAssemblyError::Cancelled(error) => {
5532 crate::fusion::FusionError::from(QueryError::Scan(error))
5533 }
5534 }
5535}
5536
5537fn accumulate_search_counters(
5538 total: &mut crate::fts::search::SearchCounters,
5539 delta: &crate::fts::search::SearchCounters,
5540) {
5541 total.docs_evaluated = total.docs_evaluated.saturating_add(delta.docs_evaluated);
5542 total.postings_decoded = total
5543 .postings_decoded
5544 .saturating_add(delta.postings_decoded);
5545 total.blocks_decoded = total.blocks_decoded.saturating_add(delta.blocks_decoded);
5546 total.blocks_skipped = total.blocks_skipped.saturating_add(delta.blocks_skipped);
5547}
5548
5549fn map_fusion_controlled_lexical_error(
5550 error: crate::fts::search::ControlledSearchError<crate::scan::ScanError>,
5551) -> crate::fusion::FusionError {
5552 match error {
5553 crate::fts::search::ControlledSearchError::Index(error) => {
5554 crate::fusion::FusionError::Leg {
5555 leg: crate::fusion::FusionLeg::Lexical,
5556 kind: crate::fusion::LegFailureKind::Lexical,
5557 detail: error.to_string(),
5558 }
5559 }
5560 crate::fts::search::ControlledSearchError::Control(error) => {
5561 crate::fusion::FusionError::from(QueryError::Scan(error))
5562 }
5563 }
5564}
5565
5566#[derive(Clone, Copy)]
5567enum PinnedLexicalQuery<'a> {
5568 Term(&'a crate::fts::search::TermQuery),
5569 Structured(&'a crate::fts::query::LexicalQuery),
5570}
5571
5572enum LexicalDocumentError {
5573 SourceOrdinalOutOfRange,
5574 RowOverflow,
5575 SealedSourceAbsent,
5576 Segment(crate::segment::SegmentError),
5577 MissingDocumentIdentity(crate::segment::SegmentId),
5578 ActiveRowAbsent,
5579}
5580
5581fn structured_lexical_document(
5582 snapshot: &PublishedSnapshot,
5583 active: &crate::ingest::ActiveSegment,
5584 sources: &[StructuredLexicalSource],
5585 doc: crate::fts::search::GlobalDocId,
5586 require_document_identity: bool,
5587) -> Result<Option<crate::ingest::DocumentVersion>, LexicalDocumentError> {
5588 let source = usize::try_from(doc.segment)
5589 .ok()
5590 .and_then(|slot| sources.get(slot))
5591 .ok_or(LexicalDocumentError::SourceOrdinalOutOfRange)?;
5592 let row = usize::try_from(doc.row).map_err(|_| LexicalDocumentError::RowOverflow)?;
5593 match source {
5594 StructuredLexicalSource::Sealed(ordinal) => {
5595 let segment = snapshot
5596 .segments()
5597 .get(*ordinal)
5598 .ok_or(LexicalDocumentError::SealedSourceAbsent)?;
5599 let document = segment
5600 .document_version(row)
5601 .map_err(LexicalDocumentError::Segment)?;
5602 if require_document_identity && document.is_none() {
5603 return Err(LexicalDocumentError::MissingDocumentIdentity(
5604 segment.meta().id,
5605 ));
5606 }
5607 Ok(document)
5608 }
5609 StructuredLexicalSource::Active => {
5610 let document = active.document(row);
5611 if require_document_identity && document.is_none() {
5612 return Err(LexicalDocumentError::ActiveRowAbsent);
5613 }
5614 Ok(document)
5615 }
5616 }
5617}
5618
5619fn map_store_lexical_document_error(
5620 error: LexicalDocumentError,
5621) -> crate::ingest::StoreLexicalError {
5622 match error {
5623 LexicalDocumentError::SourceOrdinalOutOfRange
5624 | LexicalDocumentError::RowOverflow
5625 | LexicalDocumentError::SealedSourceAbsent
5626 | LexicalDocumentError::ActiveRowAbsent => {
5627 QueryError::Store(StoreError::ActiveRowOverflow).into()
5628 }
5629 LexicalDocumentError::Segment(error) => {
5630 QueryError::Store(StoreError::Segment(error)).into()
5631 }
5632 LexicalDocumentError::MissingDocumentIdentity(segment_id) => {
5633 crate::ingest::StoreLexicalError::MissingDocumentIdentity { segment_id }
5634 }
5635 }
5636}
5637
5638fn map_fusion_lexical_document_error(
5639 error: LexicalDocumentError,
5640 invariant_kind: crate::fusion::LegFailureKind,
5641 segment_kind: crate::fusion::LegFailureKind,
5642) -> crate::fusion::FusionError {
5643 let (kind, detail) = match error {
5644 LexicalDocumentError::SourceOrdinalOutOfRange => (
5645 invariant_kind,
5646 "lexical source ordinal is out of range".to_owned(),
5647 ),
5648 LexicalDocumentError::RowOverflow => {
5649 (invariant_kind, "lexical row exceeds usize".to_owned())
5650 }
5651 LexicalDocumentError::SealedSourceAbsent => {
5652 (invariant_kind, "lexical sealed source is absent".to_owned())
5653 }
5654 LexicalDocumentError::Segment(error) => (segment_kind, error.to_string()),
5655 LexicalDocumentError::MissingDocumentIdentity(segment_id) => (
5656 invariant_kind,
5657 format!("sealed lexical segment {segment_id} has no document identity"),
5658 ),
5659 LexicalDocumentError::ActiveRowAbsent => (
5660 invariant_kind,
5661 "lexical active source row is absent".to_owned(),
5662 ),
5663 };
5664 crate::fusion::FusionError::Leg {
5665 leg: crate::fusion::FusionLeg::Lexical,
5666 kind,
5667 detail,
5668 }
5669}
5670
5671fn map_structured_leg_document_error(error: LexicalDocumentError) -> crate::fusion::FusionError {
5672 map_fusion_lexical_document_error(
5673 error,
5674 crate::fusion::LegFailureKind::Lexical,
5675 crate::fusion::LegFailureKind::Lexical,
5676 )
5677}
5678
5679fn map_term_leg_document_error(error: LexicalDocumentError) -> crate::fusion::FusionError {
5680 map_fusion_lexical_document_error(
5681 error,
5682 crate::fusion::LegFailureKind::Invariant,
5683 crate::fusion::LegFailureKind::Segment,
5684 )
5685}
5686
5687fn structured_lexical_row<'a>(
5688 snapshot: &'a PublishedSnapshot,
5689 active: &'a crate::ingest::ActiveSegment,
5690 sources: &[StructuredLexicalSource],
5691 doc: crate::fts::search::GlobalDocId,
5692) -> Result<(&'a str, crate::ingest::DocumentVersion), crate::ingest::StoreLexicalError> {
5693 let source = usize::try_from(doc.segment)
5694 .ok()
5695 .and_then(|slot| sources.get(slot))
5696 .ok_or(QueryError::Store(StoreError::ActiveRowOverflow))?;
5697 let row =
5698 usize::try_from(doc.row).map_err(|_| QueryError::Store(StoreError::ActiveRowOverflow))?;
5699 match source {
5700 StructuredLexicalSource::Sealed(ordinal) => {
5701 let segment = snapshot
5702 .segments()
5703 .get(*ordinal)
5704 .ok_or(QueryError::Store(StoreError::ActiveRowOverflow))?;
5705 let text = segment
5706 .query_stored_text()
5707 .map_err(StoreError::Segment)
5708 .map_err(QueryError::Store)?
5709 .and_then(|rows| rows.row(row).flatten())
5710 .ok_or(crate::ingest::StoreLexicalError::MissingStoredText {
5711 segment_id: segment.meta().id,
5712 row: doc.row,
5713 })?;
5714 let document = segment
5715 .document_version(row)
5716 .map_err(StoreError::Segment)
5717 .map_err(QueryError::Store)?
5718 .ok_or(crate::ingest::StoreLexicalError::MissingDocumentIdentity {
5719 segment_id: segment.meta().id,
5720 })?;
5721 Ok((text, document))
5722 }
5723 StructuredLexicalSource::Active => {
5724 let text = active
5725 .text(row)
5726 .map_err(QueryError::Store)?
5727 .ok_or(QueryError::Store(StoreError::ActiveRowOverflow))?;
5728 let document = active
5729 .document(row)
5730 .ok_or(QueryError::Store(StoreError::ActiveRowOverflow))?;
5731 Ok((text, document))
5732 }
5733 }
5734}
5735
5736type ExactLexicalLeg = (
5737 Vec<hybrid::LexicalHit>,
5738 Arc<LexicalAssembly>,
5739 crate::fts::search::SearchCounters,
5740 bool,
5741);
5742
5743fn resolve_hybrid_leg_results<Vector, Lexical>(
5744 vector: Result<Vector, crate::fusion::FusionError>,
5745 lexical: Result<Lexical, crate::fusion::FusionError>,
5746) -> Result<(Vector, Lexical), crate::fusion::FusionError> {
5747 fn control_priority(error: &crate::fusion::FusionError) -> Option<u8> {
5748 match error {
5749 crate::fusion::FusionError::ReadCancelled { .. } => Some(0),
5750 crate::fusion::FusionError::Timeout { .. } => Some(1),
5751 crate::fusion::FusionError::Cancelled { .. } => Some(2),
5752 _ => None,
5753 }
5754 }
5755
5756 match (vector, lexical) {
5757 (Ok(vector), Ok(lexical)) => Ok((vector, lexical)),
5758 (Err(vector), Ok(_)) => Err(vector),
5759 (Ok(_), Err(lexical)) => Err(lexical),
5760 (Err(vector), Err(lexical)) => {
5761 match (control_priority(&vector), control_priority(&lexical)) {
5762 (Some(vector_priority), Some(lexical_priority)) => {
5763 if vector_priority <= lexical_priority {
5764 Err(vector)
5765 } else {
5766 Err(lexical)
5767 }
5768 }
5769 (Some(_), None) => Err(vector),
5770 (None, Some(_)) => Err(lexical),
5771 (None, None) => Err(vector),
5772 }
5773 }
5774 }
5775}
5776
5777fn exact_structured_lexical_leg(
5778 inputs: LexicalInputs<'_>,
5779 _analyzer: &crate::fts::tokenizer::Analyzer,
5780 query: &crate::fts::query::LexicalQuery,
5781 bound: usize,
5782 cancellation: &QueryCancellation<'_>,
5783 preparation: &mut prepared_lexical::PreparedLexicalQuery<'_>,
5784) -> Result<ExactLexicalLeg, crate::fusion::FusionError> {
5785 let (snapshot, active) = (inputs.snapshot, inputs.active);
5786 let lexical_error = |detail: String| crate::fusion::FusionError::Leg {
5787 leg: crate::fusion::FusionLeg::Lexical,
5788 kind: crate::fusion::LegFailureKind::Lexical,
5789 detail,
5790 };
5791 let (prepared, cache_hit) = preparation.prepare_structured(inputs, cancellation)?;
5792 let assembly = &prepared.assembly;
5793 let index = &assembly.index;
5794 let sources = &assembly.sources;
5795 let Some(query_scoring) = prepared.query.as_ref() else {
5796 return Ok((
5797 Vec::new(),
5798 Arc::clone(assembly),
5799 crate::fts::search::SearchCounters::default(),
5800 cache_hit,
5801 ));
5802 };
5803 let allow_lists = assembly
5804 .alive_sets
5805 .iter()
5806 .map(|alive| alive.alive_bitmap())
5807 .collect::<Vec<_>>();
5808 let mut scratch_memory =
5809 stats::AccountedCounter::new(inputs.accounting, stats::AllocationComponent::Temporary)
5810 .map_err(QueryError::Store)?;
5811 let mut phrase = prepared_lexical::PhraseEligibility::new(
5812 index,
5813 Some(query),
5814 inputs.accounting,
5815 cancellation,
5816 )?;
5817 let result = crate::fts::prune::weighted::search(
5818 index,
5819 query_scoring,
5820 bound,
5821 &allow_lists,
5822 || {
5823 cancellation
5824 .check_graph()
5825 .map_err(QueryError::Scan)
5826 .map_err(crate::fusion::FusionError::from)
5827 },
5828 |doc| {
5829 phrase
5830 .matches(doc, cancellation)
5831 .map_err(crate::fusion::FusionError::from)
5832 },
5833 |bytes| {
5834 prepared_lexical::reserve_weighted_scratch(&mut scratch_memory, bytes)
5835 .map_err(crate::fusion::FusionError::from)
5836 },
5837 )
5838 .map_err(|error| match error {
5839 crate::fts::search::ControlledSearchError::Index(error) => lexical_error(error.to_string()),
5840 crate::fts::search::ControlledSearchError::Control(error) => error,
5841 })?;
5842 let counters = result.counters;
5843 let scored = result.hits;
5844 let mut joined = Vec::with_capacity(scored.len());
5845 for (doc, score) in scored {
5846 let document = structured_lexical_document(snapshot, active, sources, doc, false)
5847 .map_err(map_structured_leg_document_error)?;
5848 joined.push(hybrid::LexicalHit {
5849 doc,
5850 document,
5851 bm25: score,
5852 });
5853 }
5854 Ok((joined, Arc::clone(assembly), counters, cache_hit))
5855}
5856
5857fn exact_lexical_leg(
5858 inputs: LexicalInputs<'_>,
5859 bound: usize,
5860 cancellation: &QueryCancellation<'_>,
5861 preparation: &mut prepared_lexical::PreparedLexicalQuery<'_>,
5862) -> Result<ExactLexicalLeg, crate::fusion::FusionError> {
5863 let (snapshot, active) = (inputs.snapshot, inputs.active);
5864 let (prepared, cache_hit) = preparation.prepare_term(inputs, cancellation)?;
5865 let assembly = &prepared.assembly;
5866 let index = &assembly.index;
5867 let sources = &assembly.sources;
5868 if index.segments().is_empty() {
5869 return Ok((
5870 Vec::new(),
5871 Arc::clone(assembly),
5872 crate::fts::search::SearchCounters::default(),
5873 cache_hit,
5874 ));
5875 }
5876 let k = bound.min(usize::try_from(index.document_count()).unwrap_or(usize::MAX));
5877 let allow_lists = assembly
5878 .alive_sets
5879 .iter()
5880 .map(|alive| alive.alive_bitmap())
5881 .collect::<Vec<_>>();
5882 let result =
5883 exact_hybrid_lexical_search(index, prepared.scoring()?, k, &allow_lists, cancellation)?;
5884 let mut joined = Vec::with_capacity(result.hits.len());
5885 for hit in result.hits {
5886 let document = structured_lexical_document(snapshot, active, sources, hit.doc, false)
5887 .map_err(map_term_leg_document_error)?;
5888 joined.push(hybrid::LexicalHit {
5889 doc: hit.doc,
5890 document,
5891 bm25: hit.score,
5892 });
5893 }
5894 Ok((joined, Arc::clone(assembly), result.counters, cache_hit))
5895}
5896
5897fn exact_hybrid_lexical_search(
5898 index: &crate::fts::index::LexicalIndex,
5899 prepared: &crate::fts::search::PreparedTermQuery,
5900 k: usize,
5901 allow_lists: &[&crate::meta::DocBitmap],
5902 cancellation: &QueryCancellation<'_>,
5903) -> Result<crate::fts::search::SearchResult, crate::fusion::FusionError> {
5904 let allowed = allow_lists
5905 .iter()
5906 .map(|allow_list| allow_list.cardinality())
5907 .fold(0_u64, u64::saturating_add);
5908 if allowed.saturating_mul(crate::planner::LEXICAL_ALLOW_LIST_DIVISOR) <= index.document_count()
5909 {
5910 return crate::fts::search::search_allow_list_prepared_controlled(
5911 index,
5912 prepared,
5913 k,
5914 allow_lists,
5915 || cancellation.check_graph(),
5916 )
5917 .map_err(map_fusion_controlled_lexical_error);
5918 }
5919 crate::fts::prune::search_pruned_prepared_controlled(index, prepared, k, allow_lists, || {
5920 cancellation.check_graph()
5921 })
5922 .map_err(map_fusion_controlled_lexical_error)
5923}
5924
5925#[allow(clippy::too_many_arguments)]
5926fn search_pinned(
5927 pool: Option<&pool::QueryPool>,
5928 snapshot: &Arc<PublishedSnapshot>,
5929 active: &crate::ingest::ActiveSegment,
5930 accounting: &Arc<stats::Accounting>,
5931 generation: u64,
5932 epoch: Option<crate::epoch::EpochIdentity>,
5933 request: crate::ingest::SearchRequest<'_>,
5934 k: usize,
5935 options: SearchOptions,
5936 control: QueryControl,
5937 graph_bound_mode: GraphBoundMode,
5938 requested_tier: Option<SearchTier>,
5939 hybrid_scan_reason: Option<crate::planner::ScanReason>,
5940 started: std::time::Instant,
5941 clock: &dyn MonotonicClock,
5942 vector_preparation: &mut prepared::PreparedVectorQuery<'_>,
5943 #[cfg(any(test, feature = "test-support"))] vector_fault_controller: Option<
5944 &crate::scan::vector_fault::VectorFaultController,
5945 >,
5946) -> Result<crate::ingest::SearchOutcome, QueryError> {
5947 use crate::ingest::{GraphSearchStats, RowSource, SearchOutcome};
5948 use crate::quant::{QuantError, prepare_bit4_query};
5949 use crate::scan::{ScanQuery, ScanRequest, ScanRows, ScanStats};
5950
5951 let vector_started = timing_start(clock);
5952 let scan_options = options.scan();
5953 let rescore_options = options.scan_rescore();
5954 let mut rescore_counters = crate::diag::ScanRescoreCounters::default();
5955 if let Some(options) = rescore_options {
5956 options.frontier(k, 0)?;
5957 }
5958 let query = request.vector();
5959 let quantized_query_validation = if query.is_empty() {
5960 Err(QuantError::EmptyVector)
5961 } else if query.len() > crate::kernels::MAX_DOT_I8_DIMENSION {
5962 Err(QuantError::DimensionTooLarge {
5963 actual: query.len(),
5964 maximum: crate::kernels::MAX_DOT_I8_DIMENSION,
5965 })
5966 } else if let Some((index, _)) = query
5967 .iter()
5968 .enumerate()
5969 .find(|(_, value)| !value.is_finite())
5970 {
5971 Err(QuantError::NonFinite { index })
5972 } else {
5973 Ok(())
5974 };
5975 quantized_query_validation
5976 .map_err(crate::scan::ScanError::Quant)
5977 .map_err(QueryError::Scan)?;
5978 vector_preparation.validate_binding(query, epoch)?;
5979 let bit4_query = std::cell::OnceCell::new();
5980 let int8_query = std::cell::OnceCell::new();
5981 let mut candidates = Vec::new();
5982 let mut dims_touched = 0_u64;
5983 let mut bytes_read = 0_u64;
5984 let mut worst_squared_l2: Option<f64> = None;
5985 let mut worst_exhaustive = true;
5986 let mut worker_thread_ids = Vec::new();
5987 let mut graph_stats = GraphSearchStats::default();
5988 let mut plans = Vec::new();
5989 let auto_graph_options = if matches!(options.tier(), SearchTier::Auto)
5990 && snapshot.segments().iter().any(|segment| {
5991 segment
5992 .directory()
5993 .iter()
5994 .any(|entry| entry.kind == crate::segment::layout::RegionKind::GraphNodeBlocks.id())
5995 }) {
5996 Some(auto_graph_search_options(snapshot)?)
5997 } else {
5998 None
5999 };
6000
6001 if matches!(options.tier(), SearchTier::Graph(_)) {
6002 for segment in snapshot.segments() {
6003 if !segment
6004 .directory()
6005 .iter()
6006 .any(|entry| entry.kind == crate::segment::layout::RegionKind::GraphNodeBlocks.id())
6007 {
6008 return Err(QueryError::Store(StoreError::GraphUnavailable {
6009 segment_id: segment.meta().id,
6010 }));
6011 }
6012 }
6013 }
6014 let full_precision = matches!(options.tier(), SearchTier::Exact | SearchTier::Graph(_))
6015 || (matches!(options.tier(), SearchTier::Auto)
6016 && auto_uses_full_precision(snapshot, active));
6017
6018 let exact_lane = if full_precision {
6020 pool.map_or(0, |pool| {
6021 pool.reserve_exact_lane(scan_options.thread_budget)
6022 })
6023 } else {
6024 0
6025 };
6026
6027 if !active.is_empty() {
6028 let alive = active.alive().map_err(QueryError::Store)?;
6029 let mut exact_score = full_precision;
6030 let mut approximate = false;
6031 let frontier = rescore_options
6032 .map(|options| options.frontier(k, alive.live_count()))
6033 .transpose()?
6034 .unwrap_or(k);
6035 let mut exact_memory = ExactScanMemory::new(accounting)?;
6036 let mut outcome = match options.tier() {
6037 SearchTier::Auto if full_precision => {
6038 let lease = SnapshotLease::new_at(Arc::clone(snapshot), generation);
6039 scan_active_squared_l2(
6040 active,
6041 &alive,
6042 request.vector(),
6043 frontier,
6044 ExactScanDispatch {
6045 pool,
6046 accounting,
6047 control: &control,
6048 lease: &lease,
6049 thread_budget: scan_options.thread_budget,
6050 lane: exact_lane,
6051 },
6052 &mut exact_memory,
6053 #[cfg(any(test, feature = "test-support"))]
6054 vector_fault_controller,
6055 #[cfg(any(test, feature = "test-support"))]
6056 vector_fault_tier(options.tier()),
6057 )?
6058 }
6059 SearchTier::Auto | SearchTier::Scan => {
6060 let query_pool = pool.ok_or(QueryError::Store(StoreError::Synchronization {
6061 component: "active scan-tier query pool",
6062 }))?;
6063 let bit4_query = bit4_query
6064 .get_or_init(|| prepare_bit4_query(request.vector(), 0))
6065 .as_ref()
6066 .map_err(|error| crate::scan::ScanError::Quant(error.clone()))
6067 .map_err(QueryError::Scan)?;
6068 execute_store_scan(
6069 query_pool,
6070 ScanRequest {
6071 query: ScanQuery::Bit4(bit4_query),
6072 rows: ScanRows::Bit4RowMajor {
6073 codes: active.codes(),
6074 factors: active.factors(),
6075 },
6076 row_mask: Some(alive.scan_mask()),
6077 },
6078 frontier,
6079 scan_options,
6080 control.clone(),
6081 SnapshotLease::new_at(Arc::clone(snapshot), generation),
6082 #[cfg(any(test, feature = "test-support"))]
6083 vector_fault_controller,
6084 #[cfg(any(test, feature = "test-support"))]
6085 crate::scan::vector_fault::VectorRowSource::Active,
6086 #[cfg(any(test, feature = "test-support"))]
6087 vector_fault_tier(options.tier()),
6088 )?
6089 }
6090 SearchTier::Exact | SearchTier::Graph(_) => {
6091 let lease = SnapshotLease::new_at(Arc::clone(snapshot), generation);
6092 scan_active_squared_l2(
6093 active,
6094 &alive,
6095 request.vector(),
6096 frontier,
6097 ExactScanDispatch {
6098 pool,
6099 accounting,
6100 control: &control,
6101 lease: &lease,
6102 thread_budget: scan_options.thread_budget,
6103 lane: exact_lane,
6104 },
6105 &mut exact_memory,
6106 #[cfg(any(test, feature = "test-support"))]
6107 vector_fault_controller,
6108 #[cfg(any(test, feature = "test-support"))]
6109 vector_fault_tier(options.tier()),
6110 )?
6111 }
6112 };
6113 if let Some(rescore) = rescore_options {
6114 let lease = SnapshotLease::new_at(Arc::clone(snapshot), generation);
6115 let cancellation = QueryCancellation::new(&control, &lease);
6116 let (counters, exhaustive, worst) = rescored_scan::rescore(
6117 &mut outcome.candidates,
6118 &mut outcome.stats,
6119 active.vectors(),
6120 query,
6121 k,
6122 alive.live_count(),
6123 rescore,
6124 &cancellation,
6125 accounting,
6126 &mut exact_memory,
6127 )?;
6128 outcome.worst_score = worst;
6129 exact_score = true;
6130 approximate = !exhaustive;
6131 rescored_scan::accumulate(&mut rescore_counters, counters)?;
6132 }
6133 fold_worst_squared_l2(&mut worst_squared_l2, &mut worst_exhaustive, &outcome);
6134 merge_store_outcome(
6135 outcome,
6136 RowSource::Active,
6137 |row| Ok(active.document(row)),
6138 &mut candidates,
6139 &mut dims_touched,
6140 &mut bytes_read,
6141 &mut worker_thread_ids,
6142 exact_score,
6143 #[cfg(any(test, feature = "test-support"))]
6144 vector_fault_controller,
6145 #[cfg(any(test, feature = "test-support"))]
6146 vector_fault_tier(options.tier()),
6147 )?;
6148 let mut plan = crate::planner::SegmentPlan::unfiltered_scan(
6149 RowSource::Active,
6150 crate::planner::SegmentTier::ActiveScan,
6151 alive.live_count(),
6152 crate::planner::ScanReason::ActiveSegment,
6153 );
6154 plan.approximate = approximate;
6155 plans.push(plan);
6156 if matches!(graph_bound_mode, GraphBoundMode::Shared) {
6157 retain_global_top_k(&mut candidates, k);
6158 }
6159 }
6160
6161 let mut ordered_segments = snapshot.segments().iter().collect::<Vec<_>>();
6162 if matches!(graph_bound_mode, GraphBoundMode::Shared) {
6163 ordered_segments.sort_unstable_by(|left, right| {
6167 right
6168 .meta()
6169 .row_count
6170 .cmp(&left.meta().row_count)
6171 .then_with(|| left.meta().id.cmp(&right.meta().id))
6172 });
6173 }
6174
6175 for segment in ordered_segments {
6176 let alive = segment.query_alive().map_err(QueryError::Store)?;
6177 let source = RowSource::Sealed(segment.meta().id);
6178 let graph_options = match options.tier() {
6183 SearchTier::Graph(graph_options) => Some(graph_options),
6184 SearchTier::Auto
6185 if segment.directory().iter().any(|entry| {
6186 entry.kind == crate::segment::layout::RegionKind::GraphNodeBlocks.id()
6187 }) =>
6188 {
6189 auto_graph_options
6190 }
6191 SearchTier::Auto | SearchTier::Exact | SearchTier::Scan => None,
6192 };
6193 if let Some(graph_options) = graph_options {
6194 worst_exhaustive = false;
6195 let competitive_distance = global_competitive_distance(&candidates, k);
6196 let MergedSegmentGraph {
6197 plan,
6198 traversed,
6199 _result,
6200 _scratch,
6201 _lease,
6202 } = traverse_segment_graph(
6203 segment,
6204 &alive,
6205 source,
6206 graph_options,
6207 graph_bound_mode,
6208 request,
6209 k,
6210 vector_preparation,
6211 accounting,
6212 snapshot,
6213 generation,
6214 &control,
6215 competitive_distance,
6216 #[cfg(any(test, feature = "test-support"))]
6217 vector_fault_controller,
6218 options.tier(),
6219 )?
6220 .merge_into(
6221 segment,
6222 &alive,
6223 source,
6224 &mut candidates,
6225 &mut dims_touched,
6226 &mut bytes_read,
6227 &mut worker_thread_ids,
6228 &mut graph_stats,
6229 #[cfg(any(test, feature = "test-support"))]
6230 vector_fault_controller,
6231 options.tier(),
6232 )?;
6233 plans.push(plan);
6234 if traversed && matches!(graph_bound_mode, GraphBoundMode::Shared) {
6235 retain_global_top_k(&mut candidates, k);
6236 }
6237 continue;
6238 }
6239
6240 let mut exact_memory = ExactScanMemory::new(accounting)?;
6241 let frontier = rescore_options
6242 .map(|options| options.frontier(k, alive.live_count()))
6243 .transpose()?
6244 .unwrap_or(k);
6245 let mut outcome = scan_sealed_segment(
6246 pool,
6247 snapshot,
6248 segment,
6249 &alive,
6250 generation,
6251 request,
6252 frontier,
6253 scan_options,
6254 &control,
6255 &bit4_query,
6256 &int8_query,
6257 full_precision,
6258 exact_lane,
6259 accounting,
6260 &mut exact_memory,
6261 source,
6262 options.tier(),
6263 #[cfg(any(test, feature = "test-support"))]
6264 vector_fault_controller,
6265 )?;
6266 let mut exact_score = full_precision || segment.meta().scheme == 0;
6267 let mut approximate = false;
6268 if let Some(rescore) = rescore_options {
6269 let lease = SnapshotLease::new_at(Arc::clone(snapshot), generation);
6270 let cancellation = QueryCancellation::new(&control, &lease);
6271 let vectors = exact_rescore_rows_for_search(
6272 segment,
6273 #[cfg(any(test, feature = "test-support"))]
6274 vector_fault_controller,
6275 #[cfg(any(test, feature = "test-support"))]
6276 vector_fault_tier(options.tier()),
6277 )?;
6278 let (counters, exhaustive, worst) = rescored_scan::rescore(
6279 &mut outcome.candidates,
6280 &mut outcome.stats,
6281 vectors,
6282 query,
6283 k,
6284 alive.live_count(),
6285 rescore,
6286 &cancellation,
6287 accounting,
6288 &mut exact_memory,
6289 )?;
6290 outcome.worst_score = worst;
6291 exact_score = true;
6292 approximate = !exhaustive;
6293 rescored_scan::accumulate(&mut rescore_counters, counters)?;
6294 }
6295 fold_worst_squared_l2(&mut worst_squared_l2, &mut worst_exhaustive, &outcome);
6296 merge_store_outcome(
6297 outcome,
6298 source,
6299 |row| {
6300 segment
6301 .document_version(row)
6302 .map_err(StoreError::Segment)
6303 .map_err(QueryError::Store)
6304 },
6305 &mut candidates,
6306 &mut dims_touched,
6307 &mut bytes_read,
6308 &mut worker_thread_ids,
6309 exact_score,
6310 #[cfg(any(test, feature = "test-support"))]
6311 vector_fault_controller,
6312 #[cfg(any(test, feature = "test-support"))]
6313 vector_fault_tier(options.tier()),
6314 )?;
6315 let tier = if segment
6316 .directory()
6317 .iter()
6318 .any(|entry| entry.kind == crate::segment::layout::RegionKind::GraphNodeBlocks.id())
6319 {
6320 crate::planner::SegmentTier::SealedGraph
6321 } else {
6322 crate::planner::SegmentTier::SealedScan
6323 };
6324 let mut plan = crate::planner::SegmentPlan::unfiltered_scan(
6325 source,
6326 tier,
6327 alive.live_count(),
6328 sealed_scan_reason(segment, requested_tier, hybrid_scan_reason),
6329 );
6330 plan.approximate = approximate;
6331 plans.push(plan);
6332 if matches!(graph_bound_mode, GraphBoundMode::Shared) {
6333 retain_global_top_k(&mut candidates, k);
6334 }
6335 }
6336
6337 retain_global_top_k(&mut candidates, k);
6338 let stats = ScanStats {
6339 dims_touched,
6340 bytes_read,
6341 threads_used: worker_thread_ids.len(),
6342 worker_thread_ids,
6343 };
6344 accounting.record_plans(&plans).map_err(QueryError::Store)?;
6345 let vector_ceiling = Some(exact_vector_ceiling(snapshot, active, query)?);
6348 let mut diagnostics = crate::diag::QueryDiagnostics::vector(crate::diag::VectorDiagnostics {
6349 snapshot_generation: generation,
6350 indexed_through_seq: active
6351 .indexed_through_seq()
6352 .max(crate::wal::LogSeq::new(snapshot.absorbed_through())),
6353 approximate: plans.iter().any(|plan| plan.approximate),
6354 exact_rescore: candidates.iter().all(|candidate| candidate.exact_score()),
6355 requested_k: k,
6356 returned: candidates.len(),
6357 budget_exhausted: false,
6358 plan: plans,
6359 scan: stats.clone(),
6360 graph: graph_stats,
6361 epoch,
6362 elapsed: clock.now().saturating_duration_since(started),
6363 });
6364 diagnostics.scan_rescore = rescore_options.map(|_| rescore_counters);
6365 diagnostics.timings = crate::diag::QUERY_TIMING_ENABLED.then(|| crate::diag::QueryTimings {
6366 vector: timing_elapsed(clock, vector_started),
6367 ..Default::default()
6368 });
6369 Ok(SearchOutcome {
6370 candidates,
6371 vector_ceiling,
6376 stats,
6377 graph_stats,
6378 generation,
6379 epoch,
6380 diagnostics,
6381 })
6382}
6383
6384const GRAPH_WIDENING_CAP_NUMERATOR: usize = 1;
6392const GRAPH_WIDENING_CAP_DENOMINATOR: usize = 4;
6394
6395fn graph_round_is_worth_it(k: usize, corpus_rows: usize) -> bool {
6402 if corpus_rows == 0 {
6403 return true;
6404 }
6405 let ef = graph_round_ef(k);
6406 let cap =
6407 corpus_rows.saturating_mul(GRAPH_WIDENING_CAP_NUMERATOR) / GRAPH_WIDENING_CAP_DENOMINATOR;
6408 ef <= cap
6409}
6410
6411fn graph_round_ef(k: usize) -> usize {
6412 k.saturating_mul(4).max(140)
6413}
6414
6415fn sealed_scan_reason(
6416 segment: &crate::segment::reader::SegmentReader,
6417 requested_tier: Option<SearchTier>,
6418 hybrid_scan_reason: Option<crate::planner::ScanReason>,
6419) -> crate::planner::ScanReason {
6420 use crate::planner::{ExplicitScanTier, ScanReason};
6421
6422 let has_graph = segment
6423 .directory()
6424 .iter()
6425 .any(|entry| entry.kind == crate::segment::layout::RegionKind::GraphNodeBlocks.id());
6426 if has_graph && let Some(reason) = hybrid_scan_reason {
6427 return reason;
6428 }
6429 match requested_tier {
6430 Some(SearchTier::Exact) => return ScanReason::ExplicitTier(ExplicitScanTier::Exact),
6431 Some(SearchTier::Scan) => return ScanReason::ExplicitTier(ExplicitScanTier::Scan),
6432 Some(SearchTier::Auto | SearchTier::Graph(_)) | None => {}
6433 }
6434 let rows = segment.meta().row_count;
6435 let min_rows = crate::tier::thresholds::for_bucket(segment.meta().dims, segment.meta().scheme)
6436 .graph_min_rows;
6437 if rows < min_rows {
6438 ScanReason::BelowGraphThreshold { rows, min_rows }
6439 } else {
6440 ScanReason::GraphPending
6441 }
6442}
6443
6444fn snapshot_has_graph(snapshot: &PublishedSnapshot) -> bool {
6445 snapshot.segments().iter().any(|segment| {
6446 segment
6447 .directory()
6448 .iter()
6449 .any(|entry| entry.kind == crate::segment::layout::RegionKind::GraphNodeBlocks.id())
6450 })
6451}
6452
6453fn exact_vector_ceiling(
6454 snapshot: &PublishedSnapshot,
6455 active: &crate::ingest::ActiveSegment,
6456 query: &[f32],
6457) -> Result<f64, QueryError> {
6458 fold_exact_vector_ceiling(
6459 snapshot,
6460 active,
6461 query,
6462 remembered_segment_vector_ceiling_norm_range,
6463 )
6464}
6465
6466fn remembered_segment_vector_ceiling_norm_range(
6473 segment: &crate::segment::reader::SegmentReader,
6474) -> Result<crate::graph::search::GraphSegmentNormRange, QueryError> {
6475 if let Some(range) = segment.cached_vector_ceiling_norm_range() {
6476 return Ok(range);
6477 }
6478 let range = compute_segment_vector_ceiling_norm_range(segment)?;
6479 segment.remember_vector_ceiling_norm_range(range);
6480 Ok(range)
6481}
6482
6483fn fold_exact_vector_ceiling(
6489 snapshot: &PublishedSnapshot,
6490 active: &crate::ingest::ActiveSegment,
6491 query: &[f32],
6492 mut segment_range: impl FnMut(
6493 &crate::segment::reader::SegmentReader,
6494 )
6495 -> Result<crate::graph::search::GraphSegmentNormRange, QueryError>,
6496) -> Result<f64, QueryError> {
6497 let mut ceiling = 0.0_f64;
6498 if !active.is_empty() {
6499 ceiling = fold_vector_ceiling(
6502 ceiling,
6503 crate::graph::search::GraphSegmentNormRange::from_factors(active.factors()),
6504 query,
6505 )?;
6506 }
6507 for segment in snapshot.segments() {
6508 let range = segment_range(segment)?;
6509 ceiling = fold_vector_ceiling(ceiling, range, query)?;
6510 }
6511 Ok(ceiling)
6512}
6513
6514fn compute_segment_vector_ceiling_norm_range(
6516 segment: &crate::segment::reader::SegmentReader,
6517) -> Result<crate::graph::search::GraphSegmentNormRange, QueryError> {
6518 match segment.meta().scheme {
6519 4 => Ok(crate::graph::search::GraphSegmentNormRange::from_factors(
6520 segment
6521 .query_bit4_factors()
6522 .map_err(StoreError::Segment)
6523 .map_err(QueryError::Store)?,
6524 )),
6525 0 | 2 => Ok(
6526 crate::graph::search::GraphSegmentNormRange::from_exact_rows(
6527 segment
6528 .query_rescore_f32()
6529 .map_err(StoreError::Segment)
6530 .map_err(QueryError::Store)?,
6531 segment.meta().dims as usize,
6532 ),
6533 ),
6534 scheme => Err(QueryError::Store(StoreError::Segment(
6535 crate::segment::SegmentError::Geometry(format!(
6536 "vector norm enclosure does not support scheme {scheme}"
6537 )),
6538 ))),
6539 }
6540}
6541
6542fn fold_vector_ceiling(
6543 current: f64,
6544 range: crate::graph::search::GraphSegmentNormRange,
6545 query: &[f32],
6546) -> Result<f64, QueryError> {
6547 let ceiling = range.squared_l2_upper_bound(query);
6548 if !ceiling.is_finite() {
6549 return Err(QueryError::Store(StoreError::Segment(
6550 crate::segment::SegmentError::Geometry(
6551 "vector norm enclosure did not produce a finite squared-L2 ceiling".to_owned(),
6552 ),
6553 )));
6554 }
6555 Ok(current.max(ceiling))
6556}
6557
6558fn fold_worst_squared_l2(
6562 worst: &mut Option<f64>,
6563 exhaustive: &mut bool,
6564 outcome: &crate::scan::ScanOutcome,
6565) {
6566 match outcome.worst_score {
6567 Some(score) => {
6568 let squared_l2 = -f64::from(score);
6569 *worst = Some(worst.map_or(squared_l2, |current: f64| current.max(squared_l2)));
6570 }
6571 None => {
6572 if !outcome.candidates.is_empty() {
6573 *exhaustive = false;
6574 }
6575 }
6576 }
6577}
6578
6579#[allow(clippy::large_enum_variant)]
6581enum SegmentGraphResult<'a> {
6582 Pruned {
6583 plan: crate::planner::SegmentPlan,
6584 graph_stats: crate::ingest::GraphSearchStats,
6585 _lease: SnapshotLease,
6586 },
6587 Traversed {
6588 result: crate::graph::search::GraphSearchResult,
6589 plan: crate::planner::SegmentPlan,
6590 dims_touched: u64,
6591 bytes_read: u64,
6592 caller_thread: std::thread::ThreadId,
6593 graph_stats: crate::ingest::GraphSearchStats,
6594 _scratch: graph_cache::GraphScratchLease<'a>,
6595 _lease: SnapshotLease,
6596 },
6597}
6598
6599struct MergedSegmentGraph<'a> {
6600 plan: crate::planner::SegmentPlan,
6601 traversed: bool,
6602 _result: Option<crate::graph::search::GraphSearchResult>,
6603 _scratch: Option<graph_cache::GraphScratchLease<'a>>,
6604 _lease: SnapshotLease,
6605}
6606
6607impl<'a> SegmentGraphResult<'a> {
6608 #[allow(clippy::too_many_arguments)]
6609 fn merge_into(
6610 self,
6611 segment: &crate::segment::reader::SegmentReader,
6612 alive: &crate::meta::AliveSet,
6613 source: crate::ingest::RowSource,
6614 candidates: &mut Vec<crate::ingest::SearchCandidate>,
6615 dims_touched: &mut u64,
6616 bytes_read: &mut u64,
6617 worker_thread_ids: &mut Vec<std::thread::ThreadId>,
6618 graph_stats: &mut crate::ingest::GraphSearchStats,
6619 #[cfg(any(test, feature = "test-support"))] vector_fault_controller: Option<
6620 &crate::scan::vector_fault::VectorFaultController,
6621 >,
6622 _tier: SearchTier,
6623 ) -> Result<MergedSegmentGraph<'a>, QueryError> {
6624 match self {
6625 Self::Pruned {
6626 plan,
6627 graph_stats: delta,
6628 _lease,
6629 } => {
6630 add_traversal_stats(graph_stats, delta)?;
6631 Ok(MergedSegmentGraph {
6632 plan,
6633 traversed: false,
6634 _result: None,
6635 _scratch: None,
6636 _lease,
6637 })
6638 }
6639 Self::Traversed {
6640 result,
6641 plan,
6642 dims_touched: traversal_dims_touched,
6643 bytes_read: traversal_bytes_read,
6644 caller_thread,
6645 graph_stats: delta,
6646 _scratch,
6647 _lease,
6648 } => {
6649 *dims_touched = dims_touched
6650 .checked_add(traversal_dims_touched)
6651 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
6652 *bytes_read = bytes_read
6653 .checked_add(traversal_bytes_read)
6654 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
6655 if !worker_thread_ids.contains(&caller_thread) {
6656 worker_thread_ids.push(caller_thread);
6657 }
6658 add_traversal_stats(graph_stats, delta)?;
6659 reserve_global_candidates(
6660 candidates,
6661 result
6662 .candidates()
6663 .iter()
6664 .filter(|candidate| alive.is_alive(candidate.row_id()))
6665 .count(),
6666 #[cfg(any(test, feature = "test-support"))]
6667 vector_fault_controller,
6668 #[cfg(any(test, feature = "test-support"))]
6669 vector_fault_tier(_tier),
6670 )?;
6671 for candidate in result.candidates() {
6672 if !alive.is_alive(candidate.row_id()) {
6673 continue;
6674 }
6675 let score = -(candidate.distance() as f32);
6676 if !score.is_finite() {
6677 return Err(QueryError::Graph(
6678 crate::graph::search::GraphSearchError::Geometry(format!(
6679 "graph result distance for row {} is not representable as f32",
6680 candidate.row_id()
6681 )),
6682 ));
6683 }
6684 candidates.push(crate::ingest::SearchCandidate::new(
6685 crate::ingest::GlobalRowId::new(source, candidate.row_id()),
6686 segment
6687 .document_version(candidate.row_id() as usize)
6688 .map_err(StoreError::Segment)
6689 .map_err(QueryError::Store)?,
6690 score,
6691 true,
6692 ));
6693 }
6694 Ok(MergedSegmentGraph {
6695 plan,
6696 traversed: true,
6697 _result: Some(result),
6698 _scratch: Some(_scratch),
6699 _lease,
6700 })
6701 }
6702 }
6703 }
6704}
6705
6706#[allow(clippy::too_many_arguments)]
6707fn traverse_segment_graph<'a>(
6708 segment: &'a crate::segment::reader::SegmentReader,
6709 alive: &crate::meta::AliveSet,
6710 source: crate::ingest::RowSource,
6711 graph_options: GraphSearchOptions,
6712 graph_bound_mode: GraphBoundMode,
6713 request: crate::ingest::SearchRequest<'_>,
6714 k: usize,
6715 vector_preparation: &mut prepared::PreparedVectorQuery<'_>,
6716 accounting: &Arc<stats::Accounting>,
6717 snapshot: &Arc<PublishedSnapshot>,
6718 generation: u64,
6719 control: &QueryControl,
6720 competitive_distance: Option<f32>,
6721 #[cfg(any(test, feature = "test-support"))] vector_fault_controller: Option<
6722 &crate::scan::vector_fault::VectorFaultController,
6723 >,
6724 _tier: SearchTier,
6725) -> Result<SegmentGraphResult<'a>, QueryError> {
6726 let lease = SnapshotLease::new_at(Arc::clone(snapshot), generation);
6727 let cancellation = QueryCancellation::new(control, &lease);
6728 cancellation.check_graph().map_err(map_scan_error)?;
6729 let (graph, graph_validated, prepared_entry_seed_discovered, norm_range) =
6730 match graph_bound_mode {
6731 GraphBoundMode::Shared => {
6732 let prepared = segment
6733 .graph_search_cache
6734 .prepare_shared(segment, &cancellation)
6735 .map_err(map_graph_cache_error)?;
6736 (
6737 prepared.graph,
6738 prepared.graph_validated,
6739 prepared.entry_seed_discovered,
6740 Some(prepared.norm_range),
6741 )
6742 }
6743 #[cfg(test)]
6744 GraphBoundMode::Independent => {
6745 let (graph, graph_validated) = segment
6746 .graph_search_cache
6747 .bind_graph(segment)
6748 .map_err(map_graph_cache_error)?;
6749 (graph, graph_validated, false, None)
6750 }
6751 };
6752 let rescore = query_rescore_rows_for_search(
6753 segment,
6754 #[cfg(any(test, feature = "test-support"))]
6755 vector_fault_controller,
6756 #[cfg(any(test, feature = "test-support"))]
6757 vector_fault_tier(_tier),
6758 )?;
6759 let node_count = graph.node_count() as usize;
6760 let segment_k = k.min(node_count);
6761 let has_tombstones = alive.tombstone_count() != 0;
6762 let target_live_k = if has_tombstones {
6763 usize::try_from(alive.live_count())
6764 .map_err(|_| QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?
6765 .min(k)
6766 } else {
6767 segment_k
6768 };
6769 let maximum_candidate_k = if has_tombstones {
6770 let tombstone_count = usize::try_from(alive.tombstone_count())
6771 .map_err(|_| QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
6772 target_live_k
6773 .checked_add(tombstone_count)
6774 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?
6775 .min(node_count)
6776 .max(segment_k)
6777 } else {
6778 segment_k
6779 };
6780 if let Some(ef) = graph_options.ef()
6781 && ef < segment_k
6782 {
6783 return Err(QueryError::Graph(
6784 crate::graph::search::GraphSearchError::AdaptiveEf(
6785 crate::graph::search::AdaptiveEfError::ExplicitBelowK { k: segment_k, ef },
6786 ),
6787 ));
6788 }
6789 let request_for = |candidate_k| {
6790 let graph_request = crate::graph::search::GraphSearchRequest::new(
6791 request.vector(),
6792 candidate_k,
6793 graph_options.seed(),
6794 )
6795 .with_profile(graph_options.profile());
6796 graph_options.ef().map_or(graph_request, |ef| {
6801 graph_request.with_ef(ef.max(candidate_k))
6802 })
6803 };
6804 let mut candidate_k = segment_k;
6805 let mut graph_request = request_for(candidate_k);
6806 let effective_ef = graph_request
6807 .effective_ef(graph.node_count() as usize)
6808 .map_err(crate::graph::search::GraphSearchError::AdaptiveEf)
6809 .map_err(QueryError::Graph)?;
6810 let requestable_max_k = maximum_candidate_k;
6811 let scratch_ef = if has_tombstones && requestable_max_k > candidate_k {
6812 request_for(requestable_max_k)
6813 .effective_ef(node_count)
6814 .map_err(crate::graph::search::GraphSearchError::AdaptiveEf)
6815 .map_err(QueryError::Graph)?
6816 } else {
6817 effective_ef
6818 };
6819 if let (Some(norm_range), Some(competitive_distance)) = (norm_range, competitive_distance)
6820 && norm_range.squared_l2_upper_bound(request.vector()) <= f64::from(f32::MAX)
6821 && (norm_range.squared_l2_lower_bound(request.vector()) as f32) > competitive_distance
6822 {
6823 cancellation.check_graph().map_err(map_scan_error)?;
6824 let graph_stats = crate::ingest::GraphSearchStats {
6825 graph_validations: usize::from(graph_validated),
6826 entry_seed_discoveries: usize::from(prepared_entry_seed_discovered),
6827 segments_pruned_by_bound: 1,
6828 ..crate::ingest::GraphSearchStats::default()
6829 };
6830 let plan = crate::planner::SegmentPlan::unfiltered_pruned(source, alive.live_count());
6831 return Ok(SegmentGraphResult::Pruned {
6832 plan,
6833 graph_stats,
6834 _lease: lease,
6835 });
6836 }
6837 let mut scratch = segment
6838 .graph_search_cache
6839 .checkout(graph, scratch_ef, accounting, &cancellation)
6840 .map_err(map_graph_cache_error)?;
6841 let entry_seed_discovered = prepared_entry_seed_discovered || scratch.entry_seed_discovered();
6842 let entries = scratch.entries();
6843 let mut searcher = crate::graph::search::GraphSearcher::with_entry_row_ids(
6844 graph,
6845 rescore,
6846 entries,
6847 scratch.scratch_mut().map_err(map_graph_error)?,
6848 )
6849 .map_err(map_graph_error)?
6850 .with_rescore_validator(segment);
6851 #[cfg(any(test, feature = "test-support"))]
6852 if let Some(after_hops) = graph_options.cancel_after_hops() {
6853 let token = cancellation.cancel_token_for_test().ok_or_else(|| {
6854 QueryError::Graph(crate::graph::search::GraphSearchError::Geometry(
6855 "test hop cancellation requires a Cancel query control".to_owned(),
6856 ))
6857 })?;
6858 let _observed_hops = searcher.cancel_after_hops(after_hops, token);
6859 }
6860 #[cfg(any(test, feature = "test-support"))]
6861 if let Some(controller) = vector_fault_controller {
6862 searcher = searcher.with_vector_fault_controller(
6863 controller,
6864 vector_fault_source(source),
6865 crate::scan::vector_fault::VectorSearchTier::Graph,
6866 );
6867 }
6868 let mut traversal_dims_touched = 0_u64;
6869 let mut traversal_bytes_read = 0_u64;
6870 let mut traversal_epoch_clears = 0_usize;
6871 let mut traversal_candidates_scored = 0_usize;
6872 let mut traversal_candidates_rescored = 0_usize;
6873 let result = loop {
6874 let result = searcher.search_reusing(
6875 graph_request,
6876 Some(&cancellation),
6877 vector_preparation,
6878 segment,
6879 source,
6880 )?;
6881 let counters = result.counters();
6882 traversal_dims_touched = traversal_dims_touched
6883 .checked_add(counters.dims_touched())
6884 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
6885 traversal_bytes_read = traversal_bytes_read
6886 .checked_add(counters.bytes_read())
6887 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
6888 traversal_epoch_clears = traversal_epoch_clears
6889 .checked_add(usize::from(counters.visited_epoch_cleared()))
6890 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
6891 traversal_candidates_scored = traversal_candidates_scored
6892 .checked_add(counters.candidates_scored())
6893 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
6894 traversal_candidates_rescored = traversal_candidates_rescored
6895 .checked_add(counters.candidates_rescored())
6896 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
6897
6898 if !has_tombstones {
6899 break result;
6900 }
6901 let retained_live = result
6902 .candidates()
6903 .iter()
6904 .filter(|candidate| alive.is_alive(candidate.row_id()))
6905 .count();
6906 if retained_live >= target_live_k {
6907 break result;
6908 }
6909 if candidate_k >= requestable_max_k {
6910 return Err(QueryError::Graph(
6911 crate::graph::search::GraphSearchError::Geometry(format!(
6912 "graph traversal exhausted {candidate_k} candidates but retained {retained_live} live rows, expected {target_live_k}"
6913 )),
6914 ));
6915 }
6916 let widened_k = candidate_k
6917 .saturating_mul(2)
6918 .max(candidate_k.saturating_add(1))
6919 .min(requestable_max_k);
6920 if widened_k <= candidate_k {
6921 return Err(QueryError::Graph(
6922 crate::graph::search::GraphSearchError::Geometry(format!(
6923 "graph traversal could not widen beyond {candidate_k} candidates"
6924 )),
6925 ));
6926 }
6927 candidate_k = widened_k;
6928 graph_request = request_for(candidate_k);
6929 };
6930 let graph_stats = crate::ingest::GraphSearchStats {
6931 segments_traversed: 1,
6932 graph_validations: usize::from(graph_validated),
6933 entry_seed_discoveries: usize::from(entry_seed_discovered),
6934 visited_epoch_clears: traversal_epoch_clears,
6935 candidates_scored: traversal_candidates_scored,
6936 candidates_rescored: traversal_candidates_rescored,
6937 segments_pruned_by_bound: 0,
6938 };
6939 let plan = crate::planner::SegmentPlan::unfiltered_graph(
6940 source,
6941 alive.live_count(),
6942 graph_options.ef(),
6943 effective_ef,
6944 graph_options.profile(),
6945 );
6946 let caller_thread = std::thread::current().id();
6947 drop(searcher);
6948 Ok(SegmentGraphResult::Traversed {
6949 result,
6950 plan,
6951 dims_touched: traversal_dims_touched,
6952 bytes_read: traversal_bytes_read,
6953 caller_thread,
6954 graph_stats,
6955 _scratch: scratch,
6956 _lease: lease,
6957 })
6958}
6959
6960fn add_traversal_stats(
6961 stats: &mut crate::ingest::GraphSearchStats,
6962 delta: crate::ingest::GraphSearchStats,
6963) -> Result<(), QueryError> {
6964 stats.segments_traversed = stats
6965 .segments_traversed
6966 .checked_add(delta.segments_traversed)
6967 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
6968 stats.graph_validations = stats
6969 .graph_validations
6970 .checked_add(delta.graph_validations)
6971 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
6972 stats.entry_seed_discoveries = stats
6973 .entry_seed_discoveries
6974 .checked_add(delta.entry_seed_discoveries)
6975 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
6976 stats.visited_epoch_clears = stats
6977 .visited_epoch_clears
6978 .checked_add(delta.visited_epoch_clears)
6979 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
6980 stats.candidates_scored = stats
6981 .candidates_scored
6982 .checked_add(delta.candidates_scored)
6983 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
6984 stats.candidates_rescored = stats
6985 .candidates_rescored
6986 .checked_add(delta.candidates_rescored)
6987 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
6988 stats.segments_pruned_by_bound = stats
6989 .segments_pruned_by_bound
6990 .checked_add(delta.segments_pruned_by_bound)
6991 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
6992 Ok(())
6993}
6994
6995#[allow(clippy::too_many_arguments)]
6996fn scan_sealed_segment(
6997 pool: Option<&pool::QueryPool>,
6998 snapshot: &Arc<PublishedSnapshot>,
6999 segment: &crate::segment::reader::SegmentReader,
7000 alive: &crate::meta::AliveSet,
7001 generation: u64,
7002 request: crate::ingest::SearchRequest<'_>,
7003 k: usize,
7004 scan_options: crate::scan::ScanOptions,
7005 control: &QueryControl,
7006 bit4_query: &std::cell::OnceCell<Result<crate::quant::Bit4Query, crate::quant::QuantError>>,
7007 int8_query: &std::cell::OnceCell<Result<crate::quant::Int8Query, crate::quant::QuantError>>,
7008 full_precision: bool,
7009 exact_lane: usize,
7010 accounting: &Arc<stats::Accounting>,
7011 exact_memory: &mut ExactScanMemory,
7012 _source: crate::ingest::RowSource,
7013 _tier: SearchTier,
7014 #[cfg(any(test, feature = "test-support"))] vector_fault_controller: Option<
7015 &crate::scan::vector_fault::VectorFaultController,
7016 >,
7017) -> Result<crate::scan::ScanOutcome, QueryError> {
7018 use crate::scan::{ScanQuery, ScanRequest, ScanRows};
7019
7020 if full_precision {
7021 let lease = SnapshotLease::new_at(Arc::clone(snapshot), generation);
7022 let vectors = exact_rescore_rows_for_search(
7023 segment,
7024 #[cfg(any(test, feature = "test-support"))]
7025 vector_fault_controller,
7026 #[cfg(any(test, feature = "test-support"))]
7027 vector_fault_tier(_tier),
7028 )?;
7029 return scan_squared_l2(
7030 vectors,
7031 segment.meta().row_count as usize,
7032 alive,
7033 request.vector(),
7034 k,
7035 ExactScanDispatch {
7036 pool,
7037 accounting,
7038 control,
7039 lease: &lease,
7040 thread_budget: scan_options.thread_budget,
7041 lane: exact_lane,
7042 },
7043 exact_memory,
7044 #[cfg(any(test, feature = "test-support"))]
7045 vector_fault_controller,
7046 #[cfg(any(test, feature = "test-support"))]
7047 vector_fault_source(_source),
7048 #[cfg(any(test, feature = "test-support"))]
7049 vector_fault_tier(_tier),
7050 );
7051 }
7052
7053 let query_pool = pool.ok_or(QueryError::Store(StoreError::Synchronization {
7054 component: "sealed scan-tier query pool",
7055 }))?;
7056 match segment.meta().scheme {
7057 0 => execute_store_scan(
7058 query_pool,
7059 ScanRequest {
7060 query: ScanQuery::F32(request.vector()),
7061 rows: ScanRows::F32BorrowedRowMajor(
7062 segment
7063 .query_f32_codes()
7064 .map_err(StoreError::Segment)
7065 .map_err(QueryError::Store)?,
7066 ),
7067 row_mask: Some(alive.scan_mask()),
7068 },
7069 k,
7070 scan_options,
7071 control.clone(),
7072 SnapshotLease::new_at(Arc::clone(snapshot), generation),
7073 #[cfg(any(test, feature = "test-support"))]
7074 vector_fault_controller,
7075 #[cfg(any(test, feature = "test-support"))]
7076 vector_fault_source(_source),
7077 #[cfg(any(test, feature = "test-support"))]
7078 vector_fault_tier(_tier),
7079 ),
7080 4 => execute_store_scan(
7081 query_pool,
7082 ScanRequest {
7083 query: ScanQuery::Bit4(
7084 bit4_query
7085 .get_or_init(|| crate::quant::prepare_bit4_query(request.vector(), 0))
7086 .as_ref()
7087 .map_err(|error| crate::scan::ScanError::Quant(error.clone()))
7088 .map_err(QueryError::Scan)?,
7089 ),
7090 rows: ScanRows::Bit4RowMajor {
7091 codes: segment
7092 .query_bit4_codes()
7093 .map_err(StoreError::Segment)
7094 .map_err(QueryError::Store)?,
7095 factors: segment
7096 .query_bit4_factors()
7097 .map_err(StoreError::Segment)
7098 .map_err(QueryError::Store)?,
7099 },
7100 row_mask: Some(alive.scan_mask()),
7101 },
7102 k,
7103 scan_options,
7104 control.clone(),
7105 SnapshotLease::new_at(Arc::clone(snapshot), generation),
7106 #[cfg(any(test, feature = "test-support"))]
7107 vector_fault_controller,
7108 #[cfg(any(test, feature = "test-support"))]
7109 vector_fault_source(_source),
7110 #[cfg(any(test, feature = "test-support"))]
7111 vector_fault_tier(_tier),
7112 ),
7113 2 => {
7114 let factors = segment.query_int8_factors().map_err(QueryError::Store)?;
7115 execute_store_scan(
7116 query_pool,
7117 ScanRequest {
7118 query: ScanQuery::Int8(
7119 int8_query
7120 .get_or_init(|| crate::quant::prepare_int8_query(request.vector()))
7121 .as_ref()
7122 .map_err(|error| crate::scan::ScanError::Quant(error.clone()))
7123 .map_err(QueryError::Scan)?,
7124 ),
7125 rows: ScanRows::Int8RowMajor {
7126 codes: segment
7127 .query_int8_codes()
7128 .map_err(StoreError::Segment)
7129 .map_err(QueryError::Store)?,
7130 factors: factors.as_slice(),
7131 },
7132 row_mask: Some(alive.scan_mask()),
7133 },
7134 k,
7135 scan_options,
7136 control.clone(),
7137 SnapshotLease::new_at(Arc::clone(snapshot), generation),
7138 #[cfg(any(test, feature = "test-support"))]
7139 vector_fault_controller,
7140 #[cfg(any(test, feature = "test-support"))]
7141 vector_fault_source(_source),
7142 #[cfg(any(test, feature = "test-support"))]
7143 vector_fault_tier(_tier),
7144 )
7145 }
7146 scheme => Err(QueryError::Store(StoreError::Segment(
7147 crate::segment::SegmentError::Geometry(format!(
7148 "store search does not support sealed scheme {scheme}"
7149 )),
7150 ))),
7151 }
7152}
7153
7154#[allow(clippy::too_many_arguments)]
7155fn execute_store_scan(
7156 pool: &pool::QueryPool,
7157 request: crate::scan::ScanRequest<'_>,
7158 k: usize,
7159 options: crate::scan::ScanOptions,
7160 control: QueryControl,
7161 lease: SnapshotLease,
7162 #[cfg(any(test, feature = "test-support"))] controller: Option<
7163 &crate::scan::vector_fault::VectorFaultController,
7164 >,
7165 #[cfg(any(test, feature = "test-support"))] source: crate::scan::vector_fault::VectorRowSource,
7166 #[cfg(any(test, feature = "test-support"))] tier: crate::scan::vector_fault::VectorSearchTier,
7167) -> Result<crate::scan::ScanOutcome, QueryError> {
7168 #[cfg(any(test, feature = "test-support"))]
7169 if let Some(controller) = controller {
7170 let geometry = crate::scan::scan_geometry(request).map_err(QueryError::Scan)?;
7171 let cancellation = QueryCancellation::new(&control, &lease);
7172 let partition = crate::scan::scan_partition_with_vector_faults(
7173 request,
7174 k,
7175 0..geometry.row_count,
7176 Some(&cancellation),
7177 controller,
7178 source,
7179 tier,
7180 )
7181 .map_err(map_scan_error)?;
7182 return Ok(crate::scan::ScanOutcome {
7183 candidates: partition.candidates,
7184 worst_score: None,
7185 stats: crate::scan::ScanStats {
7186 dims_touched: partition.dims_touched,
7187 bytes_read: partition.bytes_read,
7188 threads_used: 1,
7189 worker_thread_ids: vec![partition.worker_thread_id],
7190 },
7191 });
7192 }
7193 pool.execute(request, k, options, control, lease)
7194}
7195
7196#[cfg(any(test, feature = "test-support"))]
7197fn vector_fault_tier(tier: SearchTier) -> crate::scan::vector_fault::VectorSearchTier {
7198 match tier {
7199 SearchTier::Exact => crate::scan::vector_fault::VectorSearchTier::Exact,
7200 SearchTier::Scan => crate::scan::vector_fault::VectorSearchTier::Scan,
7201 SearchTier::Graph(_) => crate::scan::vector_fault::VectorSearchTier::Graph,
7202 SearchTier::Auto => crate::scan::vector_fault::VectorSearchTier::Auto,
7203 }
7204}
7205
7206#[cfg(any(test, feature = "test-support"))]
7207fn vector_fault_source(
7208 source: crate::ingest::RowSource,
7209) -> crate::scan::vector_fault::VectorRowSource {
7210 match source {
7211 crate::ingest::RowSource::Active => crate::scan::vector_fault::VectorRowSource::Active,
7212 crate::ingest::RowSource::Sealed(segment) => {
7213 crate::scan::vector_fault::VectorRowSource::Sealed(*segment.as_bytes())
7214 }
7215 }
7216}
7217
7218pub(crate) fn auto_uses_full_precision(
7219 snapshot: &PublishedSnapshot,
7220 active: &crate::ingest::ActiveSegment,
7221) -> bool {
7222 let mut has_exact = false;
7223 let mut has_estimated = !active.is_empty();
7224 for segment in snapshot.segments() {
7225 if segment
7226 .directory()
7227 .iter()
7228 .any(|entry| entry.kind == crate::segment::layout::RegionKind::GraphNodeBlocks.id())
7229 {
7230 return true;
7231 }
7232 if segment.meta().row_count == 0 {
7233 continue;
7234 }
7235 match segment.meta().scheme {
7236 0 => has_exact = true,
7237 2 | 4 => has_estimated = true,
7238 _ => {}
7239 }
7240 }
7241 has_exact && has_estimated
7242}
7243
7244pub(crate) fn exact_rescore_rows(
7254 segment: &crate::segment::reader::SegmentReader,
7255) -> Result<&[f32], QueryError> {
7256 segment
7257 .query_rescore_f32()
7258 .map_err(|source| QueryError::Store(StoreError::Segment(source)))
7259}
7260
7261fn exact_rescore_rows_for_search<'a>(
7262 segment: &'a crate::segment::reader::SegmentReader,
7263 #[cfg(any(test, feature = "test-support"))] controller: Option<
7264 &crate::scan::vector_fault::VectorFaultController,
7265 >,
7266 #[cfg(any(test, feature = "test-support"))] tier: crate::scan::vector_fault::VectorSearchTier,
7267) -> Result<&'a [f32], QueryError> {
7268 #[cfg(any(test, feature = "test-support"))]
7269 if let Some(available_rows) = controller.and_then(|controller| {
7270 controller.missing_rescore_rows(
7271 crate::scan::vector_fault::VectorRowSource::Sealed(*segment.meta().id.as_bytes()),
7272 tier,
7273 crate::scan::vector_fault::MissingRescoreSite::ExactRescoreRows,
7274 segment.meta().row_count,
7275 )
7276 }) {
7277 return Err(missing_rescore_rows_error(segment, available_rows));
7278 }
7279 exact_rescore_rows(segment)
7280}
7281
7282pub(crate) fn query_rescore_rows(
7283 segment: &crate::segment::reader::SegmentReader,
7284) -> Result<&[f32], QueryError> {
7285 segment
7286 .query_rescore_f32()
7287 .map_err(|source| QueryError::Store(StoreError::Segment(source)))
7288}
7289
7290fn query_rescore_rows_for_search<'a>(
7291 segment: &'a crate::segment::reader::SegmentReader,
7292 #[cfg(any(test, feature = "test-support"))] controller: Option<
7293 &crate::scan::vector_fault::VectorFaultController,
7294 >,
7295 #[cfg(any(test, feature = "test-support"))] tier: crate::scan::vector_fault::VectorSearchTier,
7296) -> Result<&'a [f32], QueryError> {
7297 #[cfg(any(test, feature = "test-support"))]
7298 if let Some(available_rows) = controller.and_then(|controller| {
7299 controller.missing_rescore_rows(
7300 crate::scan::vector_fault::VectorRowSource::Sealed(*segment.meta().id.as_bytes()),
7301 tier,
7302 crate::scan::vector_fault::MissingRescoreSite::QueryRescoreRows,
7303 segment.meta().row_count,
7304 )
7305 }) {
7306 return Err(missing_rescore_rows_error(segment, available_rows));
7307 }
7308 query_rescore_rows(segment)
7309}
7310
7311#[cfg(any(test, feature = "test-support"))]
7312fn missing_rescore_rows_error(
7313 segment: &crate::segment::reader::SegmentReader,
7314 available_rows: u32,
7315) -> QueryError {
7316 QueryError::Store(StoreError::Segment(crate::segment::SegmentError::Geometry(
7317 format!(
7318 "exact scores unavailable for segment {}: expected {} rows, got {available_rows}",
7319 segment.meta().id,
7320 segment.meta().row_count,
7321 ),
7322 )))
7323}
7324
7325fn retain_global_top_k(candidates: &mut Vec<crate::ingest::SearchCandidate>, k: usize) {
7326 candidates.sort_unstable_by(crate::ingest::compare_search_candidates);
7327 candidates.truncate(k);
7328}
7329
7330fn global_competitive_distance(
7331 candidates: &[crate::ingest::SearchCandidate],
7332 k: usize,
7333) -> Option<f32> {
7334 if k == 0 || candidates.len() < k {
7335 return None;
7336 }
7337 candidates
7338 .get(k.saturating_sub(1))
7339 .map(|candidate| -candidate.score())
7340 .filter(|distance| distance.is_finite() && *distance >= 0.0)
7341}
7342
7343fn map_graph_cache_error(error: graph_cache::GraphCacheError) -> QueryError {
7344 match error {
7345 graph_cache::GraphCacheError::Store(error) => QueryError::Store(error),
7346 graph_cache::GraphCacheError::Search(error) => map_graph_error(error),
7347 }
7348}
7349
7350pub(crate) fn map_graph_error(error: crate::graph::search::GraphSearchError) -> QueryError {
7351 match error {
7352 crate::graph::search::GraphSearchError::ExactRescoreUnavailable(detail) => {
7353 QueryError::Store(StoreError::Segment(crate::segment::SegmentError::Geometry(
7354 detail,
7355 )))
7356 }
7357 crate::graph::search::GraphSearchError::Cancelled { partial } => {
7358 QueryError::Cancelled { partial }
7359 }
7360 crate::graph::search::GraphSearchError::Timeout { partial } => {
7361 QueryError::Timeout { partial }
7362 }
7363 crate::graph::search::GraphSearchError::ReadCancelled { partial } => {
7364 QueryError::ReadCancelled { partial }
7365 }
7366 crate::graph::search::GraphSearchError::Scan(error) => QueryError::Scan(error),
7367 error => QueryError::Graph(error),
7368 }
7369}
7370
7371#[allow(
7372 clippy::too_many_arguments,
7373 reason = "scoped dispatch and test source identity stay explicit"
7374)]
7375fn scan_active_squared_l2(
7376 active: &crate::ingest::ActiveSegment,
7377 alive: &crate::meta::AliveSet,
7378 query: &[f32],
7379 k: usize,
7380 dispatch: ExactScanDispatch<'_>,
7381 memory: &mut ExactScanMemory,
7382 #[cfg(any(test, feature = "test-support"))] controller: Option<
7383 &crate::scan::vector_fault::VectorFaultController,
7384 >,
7385 #[cfg(any(test, feature = "test-support"))] tier: crate::scan::vector_fault::VectorSearchTier,
7386) -> Result<crate::scan::ScanOutcome, QueryError> {
7387 scan_squared_l2(
7388 active.vectors(),
7389 active.row_count(),
7390 alive,
7391 query,
7392 k,
7393 dispatch,
7394 memory,
7395 #[cfg(any(test, feature = "test-support"))]
7396 controller,
7397 #[cfg(any(test, feature = "test-support"))]
7398 crate::scan::vector_fault::VectorRowSource::Active,
7399 #[cfg(any(test, feature = "test-support"))]
7400 tier,
7401 )
7402}
7403
7404#[derive(Debug)]
7405enum ExactScanError {
7406 Query(QueryError),
7407 Narrowing { row_id: usize, score: f64 },
7408}
7409
7410impl From<QueryError> for ExactScanError {
7411 fn from(error: QueryError) -> Self {
7412 Self::Query(error)
7413 }
7414}
7415
7416impl ExactScanError {
7417 fn into_query(self) -> QueryError {
7418 match self {
7419 Self::Query(error) => error,
7420 Self::Narrowing { row_id, .. } => {
7421 QueryError::Scan(crate::scan::ScanError::NonFiniteScore { row_id })
7422 }
7423 }
7424 }
7425}
7426
7427pub(crate) struct ExactScanMemory {
7428 candidates: stats::AccountedCounter,
7429 worker_ids: stats::AccountedCounter,
7430}
7431
7432impl ExactScanMemory {
7433 pub(crate) fn new(accounting: &Arc<stats::Accounting>) -> Result<Self, QueryError> {
7434 Ok(Self {
7435 candidates: stats::AccountedCounter::new(
7436 accounting,
7437 stats::AllocationComponent::Temporary,
7438 )
7439 .map_err(QueryError::Store)?,
7440 worker_ids: stats::AccountedCounter::new(
7441 accounting,
7442 stats::AllocationComponent::Temporary,
7443 )
7444 .map_err(QueryError::Store)?,
7445 })
7446 }
7447}
7448
7449struct ExactScanDispatch<'a> {
7450 pool: Option<&'a pool::QueryPool>,
7451 accounting: &'a Arc<stats::Accounting>,
7452 control: &'a QueryControl,
7453 lease: &'a SnapshotLease,
7454 thread_budget: usize,
7455 lane: usize,
7456}
7457
7458#[allow(
7459 clippy::too_many_arguments,
7460 reason = "scoped dispatch and test source identity stay explicit"
7461)]
7462fn scan_squared_l2(
7463 vectors: &[f32],
7464 row_count: usize,
7465 alive: &crate::meta::AliveSet,
7466 query: &[f32],
7467 k: usize,
7468 dispatch: ExactScanDispatch<'_>,
7469 memory: &mut ExactScanMemory,
7470 #[cfg(any(test, feature = "test-support"))] controller: Option<
7471 &crate::scan::vector_fault::VectorFaultController,
7472 >,
7473 #[cfg(any(test, feature = "test-support"))] source: crate::scan::vector_fault::VectorRowSource,
7474 #[cfg(any(test, feature = "test-support"))] tier: crate::scan::vector_fault::VectorSearchTier,
7475) -> Result<crate::scan::ScanOutcome, QueryError> {
7476 if let Some(pool) = dispatch.pool
7477 && pool.exact_workers(row_count, query.len(), dispatch.thread_budget) > 1
7478 {
7479 return pool.execute_exact(
7480 vectors,
7481 row_count,
7482 alive,
7483 query,
7484 k,
7485 dispatch.thread_budget,
7486 dispatch.lane,
7487 dispatch.control,
7488 dispatch.lease,
7489 dispatch.accounting,
7490 memory,
7491 #[cfg(any(test, feature = "test-support"))]
7492 controller,
7493 #[cfg(any(test, feature = "test-support"))]
7494 source,
7495 #[cfg(any(test, feature = "test-support"))]
7496 tier,
7497 );
7498 }
7499 let cancellation = QueryCancellation::new(dispatch.control, dispatch.lease);
7500 #[cfg(any(test, feature = "test-support"))]
7501 if controller.is_some_and(|controller| controller.has_exact_row_hooks()) {
7502 return scan_exact_partition::<false>(
7503 vectors,
7504 row_count,
7505 0..row_count,
7506 alive,
7507 query,
7508 k,
7509 &cancellation,
7510 memory,
7511 #[cfg(any(test, feature = "test-support"))]
7512 controller,
7513 #[cfg(any(test, feature = "test-support"))]
7514 source,
7515 #[cfg(any(test, feature = "test-support"))]
7516 tier,
7517 )
7518 .map_err(ExactScanError::into_query);
7519 }
7520 scan_serial_squared_l2(
7521 vectors,
7522 row_count,
7523 alive,
7524 query,
7525 k,
7526 &cancellation,
7527 &mut memory.candidates,
7528 &mut memory.worker_ids,
7529 #[cfg(any(test, feature = "test-support"))]
7530 controller,
7531 #[cfg(any(test, feature = "test-support"))]
7532 source,
7533 #[cfg(any(test, feature = "test-support"))]
7534 tier,
7535 )
7536}
7537
7538#[allow(clippy::too_many_arguments)]
7539fn scan_exact_partition<const OFFSET_ROWS: bool>(
7540 vectors: &[f32],
7541 row_count: usize,
7542 range: std::ops::Range<usize>,
7543 alive: &crate::meta::AliveSet,
7544 query: &[f32],
7545 k: usize,
7546 cancellation: &QueryCancellation<'_>,
7547 memory: &mut ExactScanMemory,
7548 #[cfg(any(test, feature = "test-support"))] controller: Option<
7549 &crate::scan::vector_fault::VectorFaultController,
7550 >,
7551 #[cfg(any(test, feature = "test-support"))] source: crate::scan::vector_fault::VectorRowSource,
7552 #[cfg(any(test, feature = "test-support"))] tier: crate::scan::vector_fault::VectorSearchTier,
7553) -> Result<crate::scan::ScanOutcome, ExactScanError> {
7554 if query.is_empty() {
7555 return Err(QueryError::Scan(crate::scan::ScanError::ZeroDimension).into());
7556 }
7557 let expected = row_count
7558 .checked_mul(query.len())
7559 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
7560 if vectors.len() != expected {
7561 return Err(QueryError::Scan(crate::scan::ScanError::RowDataLength {
7562 dimension: query.len(),
7563 actual: vectors.len(),
7564 })
7565 .into());
7566 }
7567 if range.start > range.end || range.end > row_count || (!OFFSET_ROWS && range.start != 0) {
7570 return Err(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow).into());
7571 }
7572 let mut live_rows = 0;
7573 for row in range.clone() {
7574 if row.is_multiple_of(64) {
7575 cancellation.check_graph().map_err(map_scan_error)?;
7576 }
7577 let local_row = u32::try_from(row)
7578 .map_err(|_| QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
7579 live_rows += usize::from(alive.is_alive(local_row));
7580 }
7581 cancellation.check_graph().map_err(map_scan_error)?;
7582 let mut collector = crate::scan::topk::ExactTopK::try_new(k, live_rows, &mut memory.candidates)
7583 .map_err(QueryError::Store)?;
7584 let mut worst_score: Option<f32> = None;
7585 let mut narrowing_error: Option<(usize, f64)> = None;
7586 let start = range
7587 .start
7588 .checked_mul(query.len())
7589 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
7590 let end = range
7591 .end
7592 .checked_mul(query.len())
7593 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
7594 let partition_vectors = vectors
7595 .get(start..end)
7596 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
7597 let rows = partition_vectors
7598 .chunks_exact(query.len())
7599 .enumerate()
7600 .map(|(local, vector)| {
7601 (
7602 if OFFSET_ROWS {
7603 range.start + local
7604 } else {
7605 local
7606 },
7607 vector,
7608 )
7609 })
7610 .filter(|(row, _)| u32::try_from(*row).is_ok_and(|row| alive.is_alive(row)));
7611 #[cfg(any(test, feature = "test-support"))]
7612 let mut trace =
7613 controller.and_then(|controller| controller.trace_rows(source, tier, range.clone()));
7614 let scored_rows = crate::quant::exact_squared_l2_with_sink(
7615 query,
7616 rows,
7617 |_row, is_checkpoint| {
7618 #[cfg(any(test, feature = "test-support"))]
7619 if let Some(controller) = controller {
7620 if let Some(trace) = &mut trace {
7621 trace.checked(_row);
7622 }
7623 if is_checkpoint {
7624 cancellation.check_graph().map_err(map_scan_error)?;
7625 }
7626 if controller.after_eligible_row(source, tier, _row) {
7627 return Err(QueryError::Cancelled { partial: false });
7628 }
7629 return Ok(());
7630 }
7631 if is_checkpoint {
7632 cancellation.check_graph().map_err(map_scan_error)?;
7633 }
7634 Ok(())
7635 },
7636 |row_id, exact| {
7637 let score = exact as f32;
7638 if !score.is_finite() {
7639 if narrowing_error.is_none_or(|(row, previous)| {
7643 exact.total_cmp(&previous).is_gt()
7644 || (exact.total_cmp(&previous).is_eq() && row_id < row)
7645 }) {
7646 narrowing_error = Some((row_id, exact));
7647 }
7648 return Ok(());
7649 }
7650 if worst_score.is_none_or(|worst| score.total_cmp(&worst).is_lt()) {
7651 worst_score = Some(score);
7652 }
7653 collector
7654 .try_push(
7655 crate::scan::ScanCandidate { row_id, score },
7656 &mut memory.candidates,
7657 )
7658 .map_err(QueryError::Store)
7659 },
7660 )
7661 .map_err(|error| match error {
7662 crate::quant::RescoreCheckError::Rescore(error) => map_l2_rescore_error(error),
7663 crate::quant::RescoreCheckError::Check(error) => error,
7664 })?;
7665 cancellation.check_graph().map_err(map_scan_error)?;
7666 if let Some((row_id, score)) = narrowing_error {
7667 return Err(ExactScanError::Narrowing { row_id, score });
7668 }
7669 let (candidates, _collector_capacity) = collector
7670 .try_into_sorted_with_ties(&mut memory.candidates)
7671 .map_err(QueryError::Store)?;
7672 #[cfg(any(test, feature = "test-support"))]
7673 if let Some(controller) = controller {
7674 controller.record_exact_scan(crate::scan::vector_fault::ExactScanWork {
7675 sorted_items: candidates.len(),
7676 collector_capacity: _collector_capacity,
7677 scored_rows,
7678 worst_score,
7679 ..Default::default()
7680 });
7681 }
7682 let dims = u64::try_from(query.len())
7683 .map_err(|_| QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
7684 let scored_rows = u64::try_from(scored_rows)
7685 .map_err(|_| QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
7686 let dims_touched = scored_rows
7687 .checked_mul(dims)
7688 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
7689 let bytes_read = dims_touched
7690 .checked_mul(std::mem::size_of::<f32>() as u64)
7691 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
7692 memory
7693 .worker_ids
7694 .set(std::mem::size_of::<std::thread::ThreadId>())
7695 .map_err(QueryError::Store)?;
7696 let mut worker_thread_ids = Vec::new();
7697 worker_thread_ids.try_reserve_exact(1).map_err(|_| {
7698 QueryError::Store(StoreError::AllocationFailed {
7699 needed: memory.worker_ids.bytes(),
7700 component: "exact worker IDs",
7701 })
7702 })?;
7703 worker_thread_ids.push(std::thread::current().id());
7704 Ok(crate::scan::ScanOutcome {
7705 candidates,
7706 worst_score,
7707 stats: crate::scan::ScanStats {
7708 dims_touched,
7709 bytes_read,
7710 threads_used: 1,
7711 worker_thread_ids,
7712 },
7713 })
7714}
7715
7716#[allow(clippy::too_many_arguments)]
7719fn scan_serial_squared_l2(
7720 vectors: &[f32],
7721 row_count: usize,
7722 alive: &crate::meta::AliveSet,
7723 query: &[f32],
7724 k: usize,
7725 cancellation: &QueryCancellation<'_>,
7726 memory: &mut stats::AccountedCounter,
7727 worker_memory: &mut stats::AccountedCounter,
7728 #[cfg(any(test, feature = "test-support"))] controller: Option<
7729 &crate::scan::vector_fault::VectorFaultController,
7730 >,
7731 #[cfg(any(test, feature = "test-support"))] source: crate::scan::vector_fault::VectorRowSource,
7732 #[cfg(any(test, feature = "test-support"))] tier: crate::scan::vector_fault::VectorSearchTier,
7733) -> Result<crate::scan::ScanOutcome, QueryError> {
7734 if query.is_empty() {
7735 return Err(QueryError::Scan(crate::scan::ScanError::ZeroDimension));
7736 }
7737 let expected = row_count
7738 .checked_mul(query.len())
7739 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
7740 if vectors.len() != expected {
7741 return Err(QueryError::Scan(crate::scan::ScanError::RowDataLength {
7742 dimension: query.len(),
7743 actual: vectors.len(),
7744 }));
7745 }
7746 let mut live_rows = 0;
7749 for row in 0..row_count {
7750 if row.is_multiple_of(64) {
7751 cancellation.check_graph().map_err(map_scan_error)?;
7752 }
7753 let local_row = u32::try_from(row)
7754 .map_err(|_| QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
7755 live_rows += usize::from(alive.is_alive(local_row));
7756 }
7757 cancellation.check_graph().map_err(map_scan_error)?;
7758 let mut collector =
7759 crate::scan::topk::ExactTopK::try_new(k, live_rows, memory).map_err(QueryError::Store)?;
7760 let mut worst_score: Option<f32> = None;
7761 let mut narrowing_error: Option<(usize, f64)> = None;
7762 let rows = vectors
7763 .chunks_exact(query.len())
7764 .enumerate()
7765 .filter(|(row, _)| u32::try_from(*row).is_ok_and(|row| alive.is_alive(row)));
7766 let scored_rows = crate::quant::exact_squared_l2_with_sink(
7767 query,
7768 rows,
7769 |_row, is_checkpoint| {
7770 if is_checkpoint {
7771 cancellation.check_graph().map_err(map_scan_error)?;
7772 }
7773 #[cfg(any(test, feature = "test-support"))]
7774 if controller
7775 .is_some_and(|controller| controller.after_eligible_row(source, tier, _row))
7776 {
7777 return Err(QueryError::Cancelled { partial: false });
7778 }
7779 Ok(())
7780 },
7781 |row_id, exact| {
7782 let score = exact as f32;
7783 if !score.is_finite() {
7784 if narrowing_error.is_none_or(|(row, previous)| {
7788 exact.total_cmp(&previous).is_gt()
7789 || (exact.total_cmp(&previous).is_eq() && row_id < row)
7790 }) {
7791 narrowing_error = Some((row_id, exact));
7792 }
7793 return Ok(());
7794 }
7795 if worst_score.is_none_or(|worst| score.total_cmp(&worst).is_lt()) {
7796 worst_score = Some(score);
7797 }
7798 collector
7799 .try_push(crate::scan::ScanCandidate { row_id, score }, memory)
7800 .map_err(QueryError::Store)
7801 },
7802 )
7803 .map_err(|error| match error {
7804 crate::quant::RescoreCheckError::Rescore(error) => map_l2_rescore_error(error),
7805 crate::quant::RescoreCheckError::Check(error) => error,
7806 })?;
7807 cancellation.check_graph().map_err(map_scan_error)?;
7808 if let Some((row_id, _)) = narrowing_error {
7809 return Err(QueryError::Scan(crate::scan::ScanError::NonFiniteScore {
7810 row_id,
7811 }));
7812 }
7813 let (candidates, _collector_capacity) = collector
7814 .try_into_sorted_with_ties(memory)
7815 .map_err(QueryError::Store)?;
7816 #[cfg(any(test, feature = "test-support"))]
7817 if let Some(controller) = controller {
7818 controller.record_exact_scan(crate::scan::vector_fault::ExactScanWork {
7819 sorted_items: candidates.len(),
7820 collector_capacity: _collector_capacity,
7821 scored_rows,
7822 worst_score,
7823 ..Default::default()
7824 });
7825 }
7826 let dims = u64::try_from(query.len())
7827 .map_err(|_| QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
7828 let scored_rows = u64::try_from(scored_rows)
7829 .map_err(|_| QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
7830 let dims_touched = scored_rows
7831 .checked_mul(dims)
7832 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
7833 let bytes_read = dims_touched
7834 .checked_mul(std::mem::size_of::<f32>() as u64)
7835 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
7836 worker_memory
7837 .set(std::mem::size_of::<std::thread::ThreadId>())
7838 .map_err(QueryError::Store)?;
7839 let mut worker_thread_ids = Vec::new();
7840 worker_thread_ids.try_reserve_exact(1).map_err(|_| {
7841 QueryError::Store(StoreError::AllocationFailed {
7842 needed: worker_memory.bytes(),
7843 component: "exact worker IDs",
7844 })
7845 })?;
7846 worker_thread_ids.push(std::thread::current().id());
7847 Ok(crate::scan::ScanOutcome {
7848 candidates,
7849 worst_score,
7850 stats: crate::scan::ScanStats {
7851 dims_touched,
7852 bytes_read,
7853 threads_used: 1,
7854 worker_thread_ids,
7855 },
7856 })
7857}
7858
7859fn map_l2_rescore_error(error: crate::quant::RescoreError) -> QueryError {
7860 let error = match error {
7861 crate::quant::RescoreError::ZeroDimension => crate::scan::ScanError::ZeroDimension,
7862 crate::quant::RescoreError::QueryDimension { expected, actual }
7863 | crate::quant::RescoreError::CoarseScoreCount { expected, actual }
7864 | crate::quant::RescoreError::CandidateRowCount { expected, actual } => {
7865 crate::scan::ScanError::RowDataLength {
7866 dimension: expected,
7867 actual,
7868 }
7869 }
7870 crate::quant::RescoreError::RowDataLength { dimension, actual } => {
7871 crate::scan::ScanError::RowDataLength { dimension, actual }
7872 }
7873 crate::quant::RescoreError::CandidateRowOutOfRange {
7874 row_index,
7875 row_count,
7876 ..
7877 } => crate::scan::ScanError::RowDataLength {
7878 dimension: row_count,
7879 actual: row_index,
7880 },
7881 crate::quant::RescoreError::NonFiniteCoarseScore { index }
7882 | crate::quant::RescoreError::NonFiniteExactScore { row_index: index } => {
7883 crate::scan::ScanError::NonFiniteScore { row_id: index }
7884 }
7885 crate::quant::RescoreError::ZeroK
7886 | crate::quant::RescoreError::ZeroOversample
7887 | crate::quant::RescoreError::InsufficientCandidates { .. }
7888 | crate::quant::RescoreError::ArithmeticOverflow => {
7889 crate::scan::ScanError::ArithmeticOverflow
7890 }
7891 error @ (crate::quant::RescoreError::ZeroCandidateLimit
7892 | crate::quant::RescoreError::CandidateLimitExceeded { .. }) => {
7893 crate::scan::ScanError::Rescore(error)
7894 }
7895 };
7896 QueryError::Scan(error)
7897}
7898
7899fn map_scan_error(error: crate::scan::ScanError) -> QueryError {
7900 match error {
7901 crate::scan::ScanError::Cancelled { partial } => QueryError::Cancelled { partial },
7902 crate::scan::ScanError::Timeout { partial } => QueryError::Timeout { partial },
7903 crate::scan::ScanError::ReadCancelled { partial } => QueryError::ReadCancelled { partial },
7904 error => QueryError::Scan(error),
7905 }
7906}
7907
7908#[allow(clippy::too_many_arguments)]
7909fn merge_store_outcome(
7910 outcome: crate::scan::ScanOutcome,
7911 source: crate::ingest::RowSource,
7912 document: impl Fn(usize) -> Result<Option<crate::ingest::DocumentVersion>, QueryError>,
7913 candidates: &mut Vec<crate::ingest::SearchCandidate>,
7914 dims_touched: &mut u64,
7915 bytes_read: &mut u64,
7916 worker_thread_ids: &mut Vec<std::thread::ThreadId>,
7917 exact_score: bool,
7918 #[cfg(any(test, feature = "test-support"))] vector_fault_controller: Option<
7919 &crate::scan::vector_fault::VectorFaultController,
7920 >,
7921 #[cfg(any(test, feature = "test-support"))]
7922 vector_fault_tier: crate::scan::vector_fault::VectorSearchTier,
7923) -> Result<(), QueryError> {
7924 *dims_touched = dims_touched
7925 .checked_add(outcome.stats.dims_touched)
7926 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
7927 *bytes_read = bytes_read
7928 .checked_add(outcome.stats.bytes_read)
7929 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
7930 for worker in outcome.stats.worker_thread_ids {
7931 if !worker_thread_ids.contains(&worker) {
7932 worker_thread_ids.push(worker);
7933 }
7934 }
7935 reserve_global_candidates(
7936 candidates,
7937 outcome.candidates.len(),
7938 #[cfg(any(test, feature = "test-support"))]
7939 vector_fault_controller,
7940 #[cfg(any(test, feature = "test-support"))]
7941 vector_fault_tier,
7942 )?;
7943 for candidate in outcome.candidates {
7944 let local_row = u32::try_from(candidate.row_id)
7945 .map_err(|_| QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
7946 candidates.push(crate::ingest::SearchCandidate::new(
7947 crate::ingest::GlobalRowId::new(source, local_row),
7948 document(candidate.row_id)?,
7949 candidate.score,
7950 exact_score,
7951 ));
7952 }
7953 Ok(())
7954}
7955
7956pub(crate) fn reserve_global_candidates(
7957 candidates: &mut Vec<crate::ingest::SearchCandidate>,
7958 additional: usize,
7959 #[cfg(any(test, feature = "test-support"))] vector_fault_controller: Option<
7960 &crate::scan::vector_fault::VectorFaultController,
7961 >,
7962 #[cfg(any(test, feature = "test-support"))]
7963 vector_fault_tier: crate::scan::vector_fault::VectorSearchTier,
7964) -> Result<(), QueryError> {
7965 if additional == 0 {
7966 return Ok(());
7967 }
7968 let items = u64::try_from(additional)
7969 .map_err(|_| QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
7970 let candidate_bytes = u64::try_from(std::mem::size_of::<crate::ingest::SearchCandidate>())
7971 .map_err(|_| QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
7972 let needed = items
7973 .checked_mul(candidate_bytes)
7974 .ok_or(QueryError::Scan(crate::scan::ScanError::ArithmeticOverflow))?;
7975 #[cfg(any(test, feature = "test-support"))]
7976 if vector_fault_controller.is_some_and(|controller| {
7977 controller.deny_global_candidate_allocation(vector_fault_tier, items, needed)
7978 }) {
7979 return Err(QueryError::Store(StoreError::AllocationFailed {
7980 needed,
7981 component: "vector search global candidates",
7982 }));
7983 }
7984 candidates.try_reserve_exact(additional).map_err(|_| {
7985 QueryError::Store(StoreError::AllocationFailed {
7986 needed,
7987 component: "vector search global candidates",
7988 })
7989 })
7990}
7991
7992struct AdmittedLexicalQuery<'a> {
7993 generation: u64,
7994 active: Arc<crate::ingest::ActiveSegment>,
7995 snapshot: snapshot::ReadSnapshot,
7996 active_query: ActiveQuery<'a>,
7997}
7998
7999struct AdmittedDocumentRead<'a> {
8000 generation: u64,
8001 active: Arc<crate::ingest::ActiveSegment>,
8002 snapshot: snapshot::ReadSnapshot,
8003 active_query: ActiveQuery<'a>,
8004}
8005
8006struct AdmittedVectorSearch<'a> {
8007 execution_options: SearchOptions,
8008 pool: Option<Arc<pool::QueryPool>>,
8009 snapshot: snapshot::ReadSnapshot,
8010 active_segment: Arc<crate::ingest::ActiveSegment>,
8011 generation: u64,
8012 _active_query: ActiveQuery<'a>,
8013}
8014
8015struct ActiveQuery<'a> {
8016 count: &'a AtomicU64,
8017}
8018
8019impl Drop for ActiveQuery<'_> {
8020 fn drop(&mut self) {
8021 self.count.fetch_sub(1, Ordering::Relaxed);
8022 }
8023}
8024
8025impl Drop for Store {
8026 fn drop(&mut self) {
8027 self.close_best_effort();
8028 }
8029}
8030
8031#[derive(Clone, Copy, Debug, Eq, PartialEq)]
8032struct DocumentScanRow {
8033 source: DocumentScanSource,
8034 row: u32,
8035 timestamp: i64,
8036 doc_id: crate::ingest::DocId,
8037}
8038
8039fn document_scan_predicate(
8040 predicate: Option<&crate::meta::Predicate>,
8041 timestamp_range: Option<(i64, i64)>,
8042) -> Result<Option<crate::meta::Predicate>, QueryError> {
8043 use crate::meta::{Predicate, PredicateValue, RangeBound, RangePredicate, TIMESTAMP_COLUMN};
8044
8045 let timestamp = timestamp_range
8046 .map(|(start, end)| {
8047 if start > end {
8048 return Err(QueryError::Store(StoreError::InvalidScan {
8049 detail: format!("timestamp range start {start} exceeds end {end}"),
8050 }));
8051 }
8052 Ok(Predicate::Range(RangePredicate {
8053 column: TIMESTAMP_COLUMN,
8054 lower: Some(RangeBound::inclusive(PredicateValue::I64(start))),
8055 upper: Some(RangeBound::exclusive(PredicateValue::I64(end))),
8056 }))
8057 })
8058 .transpose()?;
8059 Ok(match (predicate, timestamp) {
8060 (Some(predicate), Some(timestamp)) => {
8061 Some(Predicate::And(vec![predicate.clone(), timestamp]))
8062 }
8063 (Some(predicate), None) => Some(predicate.clone()),
8064 (None, Some(timestamp)) => Some(timestamp),
8065 (None, None) => None,
8066 })
8067}
8068
8069fn scan_control_checkpoint(
8070 control: &QueryControl,
8071 lease: &SnapshotLease,
8072 work: &mut usize,
8073) -> Result<(), QueryError> {
8074 *work = work.saturating_add(1);
8075 if (*work).is_multiple_of(64) {
8076 control.checkpoint()?;
8077 lease.check_active().map_err(QueryError::Store)?;
8078 }
8079 Ok(())
8080}
8081
8082fn active_scan_columns(
8083 active: &crate::ingest::ActiveSegment,
8084 schema: &crate::meta::Schema,
8085) -> Result<crate::meta::ColumnStore, QueryError> {
8086 let mut builder = crate::meta::ColumnStoreBuilder::new(schema.clone());
8087 for (row, timestamp) in active.timestamps().iter().copied().enumerate() {
8088 let values = active.column_values(row).map_err(QueryError::Store)?;
8089 builder
8090 .push_row(timestamp, &crate::ingest::column_inputs(&values))
8091 .map_err(|error| {
8092 QueryError::Store(StoreError::Segment(crate::segment::SegmentError::Columns(
8093 error.to_string(),
8094 )))
8095 })?;
8096 }
8097 builder.finish().map_err(|error| {
8098 QueryError::Store(StoreError::Segment(crate::segment::SegmentError::Columns(
8099 error.to_string(),
8100 )))
8101 })
8102}
8103
8104fn sealed_scan_rows(
8105 segment: &crate::segment::reader::SegmentReader,
8106 predicate: Option<&crate::meta::Predicate>,
8107) -> Result<crate::meta::DocBitmap, QueryError> {
8108 if predicate.is_some_and(|predicate| {
8109 !crate::planner::segment_may_match(segment.meta().clustering_key_range, predicate)
8110 }) {
8111 return Ok(crate::meta::DocBitmap::new());
8112 }
8113 let alive = segment.query_alive().map_err(QueryError::Store)?;
8114 match predicate {
8115 Some(predicate) => {
8116 let columns = segment.query_columns().map_err(QueryError::Store)?;
8117 crate::meta::evaluate(predicate, &columns, &alive).map_err(|error| {
8118 QueryError::Store(StoreError::Segment(crate::segment::SegmentError::Columns(
8119 error.to_string(),
8120 )))
8121 })
8122 }
8123 None => Ok(alive.alive_bitmap().clone()),
8124 }
8125}
8126
8127fn active_scan_rows(
8128 active: &crate::ingest::ActiveSegment,
8129 schema: &crate::meta::Schema,
8130 predicate: Option<&crate::meta::Predicate>,
8131) -> Result<crate::meta::DocBitmap, QueryError> {
8132 let alive = active.alive().map_err(QueryError::Store)?;
8133 match predicate {
8134 Some(predicate) => {
8135 let columns = active_scan_columns(active, schema)?;
8136 crate::meta::evaluate(predicate, &columns, &alive).map_err(|error| {
8137 QueryError::Store(StoreError::Segment(crate::segment::SegmentError::Columns(
8138 error.to_string(),
8139 )))
8140 })
8141 }
8142 None => Ok(alive.alive_bitmap().clone()),
8143 }
8144}
8145
8146fn sealed_document_scan_row(
8147 segment: &crate::segment::reader::SegmentReader,
8148 row: u32,
8149) -> Result<DocumentScanRow, QueryError> {
8150 let version = segment
8151 .document_version(row as usize)
8152 .map_err(StoreError::Segment)
8153 .map_err(QueryError::Store)?
8154 .ok_or_else(|| {
8155 QueryError::Store(StoreError::Segment(crate::segment::SegmentError::Geometry(
8156 format!("sealed document row {row} has no identity"),
8157 )))
8158 })?;
8159 let columns = segment.query_columns().map_err(QueryError::Store)?;
8160 let timestamp = columns.timestamp(row).ok_or_else(|| {
8161 QueryError::Store(StoreError::Segment(crate::segment::SegmentError::Geometry(
8162 format!("sealed document row {row} has no timestamp"),
8163 )))
8164 })?;
8165 Ok(DocumentScanRow {
8166 source: DocumentScanSource::Sealed(segment.meta().id),
8167 row,
8168 timestamp,
8169 doc_id: version.doc_id(),
8170 })
8171}
8172
8173fn active_document_scan_row(
8174 active: &crate::ingest::ActiveSegment,
8175 row: u32,
8176) -> Result<DocumentScanRow, QueryError> {
8177 let index = row as usize;
8178 let version = active
8179 .document(index)
8180 .ok_or(QueryError::Store(StoreError::ActiveRowOverflow))?;
8181 let timestamp = active
8182 .timestamps()
8183 .get(index)
8184 .copied()
8185 .ok_or(QueryError::Store(StoreError::ActiveRowOverflow))?;
8186 Ok(DocumentScanRow {
8187 source: DocumentScanSource::Active,
8188 row,
8189 timestamp,
8190 doc_id: version.doc_id(),
8191 })
8192}
8193
8194#[allow(clippy::too_many_arguments)]
8195fn collect_storage_scan_rows(
8196 snapshot: &PublishedSnapshot,
8197 active: &crate::ingest::ActiveSegment,
8198 schema: &crate::meta::Schema,
8199 predicate: Option<&crate::meta::Predicate>,
8200 cursor: Option<DocumentScanCursor>,
8201 limit: usize,
8202 control: &QueryControl,
8203 lease: &SnapshotLease,
8204 work: &mut usize,
8205) -> Result<Vec<DocumentScanRow>, QueryError> {
8206 let capacity = limit.saturating_add(1);
8207 let mut result = Vec::new();
8208 result.try_reserve_exact(capacity).map_err(|_| {
8209 QueryError::Store(StoreError::AllocationFailed {
8210 needed: allocation_bytes::<DocumentScanRow>(capacity),
8211 component: "scan positions",
8212 })
8213 })?;
8214 let mut started = cursor.is_none();
8215 'segments: for segment in snapshot.segments() {
8216 let source = DocumentScanSource::Sealed(segment.meta().id);
8217 let rows = sealed_scan_rows(segment, predicate)?;
8218 for row in rows.iter() {
8219 scan_control_checkpoint(control, lease, work)?;
8220 if !started {
8221 started =
8222 cursor.is_some_and(|cursor| cursor.source == source && cursor.next_row == row);
8223 if !started {
8224 continue;
8225 }
8226 }
8227 result.push(sealed_document_scan_row(segment, row)?);
8228 if result.len() == capacity {
8229 break 'segments;
8230 }
8231 }
8232 }
8233 if result.len() < capacity {
8234 let source = DocumentScanSource::Active;
8235 let rows = active_scan_rows(active, schema, predicate)?;
8236 for row in rows.iter() {
8237 scan_control_checkpoint(control, lease, work)?;
8238 if !started {
8239 started =
8240 cursor.is_some_and(|cursor| cursor.source == source && cursor.next_row == row);
8241 if !started {
8242 continue;
8243 }
8244 }
8245 result.push(active_document_scan_row(active, row)?);
8246 if result.len() == capacity {
8247 break;
8248 }
8249 }
8250 }
8251 if !started {
8252 return Err(QueryError::Store(StoreError::InvalidScan {
8253 detail: "scan cursor does not name a matching row".to_owned(),
8254 }));
8255 }
8256 Ok(result)
8257}
8258
8259fn document_scan_row_cmp(
8260 left: &DocumentScanRow,
8261 right: &DocumentScanRow,
8262 order: ScanOrder,
8263) -> std::cmp::Ordering {
8264 let timestamp = match order {
8265 ScanOrder::Storage | ScanOrder::TimestampAscending => left.timestamp.cmp(&right.timestamp),
8266 ScanOrder::TimestampDescending => right.timestamp.cmp(&left.timestamp),
8267 };
8268 timestamp
8269 .then_with(|| left.doc_id.cmp(&right.doc_id))
8270 .then_with(|| document_scan_source_cmp(left.source, right.source))
8271 .then_with(|| left.row.cmp(&right.row))
8272}
8273
8274fn document_scan_source_cmp(
8275 left: DocumentScanSource,
8276 right: DocumentScanSource,
8277) -> std::cmp::Ordering {
8278 match (left, right) {
8279 (DocumentScanSource::Sealed(left), DocumentScanSource::Sealed(right)) => left.cmp(&right),
8280 (DocumentScanSource::Sealed(_), DocumentScanSource::Active) => std::cmp::Ordering::Less,
8281 (DocumentScanSource::Active, DocumentScanSource::Sealed(_)) => std::cmp::Ordering::Greater,
8282 (DocumentScanSource::Active, DocumentScanSource::Active) => std::cmp::Ordering::Equal,
8283 }
8284}
8285
8286fn insert_timestamp_scan_row(
8287 rows: &mut Vec<DocumentScanRow>,
8288 candidate: DocumentScanRow,
8289 capacity: usize,
8290 order: ScanOrder,
8291) {
8292 let position = rows.partition_point(|row| {
8293 document_scan_row_cmp(row, &candidate, order) != std::cmp::Ordering::Greater
8294 });
8295 if position < capacity {
8296 rows.insert(position, candidate);
8297 if rows.len() > capacity {
8298 let _ = rows.pop();
8299 }
8300 }
8301}
8302
8303fn clustering_range_cmp(
8304 left: crate::segment::ClusteringKeyRange,
8305 right: crate::segment::ClusteringKeyRange,
8306 order: ScanOrder,
8307) -> std::cmp::Ordering {
8308 use crate::segment::ClusteringKeyRange::{Bounded, Empty, Unstamped};
8309
8310 match (left, right) {
8311 (Bounded { min_ts: left, .. }, Bounded { min_ts: right, .. })
8312 if order == ScanOrder::TimestampAscending =>
8313 {
8314 left.cmp(&right)
8315 }
8316 (Bounded { max_ts: left, .. }, Bounded { max_ts: right, .. }) => right.cmp(&left),
8317 (Unstamped, Unstamped) | (Empty, Empty) => std::cmp::Ordering::Equal,
8318 (Unstamped, _) => std::cmp::Ordering::Less,
8319 (_, Unstamped) => std::cmp::Ordering::Greater,
8320 (Empty, _) => std::cmp::Ordering::Greater,
8321 (_, Empty) => std::cmp::Ordering::Less,
8322 }
8323}
8324
8325#[allow(clippy::too_many_arguments)]
8326fn collect_timestamp_scan_rows(
8327 snapshot: &PublishedSnapshot,
8328 active: &crate::ingest::ActiveSegment,
8329 schema: &crate::meta::Schema,
8330 predicate: Option<&crate::meta::Predicate>,
8331 cursor: Option<DocumentScanCursor>,
8332 limit: usize,
8333 order: ScanOrder,
8334 control: &QueryControl,
8335 lease: &SnapshotLease,
8336 work: &mut usize,
8337) -> Result<Vec<DocumentScanRow>, QueryError> {
8338 let capacity = limit.saturating_add(1);
8339 let mut result = Vec::new();
8340 result.try_reserve_exact(capacity).map_err(|_| {
8341 QueryError::Store(StoreError::AllocationFailed {
8342 needed: allocation_bytes::<DocumentScanRow>(capacity),
8343 component: "scan positions",
8344 })
8345 })?;
8346 let cursor_row = cursor
8347 .map(|cursor| match cursor.source {
8348 DocumentScanSource::Active => active_document_scan_row(active, cursor.next_row),
8349 DocumentScanSource::Sealed(segment_id) => snapshot
8350 .segments()
8351 .iter()
8352 .find(|segment| segment.meta().id == segment_id)
8353 .ok_or_else(|| {
8354 QueryError::Store(StoreError::InvalidScan {
8355 detail: format!("scan cursor names absent segment {segment_id}"),
8356 })
8357 })
8358 .and_then(|segment| sealed_document_scan_row(segment, cursor.next_row)),
8359 })
8360 .transpose()?;
8361 let mut segments = snapshot.segments().iter().collect::<Vec<_>>();
8362 segments.sort_by(|left, right| {
8363 clustering_range_cmp(
8364 left.meta().clustering_key_range,
8365 right.meta().clustering_key_range,
8366 order,
8367 )
8368 .then_with(|| left.meta().id.cmp(&right.meta().id))
8369 });
8370 for segment in segments {
8371 let rows = sealed_scan_rows(segment, predicate)?;
8372 for row in rows.iter() {
8373 scan_control_checkpoint(control, lease, work)?;
8374 let candidate = sealed_document_scan_row(segment, row)?;
8375 if cursor_row.as_ref().is_none_or(|cursor_row| {
8376 document_scan_row_cmp(&candidate, cursor_row, order) != std::cmp::Ordering::Less
8377 }) {
8378 insert_timestamp_scan_row(&mut result, candidate, capacity, order);
8379 }
8380 }
8381 }
8382 for row in active_scan_rows(active, schema, predicate)?.iter() {
8383 scan_control_checkpoint(control, lease, work)?;
8384 let candidate = active_document_scan_row(active, row)?;
8385 if cursor_row.as_ref().is_none_or(|cursor_row| {
8386 document_scan_row_cmp(&candidate, cursor_row, order) != std::cmp::Ordering::Less
8387 }) {
8388 insert_timestamp_scan_row(&mut result, candidate, capacity, order);
8389 }
8390 }
8391 if let Some(cursor) = cursor {
8392 let starts_at_cursor = result
8393 .first()
8394 .is_some_and(|row| row.source == cursor.source && row.row == cursor.next_row);
8395 if !starts_at_cursor {
8396 return Err(QueryError::Store(StoreError::InvalidScan {
8397 detail: "scan cursor does not name a matching row".to_owned(),
8398 }));
8399 }
8400 }
8401 Ok(result)
8402}
8403
8404fn materialize_active_document(
8405 active: &crate::ingest::ActiveSegment,
8406 row: usize,
8407 fields: DocumentFields,
8408) -> Result<StoredDocument, StoreError> {
8409 let version = active.document(row).ok_or(StoreError::ActiveRowOverflow)?;
8410 let timestamp = active
8411 .timestamps()
8412 .get(row)
8413 .copied()
8414 .ok_or(StoreError::ActiveRowOverflow)?;
8415 let vector = if fields.contains(DocumentFields::VECTOR) {
8416 let dimensions = active.dims().ok_or(StoreError::ActiveRowOverflow)?;
8417 let start = row
8418 .checked_mul(dimensions)
8419 .ok_or(StoreError::ActiveRowOverflow)?;
8420 let end = start
8421 .checked_add(dimensions)
8422 .ok_or(StoreError::ActiveRowOverflow)?;
8423 let values = active
8424 .vectors()
8425 .get(start..end)
8426 .ok_or(StoreError::ActiveRowOverflow)?;
8427 Some(copy_slice(values, "get vector")?)
8428 } else {
8429 None
8430 };
8431 let text = if fields.contains(DocumentFields::TEXT) {
8432 active
8433 .text(row)?
8434 .map(|value| copy_string(value, "get text"))
8435 .transpose()?
8436 } else {
8437 None
8438 };
8439 let metadata = if fields.contains(DocumentFields::METADATA) {
8440 active
8441 .metadata(row)
8442 .map(|value| copy_slice(value, "get metadata"))
8443 .transpose()?
8444 } else {
8445 None
8446 };
8447 let attributes = if fields.contains(DocumentFields::ATTRIBUTES) {
8448 Some(active.column_values(row)?)
8449 } else {
8450 None
8451 };
8452 Ok(StoredDocument {
8453 doc_id: version.doc_id(),
8454 revision: version.revision(),
8455 timestamp,
8456 vector,
8457 text,
8458 metadata,
8459 attributes,
8460 })
8461}
8462
8463fn materialize_sealed_document(
8464 segment: &crate::segment::reader::SegmentReader,
8465 row: usize,
8466 fields: DocumentFields,
8467) -> Result<StoredDocument, StoreError> {
8468 let row_u32 = u32::try_from(row).map_err(|_| {
8469 StoreError::Segment(crate::segment::SegmentError::Geometry(
8470 "sealed document row exceeds u32".to_owned(),
8471 ))
8472 })?;
8473 let version = segment
8474 .document_version(row)
8475 .map_err(StoreError::Segment)?
8476 .ok_or_else(|| {
8477 StoreError::Segment(crate::segment::SegmentError::Geometry(format!(
8478 "sealed document row {row} has no identity"
8479 )))
8480 })?;
8481 let columns = segment.query_columns()?;
8482 let timestamp = columns.timestamp(row_u32).ok_or_else(|| {
8483 StoreError::Segment(crate::segment::SegmentError::Geometry(format!(
8484 "sealed document row {row} has no timestamp"
8485 )))
8486 })?;
8487 let vector = if fields.contains(DocumentFields::VECTOR) {
8488 segment
8489 .validate_rescore_rows(&[row_u32])
8490 .map_err(StoreError::Segment)?;
8491 let dimensions = segment.meta().dims as usize;
8492 let start = row.checked_mul(dimensions).ok_or_else(|| {
8493 StoreError::Segment(crate::segment::SegmentError::Geometry(
8494 "sealed vector row offset overflow".to_owned(),
8495 ))
8496 })?;
8497 let end = start.checked_add(dimensions).ok_or_else(|| {
8498 StoreError::Segment(crate::segment::SegmentError::Geometry(
8499 "sealed vector row end overflow".to_owned(),
8500 ))
8501 })?;
8502 let rescore = segment.query_rescore_f32().map_err(StoreError::Segment)?;
8503 let values = rescore.get(start..end).ok_or_else(|| {
8504 StoreError::Segment(crate::segment::SegmentError::Geometry(format!(
8505 "sealed vector row {row} is missing"
8506 )))
8507 })?;
8508 Some(copy_slice(values, "get vector")?)
8509 } else {
8510 None
8511 };
8512 let text = if fields.contains(DocumentFields::TEXT) {
8513 match segment.query_stored_text().map_err(StoreError::Segment)? {
8514 Some(rows) => rows
8515 .row(row)
8516 .ok_or_else(|| {
8517 StoreError::Segment(crate::segment::SegmentError::Geometry(format!(
8518 "stored-text row {row} is invalid"
8519 )))
8520 })?
8521 .map(|value| copy_string(value, "get text"))
8522 .transpose()?,
8523 None => None,
8524 }
8525 } else {
8526 None
8527 };
8528 let metadata = if fields.contains(DocumentFields::METADATA) {
8529 match segment
8530 .query_stored_metadata()
8531 .map_err(StoreError::Segment)?
8532 {
8533 Some(rows) => Some(copy_slice(
8534 rows.row(row).ok_or_else(|| {
8535 StoreError::Segment(crate::segment::SegmentError::Geometry(format!(
8536 "stored-metadata row {row} is invalid"
8537 )))
8538 })?,
8539 "get metadata",
8540 )?),
8541 None => None,
8542 }
8543 } else {
8544 None
8545 };
8546 let attributes = fields
8547 .contains(DocumentFields::ATTRIBUTES)
8548 .then(|| materialize_sealed_attributes(&columns, row_u32))
8549 .transpose()?;
8550 Ok(StoredDocument {
8551 doc_id: version.doc_id(),
8552 revision: version.revision(),
8553 timestamp,
8554 vector,
8555 text,
8556 metadata,
8557 attributes,
8558 })
8559}
8560
8561fn materialize_sealed_attributes(
8562 columns: &crate::meta::ColumnStore,
8563 row: u32,
8564) -> Result<Vec<(crate::meta::ColumnId, crate::meta::PredicateValue)>, StoreError> {
8565 let mut attributes = Vec::new();
8566 attributes
8567 .try_reserve_exact(columns.schema().user_column_count())
8568 .map_err(|_| StoreError::AllocationFailed {
8569 needed: allocation_bytes::<(crate::meta::ColumnId, crate::meta::PredicateValue)>(
8570 columns.schema().user_column_count(),
8571 ),
8572 component: "get attributes",
8573 })?;
8574 for definition in columns
8575 .schema()
8576 .columns()
8577 .iter()
8578 .filter(|definition| definition.id() != crate::meta::TIMESTAMP_COLUMN)
8579 {
8580 let column = columns.column(definition.id()).ok_or_else(|| {
8581 StoreError::Segment(crate::segment::SegmentError::Geometry(format!(
8582 "sealed attribute column {} is missing",
8583 definition.id().get()
8584 )))
8585 })?;
8586 let value = predicate_value_for_row(column, row)?;
8587 if value.is_none() && !definition.is_nullable() {
8588 return Err(StoreError::Segment(crate::segment::SegmentError::Geometry(
8589 format!(
8590 "sealed required attribute {} is null at row {row}",
8591 definition.id().get()
8592 ),
8593 )));
8594 }
8595 if let Some(value) = value {
8596 attributes.push((definition.id(), value));
8597 }
8598 }
8599 Ok(attributes)
8600}
8601
8602fn predicate_value_for_row(
8603 column: &crate::meta::Column,
8604 row: u32,
8605) -> Result<Option<crate::meta::PredicateValue>, StoreError> {
8606 use crate::meta::{Column, PredicateValue};
8607
8608 match column {
8609 Column::U64(values) => Ok(values.get(row).map(PredicateValue::U64)),
8610 Column::I64(values) => Ok(values.get(row).map(PredicateValue::I64)),
8611 Column::F64(values) => Ok(values.get(row).map(PredicateValue::F64)),
8612 Column::Bool(values) => Ok(values.get(row).map(PredicateValue::Bool)),
8613 Column::DictionaryString(values) => values
8614 .get(row)
8615 .map(|value| copy_string(value, "get attribute string").map(PredicateValue::String))
8616 .transpose(),
8617 Column::RawString(values) => values
8618 .get(row)
8619 .map(|value| copy_string(value, "get attribute string").map(PredicateValue::String))
8620 .transpose(),
8621 }
8622}
8623
8624fn copy_slice<T: Copy>(values: &[T], component: &'static str) -> Result<Vec<T>, StoreError> {
8625 let mut copied = Vec::new();
8626 copied
8627 .try_reserve_exact(values.len())
8628 .map_err(|_| StoreError::AllocationFailed {
8629 needed: allocation_bytes::<T>(values.len()),
8630 component,
8631 })?;
8632 copied.extend_from_slice(values);
8633 Ok(copied)
8634}
8635
8636fn copy_string(value: &str, component: &'static str) -> Result<String, StoreError> {
8637 let mut copied = String::new();
8638 copied
8639 .try_reserve_exact(value.len())
8640 .map_err(|_| StoreError::AllocationFailed {
8641 needed: u64::try_from(value.len()).unwrap_or(u64::MAX),
8642 component,
8643 })?;
8644 copied.push_str(value);
8645 Ok(copied)
8646}
8647
8648fn allocation_bytes<T>(count: usize) -> u64 {
8649 count
8650 .checked_mul(std::mem::size_of::<T>())
8651 .and_then(|bytes| u64::try_from(bytes).ok())
8652 .unwrap_or(u64::MAX)
8653}
8654
8655#[cfg(test)]
8656mod lexical_assembly_tests;
8657
8658#[cfg(test)]
8659mod get_tests;
8660#[cfg(test)]
8661mod scan_document_tests;
8662
8663#[cfg(test)]
8664#[allow(clippy::expect_used, clippy::panic)]
8665mod tests {
8666 use std::error::Error;
8667 use std::sync::Arc;
8668
8669 use tempfile::tempdir;
8670
8671 use crate::lifecycle::durability::DurabilityPolicy;
8672 use crate::manifest::io::commit_manifest;
8673 use crate::manifest::{Manifest, ManifestError};
8674 use crate::meta::Schema;
8675 use crate::vfs::StdVfs;
8676 use crate::vfs::crash::MemoryVfs;
8677
8678 use super::durability::{CommitTier, DurabilityMode, DurabilityPolicyError};
8679 use super::{
8680 CancelToken, GraphSearchOptions, ManualMonotonicClock, OpenOptions, QueryControl,
8681 SearchOptions, SearchTier, Store, StoreError, StoreTestDependencies,
8682 resolve_hybrid_leg_results,
8683 };
8684
8685 #[test]
8686 fn astra_12_vocabulary_charge_survives_eviction_and_close() {
8687 use crate::fts::{index::DEFAULT_FIELD, query::LexicalQuery};
8688 use crate::ingest::{DocId, DocumentVersion, IngestBatch, IngestDocument, Revision};
8689 let directory = tempdir().expect("directory");
8690 let store = Store::open(directory.path(), OpenOptions::default()).expect("open");
8691 store
8692 .ingest(IngestBatch::new(vec![
8693 IngestDocument::new(
8694 DocumentVersion::new(DocId::new(1), Revision::new(1)),
8695 vec![1.0, 0.0],
8696 )
8697 .with_text("alpha alphabet beta"),
8698 ]))
8699 .expect("ingest");
8700 store
8701 .search_lexical_structured(
8702 &LexicalQuery::prefix(b"alph".to_vec(), DEFAULT_FIELD),
8703 1,
8704 64,
8705 QueryControl::Cancel(CancelToken::new()),
8706 )
8707 .expect("build cache");
8708 let dictionary = {
8709 let entry = store.lexical_index_cache.entry.lock().expect("assembly");
8710 let vocabulary = entry
8711 .as_ref()
8712 .expect("assembly")
8713 .assembly
8714 .vocabulary
8715 .lock()
8716 .expect("vocabulary");
8717 Arc::clone(vocabulary.as_ref().expect("cached"))
8718 };
8719 let bytes = dictionary._memory.bytes();
8720 assert!(bytes > 17);
8721 let second_owner = Arc::clone(&dictionary);
8722 drop(
8723 store
8724 .lexical_index_cache
8725 .entry
8726 .lock()
8727 .expect("evict")
8728 .take(),
8729 );
8730 assert_eq!(store.stats().expect("held bytes").cache_bytes, bytes);
8731 store
8732 .search_lexical_structured(
8733 &LexicalQuery::prefix(b"alph".to_vec(), DEFAULT_FIELD),
8734 1,
8735 64,
8736 QueryControl::Cancel(CancelToken::new()),
8737 )
8738 .expect("replacement cache");
8739 assert!(store.stats().expect("both dictionaries").cache_bytes > bytes);
8740 store.close().expect("close");
8741 assert_eq!(
8742 store
8743 .accounting
8744 .audit()
8745 .expect("closed accounting")
8746 .cache_bytes,
8747 bytes
8748 );
8749 drop(dictionary);
8750 assert_eq!(
8751 store.accounting.audit().expect("second owner").cache_bytes,
8752 bytes
8753 );
8754 assert_eq!(
8755 second_owner.view.iter().collect::<Vec<_>>(),
8756 vec![b"alpha".as_slice(), b"alphabet", b"beta"]
8757 );
8758 drop(second_owner);
8759 assert_eq!(store.accounting.audit().expect("last owner").cache_bytes, 0);
8760 assert_eq!(
8761 store.accounting.audit().expect("scratch").temporary_bytes,
8762 0
8763 );
8764 }
8765
8766 #[test]
8767 fn astra_12_refused_vocabulary_is_not_published() {
8768 use crate::fts::{index::DEFAULT_FIELD, query::LexicalQuery};
8769 use crate::ingest::{
8770 DocId, DocumentVersion, IngestBatch, IngestDocument, Revision, StoreLexicalError,
8771 };
8772 let directory = tempdir().expect("directory");
8773 let writer = Store::open(directory.path(), OpenOptions::default()).expect("open");
8774 writer
8775 .ingest(IngestBatch::new(vec![
8776 IngestDocument::new(
8777 DocumentVersion::new(DocId::new(1), Revision::new(1)),
8778 vec![1.0, 0.0],
8779 )
8780 .with_text("alpha alphabet beta"),
8781 ]))
8782 .expect("ingest");
8783 writer.seal().expect("seal");
8784 writer.close().expect("close writer");
8785 let store = Store::open(
8786 directory.path(),
8787 OpenOptions::default().with_max_temp_bytes(1),
8788 )
8789 .expect("small budget");
8790 for _ in 0..2 {
8791 let result = store.search_lexical_structured(
8792 &LexicalQuery::prefix(b"alph".to_vec(), DEFAULT_FIELD),
8793 1,
8794 64,
8795 QueryControl::Cancel(CancelToken::new()),
8796 );
8797 assert!(
8798 matches!(
8799 result,
8800 Err(StoreLexicalError::Query(super::QueryError::Store(
8801 StoreError::BudgetExceeded {
8802 component: "temporary",
8803 ..
8804 }
8805 )))
8806 ),
8807 "{result:?}"
8808 );
8809 assert_eq!(store.stats().expect("released").temporary_bytes, 0);
8810 let entry = store.lexical_index_cache.entry.lock().expect("assembly");
8811 assert!(
8812 entry
8813 .as_ref()
8814 .expect("assembled")
8815 .assembly
8816 .vocabulary
8817 .lock()
8818 .expect("dictionary")
8819 .is_none()
8820 );
8821 }
8822 store.close().expect("close");
8823 assert_eq!(
8824 store
8825 .accounting
8826 .audit()
8827 .expect("released cache")
8828 .cache_bytes,
8829 0
8830 );
8831 }
8832
8833 #[test]
8834 fn astra_03_lexical_assembly_charge_survives_cache_eviction() {
8835 use crate::fts::index::DEFAULT_FIELD;
8836 use crate::fts::search::TermQuery;
8837 use crate::ingest::{DocId, DocumentVersion, IngestBatch, IngestDocument, Revision};
8838 let directory = tempdir().expect("directory");
8839 let store = Store::open(directory.path(), OpenOptions::default()).expect("open");
8840 store
8841 .ingest(IngestBatch::new(vec![
8842 IngestDocument::new(
8843 DocumentVersion::new(DocId::new(1), Revision::new(1)),
8844 vec![1.0, 0.0],
8845 )
8846 .with_text("alpha beta"),
8847 ]))
8848 .expect("ingest");
8849 store
8850 .search_lexical(
8851 &TermQuery::flat(vec![b"alpha".to_vec()], &[DEFAULT_FIELD]),
8852 1,
8853 QueryControl::Cancel(CancelToken::new()),
8854 )
8855 .expect("build lexical cache");
8856 let (first, second) = {
8857 let entry = store.lexical_index_cache.entry.lock().expect("cache lock");
8858 let assembly = &entry.as_ref().expect("cached assembly").assembly;
8859 (Arc::clone(assembly), Arc::clone(assembly))
8860 };
8861 let charged = store.stats().expect("charged stats").cache_bytes;
8862 let evicted = store
8863 .lexical_index_cache
8864 .entry
8865 .lock()
8866 .expect("cache lock")
8867 .take();
8868 drop(evicted);
8869 assert_eq!(
8870 first.index.document_count(),
8871 1,
8872 "eviction left the query data alive"
8873 );
8874 assert_eq!(
8875 store.stats().expect("held stats").cache_bytes,
8876 charged,
8877 "cache eviction must not release a query-owned assembly's charge"
8878 );
8879 drop(first);
8880 assert_eq!(
8881 store.stats().expect("second owner stats").cache_bytes,
8882 charged
8883 );
8884 drop(second);
8885 let released = store.stats().expect("released stats").cache_bytes;
8886 assert_eq!(
8887 released, 0,
8888 "the final owner must release the full assembly charge"
8889 );
8890 println!("assembly cache bytes: retained={charged}, after final owner={released}");
8891 store.close().expect("close");
8892 }
8893
8894 const CEILING_ROWS: [[f32; 4]; 3] = [
8896 [0.5, -0.25, 0.75, 0.125],
8897 [3.0, 1.5, -2.5, 0.25],
8898 [0.0, 0.0, 0.0, 0.0],
8899 ];
8900
8901 fn publish_f32_segment(directory: &std::path::Path) {
8903 let created = Store::open(directory, OpenOptions::default()).expect("create f32 store");
8904 created.close().expect("close f32 store");
8905
8906 let schema = Schema::timestamp_only();
8907 let mut columns = crate::meta::ColumnStoreBuilder::new(schema.clone());
8908 for row in 0..CEILING_ROWS.len() {
8909 columns
8910 .push_row(i64::try_from(row).expect("row timestamp"), &[])
8911 .expect("timestamp row");
8912 }
8913 let columns = columns.finish().expect("columns");
8914 let alive =
8915 crate::meta::AliveSet::new(u32::try_from(CEILING_ROWS.len()).expect("row count"));
8916 let rescore = CEILING_ROWS.iter().flatten().copied().collect::<Vec<f32>>();
8917 let codes = rescore
8918 .iter()
8919 .flat_map(|value| value.to_bits().to_le_bytes())
8920 .collect::<Vec<u8>>();
8921 let policy = DurabilityPolicy::new(DurabilityMode::Derived, CommitTier::None)
8922 .expect("f32 fixture durability");
8923 let meta = crate::segment::writer::write_segment(
8924 &StdVfs,
8925 directory,
8926 crate::segment::writer::SegmentBuild {
8927 id: crate::segment::SegmentId::new(11, [7; 10]),
8928 scheme: 0,
8929 dims: 4,
8930 codes: &codes,
8931 factors: crate::segment::writer::SegmentFactors::F32,
8932 rescore: &rescore,
8933 columns: &columns,
8934 alive: &alive,
8935 },
8936 policy,
8937 )
8938 .expect("write scheme-0 segment");
8939 commit_manifest(
8940 &StdVfs,
8941 directory,
8942 &Manifest {
8943 generation: 1,
8944 log_seq: 0,
8945 segments: vec![meta],
8946 epochs: Vec::new(),
8947 epoch_alias: None,
8948 schema,
8949 },
8950 policy,
8951 )
8952 .expect("publish scheme-0 segment");
8953 }
8954
8955 fn ingest_ceiling_rows(store: &Store, keep_active: usize, one_segment_per_row: bool) {
8958 use crate::ingest::{DocId, DocumentVersion, IngestBatch, IngestDocument, Revision};
8959
8960 let sealed = CEILING_ROWS.len() - keep_active;
8961 for (index, row) in CEILING_ROWS.iter().enumerate() {
8962 let version = DocumentVersion::new(
8963 DocId::new(u128::try_from(index).expect("document id")),
8964 Revision::new(1),
8965 );
8966 store
8967 .ingest(IngestBatch::new(vec![IngestDocument::new(
8968 version,
8969 row.to_vec(),
8970 )]))
8971 .expect("ingest ceiling row");
8972 let seal_here = if one_segment_per_row {
8973 index < sealed
8974 } else {
8975 index + 1 == sealed
8976 };
8977 if seal_here {
8978 store.seal().expect("seal ceiling rows");
8979 }
8980 }
8981 }
8982
8983 fn assert_remembered_ceiling_is_exact(store: &Store, expected_segments: usize) {
8986 let queries: [Vec<f32>; 3] = [
8987 vec![0.0, 0.0, 0.0, 0.0],
8988 vec![1.0, -1.0, 0.5, 2.0],
8989 vec![7.5, 0.25, -3.5, 0.125],
8990 ];
8991 let admitted = store
8992 .admit_vector_search(SearchOptions::default())
8993 .expect("admit ceiling query");
8994 let snapshot = &admitted.snapshot;
8995 let active = &admitted.active_segment;
8996 assert_eq!(snapshot.segments().len(), expected_segments);
8997 for segment in snapshot.segments() {
8998 assert!(
8999 segment.cached_vector_ceiling_norm_range().is_none(),
9000 "a segment remembered an enclosure before any query walked it"
9001 );
9002 }
9003
9004 for query in &queries {
9005 let cold = super::exact_vector_ceiling(snapshot, active, query).expect("cold ceiling");
9006 for segment in snapshot.segments() {
9007 let remembered = segment
9008 .cached_vector_ceiling_norm_range()
9009 .expect("a walked segment must remember its enclosure");
9010 let recomputed = super::compute_segment_vector_ceiling_norm_range(segment)
9011 .expect("recomputed enclosure");
9012 assert_eq!(
9013 remembered.endpoint_bits(),
9014 recomputed.endpoint_bits(),
9015 "remembered enclosure differs from the recomputed one"
9016 );
9017 }
9018 let warm = super::exact_vector_ceiling(snapshot, active, query).expect("warm ceiling");
9019 let recomputed = super::fold_exact_vector_ceiling(
9020 snapshot,
9021 active,
9022 query,
9023 super::compute_segment_vector_ceiling_norm_range,
9024 )
9025 .expect("recomputed ceiling");
9026 assert_eq!(
9027 cold.to_bits(),
9028 warm.to_bits(),
9029 "the remembered ceiling changed between queries"
9030 );
9031 assert_eq!(
9032 cold.to_bits(),
9033 recomputed.to_bits(),
9034 "the remembered ceiling is not the recomputed ceiling"
9035 );
9036 }
9037 }
9038
9039 #[test]
9040 fn astra_00_elapsed_includes_ceiling_work() {
9041 use std::time::{Duration, Instant};
9042 struct CeilingClock {
9043 base: Instant,
9044 snapshot: Arc<super::PublishedSnapshot>,
9045 }
9046 impl super::MonotonicClock for CeilingClock {
9047 fn now(&self) -> Instant {
9048 if self
9051 .snapshot
9052 .segments()
9053 .iter()
9054 .all(|segment| segment.cached_vector_ceiling_norm_range().is_some())
9055 {
9056 self.base + Duration::from_secs(7)
9057 } else {
9058 self.base
9059 }
9060 }
9061 }
9062 let directory = tempdir().expect("ceiling timer store");
9063 let mut store = Store::open(directory.path(), OpenOptions::default()).expect("open");
9064 ingest_ceiling_rows(&store, 0, false);
9065 let admitted = store
9066 .admit_vector_search(SearchOptions::default())
9067 .expect("snapshot");
9068 assert_eq!(admitted.snapshot.segments().len(), 1);
9069 assert!(
9070 admitted
9071 .snapshot
9072 .segments()
9073 .iter()
9074 .all(|segment| { segment.cached_vector_ceiling_norm_range().is_none() })
9075 );
9076 let snapshot = Arc::clone(&admitted.snapshot);
9077 drop(admitted);
9078 store.clock = Arc::new(CeilingClock {
9079 base: Instant::now(),
9080 snapshot,
9081 });
9082 let outcome = store
9083 .search(
9084 crate::ingest::SearchRequest::new(&[1.0, 0.0, 0.0, 0.0]),
9085 1,
9086 SearchOptions::default().with_tier(SearchTier::Exact),
9087 QueryControl::Cancel(CancelToken::new()),
9088 )
9089 .expect("exact query");
9090 assert_eq!(outcome.diagnostics.elapsed, Duration::from_secs(7));
9091 store.clock = Arc::new(super::SystemMonotonicClock);
9093 store.close().expect("close");
9094 }
9095
9096 #[test]
9097 fn the_remembered_vector_ceiling_is_the_recomputed_ceiling_for_every_scheme() {
9098 let bit4_directory = tempdir().expect("Bit4 ceiling store directory");
9099 let bit4 = Store::open(bit4_directory.path(), OpenOptions::default())
9100 .expect("open Bit4 ceiling store");
9101 ingest_ceiling_rows(&bit4, 0, false);
9102 assert_remembered_ceiling_is_exact(&bit4, 1);
9103 bit4.close().expect("close Bit4 ceiling store");
9104
9105 let split_directory = tempdir().expect("multi-segment ceiling store directory");
9106 let split = Store::open(split_directory.path(), OpenOptions::default())
9107 .expect("open multi-segment ceiling store");
9108 ingest_ceiling_rows(&split, 0, true);
9109 assert_remembered_ceiling_is_exact(&split, CEILING_ROWS.len());
9110 split.close().expect("close multi-segment ceiling store");
9111
9112 let int8_directory = tempdir().expect("Int8 ceiling store directory");
9113 let int8 = Store::open_with_test_dependencies(
9114 int8_directory.path(),
9115 OpenOptions::default(),
9116 StoreTestDependencies::new(Arc::new(StdVfs), Arc::new(super::SystemMonotonicClock))
9117 .with_vector_seal_scheme(crate::quant::QuantScheme::Int8),
9118 )
9119 .expect("open Int8 ceiling store");
9120 ingest_ceiling_rows(&int8, 0, false);
9121 assert_remembered_ceiling_is_exact(&int8, 1);
9122 int8.close().expect("close Int8 ceiling store");
9123
9124 let f32_directory = tempdir().expect("F32 ceiling store directory");
9125 publish_f32_segment(f32_directory.path());
9126 let f32_store = Store::open(f32_directory.path(), OpenOptions::default())
9127 .expect("open F32 ceiling store");
9128 assert_remembered_ceiling_is_exact(&f32_store, 1);
9129 f32_store.close().expect("close F32 ceiling store");
9130
9131 let active_directory = tempdir().expect("active ceiling store directory");
9132 let active = Store::open(active_directory.path(), OpenOptions::default())
9133 .expect("open active ceiling store");
9134 ingest_ceiling_rows(&active, 2, false);
9135 assert!(
9136 !active
9137 .admit_vector_search(SearchOptions::default())
9138 .expect("admit active ceiling query")
9139 .active_segment
9140 .is_empty(),
9141 "the active-segment case sealed every row"
9142 );
9143 assert_remembered_ceiling_is_exact(&active, 1);
9144 active.close().expect("close active ceiling store");
9145 }
9146
9147 #[test]
9148 fn exact_and_graph_tiers_preserve_quantizer_validation_errors() {
9149 let directory = tempdir().expect("store directory");
9150 let store = Store::open(directory.path(), OpenOptions::default()).expect("open");
9151 let query = vec![0.0_f32; crate::kernels::MAX_DOT_I8_DIMENSION + 1];
9152
9153 for tier in [
9154 SearchTier::Exact,
9155 SearchTier::Graph(GraphSearchOptions::default()),
9156 ] {
9157 let error = store
9158 .search(
9159 crate::ingest::SearchRequest::new(&query),
9160 1,
9161 SearchOptions::default().with_tier(tier),
9162 QueryControl::Cancel(CancelToken::new()),
9163 )
9164 .expect_err("oversized quantized query must retain its error");
9165 assert!(matches!(
9166 error,
9167 super::QueryError::Scan(crate::scan::ScanError::Quant(
9168 crate::quant::QuantError::DimensionTooLarge { actual, maximum }
9169 )) if actual == query.len()
9170 && maximum == crate::kernels::MAX_DOT_I8_DIMENSION
9171 ));
9172 }
9173 }
9174
9175 #[cfg(feature = "allocation-audit")]
9176 #[test]
9177 fn exact_tier_query_does_not_prepare_quantized_forms() {
9178 let directory = tempdir().expect("store directory");
9179 let store = Store::open(directory.path(), OpenOptions::default()).expect("open");
9180 let query = [0.25_f32; 128];
9181 let options = SearchOptions::default().with_tier(SearchTier::Exact);
9182 let warm = store
9183 .search(
9184 crate::ingest::SearchRequest::new(&query),
9185 1,
9186 options,
9187 QueryControl::Cancel(CancelToken::new()),
9188 )
9189 .expect("warm exact-tier search");
9190 assert!(warm.candidates.is_empty());
9191 let control = QueryControl::Cancel(CancelToken::new());
9192
9193 let (outcome, report) = crate::allocation_audit::audit_engine_path(|| {
9194 store.search(
9195 crate::ingest::SearchRequest::new(&query),
9196 1,
9197 options,
9198 control,
9199 )
9200 });
9201 let outcome = outcome.expect("audited exact-tier search");
9202 assert!(outcome.candidates.is_empty());
9203 assert_eq!(report.allocations, 0, "Exact tier prepared quantized forms");
9204 assert_eq!(report.attributed_bytes, 0);
9205 assert_eq!(report.unattributed_bytes, 0);
9206 }
9207
9208 #[test]
9209 fn astra_18_scoped_work_is_joined_before_cancel_returns() {
9210 use std::sync::atomic::{AtomicBool, Ordering};
9211 use std::sync::mpsc;
9212 let worker = super::pool::LexicalWorker::start().expect("worker");
9213 let token = CancelToken::new();
9214 let work_token = token.clone();
9215 let finished = AtomicBool::new(false);
9216 let mut borrowed_rows = Vec::new();
9217 let (started_tx, started_rx) = mpsc::sync_channel(0);
9218 let (release_tx, release_rx) = mpsc::sync_channel(0);
9219 let worker_rows = &mut borrowed_rows;
9220 let worker_finished = &finished;
9221 let (caller, lexical) = worker
9222 .run_scoped(
9223 move || {
9224 started_tx.send(()).expect("entered borrowed job");
9225 release_rx.recv().expect("caller has cancelled");
9226 assert!(work_token.is_cancelled());
9227 worker_rows.push(42_u32);
9228 worker_finished.store(true, Ordering::Release);
9229 Err::<(), _>(crate::fusion::FusionError::Cancelled { partial: false })
9230 },
9231 || {
9232 started_rx.recv().expect("worker entered");
9233 token.cancel();
9234 assert!(!finished.load(Ordering::Acquire));
9235 release_tx.send(()).expect("release borrowed job");
9236 Err::<(), _>(crate::fusion::FusionError::Cancelled { partial: false })
9237 },
9238 )
9239 .expect("scoped dispatch");
9240 assert!(finished.load(Ordering::Acquire));
9241 assert_eq!(borrowed_rows, [42]);
9242 assert_eq!(
9243 resolve_hybrid_leg_results(caller, lexical.expect("worker returned")),
9244 Err(crate::fusion::FusionError::Cancelled { partial: false }),
9245 );
9246 worker.stop_and_join().expect("join worker");
9247 }
9248
9249 #[test]
9250 fn astra_18_control_error_precedence_is_completion_order_independent() {
9251 use crate::fusion::{FusionError, FusionLeg, LegFailureKind};
9252 use std::sync::mpsc;
9253 let failures = [
9254 FusionError::Leg {
9255 leg: FusionLeg::Vector,
9256 kind: LegFailureKind::Scan,
9257 detail: "scan".to_owned(),
9258 },
9259 FusionError::Cancelled { partial: false },
9260 FusionError::Timeout { partial: false },
9261 FusionError::ReadCancelled { partial: false },
9262 ];
9263 let worker = super::pool::LexicalWorker::start().expect("worker");
9264 for (a, vector) in failures.iter().enumerate() {
9265 for (b, lexical) in failures.iter().enumerate() {
9266 for lexical_first in [false, true] {
9267 let (lexical_tx, lexical_rx) = mpsc::sync_channel(1);
9268 let (vector_tx, vector_rx) = mpsc::sync_channel(1);
9269 let (vector_result, lexical_result) = worker
9270 .run_scoped(
9271 move || {
9272 if !lexical_first {
9273 vector_rx.recv().expect("vector completed");
9274 }
9275 if lexical_first {
9276 lexical_tx.send(()).expect("lexical completed");
9277 }
9278 Err::<(), _>(lexical.clone())
9279 },
9280 || {
9281 if lexical_first {
9282 lexical_rx.recv().expect("lexical completed");
9283 }
9284 if !lexical_first {
9285 vector_tx.send(()).expect("vector completed");
9286 }
9287 Err::<(), _>(vector.clone())
9288 },
9289 )
9290 .expect("dispatch");
9291 assert_eq!(
9292 resolve_hybrid_leg_results(vector_result, lexical_result.expect("joined")),
9293 Err(failures.get(a.max(b)).expect("known failure index").clone()),
9294 "vector={a} lexical={b} lexical_first={lexical_first}",
9295 );
9296 }
9297 }
9298 }
9299 worker.stop_and_join().expect("join worker");
9300 }
9301
9302 #[test]
9303 fn hybrid_dual_failure_precedence_is_independent_of_leg_completion_order() {
9304 use crate::fusion::{FusionError, FusionLeg, LegFailureKind};
9305
9306 let vector_failure = FusionError::Leg {
9307 leg: FusionLeg::Vector,
9308 kind: LegFailureKind::Scan,
9309 detail: "vector failure".to_owned(),
9310 };
9311 let lexical_failure = FusionError::Leg {
9312 leg: FusionLeg::Lexical,
9313 kind: LegFailureKind::Lexical,
9314 detail: "lexical failure".to_owned(),
9315 };
9316 assert_eq!(
9317 resolve_hybrid_leg_results::<(), ()>(
9318 Err(vector_failure.clone()),
9319 Err(lexical_failure.clone()),
9320 ),
9321 Err(vector_failure.clone()),
9322 );
9323 assert_eq!(
9324 resolve_hybrid_leg_results::<(), ()>(
9325 Err(vector_failure.clone()),
9326 Err(FusionError::Cancelled { partial: false }),
9327 ),
9328 Err(FusionError::Cancelled { partial: false }),
9329 );
9330 assert_eq!(
9331 resolve_hybrid_leg_results::<(), ()>(
9332 Err(FusionError::Timeout { partial: false }),
9333 Err(FusionError::ReadCancelled { partial: false }),
9334 ),
9335 Err(FusionError::ReadCancelled { partial: false }),
9336 );
9337 assert_eq!(
9338 resolve_hybrid_leg_results::<(), ()>(
9339 Err(FusionError::Cancelled { partial: false }),
9340 Err(FusionError::Timeout { partial: false }),
9341 ),
9342 Err(FusionError::Timeout { partial: false }),
9343 );
9344 assert_eq!(
9345 resolve_hybrid_leg_results::<(), ()>(
9346 Err(FusionError::ReadCancelled { partial: false }),
9347 Err(lexical_failure),
9348 ),
9349 Err(FusionError::ReadCancelled { partial: false }),
9350 );
9351 }
9352
9353 #[test]
9354 fn accounting_errors_report_the_typed_context() {
9355 let statistics = StoreError::Statistics {
9356 component: "mapped resident bytes",
9357 source: std::io::Error::other("mincore failed"),
9358 };
9359 assert_eq!(
9360 statistics.to_string(),
9361 "store statistics mapped resident bytes: mincore failed"
9362 );
9363 assert_eq!(
9364 statistics.source().map(ToString::to_string),
9365 Some("mincore failed".to_owned())
9366 );
9367
9368 let budget = StoreError::BudgetExceeded {
9369 needed: 65,
9370 budget: 64,
9371 component: "wal",
9372 };
9373 assert_eq!(
9374 budget.to_string(),
9375 "store wal allocation needs 65 bytes, budget is 64 bytes"
9376 );
9377 assert!(budget.source().is_none());
9378
9379 let allocation = StoreError::AllocationFailed {
9380 needed: 23,
9381 component: "snapshot",
9382 };
9383 assert_eq!(
9384 allocation.to_string(),
9385 "store snapshot allocator rejected 23 bytes"
9386 );
9387 assert!(allocation.source().is_none());
9388 }
9389
9390 #[test]
9391 fn open_options_resolve_durability_before_filesystem_mutation() {
9392 let parent = tempdir().expect("parent directory");
9393 let store_path = parent.path().join("not-created");
9394 let options =
9395 OpenOptions::new().with_durability(DurabilityMode::Attached, CommitTier::Durable);
9396
9397 let error = match Store::open(&store_path, options) {
9398 Ok(store) => {
9399 drop(store);
9400 panic!("attached durability unexpectedly opened")
9401 }
9402 Err(error) => error,
9403 };
9404
9405 assert!(matches!(
9406 error,
9407 StoreError::Durability(DurabilityPolicyError::AttachedNotYetSupported)
9408 ));
9409 assert!(!store_path.exists(), "rejected open created store files");
9410 }
9411
9412 #[test]
9413 fn memory_vfs_store_open_does_not_create_host_directory() {
9414 let parent = tempdir().expect("parent directory");
9415 let store_path = parent.path().join("memory-only-store");
9416 let dependencies = StoreTestDependencies::new(
9417 Arc::new(MemoryVfs::new()),
9418 Arc::new(ManualMonotonicClock::new()),
9419 );
9420
9421 let store =
9422 Store::open_with_test_dependencies(&store_path, OpenOptions::read_only(), dependencies)
9423 .expect("MemoryVfs-backed store open");
9424 store.close().expect("MemoryVfs-backed store close");
9425
9426 assert!(
9427 !store_path.exists(),
9428 "MemoryVfs-backed Store touched the host filesystem"
9429 );
9430 }
9431
9432 #[test]
9433 fn open_refuses_a_manifest_ahead_of_the_checked_wal() {
9434 let directory = tempdir().expect("store directory");
9435 let policy = DurabilityPolicy::new(DurabilityMode::Derived, CommitTier::Ordered)
9436 .expect("derived policy");
9437 commit_manifest(
9438 &StdVfs,
9439 directory.path(),
9440 &Manifest {
9441 generation: 1,
9442 log_seq: 1,
9443 segments: Vec::new(),
9444 epochs: Vec::new(),
9445 epoch_alias: None,
9446 schema: Schema::new(Vec::new()).expect("schema"),
9447 },
9448 policy,
9449 )
9450 .expect("manifest");
9451
9452 let error = match Store::open(directory.path(), OpenOptions::default()) {
9453 Ok(store) => {
9454 drop(store);
9455 panic!("manifest ahead of missing WAL unexpectedly opened")
9456 }
9457 Err(error) => error,
9458 };
9459
9460 assert!(matches!(
9461 error,
9462 StoreError::Manifest(ManifestError::AheadOfLog {
9463 snapshot: 1,
9464 durable: 0
9465 })
9466 ));
9467 }
9468}