1use std::sync::Arc;
4use std::sync::atomic::{AtomicU64, Ordering};
5
6use crate::graph::search::{FilteredGraphSearchOutcome, GraphSearchCounters, GraphSearchRequest};
7use crate::ingest::{RowSource, SearchCandidate, SearchRequest};
8use crate::lifecycle::{
9 ExactScanMemory, GraphSearchOptions, PublishedSnapshot, QueryCancellation, QueryControl,
10 QueryError, SearchOptions, SearchTier, SnapshotLease, Store, StoreError, StoreState,
11 auto_graph_search_options, auto_uses_full_precision, exact_rescore_rows, rescored_scan,
12};
13use crate::meta::{ColumnStore, ColumnStoreBuilder, EvalError, Predicate, evaluate};
14use crate::quant::{Bit4Query, Int8Query, prepare_bit4_query, prepare_int8_query};
15use crate::scan::{
16 Int8Factors, ScanError, ScanQuery, ScanRequest as VectorScanRequest, ScanRows, ScanStats,
17 gather_top_k, scan_partition,
18};
19use crate::segment::layout::RegionKind;
20
21use super::{
22 ALLOW_LIST_ROWS_THRESHOLD, PlanError, PlanFallback, PlanNode, SegmentBranch, SegmentPlan,
23 SegmentTier, choose_scan_branch, segment_may_match, validate_predicate,
24};
25#[cfg(any(test, feature = "test-seams"))]
26use super::{
27 MetadataControllerError, MetadataExecutionReceipt, MetadataQueryContext, MetadataTestController,
28};
29
30const FILTERED_VISITED_BUDGET_MULTIPLIER: usize = 2;
34
35#[derive(Clone, Debug, PartialEq)]
37pub struct FilteredSearchOutcome {
38 pub candidates: Vec<SearchCandidate>,
40 pub stats: ScanStats,
42 pub generation: u64,
44 pub plans: Vec<SegmentPlan>,
46 pub diagnostics: crate::diag::QueryDiagnostics,
48}
49
50#[derive(Debug)]
52pub enum FilteredSearchError {
53 Plan(PlanError),
55 Query(QueryError),
57 ActiveMetadata(String),
59 PlanReportMismatch {
61 reported: SegmentBranch,
63 executed: SegmentBranch,
65 },
66 InvalidPlanNode(&'static str),
68 #[cfg(any(test, feature = "test-seams"))]
70 TestController(MetadataControllerError),
71}
72
73impl std::fmt::Display for FilteredSearchError {
74 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
75 match self {
76 Self::Plan(error) => error.fmt(formatter),
77 Self::Query(error) => error.fmt(formatter),
78 Self::ActiveMetadata(detail) => {
79 write!(formatter, "active metadata construction failed: {detail}")
80 }
81 Self::PlanReportMismatch { reported, executed } => write!(
82 formatter,
83 "reported branch {reported:?} differs from executed branch {executed:?}"
84 ),
85 Self::InvalidPlanNode(detail) => {
86 write!(formatter, "invalid vector plan node: {detail}")
87 }
88 #[cfg(any(test, feature = "test-seams"))]
89 Self::TestController(error) => error.fmt(formatter),
90 }
91 }
92}
93
94impl std::error::Error for FilteredSearchError {
95 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
96 match self {
97 Self::Plan(error) => Some(error),
98 Self::Query(error) => Some(error),
99 Self::ActiveMetadata(_)
100 | Self::PlanReportMismatch { .. }
101 | Self::InvalidPlanNode(_) => None,
102 #[cfg(any(test, feature = "test-seams"))]
103 Self::TestController(error) => Some(error),
104 }
105 }
106}
107
108impl From<PlanError> for FilteredSearchError {
109 fn from(error: PlanError) -> Self {
110 Self::Plan(error)
111 }
112}
113
114impl From<QueryError> for FilteredSearchError {
115 fn from(error: QueryError) -> Self {
116 Self::Query(error)
117 }
118}
119
120#[cfg(any(test, feature = "test-seams"))]
121impl From<MetadataControllerError> for FilteredSearchError {
122 fn from(error: MetadataControllerError) -> Self {
123 Self::TestController(error)
124 }
125}
126
127impl Store {
128 pub fn search_filtered(
134 &self,
135 request: SearchRequest<'_>,
136 predicate: &Predicate,
137 k: usize,
138 options: impl Into<SearchOptions>,
139 control: QueryControl,
140 ) -> Result<FilteredSearchOutcome, FilteredSearchError> {
141 execute_store(
142 self,
143 request,
144 predicate,
145 k,
146 options.into(),
147 control.with_clock(Arc::clone(&self.clock)),
148 )
149 }
150}
151
152fn execute_store(
153 store: &Store,
154 request: SearchRequest<'_>,
155 predicate: &Predicate,
156 k: usize,
157 options: SearchOptions,
158 control: QueryControl,
159) -> Result<FilteredSearchOutcome, FilteredSearchError> {
160 let started = std::time::Instant::now();
161 let state = store
162 .state
163 .lock()
164 .map_err(|_| QueryError::Store(StoreError::Synchronization { component: "state" }))?;
165 match *state {
166 StoreState::Open => {}
167 StoreState::Closing => return Err(QueryError::Store(StoreError::Closing).into()),
168 StoreState::Closed => return Err(QueryError::Store(StoreError::Closed).into()),
169 }
170 let active_guard = store.active.lock().map_err(|_| {
171 QueryError::Store(StoreError::Synchronization {
172 component: "active segment",
173 })
174 })?;
175 let active_state = active_guard
176 .as_ref()
177 .ok_or(QueryError::Store(StoreError::Closed))?;
178 let generation = active_state.generation;
179 let active = Arc::clone(&active_state.segment);
180 let snapshot = store
181 .snapshot
182 .read()
183 .map_err(|_| {
184 QueryError::Store(StoreError::Synchronization {
185 component: "published snapshot",
186 })
187 })?
188 .as_ref()
189 .cloned()
190 .ok_or(QueryError::Store(StoreError::Closed))?;
191
192 let schema = store.schema.clone();
193 validate_predicate(predicate, &schema)?;
194 if matches!(options.tier(), SearchTier::Graph(_)) {
195 for segment in snapshot.segments() {
196 if !has_graph(segment) {
197 return Err(QueryError::Store(StoreError::GraphUnavailable {
198 segment_id: segment.meta().id,
199 })
200 .into());
201 }
202 }
203 }
204
205 store.active_queries.fetch_add(1, Ordering::Relaxed);
206 let query_guard = ActiveFilteredQuery {
207 count: &store.active_queries,
208 };
209 drop(active_guard);
210 drop(state);
211
212 let lease = SnapshotLease::new_at(Arc::clone(&snapshot), generation);
213 let cancellation = QueryCancellation::new(&control, &lease);
214 cancellation.check_graph().map_err(map_scan_error)?;
215 #[cfg(any(test, feature = "test-seams"))]
216 let metadata_controller = store.metadata_test_controller.as_deref();
217 #[cfg(any(test, feature = "test-seams"))]
218 let metadata_query = metadata_controller
219 .map(MetadataTestController::begin_query)
220 .transpose()?
221 .flatten();
222 let score = || {
223 execute_pinned(
224 &snapshot,
225 &active,
226 &schema,
227 &store.accounting,
228 generation,
229 store.epoch_identity(),
230 request,
231 predicate,
232 k,
233 options,
234 &cancellation,
235 started,
236 #[cfg(any(test, feature = "test-seams"))]
237 metadata_controller,
238 #[cfg(any(test, feature = "test-seams"))]
239 metadata_query.as_ref(),
240 #[cfg(any(test, feature = "test-seams"))]
241 store.vector_fault_controller.as_ref(),
242 )
243 };
244 #[cfg(any(test, feature = "test-seams"))]
245 let result = crate::kernels::vector_fault::run_store_scoring(
246 store.kernel_fault_controller.as_ref(),
247 score,
248 );
249 #[cfg(not(any(test, feature = "test-seams")))]
250 let result = score();
251 #[cfg(any(test, feature = "test-seams"))]
252 {
253 if let Some(controller) = store.vector_fault_controller.as_ref() {
254 controller.finalize_search(result.is_ok());
255 }
256 }
257 drop(query_guard);
258 result
259}
260
261#[allow(clippy::too_many_arguments)]
262fn execute_pinned(
263 snapshot: &PublishedSnapshot,
264 active: &crate::ingest::ActiveSegment,
265 schema: &crate::meta::Schema,
266 accounting: &Arc<crate::lifecycle::stats::Accounting>,
267 generation: u64,
268 epoch: Option<crate::epoch::EpochIdentity>,
269 request: SearchRequest<'_>,
270 predicate: &Predicate,
271 k: usize,
272 options: SearchOptions,
273 cancellation: &QueryCancellation<'_>,
274 started: std::time::Instant,
275 #[cfg(any(test, feature = "test-seams"))] metadata_controller: Option<&MetadataTestController>,
276 #[cfg(any(test, feature = "test-seams"))] metadata_query: Option<&MetadataQueryContext>,
277 #[cfg(any(test, feature = "test-seams"))] vector_fault_controller: Option<
278 &crate::scan::vector_fault::VectorFaultController,
279 >,
280) -> Result<FilteredSearchOutcome, FilteredSearchError> {
281 let mut candidates = Vec::new();
282 let mut rescore_counters = crate::diag::ScanRescoreCounters::default();
283 if let Some(options) = options.scan_rescore() {
284 options.frontier(k, 0)?;
285 }
286 let mut plans = Vec::new();
287 let mut stats = MutableStats::default();
288 let mut graph_stats = crate::ingest::GraphSearchStats::default();
289 let bit4_query = prepare_bit4_query(request.vector(), 0)
290 .map_err(ScanError::Quant)
291 .map_err(QueryError::Scan)?;
292 let int8_query = prepare_int8_query(request.vector())
293 .map_err(ScanError::Quant)
294 .map_err(QueryError::Scan)?;
295 let auto_uses_graph =
296 matches!(options.tier(), SearchTier::Auto) && snapshot.segments().iter().any(has_graph);
297 let auto_graph_options = if auto_uses_graph {
298 Some(auto_graph_search_options(snapshot)?)
299 } else {
300 None
301 };
302 let full_precision = matches!(options.tier(), SearchTier::Exact | SearchTier::Graph(_))
303 || (matches!(options.tier(), SearchTier::Auto)
304 && auto_uses_full_precision(snapshot, active));
305
306 if !active.is_empty() {
307 execute_active_filtered(
308 active,
309 schema,
310 request,
311 predicate,
312 k,
313 options,
314 cancellation,
315 full_precision,
316 &bit4_query,
317 &mut candidates,
318 &mut stats,
319 &mut plans,
320 accounting,
321 &mut rescore_counters,
322 #[cfg(any(test, feature = "test-seams"))]
323 metadata_controller,
324 #[cfg(any(test, feature = "test-seams"))]
325 metadata_query,
326 #[cfg(any(test, feature = "test-seams"))]
327 vector_fault_controller,
328 )?;
329 }
330
331 for segment in ordered_filtered_segments(snapshot, options, auto_uses_graph) {
332 let source = RowSource::Sealed(segment.meta().id);
333 let graph_selected =
334 has_graph(segment) && matches!(options.tier(), SearchTier::Auto | SearchTier::Graph(_));
335 let tier = if graph_selected {
336 SegmentTier::SealedGraph
337 } else {
338 SegmentTier::SealedScan
339 };
340 if !segment_may_match(segment.meta().clustering_key_range, predicate) {
341 #[cfg(any(test, feature = "test-seams"))]
342 record_pruned_receipt(
343 metadata_controller,
344 metadata_query,
345 source,
346 segment.meta().row_count,
347 0,
348 )?;
349 plans.push(
350 SegmentPlan::exact(source, tier, SegmentBranch::Pruned, 0, None)
351 .with_predicate(predicate),
352 );
353 continue;
354 }
355 let columns = match segment.query_columns() {
356 Ok(columns) => columns,
357 Err(error) => {
358 #[cfg(any(test, feature = "test-seams"))]
359 record_column_refusal(metadata_controller, metadata_query, source, &error)?;
360 return Err(QueryError::Store(error).into());
361 }
362 };
363 let alive = match segment.query_alive() {
364 Ok(alive) => alive,
365 Err(error) => {
366 #[cfg(any(test, feature = "test-seams"))]
367 record_alive_refusal(metadata_controller, metadata_query, source, &error)?;
368 return Err(QueryError::Store(error).into());
369 }
370 };
371 let allow_list = evaluate_for_plan(predicate, &columns, &alive)?;
372 let row_count = segment.meta().row_count as usize;
373 let graph_options = match options.tier() {
374 SearchTier::Graph(graph_options) if graph_selected => Some(graph_options),
375 SearchTier::Auto if graph_selected => auto_graph_options,
376 SearchTier::Auto | SearchTier::Exact | SearchTier::Scan | SearchTier::Graph(_) => None,
377 };
378 let graph_outcome = sealed_graph_filtered(
379 segment,
380 request.vector(),
381 predicate,
382 &allow_list,
383 k,
384 graph_options,
385 cancellation,
386 accounting,
387 source,
388 row_count,
389 options,
390 &mut candidates,
391 &mut stats,
392 &mut graph_stats,
393 &mut plans,
394 #[cfg(any(test, feature = "test-seams"))]
395 metadata_controller,
396 #[cfg(any(test, feature = "test-seams"))]
397 metadata_query,
398 #[cfg(any(test, feature = "test-seams"))]
399 vector_fault_controller,
400 )?;
401 if !matches!(graph_outcome, SealedGraphOutcome::Scan) {
402 continue;
403 }
404 let branch = choose_scan_branch(allow_list.cardinality());
405 #[cfg(any(test, feature = "test-seams"))]
406 record_selectivity_decision(
407 metadata_controller,
408 metadata_query,
409 source,
410 allow_list.cardinality(),
411 branch,
412 )?;
413 let scan_reason = (tier == SegmentTier::SealedScan)
414 .then(|| sealed_scan_reason(segment, options.explicit_tier()));
415 let mut plan =
416 SegmentPlan::exact(source, tier, branch, allow_list.cardinality(), scan_reason)
417 .with_predicate(predicate);
418 #[cfg(any(test, feature = "test-seams"))]
419 let receipt_context = metadata_execution_receipt_context(
420 metadata_controller,
421 metadata_query,
422 source,
423 row_count,
424 allow_list.cardinality(),
425 request.vector().len(),
426 true,
427 )?;
428 let frontier = options
429 .scan_rescore()
430 .map(|options| options.frontier(k, allow_list.cardinality()))
431 .transpose()?
432 .unwrap_or(k);
433 let (mut local, mut local_stats, execution) = scan_sealed_filtered(
434 segment,
435 &plan,
436 row_count,
437 request.vector(),
438 &allow_list,
439 frontier,
440 cancellation,
441 full_precision,
442 &bit4_query,
443 &int8_query,
444 #[cfg(any(test, feature = "test-seams"))]
445 receipt_context,
446 )?;
447 let _rescore_memory = if let Some(rescore) = options.scan_rescore() {
448 let mut memory = ExactScanMemory::new(accounting)?;
449 let (counters, exhaustive, _) = rescored_scan::rescore(
450 &mut local,
451 &mut local_stats,
452 exact_rescore_rows(segment)?,
453 request.vector(),
454 k,
455 allow_list.cardinality(),
456 rescore,
457 cancellation,
458 accounting,
459 &mut memory,
460 )?;
461 plan.approximate = !exhaustive;
462 rescored_scan::accumulate(&mut rescore_counters, counters)?;
463 Some(memory)
464 } else {
465 None
466 };
467 #[cfg(any(test, feature = "test-seams"))]
468 let reported_branch =
469 metadata_query.map_or(plan.branch, |query| query.reported_branch(plan.branch));
470 #[cfg(not(any(test, feature = "test-seams")))]
471 let reported_branch = plan.branch;
472 verify_execution_branch(reported_branch, execution.branch)?;
473 append_candidates(
474 local,
475 source,
476 |row| {
477 segment
478 .document_version(row)
479 .map_err(StoreError::Segment)
480 .map_err(QueryError::Store)
481 },
482 &mut candidates,
483 full_precision || segment.meta().scheme == 0 || options.scan_rescore().is_some(),
484 #[cfg(any(test, feature = "test-seams"))]
485 vector_fault_controller,
486 #[cfg(any(test, feature = "test-seams"))]
487 vector_fault_tier(options.tier()),
488 )?;
489 stats.add(local_stats)?;
490 plans.push(plan);
491 retain_global_top_k(&mut candidates, k);
492 }
493
494 candidates.sort_unstable_by(crate::ingest::compare_search_candidates);
495 candidates.truncate(k);
496 cancellation.check_graph().map_err(map_scan_error)?;
497 let stats = stats.finish();
498 accounting.record_plans(&plans).map_err(QueryError::Store)?;
499 let mut diagnostics = crate::diag::QueryDiagnostics::vector(crate::diag::VectorDiagnostics {
500 snapshot_generation: generation,
501 indexed_through_seq: active
502 .indexed_through_seq()
503 .max(crate::wal::LogSeq::new(snapshot.absorbed_through())),
504 approximate: plans.iter().any(|plan| plan.approximate),
505 exact_rescore: candidates.iter().all(|candidate| candidate.exact_score()),
506 requested_k: k,
507 returned: candidates.len(),
508 budget_exhausted: plans.iter().any(|plan| {
509 matches!(
510 plan.fallback,
511 PlanFallback::VisitedBudget | PlanFallback::CandidateShortfall
512 )
513 }),
514 plan: plans.clone(),
515 scan: stats.clone(),
516 graph: graph_stats,
517 epoch,
518 elapsed: started.elapsed(),
519 });
520 diagnostics.scan_rescore = options.scan_rescore().map(|_| rescore_counters);
521 Ok(FilteredSearchOutcome {
522 candidates,
523 stats,
524 generation,
525 plans,
526 diagnostics,
527 })
528}
529
530fn ordered_filtered_segments(
531 snapshot: &PublishedSnapshot,
532 options: SearchOptions,
533 auto_uses_graph: bool,
534) -> Vec<&crate::segment::reader::SegmentReader> {
535 let mut segments = snapshot.segments().iter().collect::<Vec<_>>();
536 if auto_uses_graph || matches!(options.tier(), SearchTier::Graph(_)) {
537 segments.sort_unstable_by(|left, right| {
538 right
539 .meta()
540 .row_count
541 .cmp(&left.meta().row_count)
542 .then_with(|| left.meta().id.cmp(&right.meta().id))
543 });
544 }
545 segments
546}
547
548#[allow(clippy::too_many_arguments)]
549fn execute_active_filtered(
550 active: &crate::ingest::ActiveSegment,
551 schema: &crate::meta::Schema,
552 request: SearchRequest<'_>,
553 predicate: &Predicate,
554 k: usize,
555 options: SearchOptions,
556 cancellation: &QueryCancellation<'_>,
557 full_precision: bool,
558 bit4_query: &Bit4Query,
559 candidates: &mut Vec<SearchCandidate>,
560 stats: &mut MutableStats,
561 plans: &mut Vec<SegmentPlan>,
562 accounting: &Arc<crate::lifecycle::stats::Accounting>,
563 rescore_counters: &mut crate::diag::ScanRescoreCounters,
564 #[cfg(any(test, feature = "test-seams"))] metadata_controller: Option<&MetadataTestController>,
565 #[cfg(any(test, feature = "test-seams"))] metadata_query: Option<&MetadataQueryContext>,
566 #[cfg(any(test, feature = "test-seams"))] vector_fault_controller: Option<
567 &crate::scan::vector_fault::VectorFaultController,
568 >,
569) -> Result<(), FilteredSearchError> {
570 let alive = active.alive().map_err(QueryError::Store)?;
571 let columns = active_columns(active, schema)?;
572 #[cfg(any(test, feature = "test-seams"))]
573 let alive = match metadata_query {
574 Some(query) => query.evaluator_alive(columns.row_count(), alive)?,
575 None => alive,
576 };
577 let allow_list = evaluate_for_plan(predicate, &columns, &alive)?;
578 let branch = choose_scan_branch(allow_list.cardinality());
579 let source = RowSource::Active;
580 #[cfg(any(test, feature = "test-seams"))]
581 record_selectivity_decision(
582 metadata_controller,
583 metadata_query,
584 source,
585 allow_list.cardinality(),
586 branch,
587 )?;
588 let mut plan = SegmentPlan::exact(
589 source,
590 SegmentTier::ActiveScan,
591 branch,
592 allow_list.cardinality(),
593 Some(crate::planner::ScanReason::ActiveSegment),
594 )
595 .with_predicate(predicate);
596 #[cfg(any(test, feature = "test-seams"))]
597 let receipt_context = metadata_execution_receipt_context(
598 metadata_controller,
599 metadata_query,
600 source,
601 active.row_count(),
602 allow_list.cardinality(),
603 request.vector().len(),
604 false,
605 )?;
606 let frontier = options
607 .scan_rescore()
608 .map(|options| options.frontier(k, allow_list.cardinality()))
609 .transpose()?
610 .unwrap_or(k);
611 let (mut local, mut local_stats, execution) = if full_precision {
612 execute_squared_l2_plan(
613 &plan.node,
614 active.vectors(),
615 active.row_count(),
616 request.vector(),
617 &allow_list,
618 frontier,
619 cancellation,
620 #[cfg(any(test, feature = "test-seams"))]
621 receipt_context,
622 )?
623 } else {
624 execute_scan_plan(
625 &plan.node,
626 VectorScanRequest {
627 query: ScanQuery::Bit4(bit4_query),
628 rows: ScanRows::Bit4RowMajor {
629 codes: active.codes(),
630 factors: active.factors(),
631 },
632 row_mask: Some(allow_list.as_roaring()),
633 },
634 active.row_count(),
635 frontier,
636 cancellation,
637 #[cfg(any(test, feature = "test-seams"))]
638 receipt_context,
639 )?
640 };
641 let _rescore_memory = if let Some(rescore) = options.scan_rescore() {
642 let mut memory = ExactScanMemory::new(accounting)?;
643 let (counters, exhaustive, _) = rescored_scan::rescore(
644 &mut local,
645 &mut local_stats,
646 active.vectors(),
647 request.vector(),
648 k,
649 allow_list.cardinality(),
650 rescore,
651 cancellation,
652 accounting,
653 &mut memory,
654 )?;
655 plan.approximate = !exhaustive;
656 rescored_scan::accumulate(rescore_counters, counters)?;
657 Some(memory)
658 } else {
659 None
660 };
661 #[cfg(any(test, feature = "test-seams"))]
662 let reported_branch =
663 metadata_query.map_or(plan.branch, |query| query.reported_branch(plan.branch));
664 #[cfg(not(any(test, feature = "test-seams")))]
665 let reported_branch = plan.branch;
666 verify_execution_branch(reported_branch, execution.branch)?;
667 append_candidates(
668 local,
669 source,
670 |row| Ok(active.document(row)),
671 candidates,
672 full_precision || options.scan_rescore().is_some(),
673 #[cfg(any(test, feature = "test-seams"))]
674 vector_fault_controller,
675 #[cfg(any(test, feature = "test-seams"))]
676 vector_fault_tier(options.tier()),
677 )?;
678 stats.add(local_stats)?;
679 plans.push(plan);
680 retain_global_top_k(candidates, k);
681 Ok(())
682}
683
684#[allow(clippy::too_many_arguments)]
685fn scan_sealed_filtered(
686 segment: &crate::segment::reader::SegmentReader,
687 plan: &SegmentPlan,
688 row_count: usize,
689 query: &[f32],
690 allow_list: &crate::meta::DocBitmap,
691 k: usize,
692 cancellation: &QueryCancellation<'_>,
693 full_precision: bool,
694 bit4_query: &Bit4Query,
695 int8_query: &Int8Query,
696 #[cfg(any(test, feature = "test-seams"))] receipt_context: Option<
697 MetadataExecutionReceiptContext<'_>,
698 >,
699) -> Result<(Vec<LocalCandidate>, ScanStats, ScanExecutionFacts), FilteredSearchError> {
700 if full_precision {
701 let vectors = exact_rescore_rows(segment)?;
702 return execute_squared_l2_plan(
703 &plan.node,
704 vectors,
705 row_count,
706 query,
707 allow_list,
708 k,
709 cancellation,
710 #[cfg(any(test, feature = "test-seams"))]
711 receipt_context,
712 );
713 }
714 match segment.meta().scheme {
715 0 => execute_scan_plan(
716 &plan.node,
717 VectorScanRequest {
718 query: ScanQuery::F32(query),
719 rows: ScanRows::F32BorrowedRowMajor(
720 segment
721 .query_f32_codes()
722 .map_err(StoreError::Segment)
723 .map_err(QueryError::Store)?,
724 ),
725 row_mask: Some(allow_list.as_roaring()),
726 },
727 row_count,
728 k,
729 cancellation,
730 #[cfg(any(test, feature = "test-seams"))]
731 receipt_context,
732 ),
733 2 => {
734 let factors = segment
735 .query_int8_factors_slice()
736 .map_err(StoreError::Segment)
737 .map_err(QueryError::Store)?
738 .iter()
739 .map(|factor| Int8Factors::new(factor.scale, factor.offset))
740 .collect::<Result<Vec<_>, _>>()
741 .map_err(|_| {
742 QueryError::Store(StoreError::Segment(crate::segment::SegmentError::Geometry(
743 "Int8 factor is not finite and non-negative".to_owned(),
744 )))
745 })?;
746 execute_scan_plan(
747 &plan.node,
748 VectorScanRequest {
749 query: ScanQuery::Int8(int8_query),
750 rows: ScanRows::Int8RowMajor {
751 codes: segment
752 .query_int8_codes()
753 .map_err(StoreError::Segment)
754 .map_err(QueryError::Store)?,
755 factors: &factors,
756 },
757 row_mask: Some(allow_list.as_roaring()),
758 },
759 row_count,
760 k,
761 cancellation,
762 #[cfg(any(test, feature = "test-seams"))]
763 receipt_context,
764 )
765 }
766 4 => execute_scan_plan(
767 &plan.node,
768 VectorScanRequest {
769 query: ScanQuery::Bit4(bit4_query),
770 rows: ScanRows::Bit4RowMajor {
771 codes: segment
772 .query_bit4_codes()
773 .map_err(StoreError::Segment)
774 .map_err(QueryError::Store)?,
775 factors: segment
776 .query_bit4_factors()
777 .map_err(StoreError::Segment)
778 .map_err(QueryError::Store)?,
779 },
780 row_mask: Some(allow_list.as_roaring()),
781 },
782 row_count,
783 k,
784 cancellation,
785 #[cfg(any(test, feature = "test-seams"))]
786 receipt_context,
787 ),
788 scheme => Err(QueryError::Store(StoreError::Segment(
789 crate::segment::SegmentError::Geometry(format!(
790 "filtered store search does not support sealed scheme {scheme}"
791 )),
792 ))
793 .into()),
794 }
795}
796
797#[derive(Clone, Copy, Debug, Eq, PartialEq)]
798enum SealedGraphOutcome {
799 Scan,
800 Pruned,
801 Executed,
802}
803
804#[allow(clippy::too_many_arguments)]
805fn sealed_graph_filtered(
806 segment: &crate::segment::reader::SegmentReader,
807 query: &[f32],
808 predicate: &Predicate,
809 allow_list: &crate::meta::DocBitmap,
810 k: usize,
811 graph_options: Option<GraphSearchOptions>,
812 cancellation: &QueryCancellation<'_>,
813 accounting: &Arc<crate::lifecycle::stats::Accounting>,
814 source: RowSource,
815 row_count: usize,
816 _options: SearchOptions,
817 candidates: &mut Vec<SearchCandidate>,
818 stats: &mut MutableStats,
819 graph_stats: &mut crate::ingest::GraphSearchStats,
820 plans: &mut Vec<SegmentPlan>,
821 #[cfg(any(test, feature = "test-seams"))] metadata_controller: Option<&MetadataTestController>,
822 #[cfg(any(test, feature = "test-seams"))] metadata_query: Option<&MetadataQueryContext>,
823 #[cfg(any(test, feature = "test-seams"))] vector_fault_controller: Option<
824 &crate::scan::vector_fault::VectorFaultController,
825 >,
826) -> Result<SealedGraphOutcome, FilteredSearchError> {
827 let Some(graph_options) = graph_options else {
828 return Ok(SealedGraphOutcome::Scan);
829 };
830 validate_filtered_explicit_ef(graph_options, k.min(row_count), row_count)?;
831 if allow_list.cardinality() <= ALLOW_LIST_ROWS_THRESHOLD {
832 return Ok(SealedGraphOutcome::Scan);
833 }
834 let prepared = segment
835 .graph_search_cache
836 .prepare_shared(segment, cancellation)
837 .map_err(map_graph_cache_error)?;
838 let norm_range = prepared.norm_range;
839 if let Some(competitive_distance) = global_competitive_distance(candidates, k)
840 && norm_range.squared_l2_upper_bound(query) <= f64::from(f32::MAX)
841 && (norm_range.squared_l2_lower_bound(query) as f32) > competitive_distance
842 {
843 cancellation.check_graph().map_err(map_scan_error)?;
844 #[cfg(any(test, feature = "test-seams"))]
845 record_pruned_receipt(
846 metadata_controller,
847 metadata_query,
848 source,
849 segment.meta().row_count,
850 allow_list.cardinality(),
851 )?;
852 plans.push(
853 SegmentPlan::exact(
854 source,
855 SegmentTier::SealedGraph,
856 SegmentBranch::Pruned,
857 allow_list.cardinality(),
858 None,
859 )
860 .with_predicate(predicate),
861 );
862 return Ok(SealedGraphOutcome::Pruned);
863 }
864 let execution = execute_filtered_graph(
865 segment,
866 query,
867 allow_list,
868 k,
869 graph_options,
870 cancellation,
871 accounting,
872 prepared.graph_validated,
873 prepared.entry_seed_discovered,
874 #[cfg(any(test, feature = "test-seams"))]
875 metadata_query,
876 #[cfg(any(test, feature = "test-seams"))]
877 metadata_execution_receipt_context(
878 metadata_controller,
879 metadata_query,
880 source,
881 row_count,
882 allow_list.cardinality(),
883 query.len(),
884 true,
885 )?,
886 #[cfg(any(test, feature = "test-seams"))]
887 vector_fault_controller,
888 )?;
889 let plan = SegmentPlan::filtered_graph(
890 source,
891 allow_list.cardinality(),
892 execution.ef_requested,
893 execution.ef_effective,
894 execution.branch,
895 execution.fallback,
896 graph_options.profile(),
897 )
898 .with_predicate(predicate);
899 #[cfg(any(test, feature = "test-seams"))]
900 let reported_branch =
901 metadata_query.map_or(plan.branch, |query| query.reported_branch(plan.branch));
902 #[cfg(not(any(test, feature = "test-seams")))]
903 let reported_branch = plan.branch;
904 verify_execution_branch(reported_branch, execution.branch)?;
905 #[cfg(any(test, feature = "test-seams"))]
906 record_visited_fallback_feature(
907 metadata_controller,
908 metadata_query,
909 source,
910 allow_list.cardinality(),
911 &execution,
912 )?;
913 append_candidates(
914 execution.candidates,
915 source,
916 |row| {
917 segment
918 .document_version(row)
919 .map_err(StoreError::Segment)
920 .map_err(QueryError::Store)
921 },
922 candidates,
923 true,
924 #[cfg(any(test, feature = "test-seams"))]
925 vector_fault_controller,
926 #[cfg(any(test, feature = "test-seams"))]
927 vector_fault_tier(_options.tier()),
928 )?;
929 stats.add(execution.stats)?;
930 add_graph_stats(graph_stats, execution.graph_stats)?;
931 plans.push(plan);
932 retain_global_top_k(candidates, k);
933 Ok(SealedGraphOutcome::Executed)
934}
935
936fn sealed_scan_reason(
937 segment: &crate::segment::reader::SegmentReader,
938 requested_tier: Option<SearchTier>,
939) -> crate::planner::ScanReason {
940 use crate::planner::{ExplicitScanTier, ScanReason};
941
942 match requested_tier {
943 Some(SearchTier::Exact) => return ScanReason::ExplicitTier(ExplicitScanTier::Exact),
944 Some(SearchTier::Scan) => return ScanReason::ExplicitTier(ExplicitScanTier::Scan),
945 Some(SearchTier::Auto | SearchTier::Graph(_)) | None => {}
946 }
947 let rows = segment.meta().row_count;
948 let min_rows = crate::tier::thresholds::for_bucket(segment.meta().dims, segment.meta().scheme)
949 .graph_min_rows;
950 if rows < min_rows {
951 ScanReason::BelowGraphThreshold { rows, min_rows }
952 } else {
953 ScanReason::GraphPending
954 }
955}
956
957struct FilteredGraphExecution {
958 candidates: Vec<LocalCandidate>,
959 stats: ScanStats,
960 branch: SegmentBranch,
961 fallback: PlanFallback,
962 ef_requested: Option<usize>,
963 ef_effective: usize,
964 graph_stats: crate::ingest::GraphSearchStats,
965 #[cfg(any(test, feature = "test-seams"))]
966 graph_nodes_visited: usize,
967 #[cfg(any(test, feature = "test-seams"))]
968 exact_fallback_rows_examined: u64,
969 #[cfg(any(test, feature = "test-seams"))]
970 visited_budget: usize,
971 #[cfg(any(test, feature = "test-seams"))]
972 visited_budget_guard: Option<(usize, usize)>,
973}
974
975fn validate_filtered_explicit_ef(
976 options: GraphSearchOptions,
977 segment_k: usize,
978 rows: usize,
979) -> Result<(), FilteredSearchError> {
980 let Some(ef) = options.ef() else {
981 return Ok(());
982 };
983 if ef < segment_k {
984 return Err(
985 QueryError::Graph(crate::graph::search::GraphSearchError::AdaptiveEf(
986 crate::graph::search::AdaptiveEfError::ExplicitBelowK { k: segment_k, ef },
987 ))
988 .into(),
989 );
990 }
991 if ef > rows {
992 return Err(
993 QueryError::Graph(crate::graph::search::GraphSearchError::AdaptiveEf(
994 crate::graph::search::AdaptiveEfError::ExplicitExceedsRows { ef, rows },
995 ))
996 .into(),
997 );
998 }
999 Ok(())
1000}
1001
1002struct FilteredGraphBudget {
1003 target_k: usize,
1004 ef_effective: usize,
1005 filtered_visited_budget: usize,
1006 scratch_capacity: usize,
1007 maximum_corrective_ef: usize,
1008}
1009
1010#[allow(clippy::too_many_arguments)]
1011fn filtered_graph_budget(
1012 graph: crate::graph::block::GraphNodeBlocks<'_>,
1013 query: &[f32],
1014 options: GraphSearchOptions,
1015 allow_count: usize,
1016 node_count: usize,
1017 k: usize,
1018 #[cfg(any(test, feature = "test-seams"))] metadata_query: Option<&MetadataQueryContext>,
1019) -> Result<FilteredGraphBudget, FilteredSearchError> {
1020 let target_k = k.min(node_count).min(allow_count);
1021 if target_k == 0 {
1022 return Err(
1023 QueryError::Graph(crate::graph::search::GraphSearchError::Geometry(
1024 "filtered graph traversal was selected with an empty effective mask".to_owned(),
1025 ))
1026 .into(),
1027 );
1028 }
1029 let base_request =
1030 GraphSearchRequest::new(query, target_k, options.seed()).with_profile(options.profile());
1031 let base_request = options
1032 .ef()
1033 .map_or(base_request, |ef| base_request.with_ef(ef));
1034 let base_ef = base_request
1035 .effective_ef(node_count)
1036 .map_err(crate::graph::search::GraphSearchError::AdaptiveEf)
1037 .map_err(QueryError::Graph)?;
1038 let selectivity_scaled = base_ef
1039 .checked_mul(node_count)
1040 .map(|value| value.div_ceil(allow_count))
1041 .ok_or(QueryError::Scan(ScanError::ArithmeticOverflow))?;
1042 let ef_effective = selectivity_scaled.min(allow_count).max(target_k);
1043 let max_degree = usize::from(graph.layout().max_degree());
1044 let computed_visited_budget = ef_effective
1045 .checked_mul(max_degree)
1046 .and_then(|value| value.checked_mul(FILTERED_VISITED_BUDGET_MULTIPLIER))
1047 .ok_or(QueryError::Scan(ScanError::ArithmeticOverflow))?
1048 .min(node_count);
1049 #[cfg(any(test, feature = "test-seams"))]
1050 let filtered_visited_budget = metadata_query
1051 .and_then(MetadataQueryContext::visited_budget_override)
1052 .unwrap_or(computed_visited_budget)
1053 .min(node_count);
1054 #[cfg(not(any(test, feature = "test-seams")))]
1055 let filtered_visited_budget = computed_visited_budget;
1056 let scratch_capacity = ef_effective.max(filtered_visited_budget);
1057 let maximum_corrective_ef = allow_count.min(scratch_capacity);
1058 Ok(FilteredGraphBudget {
1059 target_k,
1060 ef_effective,
1061 filtered_visited_budget,
1062 scratch_capacity,
1063 maximum_corrective_ef,
1064 })
1065}
1066
1067struct FallbackContext<'a, 'cancel> {
1068 rescore: &'a [f32],
1069 node_count: usize,
1070 query: &'a [f32],
1071 allow_list: &'a crate::meta::DocBitmap,
1072 k: usize,
1073 cancellation: &'a QueryCancellation<'cancel>,
1074 ef_requested: Option<usize>,
1075 #[cfg(any(test, feature = "test-seams"))]
1076 filtered_visited_budget: usize,
1077 #[cfg(any(test, feature = "test-seams"))]
1078 receipt_context: Option<MetadataExecutionReceiptContext<'a>>,
1079}
1080
1081#[allow(clippy::too_many_arguments)]
1082fn execute_filtered_graph(
1083 segment: &crate::segment::reader::SegmentReader,
1084 query: &[f32],
1085 allow_list: &crate::meta::DocBitmap,
1086 k: usize,
1087 options: GraphSearchOptions,
1088 cancellation: &QueryCancellation<'_>,
1089 accounting: &Arc<crate::lifecycle::stats::Accounting>,
1090 graph_validated_before: bool,
1091 entry_seed_discovered_before: bool,
1092 #[cfg(any(test, feature = "test-seams"))] metadata_query: Option<&MetadataQueryContext>,
1093 #[cfg(any(test, feature = "test-seams"))] receipt_context: Option<
1094 MetadataExecutionReceiptContext<'_>,
1095 >,
1096 #[cfg(any(test, feature = "test-seams"))] vector_fault_controller: Option<
1097 &crate::scan::vector_fault::VectorFaultController,
1098 >,
1099) -> Result<FilteredGraphExecution, FilteredSearchError> {
1100 cancellation.check_graph().map_err(map_scan_error)?;
1101 let prepared = segment
1102 .graph_search_cache
1103 .prepare_shared(segment, cancellation)
1104 .map_err(map_graph_cache_error)?;
1105 let graph = prepared.graph;
1106 let rescore = crate::lifecycle::query_rescore_rows(segment)?;
1107 let node_count = graph.node_count() as usize;
1108 let allow_count = usize::try_from(allow_list.cardinality())
1109 .map_err(|_| QueryError::Scan(ScanError::ArithmeticOverflow))?;
1110 let FilteredGraphBudget {
1111 target_k,
1112 ef_effective,
1113 filtered_visited_budget,
1114 scratch_capacity,
1115 maximum_corrective_ef,
1116 } = filtered_graph_budget(
1117 graph,
1118 query,
1119 options,
1120 allow_count,
1121 node_count,
1122 k,
1123 #[cfg(any(test, feature = "test-seams"))]
1124 metadata_query,
1125 )?;
1126 let mut scratch = segment
1127 .graph_search_cache
1128 .checkout(graph, scratch_capacity, accounting, cancellation)
1129 .map_err(map_graph_cache_error)?;
1130 let entry_seed_discovered = entry_seed_discovered_before
1131 || prepared.entry_seed_discovered
1132 || scratch.entry_seed_discovered();
1133 let entries = scratch.entries();
1134 let mut searcher = crate::graph::search::GraphSearcher::with_entry_row_ids(
1135 graph,
1136 rescore,
1137 entries,
1138 scratch.scratch_mut().map_err(map_graph_error)?,
1139 )
1140 .map_err(map_graph_error)?
1141 .with_rescore_validator(segment);
1142 #[cfg(any(test, feature = "test-seams"))]
1143 if let Some(controller) = vector_fault_controller {
1144 searcher = searcher.with_vector_fault_controller(
1145 controller,
1146 crate::scan::vector_fault::VectorRowSource::Sealed(*segment.meta().id.as_bytes()),
1147 crate::scan::vector_fault::VectorSearchTier::Graph,
1148 );
1149 }
1150 let ef_requested = options.ef();
1151 let mut current_ef = ef_effective;
1152 let mut traversal_stats = MutableStats::default();
1153 let mut graph_stats = crate::ingest::GraphSearchStats {
1154 graph_validations: usize::from(graph_validated_before || prepared.graph_validated),
1155 entry_seed_discoveries: usize::from(entry_seed_discovered),
1156 ..crate::ingest::GraphSearchStats::default()
1157 };
1158 #[cfg(any(test, feature = "test-seams"))]
1159 let mut graph_nodes_visited = 0_usize;
1160 let fallback_context = FallbackContext {
1161 rescore,
1162 node_count,
1163 query,
1164 allow_list,
1165 k,
1166 cancellation,
1167 ef_requested,
1168 #[cfg(any(test, feature = "test-seams"))]
1169 filtered_visited_budget,
1170 #[cfg(any(test, feature = "test-seams"))]
1171 receipt_context,
1172 };
1173 loop {
1174 let request = GraphSearchRequest::new(query, target_k, options.seed())
1175 .with_profile(options.profile())
1176 .with_ef(current_ef);
1177 let outcome = match searcher.search_filtered(
1178 request,
1179 allow_list,
1180 filtered_visited_budget,
1181 Some(cancellation),
1182 ) {
1183 Ok(outcome) => outcome,
1184 Err(crate::graph::search::GraphSearchError::FilteredCandidateShortfall {
1185 counters,
1186 ..
1187 }) => {
1188 #[cfg(any(test, feature = "test-seams"))]
1189 {
1190 graph_nodes_visited = graph_nodes_visited
1191 .checked_add(counters.visited())
1192 .ok_or(QueryError::Scan(ScanError::ArithmeticOverflow))?;
1193 }
1194 add_graph_counters(&mut graph_stats, counters)?;
1195 traversal_stats.add(graph_scan_stats(counters)?)?;
1196 return exact_filtered_graph_fallback(
1197 &fallback_context,
1198 traversal_stats,
1199 PlanFallback::CandidateShortfall,
1200 current_ef,
1201 graph_stats,
1202 #[cfg(any(test, feature = "test-seams"))]
1203 graph_nodes_visited,
1204 #[cfg(any(test, feature = "test-seams"))]
1205 None,
1206 );
1207 }
1208 Err(error) => return Err(map_graph_error(error)),
1209 };
1210 match outcome {
1211 FilteredGraphSearchOutcome::Traversed(result) => {
1212 #[cfg(any(test, feature = "test-seams"))]
1213 {
1214 graph_nodes_visited = graph_nodes_visited
1215 .checked_add(result.counters().visited())
1216 .ok_or(QueryError::Scan(ScanError::ArithmeticOverflow))?;
1217 }
1218 add_graph_counters(&mut graph_stats, result.counters())?;
1219 traversal_stats.add(graph_scan_stats(result.counters())?)?;
1220 if result.candidates().len() >= target_k {
1221 let candidates = result
1222 .candidates()
1223 .iter()
1224 .map(|candidate| LocalCandidate {
1225 row: candidate.row_id() as usize,
1226 score: -(candidate.distance() as f32),
1227 })
1228 .collect();
1229 let execution = FilteredGraphExecution {
1230 candidates,
1231 stats: traversal_stats.finish(),
1232 branch: SegmentBranch::FilteredGraph,
1233 fallback: if ef_requested.is_some_and(|requested| current_ef > requested) {
1234 PlanFallback::EfWidened
1235 } else {
1236 PlanFallback::None
1237 },
1238 ef_requested,
1239 ef_effective: current_ef,
1240 graph_stats: with_traversed_segment(graph_stats)?,
1241 #[cfg(any(test, feature = "test-seams"))]
1242 graph_nodes_visited,
1243 #[cfg(any(test, feature = "test-seams"))]
1244 exact_fallback_rows_examined: 0,
1245 #[cfg(any(test, feature = "test-seams"))]
1246 visited_budget: filtered_visited_budget,
1247 #[cfg(any(test, feature = "test-seams"))]
1248 visited_budget_guard: None,
1249 };
1250 #[cfg(any(test, feature = "test-seams"))]
1251 if let Some(context) = fallback_context.receipt_context {
1252 context.record_graph(&execution)?;
1253 }
1254 return Ok(execution);
1255 }
1256 if current_ef < maximum_corrective_ef {
1257 current_ef = current_ef
1258 .saturating_mul(2)
1259 .max(current_ef.saturating_add(1))
1260 .min(maximum_corrective_ef);
1261 continue;
1262 }
1263 return exact_filtered_graph_fallback(
1264 &fallback_context,
1265 traversal_stats,
1266 PlanFallback::CandidateShortfall,
1267 current_ef,
1268 graph_stats,
1269 #[cfg(any(test, feature = "test-seams"))]
1270 graph_nodes_visited,
1271 #[cfg(any(test, feature = "test-seams"))]
1272 None,
1273 );
1274 }
1275 FilteredGraphSearchOutcome::VisitedBudgetExceeded {
1276 visited,
1277 budget,
1278 counters,
1279 } => {
1280 #[cfg(not(any(test, feature = "test-seams")))]
1281 let _ = (visited, budget);
1282 #[cfg(any(test, feature = "test-seams"))]
1283 {
1284 graph_nodes_visited = graph_nodes_visited
1285 .checked_add(counters.visited())
1286 .ok_or(QueryError::Scan(ScanError::ArithmeticOverflow))?;
1287 }
1288 add_graph_counters(&mut graph_stats, counters)?;
1289 traversal_stats.add(graph_scan_stats(counters)?)?;
1290 return exact_filtered_graph_fallback(
1291 &fallback_context,
1292 traversal_stats,
1293 PlanFallback::VisitedBudget,
1294 current_ef,
1295 graph_stats,
1296 #[cfg(any(test, feature = "test-seams"))]
1297 graph_nodes_visited,
1298 #[cfg(any(test, feature = "test-seams"))]
1299 Some((visited, budget)),
1300 );
1301 }
1302 }
1303 }
1304}
1305
1306fn exact_filtered_graph_fallback(
1307 context: &FallbackContext<'_, '_>,
1308 mut traversal_stats: MutableStats,
1309 fallback: PlanFallback,
1310 ef_effective: usize,
1311 graph_stats: crate::ingest::GraphSearchStats,
1312 #[cfg(any(test, feature = "test-seams"))] graph_nodes_visited: usize,
1313 #[cfg(any(test, feature = "test-seams"))] visited_budget_guard: Option<(usize, usize)>,
1314) -> Result<FilteredGraphExecution, FilteredSearchError> {
1315 let (candidates, exact_stats, _) = execute_squared_l2_rows(
1316 context.rescore,
1317 context.node_count,
1318 context.query,
1319 context.allow_list,
1320 SegmentBranch::ExactAllowList,
1321 context.k,
1322 context.cancellation,
1323 #[cfg(any(test, feature = "test-seams"))]
1324 None,
1325 )?;
1326 traversal_stats.add(exact_stats)?;
1327 #[cfg(any(test, feature = "test-seams"))]
1328 let exact_fallback_rows_examined = context.allow_list.cardinality();
1329 let execution = FilteredGraphExecution {
1330 candidates,
1331 stats: traversal_stats.finish(),
1332 branch: SegmentBranch::GraphExactFallback,
1333 fallback,
1334 ef_requested: context.ef_requested,
1335 ef_effective,
1336 graph_stats: with_traversed_segment(graph_stats)?,
1337 #[cfg(any(test, feature = "test-seams"))]
1338 graph_nodes_visited,
1339 #[cfg(any(test, feature = "test-seams"))]
1340 exact_fallback_rows_examined,
1341 #[cfg(any(test, feature = "test-seams"))]
1342 visited_budget: context.filtered_visited_budget,
1343 #[cfg(any(test, feature = "test-seams"))]
1344 visited_budget_guard,
1345 };
1346 #[cfg(any(test, feature = "test-seams"))]
1347 if let Some(receipt_context) = context.receipt_context {
1348 receipt_context.record_graph(&execution)?;
1349 }
1350 Ok(execution)
1351}
1352
1353fn add_graph_counters(
1354 stats: &mut crate::ingest::GraphSearchStats,
1355 counters: GraphSearchCounters,
1356) -> Result<(), FilteredSearchError> {
1357 stats.visited_epoch_clears = stats
1358 .visited_epoch_clears
1359 .checked_add(usize::from(counters.visited_epoch_cleared()))
1360 .ok_or(QueryError::Scan(ScanError::ArithmeticOverflow))?;
1361 stats.candidates_scored = stats
1362 .candidates_scored
1363 .checked_add(counters.candidates_scored())
1364 .ok_or(QueryError::Scan(ScanError::ArithmeticOverflow))?;
1365 stats.candidates_rescored = stats
1366 .candidates_rescored
1367 .checked_add(counters.candidates_rescored())
1368 .ok_or(QueryError::Scan(ScanError::ArithmeticOverflow))?;
1369 Ok(())
1370}
1371
1372fn with_traversed_segment(
1373 mut stats: crate::ingest::GraphSearchStats,
1374) -> Result<crate::ingest::GraphSearchStats, FilteredSearchError> {
1375 stats.segments_traversed = stats
1376 .segments_traversed
1377 .checked_add(1)
1378 .ok_or(QueryError::Scan(ScanError::ArithmeticOverflow))?;
1379 Ok(stats)
1380}
1381
1382fn add_graph_stats(
1383 total: &mut crate::ingest::GraphSearchStats,
1384 delta: crate::ingest::GraphSearchStats,
1385) -> Result<(), FilteredSearchError> {
1386 total.segments_traversed = total
1387 .segments_traversed
1388 .checked_add(delta.segments_traversed)
1389 .ok_or(QueryError::Scan(ScanError::ArithmeticOverflow))?;
1390 total.graph_validations = total
1391 .graph_validations
1392 .checked_add(delta.graph_validations)
1393 .ok_or(QueryError::Scan(ScanError::ArithmeticOverflow))?;
1394 total.entry_seed_discoveries = total
1395 .entry_seed_discoveries
1396 .checked_add(delta.entry_seed_discoveries)
1397 .ok_or(QueryError::Scan(ScanError::ArithmeticOverflow))?;
1398 total.visited_epoch_clears = total
1399 .visited_epoch_clears
1400 .checked_add(delta.visited_epoch_clears)
1401 .ok_or(QueryError::Scan(ScanError::ArithmeticOverflow))?;
1402 total.candidates_scored = total
1403 .candidates_scored
1404 .checked_add(delta.candidates_scored)
1405 .ok_or(QueryError::Scan(ScanError::ArithmeticOverflow))?;
1406 total.candidates_rescored = total
1407 .candidates_rescored
1408 .checked_add(delta.candidates_rescored)
1409 .ok_or(QueryError::Scan(ScanError::ArithmeticOverflow))?;
1410 total.segments_pruned_by_bound = total
1411 .segments_pruned_by_bound
1412 .checked_add(delta.segments_pruned_by_bound)
1413 .ok_or(QueryError::Scan(ScanError::ArithmeticOverflow))?;
1414 Ok(())
1415}
1416
1417fn graph_scan_stats(counters: GraphSearchCounters) -> Result<ScanStats, FilteredSearchError> {
1418 let did_work = counters.candidates_scored() != 0 || counters.candidates_rescored() != 0;
1419 Ok(ScanStats {
1420 dims_touched: counters.dims_touched(),
1421 bytes_read: counters.bytes_read(),
1422 threads_used: usize::from(did_work),
1423 worker_thread_ids: did_work
1424 .then(|| std::thread::current().id())
1425 .into_iter()
1426 .collect(),
1427 })
1428}
1429
1430fn map_graph_cache_error(
1431 error: crate::lifecycle::graph_cache::GraphCacheError,
1432) -> FilteredSearchError {
1433 match error {
1434 crate::lifecycle::graph_cache::GraphCacheError::Store(error) => {
1435 QueryError::Store(error).into()
1436 }
1437 crate::lifecycle::graph_cache::GraphCacheError::Search(error) => map_graph_error(error),
1438 }
1439}
1440
1441fn map_graph_error(error: crate::graph::search::GraphSearchError) -> FilteredSearchError {
1442 let query = match error {
1443 crate::graph::search::GraphSearchError::ExactRescoreUnavailable(detail) => {
1444 QueryError::Store(StoreError::Segment(crate::segment::SegmentError::Geometry(
1445 detail,
1446 )))
1447 }
1448 crate::graph::search::GraphSearchError::Cancelled { partial } => {
1449 QueryError::Cancelled { partial }
1450 }
1451 crate::graph::search::GraphSearchError::Timeout { partial } => {
1452 QueryError::Timeout { partial }
1453 }
1454 crate::graph::search::GraphSearchError::ReadCancelled { partial } => {
1455 QueryError::ReadCancelled { partial }
1456 }
1457 error => QueryError::Graph(error),
1458 };
1459 query.into()
1460}
1461
1462fn active_columns(
1463 active: &crate::ingest::ActiveSegment,
1464 schema: &crate::meta::Schema,
1465) -> Result<ColumnStore, FilteredSearchError> {
1466 let mut builder = ColumnStoreBuilder::new(schema.clone());
1467 for (row, timestamp) in active.timestamps().iter().enumerate() {
1468 let values = active.column_values(row).map_err(QueryError::Store)?;
1469 builder
1470 .push_row(*timestamp, &crate::ingest::column_inputs(&values))
1471 .map_err(|error| FilteredSearchError::ActiveMetadata(error.to_string()))?;
1472 }
1473 builder
1474 .finish()
1475 .map_err(|error| FilteredSearchError::ActiveMetadata(error.to_string()))
1476}
1477
1478#[cfg(any(test, feature = "test-seams"))]
1479fn record_selectivity_decision(
1480 controller: Option<&MetadataTestController>,
1481 query: Option<&MetadataQueryContext>,
1482 source: RowSource,
1483 cardinality: u64,
1484 branch: SegmentBranch,
1485) -> Result<(), FilteredSearchError> {
1486 if let (Some(controller), Some(query)) = (controller, query) {
1487 query.record_selectivity(
1488 controller,
1489 source,
1490 cardinality,
1491 ALLOW_LIST_ROWS_THRESHOLD,
1492 branch,
1493 )?;
1494 }
1495 Ok(())
1496}
1497
1498#[cfg(any(test, feature = "test-seams"))]
1499fn record_pruned_receipt(
1500 controller: Option<&MetadataTestController>,
1501 query: Option<&MetadataQueryContext>,
1502 source: RowSource,
1503 row_count: u32,
1504 filter_cardinality: u64,
1505) -> Result<(), FilteredSearchError> {
1506 if let (Some(controller), Some(query)) = (controller, query) {
1507 controller.record_execution(MetadataExecutionReceipt {
1508 query_id: query.query_id(),
1509 source,
1510 row_count: u64::from(row_count),
1511 filter_cardinality,
1512 branch: SegmentBranch::Pruned,
1513 fallback: PlanFallback::None,
1514 rows_examined: 0,
1515 allowed_rows_examined: 0,
1516 vectors_scored: 0,
1517 graph_nodes_visited: 0,
1518 exact_fallback_rows_examined: 0,
1519 returned_candidates: 0,
1520 ef_effective: None,
1521 visited_budget: None,
1522 sealed: true,
1523 })?;
1524 }
1525 Ok(())
1526}
1527
1528#[cfg(any(test, feature = "test-seams"))]
1529fn record_column_refusal(
1530 controller: Option<&MetadataTestController>,
1531 query: Option<&MetadataQueryContext>,
1532 source: RowSource,
1533 error: &StoreError,
1534) -> Result<(), FilteredSearchError> {
1535 if let (Some(controller), Some(query)) = (controller, query) {
1536 query.record_column_refusal(controller, source, error)?;
1537 }
1538 Ok(())
1539}
1540
1541#[cfg(any(test, feature = "test-seams"))]
1542fn record_alive_refusal(
1543 controller: Option<&MetadataTestController>,
1544 query: Option<&MetadataQueryContext>,
1545 source: RowSource,
1546 error: &StoreError,
1547) -> Result<(), FilteredSearchError> {
1548 if let (Some(controller), Some(query)) = (controller, query) {
1549 query.record_alive_refusal(controller, source, error)?;
1550 }
1551 Ok(())
1552}
1553
1554#[cfg(any(test, feature = "test-seams"))]
1555fn record_visited_fallback_feature(
1556 controller: Option<&MetadataTestController>,
1557 query: Option<&MetadataQueryContext>,
1558 source: RowSource,
1559 filter_cardinality: u64,
1560 execution: &FilteredGraphExecution,
1561) -> Result<(), FilteredSearchError> {
1562 let (Some(controller), Some(query), Some((visited, budget))) =
1563 (controller, query, execution.visited_budget_guard)
1564 else {
1565 return Ok(());
1566 };
1567 query.record_visited_fallback(
1568 controller,
1569 source,
1570 visited,
1571 budget,
1572 filter_cardinality,
1573 execution.exact_fallback_rows_examined,
1574 u64::try_from(execution.candidates.len())
1575 .map_err(|_| QueryError::Scan(ScanError::ArithmeticOverflow))?,
1576 execution.fallback,
1577 )?;
1578 Ok(())
1579}
1580
1581#[derive(Clone, Copy, Debug)]
1582struct LocalCandidate {
1583 row: usize,
1584 score: f32,
1585}
1586
1587impl rescored_scan::RescoreCandidate for LocalCandidate {
1588 fn row_id(&self) -> usize {
1589 self.row
1590 }
1591 fn score(&self) -> f32 {
1592 self.score
1593 }
1594 fn set_score(&mut self, score: f32) {
1595 self.score = score;
1596 }
1597}
1598
1599#[derive(Clone, Copy, Debug)]
1602struct ScanExecutionFacts {
1603 branch: SegmentBranch,
1604 #[cfg(any(test, feature = "test-seams"))]
1605 rows_examined: u64,
1606 #[cfg(any(test, feature = "test-seams"))]
1607 allowed_rows_examined: u64,
1608}
1609
1610#[cfg(any(test, feature = "test-seams"))]
1611#[derive(Clone, Copy)]
1612struct MetadataExecutionReceiptContext<'a> {
1613 controller: &'a MetadataTestController,
1614 query: &'a MetadataQueryContext,
1615 source: RowSource,
1616 row_count: u64,
1617 filter_cardinality: u64,
1618 dimensions: u64,
1619 sealed: bool,
1620}
1621
1622#[cfg(any(test, feature = "test-seams"))]
1623impl MetadataExecutionReceiptContext<'_> {
1624 fn record_scan(
1625 self,
1626 execution: ScanExecutionFacts,
1627 stats: &ScanStats,
1628 returned_candidates: usize,
1629 ) -> Result<(), FilteredSearchError> {
1630 let vectors_scored = if self.dimensions == 0 {
1631 0
1632 } else {
1633 stats.dims_touched / self.dimensions
1634 };
1635 self.controller.record_execution(MetadataExecutionReceipt {
1636 query_id: self.query.query_id(),
1637 source: self.source,
1638 row_count: self.row_count,
1639 filter_cardinality: self.filter_cardinality,
1640 branch: execution.branch,
1641 fallback: PlanFallback::None,
1642 rows_examined: execution.rows_examined,
1643 allowed_rows_examined: execution.allowed_rows_examined,
1644 vectors_scored,
1645 graph_nodes_visited: 0,
1646 exact_fallback_rows_examined: 0,
1647 returned_candidates: u64::try_from(returned_candidates)
1648 .map_err(|_| QueryError::Scan(ScanError::ArithmeticOverflow))?,
1649 ef_effective: None,
1650 visited_budget: None,
1651 sealed: self.sealed,
1652 })?;
1653 Ok(())
1654 }
1655
1656 fn record_graph(self, execution: &FilteredGraphExecution) -> Result<(), FilteredSearchError> {
1657 let graph_nodes_visited = u64::try_from(execution.graph_nodes_visited)
1658 .map_err(|_| QueryError::Scan(ScanError::ArithmeticOverflow))?;
1659 self.controller.record_execution(MetadataExecutionReceipt {
1660 query_id: self.query.query_id(),
1661 source: self.source,
1662 row_count: self.row_count,
1663 filter_cardinality: self.filter_cardinality,
1664 branch: execution.branch,
1665 fallback: execution.fallback,
1666 rows_examined: graph_nodes_visited,
1667 allowed_rows_examined: self.filter_cardinality,
1668 vectors_scored: u64::try_from(execution.graph_stats.candidates_scored)
1669 .map_err(|_| QueryError::Scan(ScanError::ArithmeticOverflow))?,
1670 graph_nodes_visited,
1671 exact_fallback_rows_examined: execution.exact_fallback_rows_examined,
1672 returned_candidates: u64::try_from(execution.candidates.len())
1673 .map_err(|_| QueryError::Scan(ScanError::ArithmeticOverflow))?,
1674 ef_effective: Some(execution.ef_effective),
1675 visited_budget: Some(execution.visited_budget),
1676 sealed: self.sealed,
1677 })?;
1678 Ok(())
1679 }
1680}
1681
1682#[cfg(any(test, feature = "test-seams"))]
1683#[allow(clippy::too_many_arguments)]
1684fn metadata_execution_receipt_context<'a>(
1685 controller: Option<&'a MetadataTestController>,
1686 query: Option<&'a MetadataQueryContext>,
1687 source: RowSource,
1688 row_count: usize,
1689 filter_cardinality: u64,
1690 dimensions: usize,
1691 sealed: bool,
1692) -> Result<Option<MetadataExecutionReceiptContext<'a>>, FilteredSearchError> {
1693 let (Some(controller), Some(query)) = (controller, query) else {
1694 return Ok(None);
1695 };
1696 Ok(Some(MetadataExecutionReceiptContext {
1697 controller,
1698 query,
1699 source,
1700 row_count: u64::try_from(row_count)
1701 .map_err(|_| QueryError::Scan(ScanError::ArithmeticOverflow))?,
1702 filter_cardinality,
1703 dimensions: u64::try_from(dimensions)
1704 .map_err(|_| QueryError::Scan(ScanError::ArithmeticOverflow))?,
1705 sealed,
1706 }))
1707}
1708
1709fn execute_scan_request(
1710 request: VectorScanRequest<'_>,
1711 row_count: usize,
1712 branch: SegmentBranch,
1713 k: usize,
1714 cancellation: &QueryCancellation<'_>,
1715 #[cfg(any(test, feature = "test-seams"))] receipt_context: Option<
1716 MetadataExecutionReceiptContext<'_>,
1717 >,
1718) -> Result<(Vec<LocalCandidate>, ScanStats, ScanExecutionFacts), FilteredSearchError> {
1719 match branch {
1720 SegmentBranch::FilteredGraph | SegmentBranch::Graph => {
1721 Err(FilteredSearchError::InvalidPlanNode(
1722 "filtered graph branch reached exact scan request",
1723 ))
1724 }
1725 SegmentBranch::Pruned => finish_scan_request(
1726 crate::scan::ScanOutcome {
1727 candidates: Vec::new(),
1728 worst_score: None,
1729 stats: ScanStats {
1730 dims_touched: 0,
1731 bytes_read: 0,
1732 threads_used: 0,
1733 worker_thread_ids: Vec::new(),
1734 },
1735 },
1736 ScanExecutionFacts {
1737 branch: SegmentBranch::Pruned,
1738 #[cfg(any(test, feature = "test-seams"))]
1739 rows_examined: 0,
1740 #[cfg(any(test, feature = "test-seams"))]
1741 allowed_rows_examined: 0,
1742 },
1743 #[cfg(any(test, feature = "test-seams"))]
1744 receipt_context,
1745 ),
1746 SegmentBranch::ExactAllowList => {
1747 let outcome = gather_top_k(request, k, Some(cancellation)).map_err(map_scan_error)?;
1748 finish_scan_request(
1749 outcome,
1750 ScanExecutionFacts {
1751 branch: SegmentBranch::ExactAllowList,
1752 #[cfg(any(test, feature = "test-seams"))]
1753 rows_examined: request.row_mask.map_or(0, roaring::RoaringBitmap::len),
1754 #[cfg(any(test, feature = "test-seams"))]
1755 allowed_rows_examined: request.row_mask.map_or(0, roaring::RoaringBitmap::len),
1756 },
1757 #[cfg(any(test, feature = "test-seams"))]
1758 receipt_context,
1759 )
1760 }
1761 SegmentBranch::MaskedScan | SegmentBranch::GraphExactFallback => {
1762 let partition = scan_partition(request, k, 0..row_count, Some(cancellation))
1763 .map_err(map_scan_error)?;
1764 finish_scan_request(
1765 crate::scan::ScanOutcome {
1766 candidates: partition.candidates,
1767 worst_score: None,
1768 stats: ScanStats {
1769 dims_touched: partition.dims_touched,
1770 bytes_read: partition.bytes_read,
1771 threads_used: 1,
1772 worker_thread_ids: vec![partition.worker_thread_id],
1773 },
1774 },
1775 ScanExecutionFacts {
1776 branch,
1777 #[cfg(any(test, feature = "test-seams"))]
1778 rows_examined: u64::try_from(row_count)
1779 .map_err(|_| QueryError::Scan(ScanError::ArithmeticOverflow))?,
1780 #[cfg(any(test, feature = "test-seams"))]
1781 allowed_rows_examined: request
1782 .row_mask
1783 .map_or_else(|| u64::try_from(row_count), |mask| Ok(mask.len()))
1784 .map_err(|_| QueryError::Scan(ScanError::ArithmeticOverflow))?,
1785 },
1786 #[cfg(any(test, feature = "test-seams"))]
1787 receipt_context,
1788 )
1789 }
1790 }
1791}
1792
1793fn finish_scan_request(
1794 outcome: crate::scan::ScanOutcome,
1795 execution: ScanExecutionFacts,
1796 #[cfg(any(test, feature = "test-seams"))] receipt_context: Option<
1797 MetadataExecutionReceiptContext<'_>,
1798 >,
1799) -> Result<(Vec<LocalCandidate>, ScanStats, ScanExecutionFacts), FilteredSearchError> {
1800 #[cfg(any(test, feature = "test-seams"))]
1801 if let Some(context) = receipt_context {
1802 context.record_scan(execution, &outcome.stats, outcome.candidates.len())?;
1803 }
1804 Ok((
1805 outcome
1806 .candidates
1807 .into_iter()
1808 .map(|candidate| LocalCandidate {
1809 row: candidate.row_id,
1810 score: candidate.score,
1811 })
1812 .collect(),
1813 outcome.stats,
1814 execution,
1815 ))
1816}
1817
1818fn execute_scan_plan(
1819 node: &PlanNode,
1820 request: VectorScanRequest<'_>,
1821 row_count: usize,
1822 k: usize,
1823 cancellation: &QueryCancellation<'_>,
1824 #[cfg(any(test, feature = "test-seams"))] receipt_context: Option<
1825 MetadataExecutionReceiptContext<'_>,
1826 >,
1827) -> Result<(Vec<LocalCandidate>, ScanStats, ScanExecutionFacts), FilteredSearchError> {
1828 match node {
1829 PlanNode::BitmapIntersect { input, .. } => execute_scan_plan(
1830 input,
1831 request,
1832 row_count,
1833 k,
1834 cancellation,
1835 #[cfg(any(test, feature = "test-seams"))]
1836 receipt_context,
1837 ),
1838 PlanNode::Scan { branch, .. } => execute_scan_request(
1839 request,
1840 row_count,
1841 *branch,
1842 k,
1843 cancellation,
1844 #[cfg(any(test, feature = "test-seams"))]
1845 receipt_context,
1846 ),
1847 PlanNode::Graph {
1848 fallback: Some(fallback),
1849 ..
1850 } => {
1851 let (candidates, stats, mut execution) = execute_scan_plan(
1852 fallback,
1853 request,
1854 row_count,
1855 k,
1856 cancellation,
1857 #[cfg(any(test, feature = "test-seams"))]
1858 None,
1859 )?;
1860 execution.branch = SegmentBranch::GraphExactFallback;
1861 #[cfg(any(test, feature = "test-seams"))]
1862 if let Some(context) = receipt_context {
1863 context.record_scan(execution, &stats, candidates.len())?;
1864 }
1865 Ok((candidates, stats, execution))
1866 }
1867 PlanNode::Graph { fallback: None, .. } => Err(FilteredSearchError::InvalidPlanNode(
1868 "graph node has no exact fallback before task 19-M6",
1869 )),
1870 PlanNode::Lexical { .. } => Err(FilteredSearchError::InvalidPlanNode(
1871 "lexical node reached vector executor",
1872 )),
1873 PlanNode::Fusion { .. } => Err(FilteredSearchError::InvalidPlanNode(
1874 "fusion node reached vector executor before task 17",
1875 )),
1876 }
1877}
1878
1879#[allow(clippy::too_many_arguments)]
1880fn execute_squared_l2_plan(
1881 node: &PlanNode,
1882 vectors: &[f32],
1883 row_count: usize,
1884 query: &[f32],
1885 allow_list: &crate::meta::DocBitmap,
1886 k: usize,
1887 cancellation: &QueryCancellation<'_>,
1888 #[cfg(any(test, feature = "test-seams"))] receipt_context: Option<
1889 MetadataExecutionReceiptContext<'_>,
1890 >,
1891) -> Result<(Vec<LocalCandidate>, ScanStats, ScanExecutionFacts), FilteredSearchError> {
1892 match node {
1893 PlanNode::BitmapIntersect { input, .. } => execute_squared_l2_plan(
1894 input,
1895 vectors,
1896 row_count,
1897 query,
1898 allow_list,
1899 k,
1900 cancellation,
1901 #[cfg(any(test, feature = "test-seams"))]
1902 receipt_context,
1903 ),
1904 PlanNode::Scan { branch, .. } => execute_squared_l2_rows(
1905 vectors,
1906 row_count,
1907 query,
1908 allow_list,
1909 *branch,
1910 k,
1911 cancellation,
1912 #[cfg(any(test, feature = "test-seams"))]
1913 receipt_context,
1914 ),
1915 PlanNode::Graph {
1916 fallback: Some(fallback),
1917 ..
1918 } => {
1919 let (candidates, stats, mut execution) = execute_squared_l2_plan(
1920 fallback,
1921 vectors,
1922 row_count,
1923 query,
1924 allow_list,
1925 k,
1926 cancellation,
1927 #[cfg(any(test, feature = "test-seams"))]
1928 None,
1929 )?;
1930 execution.branch = SegmentBranch::GraphExactFallback;
1931 #[cfg(any(test, feature = "test-seams"))]
1932 if let Some(context) = receipt_context {
1933 context.record_scan(execution, &stats, candidates.len())?;
1934 }
1935 Ok((candidates, stats, execution))
1936 }
1937 PlanNode::Graph { fallback: None, .. } => Err(FilteredSearchError::InvalidPlanNode(
1938 "graph node has no exact fallback before task 19-M6",
1939 )),
1940 PlanNode::Lexical { .. } => Err(FilteredSearchError::InvalidPlanNode(
1941 "lexical node reached vector executor",
1942 )),
1943 PlanNode::Fusion { .. } => Err(FilteredSearchError::InvalidPlanNode(
1944 "fusion node reached vector executor before task 17",
1945 )),
1946 }
1947}
1948
1949#[allow(clippy::too_many_arguments)]
1950fn execute_squared_l2_rows(
1951 vectors: &[f32],
1952 row_count: usize,
1953 query: &[f32],
1954 allow_list: &crate::meta::DocBitmap,
1955 branch: SegmentBranch,
1956 k: usize,
1957 cancellation: &QueryCancellation<'_>,
1958 #[cfg(any(test, feature = "test-seams"))] receipt_context: Option<
1959 MetadataExecutionReceiptContext<'_>,
1960 >,
1961) -> Result<(Vec<LocalCandidate>, ScanStats, ScanExecutionFacts), FilteredSearchError> {
1962 if query.is_empty() {
1963 return Err(QueryError::Scan(ScanError::ZeroDimension).into());
1964 }
1965 let expected = row_count
1966 .checked_mul(query.len())
1967 .ok_or(QueryError::Scan(ScanError::ArithmeticOverflow))?;
1968 if vectors.len() != expected {
1969 return Err(QueryError::Scan(ScanError::RowDataLength {
1970 dimension: query.len(),
1971 actual: vectors.len(),
1972 })
1973 .into());
1974 }
1975 let mut candidates = Vec::new();
1976 let mut scored_rows = 0_u64;
1977 #[cfg(any(test, feature = "test-seams"))]
1978 let mut rows_examined = 0_u64;
1979 #[cfg(any(test, feature = "test-seams"))]
1980 let mut allowed_rows_examined = 0_u64;
1981 match branch {
1982 SegmentBranch::FilteredGraph | SegmentBranch::Graph => {
1983 Err(FilteredSearchError::InvalidPlanNode(
1984 "filtered graph branch reached exact squared-L2 rows",
1985 ))
1986 }
1987 SegmentBranch::Pruned => finish_squared_l2_rows(
1988 candidates,
1989 scored_rows,
1990 query.len(),
1991 SegmentBranch::Pruned,
1992 k,
1993 #[cfg(any(test, feature = "test-seams"))]
1994 rows_examined,
1995 #[cfg(any(test, feature = "test-seams"))]
1996 allowed_rows_examined,
1997 #[cfg(any(test, feature = "test-seams"))]
1998 receipt_context,
1999 ),
2000 SegmentBranch::ExactAllowList => {
2001 for row in allow_list.iter() {
2002 #[cfg(any(test, feature = "test-seams"))]
2003 {
2004 rows_examined = rows_examined
2005 .checked_add(1)
2006 .ok_or(QueryError::Scan(ScanError::ArithmeticOverflow))?;
2007 allowed_rows_examined = allowed_rows_examined
2008 .checked_add(1)
2009 .ok_or(QueryError::Scan(ScanError::ArithmeticOverflow))?;
2010 }
2011 let row = usize::try_from(row)
2012 .map_err(|_| QueryError::Scan(ScanError::ArithmeticOverflow))?;
2013 score_one(
2014 vectors,
2015 row_count,
2016 query,
2017 row,
2018 &mut candidates,
2019 &mut scored_rows,
2020 cancellation,
2021 )?;
2022 }
2023 finish_squared_l2_rows(
2024 candidates,
2025 scored_rows,
2026 query.len(),
2027 SegmentBranch::ExactAllowList,
2028 k,
2029 #[cfg(any(test, feature = "test-seams"))]
2030 rows_examined,
2031 #[cfg(any(test, feature = "test-seams"))]
2032 allowed_rows_examined,
2033 #[cfg(any(test, feature = "test-seams"))]
2034 receipt_context,
2035 )
2036 }
2037 SegmentBranch::MaskedScan | SegmentBranch::GraphExactFallback => {
2038 for row in 0..row_count {
2039 #[cfg(any(test, feature = "test-seams"))]
2040 {
2041 rows_examined = rows_examined
2042 .checked_add(1)
2043 .ok_or(QueryError::Scan(ScanError::ArithmeticOverflow))?;
2044 }
2045 if row.is_multiple_of(64) {
2046 cancellation.check_graph().map_err(map_scan_error)?;
2047 }
2048 let local_row = u32::try_from(row)
2049 .map_err(|_| QueryError::Scan(ScanError::ArithmeticOverflow))?;
2050 if allow_list.contains(local_row) {
2051 #[cfg(any(test, feature = "test-seams"))]
2052 {
2053 allowed_rows_examined = allowed_rows_examined
2054 .checked_add(1)
2055 .ok_or(QueryError::Scan(ScanError::ArithmeticOverflow))?;
2056 }
2057 score_one(
2058 vectors,
2059 row_count,
2060 query,
2061 row,
2062 &mut candidates,
2063 &mut scored_rows,
2064 cancellation,
2065 )?;
2066 }
2067 }
2068 finish_squared_l2_rows(
2069 candidates,
2070 scored_rows,
2071 query.len(),
2072 branch,
2073 k,
2074 #[cfg(any(test, feature = "test-seams"))]
2075 rows_examined,
2076 #[cfg(any(test, feature = "test-seams"))]
2077 allowed_rows_examined,
2078 #[cfg(any(test, feature = "test-seams"))]
2079 receipt_context,
2080 )
2081 }
2082 }
2083}
2084
2085#[allow(clippy::too_many_arguments)]
2086fn finish_squared_l2_rows(
2087 mut candidates: Vec<LocalCandidate>,
2088 scored_rows: u64,
2089 dimensions: usize,
2090 executed_branch: SegmentBranch,
2091 k: usize,
2092 #[cfg(any(test, feature = "test-seams"))] rows_examined: u64,
2093 #[cfg(any(test, feature = "test-seams"))] allowed_rows_examined: u64,
2094 #[cfg(any(test, feature = "test-seams"))] receipt_context: Option<
2095 MetadataExecutionReceiptContext<'_>,
2096 >,
2097) -> Result<(Vec<LocalCandidate>, ScanStats, ScanExecutionFacts), FilteredSearchError> {
2098 candidates.sort_unstable_by(|left, right| {
2099 right
2100 .score
2101 .total_cmp(&left.score)
2102 .then_with(|| left.row.cmp(&right.row))
2103 });
2104 crate::scan::truncate_to_k_with_score_ties(&mut candidates, k, |candidate| candidate.score);
2105 let dimensions =
2106 u64::try_from(dimensions).map_err(|_| QueryError::Scan(ScanError::ArithmeticOverflow))?;
2107 let dims_touched = scored_rows
2108 .checked_mul(dimensions)
2109 .ok_or(QueryError::Scan(ScanError::ArithmeticOverflow))?;
2110 let bytes_read = dims_touched
2111 .checked_mul(std::mem::size_of::<f32>() as u64)
2112 .ok_or(QueryError::Scan(ScanError::ArithmeticOverflow))?;
2113 let stats = ScanStats {
2114 dims_touched,
2115 bytes_read,
2116 threads_used: usize::from(scored_rows != 0),
2117 worker_thread_ids: (scored_rows != 0)
2118 .then(|| std::thread::current().id())
2119 .into_iter()
2120 .collect(),
2121 };
2122 let execution = ScanExecutionFacts {
2123 branch: executed_branch,
2124 #[cfg(any(test, feature = "test-seams"))]
2125 rows_examined,
2126 #[cfg(any(test, feature = "test-seams"))]
2127 allowed_rows_examined,
2128 };
2129 #[cfg(any(test, feature = "test-seams"))]
2130 if let Some(context) = receipt_context {
2131 context.record_scan(execution, &stats, candidates.len())?;
2132 }
2133 Ok((candidates, stats, execution))
2134}
2135
2136#[allow(clippy::too_many_arguments)]
2137fn score_one(
2138 vectors: &[f32],
2139 row_count: usize,
2140 query: &[f32],
2141 row: usize,
2142 candidates: &mut Vec<LocalCandidate>,
2143 scored_rows: &mut u64,
2144 cancellation: &QueryCancellation<'_>,
2145) -> Result<(), FilteredSearchError> {
2146 if row >= row_count {
2147 return Err(QueryError::Scan(ScanError::ArithmeticOverflow).into());
2148 }
2149 if row.is_multiple_of(64) {
2150 cancellation.check_graph().map_err(map_scan_error)?;
2151 }
2152 let start = row
2153 .checked_mul(query.len())
2154 .ok_or(QueryError::Scan(ScanError::ArithmeticOverflow))?;
2155 let end = start
2156 .checked_add(query.len())
2157 .ok_or(QueryError::Scan(ScanError::ArithmeticOverflow))?;
2158 let vector = vectors
2159 .get(start..end)
2160 .ok_or(QueryError::Scan(ScanError::RowDataLength {
2161 dimension: query.len(),
2162 actual: vectors.len(),
2163 }))?;
2164 let distance = vector
2165 .iter()
2166 .zip(query)
2167 .map(|(left, right)| {
2168 let delta = f64::from(*left) - f64::from(*right);
2169 delta * delta
2170 })
2171 .sum::<f64>();
2172 let score = -(distance as f32);
2173 if !score.is_finite() {
2174 return Err(QueryError::Scan(ScanError::NonFiniteScore { row_id: row }).into());
2175 }
2176 candidates.push(LocalCandidate { row, score });
2177 *scored_rows = scored_rows
2178 .checked_add(1)
2179 .ok_or(QueryError::Scan(ScanError::ArithmeticOverflow))?;
2180 Ok(())
2181}
2182
2183fn append_candidates(
2184 local: Vec<LocalCandidate>,
2185 source: RowSource,
2186 document: impl Fn(usize) -> Result<Option<crate::ingest::DocumentVersion>, QueryError>,
2187 candidates: &mut Vec<SearchCandidate>,
2188 exact_score: bool,
2189 #[cfg(any(test, feature = "test-seams"))] vector_fault_controller: Option<
2190 &crate::scan::vector_fault::VectorFaultController,
2191 >,
2192 #[cfg(any(test, feature = "test-seams"))]
2193 vector_fault_tier: crate::scan::vector_fault::VectorSearchTier,
2194) -> Result<(), FilteredSearchError> {
2195 crate::lifecycle::reserve_global_candidates(
2196 candidates,
2197 local.len(),
2198 #[cfg(any(test, feature = "test-seams"))]
2199 vector_fault_controller,
2200 #[cfg(any(test, feature = "test-seams"))]
2201 vector_fault_tier,
2202 )?;
2203 for candidate in local {
2204 let local_row = u32::try_from(candidate.row)
2205 .map_err(|_| QueryError::Scan(ScanError::ArithmeticOverflow))?;
2206 candidates.push(SearchCandidate::new(
2207 crate::ingest::GlobalRowId::new(source, local_row),
2208 document(candidate.row)?,
2209 candidate.score,
2210 exact_score,
2211 ));
2212 }
2213 Ok(())
2214}
2215
2216#[cfg(any(test, feature = "test-seams"))]
2217fn vector_fault_tier(tier: SearchTier) -> crate::scan::vector_fault::VectorSearchTier {
2218 match tier {
2219 SearchTier::Exact => crate::scan::vector_fault::VectorSearchTier::Exact,
2220 SearchTier::Scan => crate::scan::vector_fault::VectorSearchTier::Scan,
2221 SearchTier::Graph(_) => crate::scan::vector_fault::VectorSearchTier::Graph,
2222 SearchTier::Auto => crate::scan::vector_fault::VectorSearchTier::Auto,
2223 }
2224}
2225
2226fn retain_global_top_k(candidates: &mut Vec<SearchCandidate>, k: usize) {
2227 candidates.sort_unstable_by(crate::ingest::compare_search_candidates);
2228 candidates.truncate(k);
2229}
2230
2231fn global_competitive_distance(candidates: &[SearchCandidate], k: usize) -> Option<f32> {
2232 if k == 0 || candidates.len() < k {
2233 return None;
2234 }
2235 candidates
2236 .get(k - 1)
2237 .map(|candidate| (-candidate.score()).max(0.0))
2238 .filter(|distance| distance.is_finite())
2239}
2240
2241fn has_graph(segment: &crate::segment::reader::SegmentReader) -> bool {
2242 segment
2243 .directory()
2244 .iter()
2245 .any(|entry| entry.kind == RegionKind::GraphNodeBlocks.id())
2246}
2247
2248fn map_eval_error(error: EvalError) -> FilteredSearchError {
2249 let plan = match error {
2250 EvalError::UnknownColumn(column) => PlanError::UnknownColumn(column),
2251 EvalError::TypeMismatch {
2252 column,
2253 expected,
2254 actual,
2255 } => PlanError::TypeMismatch {
2256 column,
2257 expected,
2258 actual,
2259 },
2260 EvalError::RangeRequiresNumericColumn(column) => {
2261 PlanError::RangeRequiresNumericColumn(column)
2262 }
2263 EvalError::RowCountMismatch { columns, alive } => {
2264 PlanError::RowCountMismatch { columns, alive }
2265 }
2266 };
2267 FilteredSearchError::Plan(plan)
2268}
2269
2270fn evaluate_for_plan(
2271 predicate: &Predicate,
2272 columns: &ColumnStore,
2273 alive: &crate::meta::AliveSet,
2274) -> Result<crate::meta::DocBitmap, FilteredSearchError> {
2275 evaluate(predicate, columns, alive).map_err(map_eval_error)
2276}
2277
2278fn map_scan_error(error: ScanError) -> FilteredSearchError {
2279 let query = match error {
2280 ScanError::Cancelled { partial } => QueryError::Cancelled { partial },
2281 ScanError::Timeout { partial } => QueryError::Timeout { partial },
2282 ScanError::ReadCancelled { partial } => QueryError::ReadCancelled { partial },
2283 error => QueryError::Scan(error),
2284 };
2285 FilteredSearchError::Query(query)
2286}
2287
2288fn verify_execution_branch(
2289 reported: SegmentBranch,
2290 executed: SegmentBranch,
2291) -> Result<(), FilteredSearchError> {
2292 if reported == executed {
2293 Ok(())
2294 } else {
2295 Err(FilteredSearchError::PlanReportMismatch { reported, executed })
2296 }
2297}
2298
2299#[derive(Default)]
2300struct MutableStats {
2301 dims_touched: u64,
2302 bytes_read: u64,
2303 worker_thread_ids: Vec<std::thread::ThreadId>,
2304}
2305
2306impl MutableStats {
2307 fn add(&mut self, stats: ScanStats) -> Result<(), FilteredSearchError> {
2308 self.dims_touched = self
2309 .dims_touched
2310 .checked_add(stats.dims_touched)
2311 .ok_or(QueryError::Scan(ScanError::ArithmeticOverflow))?;
2312 self.bytes_read = self
2313 .bytes_read
2314 .checked_add(stats.bytes_read)
2315 .ok_or(QueryError::Scan(ScanError::ArithmeticOverflow))?;
2316 for thread in stats.worker_thread_ids {
2317 if !self.worker_thread_ids.contains(&thread) {
2318 self.worker_thread_ids.push(thread);
2319 }
2320 }
2321 Ok(())
2322 }
2323
2324 fn finish(self) -> ScanStats {
2325 ScanStats {
2326 dims_touched: self.dims_touched,
2327 bytes_read: self.bytes_read,
2328 threads_used: self.worker_thread_ids.len(),
2329 worker_thread_ids: self.worker_thread_ids,
2330 }
2331 }
2332}
2333
2334struct ActiveFilteredQuery<'a> {
2335 count: &'a AtomicU64,
2336}
2337
2338impl Drop for ActiveFilteredQuery<'_> {
2339 fn drop(&mut self) {
2340 self.count.fetch_sub(1, Ordering::Relaxed);
2341 }
2342}
2343
2344#[allow(clippy::expect_used, clippy::panic)]
2345#[cfg(test)]
2346mod tests {
2347 use super::*;
2348 use crate::meta::{AliveSet, ColumnStoreBuilder, Schema};
2349
2350 #[test]
2351 fn metadata_execution_receipt_report_mismatch_is_typed() {
2352 assert!(matches!(
2353 verify_execution_branch(SegmentBranch::MaskedScan, SegmentBranch::ExactAllowList),
2354 Err(FilteredSearchError::PlanReportMismatch {
2355 reported: SegmentBranch::MaskedScan,
2356 executed: SegmentBranch::ExactAllowList,
2357 })
2358 ));
2359 }
2360
2361 #[test]
2362 fn metadata_execution_receipt_row_count_mismatch_uses_real_evaluator_inputs() {
2363 let mut builder = ColumnStoreBuilder::new(Schema::timestamp_only());
2364 builder.push_row(1, &[]).expect("one real metadata row");
2365 let columns = builder.finish().expect("real metadata columns");
2366 let alive = AliveSet::new(2);
2367
2368 assert!(matches!(
2369 evaluate_for_plan(&Predicate::And(Vec::new()), &columns, &alive),
2370 Err(FilteredSearchError::Plan(PlanError::RowCountMismatch {
2371 columns: 1,
2372 alive: 2,
2373 }))
2374 ));
2375 }
2376}