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