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