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