1use crate::graph::block::{CheckedNodeId, GraphNodeBlocks, GraphNodeError};
4use crate::kernels::{Bit4Rows4, GatherShapeError, score_bit4_ptrs};
5use crate::lifecycle::QueryCancellation;
6use crate::meta::DocBitmap;
7use crate::quant::{
8 Bit4Query, QuantError, RescoreError, RescoreMetric, RescorePool, prepare_bit4_query,
9 rescore_top_k,
10};
11use crate::scan::ScanError;
12
13#[cfg(any(test, feature = "test-seams"))]
16#[derive(Default)]
17#[doc(hidden)]
18pub struct GraphQueryTestObservations {
19 pub traversals: usize,
21 pub preparations: Vec<(usize, usize, u64)>,
23 pub exact_rows: Vec<usize>,
25 pub validated_pools: Vec<Vec<u32>>,
27}
28
29#[cfg(any(test, feature = "test-seams"))]
30pub(crate) fn observe_exact_score(row: usize) {
31 GRAPH_QUERY_TEST_OBSERVATIONS.with(|observations| {
32 if let Some(observations) = observations.borrow_mut().as_mut() {
33 observations.exact_rows.push(row);
34 }
35 });
36}
37
38#[cfg(any(test, feature = "test-seams"))]
39thread_local! {
40 static GRAPH_QUERY_TEST_OBSERVATIONS: std::cell::RefCell<Option<GraphQueryTestObservations>> =
41 const { std::cell::RefCell::new(None) };
42}
43
44#[cfg(any(test, feature = "test-seams"))]
46#[doc(hidden)]
47pub fn begin_graph_query_test_observations() {
48 GRAPH_QUERY_TEST_OBSERVATIONS.with(|observations| {
49 *observations.borrow_mut() = Some(GraphQueryTestObservations::default());
50 });
51}
52
53#[cfg(any(test, feature = "test-seams"))]
55#[doc(hidden)]
56pub fn take_graph_query_test_observations() -> GraphQueryTestObservations {
57 GRAPH_QUERY_TEST_OBSERVATIONS.with(|observations| observations.take().unwrap_or_default())
58}
59
60pub(crate) struct PreparedGraphQuery {
62 query: Bit4Query,
63 norm: f64,
64}
65
66impl PreparedGraphQuery {
67 pub(crate) fn new(
68 query: &[f32],
69 padded_dimensions: usize,
70 seed: u64,
71 ) -> Result<Self, GraphSearchError> {
72 if padded_dimensions < query.len() {
73 return Err(GraphSearchError::Geometry(
74 "query padding is smaller than its original dimensions".to_owned(),
75 ));
76 }
77 let padded = if padded_dimensions == query.len() {
78 None
79 } else {
80 let mut values = Vec::with_capacity(padded_dimensions);
81 values.extend_from_slice(query);
82 values.resize(padded_dimensions, 0.0);
83 Some(values)
84 };
85 let prepared = prepare_bit4_query(padded.as_deref().unwrap_or(query), seed)?;
86 #[cfg(any(test, feature = "test-seams"))]
87 GRAPH_QUERY_TEST_OBSERVATIONS.with(|observations| {
88 if let Some(observations) = observations.borrow_mut().as_mut() {
89 observations
90 .preparations
91 .push((query.len(), padded_dimensions, seed));
92 }
93 });
94 Ok(Self {
95 query: prepared,
96 norm: squared_norm(query),
97 })
98 }
99
100 pub(crate) fn resident_bytes(&self) -> usize {
101 self.query.resident_bytes()
102 }
103
104 pub(crate) fn bit4(&self) -> &Bit4Query {
105 &self.query
106 }
107}
108
109const SIFT_RESEARCH_FLOOR_TENTHS: usize = 14;
111const EF_TENTHS_DENOMINATOR: usize = 10;
113const SIFT_RESEARCH_ABSOLUTE_FLOOR: usize = 140;
117const SIFT_SHIPPED_EF_PER_K: usize = 2;
121const ANGULAR_EF_PER_K: usize = 4;
126pub const UNIT_NORM_SQUARED_TOLERANCE: f64 = 1.0e-3;
128
129pub(crate) fn non_unit_squared_norm(vector: &[f32]) -> Option<f64> {
130 match non_unit_squared_norm_controlled(vector, &mut |_| Ok::<_, std::convert::Infallible>(())) {
131 Ok(value) => value,
132 Err(never) => match never {},
133 }
134}
135
136pub(crate) fn non_unit_squared_norm_controlled<E>(
137 vector: &[f32],
138 control: &mut impl FnMut(u64) -> Result<(), E>,
139) -> Result<Option<f64>, E> {
140 if vector.is_empty() {
141 return Ok(None);
142 }
143 let mut squared_norm = 0.0_f64;
144 for chunk in vector.chunks(256) {
145 control(chunk.len() as u64)?;
146 for value in chunk {
147 if !value.is_finite() {
148 return Ok(None);
149 }
150 squared_norm += f64::from(*value) * f64::from(*value);
151 }
152 }
153 Ok(((squared_norm - 1.0).abs() > UNIT_NORM_SQUARED_TOLERANCE).then_some(squared_norm))
154}
155
156#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
158pub enum GraphSearchProfile {
159 #[default]
161 SiftClass,
162 Angular,
164}
165
166#[derive(Clone, Copy, Debug, Eq, PartialEq)]
168pub enum GraphDistanceMetric {
169 SquaredL2,
171}
172
173impl std::fmt::Display for GraphDistanceMetric {
174 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
175 match self {
176 Self::SquaredL2 => formatter.write_str("squared-l2"),
177 }
178 }
179}
180
181#[derive(Clone, Copy, Debug, Eq, PartialEq)]
183pub struct GraphProfileShape {
184 pub document_dims: u32,
186 pub document_normalization: crate::epoch::Normalization,
188 pub query_dims: u32,
190 pub query_normalization: crate::epoch::Normalization,
192 pub metric: GraphDistanceMetric,
194}
195
196#[derive(Clone, Copy, Debug, Eq, PartialEq)]
198pub enum EpochGraphProfile {
199 SiftClass,
201 AngularClass,
203}
204
205impl EpochGraphProfile {
206 #[must_use]
208 pub const fn search_profile(self) -> GraphSearchProfile {
209 match self {
210 Self::SiftClass => GraphSearchProfile::SiftClass,
211 Self::AngularClass => GraphSearchProfile::Angular,
212 }
213 }
214
215 #[must_use]
217 pub const fn build_params(self) -> crate::graph::GraphParams {
218 match self {
219 Self::SiftClass => crate::graph::GraphParams::sift_1m(),
220 Self::AngularClass => crate::graph::GraphParams::angular(),
221 }
222 }
223}
224
225#[derive(Clone, Copy, Debug, Eq, PartialEq)]
227pub enum GraphProfileError {
228 EpochUnstamped,
230 EpochNotRegistered {
232 epoch: crate::epoch::EpochIdentity,
234 },
235 UnrecognizedEpochShape {
237 epoch: crate::epoch::EpochIdentity,
239 shape: GraphProfileShape,
241 },
242}
243
244impl std::fmt::Display for GraphProfileError {
245 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
246 match self {
247 Self::EpochUnstamped => formatter.write_str(
248 "graph profile requires a persisted epoch registry, but the store is unstamped",
249 ),
250 Self::EpochNotRegistered { epoch } => write!(
251 formatter,
252 "graph profile epoch ({}, {}) is absent from the pinned registry",
253 epoch.embedding, epoch.tokenizer
254 ),
255 Self::UnrecognizedEpochShape { epoch, shape } => write!(
256 formatter,
257 "graph profile epoch ({}, {}) has unrecognized shape: document dims {}, document normalization {:?}, query dims {}, query normalization {:?}, metric {}",
258 epoch.embedding,
259 epoch.tokenizer,
260 shape.document_dims,
261 shape.document_normalization,
262 shape.query_dims,
263 shape.query_normalization,
264 shape.metric
265 ),
266 }
267 }
268}
269
270impl std::error::Error for GraphProfileError {}
271
272pub fn select_epoch_graph_profile(
274 epoch: &crate::manifest::EpochMeta,
275 metric: GraphDistanceMetric,
276) -> Result<EpochGraphProfile, GraphProfileError> {
277 let identity = crate::epoch::EpochIdentity {
278 embedding: epoch.id,
279 tokenizer: epoch.tokenizer,
280 };
281 let shape = GraphProfileShape {
282 document_dims: epoch.embedding.document.dims,
283 document_normalization: epoch.embedding.document.normalization,
284 query_dims: epoch.embedding.query.dims,
285 query_normalization: epoch.embedding.query.normalization,
286 metric,
287 };
288 if shape.document_dims == shape.query_dims && shape.metric == GraphDistanceMetric::SquaredL2 {
301 match (shape.document_normalization, shape.query_normalization) {
302 (crate::epoch::Normalization::None, crate::epoch::Normalization::None) => {
303 return Ok(EpochGraphProfile::SiftClass);
304 }
305 (crate::epoch::Normalization::L2, crate::epoch::Normalization::L2) => {
306 return Ok(EpochGraphProfile::AngularClass);
309 }
310 _ => {}
311 }
312 }
313 Err(GraphProfileError::UnrecognizedEpochShape {
314 epoch: identity,
315 shape,
316 })
317}
318
319#[derive(Clone, Copy, Debug, Eq, PartialEq)]
321pub enum AdaptiveEfError {
322 ZeroK,
324 KExceedsRows {
326 k: usize,
328 rows: usize,
330 },
331 ExplicitBelowK {
333 k: usize,
335 ef: usize,
337 },
338 ExplicitExceedsRows {
340 ef: usize,
342 rows: usize,
344 },
345 ArithmeticOverflow,
347}
348
349impl std::fmt::Display for AdaptiveEfError {
350 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
351 match self {
352 Self::ZeroK => formatter.write_str("adaptive graph top-k must be positive"),
353 Self::KExceedsRows { k, rows } => {
354 write!(formatter, "graph top-k {k} exceeds graph row count {rows}")
355 }
356 Self::ExplicitBelowK { k, ef } => {
357 write!(formatter, "explicit graph ef {ef} is below top-k {k}")
358 }
359 Self::ExplicitExceedsRows { ef, rows } => {
360 write!(
361 formatter,
362 "explicit graph ef {ef} exceeds graph row count {rows}"
363 )
364 }
365 Self::ArithmeticOverflow => formatter.write_str("adaptive graph ef overflowed usize"),
366 }
367 }
368}
369
370impl std::error::Error for AdaptiveEfError {}
371
372#[derive(Clone, Copy, Debug)]
374pub struct GraphSearchRequest<'a> {
375 query: &'a [f32],
376 k: usize,
377 ef: Option<usize>,
378 profile: GraphSearchProfile,
379 seed: u64,
380 trace_candidates: bool,
381 observed_core_class: QueryCoreClass,
382 prefetch: TraversalPrefetch,
383}
384
385impl<'a> GraphSearchRequest<'a> {
386 #[must_use]
388 pub const fn new(query: &'a [f32], k: usize, seed: u64) -> Self {
389 Self {
390 query,
391 k,
392 ef: None,
393 profile: GraphSearchProfile::SiftClass,
394 seed,
395 trace_candidates: false,
396 observed_core_class: QueryCoreClass::Unverified,
397 prefetch: TraversalPrefetch::Enabled,
398 }
399 }
400
401 #[must_use]
403 pub const fn with_profile(mut self, profile: GraphSearchProfile) -> Self {
404 self.profile = profile;
405 self
406 }
407
408 #[must_use]
410 pub const fn with_ef(mut self, ef: usize) -> Self {
411 self.ef = Some(ef);
412 self
413 }
414
415 pub fn research_floor_ef(self) -> Result<usize, AdaptiveEfError> {
421 if self.k == 0 {
422 return Err(AdaptiveEfError::ZeroK);
423 }
424 match self.profile {
425 GraphSearchProfile::SiftClass => self
426 .k
427 .checked_mul(SIFT_RESEARCH_FLOOR_TENTHS)
428 .map(|value| value.div_ceil(EF_TENTHS_DENOMINATOR))
429 .map(|value| value.max(SIFT_RESEARCH_ABSOLUTE_FLOOR))
430 .ok_or(AdaptiveEfError::ArithmeticOverflow),
431 GraphSearchProfile::Angular => self
432 .k
433 .checked_mul(ANGULAR_EF_PER_K)
434 .ok_or(AdaptiveEfError::ArithmeticOverflow),
435 }
436 }
437
438 pub fn effective_ef(self, rows: usize) -> Result<usize, AdaptiveEfError> {
445 if self.k == 0 {
446 return Err(AdaptiveEfError::ZeroK);
447 }
448 if self.k > rows {
449 return Err(AdaptiveEfError::KExceedsRows { k: self.k, rows });
450 }
451 if let Some(ef) = self.ef {
452 if ef < self.k {
453 return Err(AdaptiveEfError::ExplicitBelowK { k: self.k, ef });
454 }
455 if ef > rows {
456 return Err(AdaptiveEfError::ExplicitExceedsRows { ef, rows });
457 }
458 return Ok(ef);
459 }
460 let floor = self.research_floor_ef()?;
461 let target = match self.profile {
462 GraphSearchProfile::SiftClass => self
463 .k
464 .checked_mul(SIFT_SHIPPED_EF_PER_K)
465 .map(|ef| ef.max(floor))
466 .ok_or(AdaptiveEfError::ArithmeticOverflow)?,
467 GraphSearchProfile::Angular => floor,
468 };
469 Ok(target.min(rows))
470 }
471
472 #[must_use]
474 pub const fn with_candidate_trace(mut self) -> Self {
475 self.trace_candidates = true;
476 self
477 }
478
479 #[must_use]
481 pub const fn with_observed_core_class(mut self, core_class: QueryCoreClass) -> Self {
482 self.observed_core_class = core_class;
483 self
484 }
485
486 #[must_use]
488 pub const fn with_prefetch(mut self, prefetch: TraversalPrefetch) -> Self {
489 self.prefetch = prefetch;
490 self
491 }
492}
493
494#[derive(Clone, Copy, Debug, Eq, PartialEq)]
496pub enum TraversalPrefetch {
497 Enabled,
499 Disabled,
501}
502
503impl TraversalPrefetch {
504 const fn is_enabled(self) -> bool {
505 matches!(self, Self::Enabled)
506 }
507}
508
509#[derive(Clone, Copy, Debug, PartialEq)]
511pub struct GraphSearchCandidate {
512 row_id: u32,
513 distance: f64,
514}
515
516impl GraphSearchCandidate {
517 #[must_use]
519 pub const fn row_id(self) -> u32 {
520 self.row_id
521 }
522
523 #[must_use]
525 pub const fn distance(self) -> f64 {
526 self.distance
527 }
528}
529
530#[derive(Clone, Copy, Debug, Eq, PartialEq)]
532pub enum QueryQosClass {
533 UserInteractive,
535 UserInitiated,
537 Default,
539 Utility,
541 Background,
543 Unspecified,
545 Unavailable,
547 Unknown,
549}
550
551#[derive(Clone, Copy, Debug, Eq, PartialEq)]
553pub enum QueryCoreClass {
554 Performance,
556 Efficiency,
558 Unverified,
560}
561
562#[derive(Clone, Copy, Debug, Eq, PartialEq)]
564pub struct GraphSearchCounters {
565 effective_ef: usize,
566 hops: usize,
567 candidates_scored: usize,
568 candidates_rescored: usize,
569 candidate_window_peak: usize,
570 pushes: usize,
571 visited: usize,
572 visited_epoch_cleared: bool,
573 dims_touched: u64,
574 bytes_read: u64,
575 qos_class: QueryQosClass,
576 qos_relative_priority: i32,
577 core_class: QueryCoreClass,
578}
579
580#[derive(Clone, Copy, Debug, Eq, PartialEq)]
582pub struct DeterministicGraphWork {
583 pub effective_ef: usize,
585 pub hops: usize,
587 pub candidates_scored: usize,
589 pub candidates_rescored: usize,
591 pub pushes: usize,
593 pub visited: usize,
595}
596
597impl GraphSearchCounters {
598 pub const fn candidate_window_peak(self) -> usize {
600 self.candidate_window_peak
601 }
602
603 #[must_use]
605 pub const fn deterministic_work(self) -> DeterministicGraphWork {
606 DeterministicGraphWork {
607 effective_ef: self.effective_ef,
608 hops: self.hops,
609 candidates_scored: self.candidates_scored,
610 candidates_rescored: self.candidates_rescored,
611 pushes: self.pushes,
612 visited: self.visited,
613 }
614 }
615
616 #[must_use]
618 pub const fn effective_ef(self) -> usize {
619 self.effective_ef
620 }
621
622 #[must_use]
624 pub const fn hops(self) -> usize {
625 self.hops
626 }
627
628 #[must_use]
630 pub const fn candidates_scored(self) -> usize {
631 self.candidates_scored
632 }
633
634 #[must_use]
636 pub const fn candidates_rescored(self) -> usize {
637 self.candidates_rescored
638 }
639
640 #[must_use]
642 pub const fn pushes(self) -> usize {
643 self.pushes
644 }
645
646 #[must_use]
648 pub const fn visited(self) -> usize {
649 self.visited
650 }
651
652 #[must_use]
654 pub const fn visited_epoch_cleared(self) -> bool {
655 self.visited_epoch_cleared
656 }
657
658 #[must_use]
660 pub const fn dims_touched(self) -> u64 {
661 self.dims_touched
662 }
663
664 #[must_use]
666 pub const fn bytes_read(self) -> u64 {
667 self.bytes_read
668 }
669
670 #[must_use]
672 pub const fn qos_class(self) -> QueryQosClass {
673 self.qos_class
674 }
675
676 #[must_use]
678 pub const fn qos_relative_priority(self) -> i32 {
679 self.qos_relative_priority
680 }
681
682 #[must_use]
684 pub const fn core_class(self) -> QueryCoreClass {
685 self.core_class
686 }
687}
688
689#[derive(Clone, Debug, PartialEq)]
691pub struct GraphSearchResult {
692 candidates: Vec<GraphSearchCandidate>,
693 counters: GraphSearchCounters,
694 candidate_sequence: Option<Vec<u32>>,
695}
696
697impl GraphSearchResult {
698 #[must_use]
700 pub fn candidates(&self) -> &[GraphSearchCandidate] {
701 &self.candidates
702 }
703
704 #[must_use]
706 pub const fn counters(&self) -> GraphSearchCounters {
707 self.counters
708 }
709
710 #[must_use]
712 pub fn candidate_sequence(&self) -> Option<&[u32]> {
713 self.candidate_sequence.as_deref()
714 }
715
716 #[must_use]
718 pub fn recall_against(&self, truth: &[u32]) -> usize {
719 self.candidates
720 .iter()
721 .filter(|candidate| truth.contains(&candidate.row_id))
722 .count()
723 }
724}
725
726#[derive(Clone, Debug, PartialEq)]
728pub enum FilteredGraphSearchOutcome {
729 Traversed(GraphSearchResult),
731 VisitedBudgetExceeded {
733 visited: usize,
735 budget: usize,
737 counters: GraphSearchCounters,
739 },
740}
741
742impl FilteredGraphSearchOutcome {
743 #[must_use]
745 pub const fn counters(&self) -> GraphSearchCounters {
746 match self {
747 Self::Traversed(result) => result.counters,
748 Self::VisitedBudgetExceeded { counters, .. } => *counters,
749 }
750 }
751}
752
753enum TraversalOutcome {
754 Retained {
755 counters: GraphSearchCounters,
756 candidate_sequence: Option<Vec<u32>>,
757 },
758 Budget(SearchInnerOutcome),
759}
760
761enum SearchInnerOutcome {
762 Traversed(GraphSearchResult),
763 VisitedBudgetExceeded {
764 visited: usize,
765 budget: usize,
766 counters: GraphSearchCounters,
767 },
768}
769
770#[derive(Clone, Debug, PartialEq)]
772pub enum GraphSearchError {
773 Geometry(String),
775 Profile(GraphProfileError),
777 Graph(GraphNodeError),
779 Quant(QuantError),
781 Gather(GatherShapeError),
783 AdaptiveEf(AdaptiveEfError),
785 Rescore(RescoreError),
787 FilteredCandidateShortfall {
789 k: usize,
791 candidates: usize,
793 counters: GraphSearchCounters,
795 },
796 ExactRescoreUnavailable(String),
798 VisitedCapExceeded {
800 visited: usize,
802 cap: usize,
804 },
805 Cancelled {
807 partial: bool,
809 },
810 Timeout {
812 partial: bool,
814 },
815 ReadCancelled {
817 partial: bool,
819 },
820 Scan(ScanError),
822}
823
824impl std::fmt::Display for GraphSearchError {
825 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
826 match self {
827 Self::Geometry(detail) => {
828 write!(formatter, "graph search geometry is invalid: {detail}")
829 }
830 Self::Profile(error) => error.fmt(formatter),
831 Self::Graph(error) => error.fmt(formatter),
832 Self::Quant(error) => error.fmt(formatter),
833 Self::Gather(error) => error.fmt(formatter),
834 Self::AdaptiveEf(error) => error.fmt(formatter),
835 Self::Rescore(error) => error.fmt(formatter),
836 Self::FilteredCandidateShortfall { k, candidates, .. } => write!(
837 formatter,
838 "filtered graph search retained {candidates} candidates for requested k={k}"
839 ),
840 Self::ExactRescoreUnavailable(detail) => formatter.write_str(detail),
841 Self::VisitedCapExceeded { visited, cap } => write!(
842 formatter,
843 "graph search visited {visited} nodes, exceeding the fail-closed cap {cap}"
844 ),
845 Self::Cancelled { partial } => {
846 write!(formatter, "graph search was cancelled (partial={partial})")
847 }
848 Self::Timeout { partial } => {
849 write!(
850 formatter,
851 "graph search deadline expired (partial={partial})"
852 )
853 }
854 Self::ReadCancelled { partial } => write!(
855 formatter,
856 "store close cancelled graph search (partial={partial})"
857 ),
858 Self::Scan(error) => error.fmt(formatter),
859 }
860 }
861}
862
863impl std::error::Error for GraphSearchError {
864 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
865 match self {
866 Self::Profile(error) => Some(error),
867 Self::Graph(error) => Some(error),
868 Self::Quant(error) => Some(error),
869 Self::Gather(error) => Some(error),
870 Self::AdaptiveEf(error) => Some(error),
871 Self::Rescore(error) => Some(error),
872 Self::Scan(error) => Some(error),
873 Self::Geometry(_)
874 | Self::ExactRescoreUnavailable(_)
875 | Self::FilteredCandidateShortfall { .. }
876 | Self::VisitedCapExceeded { .. }
877 | Self::Cancelled { .. }
878 | Self::Timeout { .. }
879 | Self::ReadCancelled { .. } => None,
880 }
881 }
882}
883
884impl From<GraphNodeError> for GraphSearchError {
885 fn from(error: GraphNodeError) -> Self {
886 Self::Graph(error)
887 }
888}
889
890impl From<QuantError> for GraphSearchError {
891 fn from(error: QuantError) -> Self {
892 Self::Quant(error)
893 }
894}
895
896impl From<GatherShapeError> for GraphSearchError {
897 fn from(error: GatherShapeError) -> Self {
898 Self::Gather(error)
899 }
900}
901
902impl From<AdaptiveEfError> for GraphSearchError {
903 fn from(error: AdaptiveEfError) -> Self {
904 Self::AdaptiveEf(error)
905 }
906}
907
908impl From<RescoreError> for GraphSearchError {
909 fn from(error: RescoreError) -> Self {
910 Self::Rescore(error)
911 }
912}
913
914#[derive(Clone, Copy, Debug)]
915struct ScoredNode {
916 row_id: CheckedNodeId,
917 distance: f64,
918}
919
920#[derive(Clone, Copy, Debug)]
926pub(crate) struct GraphSegmentNormRange {
927 minimum: f64,
928 maximum: f64,
929}
930
931impl GraphSegmentNormRange {
932 pub(crate) fn from_factors(factors: &[crate::quant::Bit4Factors]) -> Self {
933 let mut minimum = f64::INFINITY;
934 let mut maximum = 0.0_f64;
935 for factors in factors {
936 let [scale, normalized_norm, _] = factors.persisted_fields();
937 if scale < 0.0 || normalized_norm < 0.0 {
938 return Self::unbounded();
939 }
940 let center = f64::from(scale) * f64::from(normalized_norm);
941 let normalized_ulp = f64::from(next_up_f32(normalized_norm) - normalized_norm);
942 let uncertainty = f64::from(scale) * normalized_ulp;
943 minimum = minimum.min((center - uncertainty).max(0.0));
944 maximum = maximum.max(center + uncertainty);
945 }
946 if !minimum.is_finite() || !maximum.is_finite() {
947 return Self::unbounded();
948 }
949 Self { minimum, maximum }
950 }
951
952 pub(crate) fn from_exact_rows(rows: &[f32], dimensions: usize) -> Self {
953 if dimensions == 0 || rows.is_empty() || !rows.len().is_multiple_of(dimensions) {
954 return Self::unbounded();
955 }
956 let relative_error = accumulation_relative_error(dimensions);
957 let mut maximum = 0.0_f64;
958 for row in rows.chunks_exact(dimensions) {
959 let squared_norm = row
960 .iter()
961 .map(|value| {
962 let value = f64::from(*value);
963 value * value
964 })
965 .sum::<f64>();
966 let norm = next_up_f64((squared_norm / (1.0 - relative_error)).max(0.0).sqrt());
967 maximum = maximum.max(norm);
968 }
969 if !maximum.is_finite() {
970 return Self::unbounded();
971 }
972 Self {
973 minimum: 0.0,
974 maximum,
975 }
976 }
977
978 pub(crate) fn from_graph(
979 graph: GraphNodeBlocks<'_>,
980 cancellation: Option<&QueryCancellation<'_>>,
981 ) -> Result<Self, GraphSearchError> {
982 let mut minimum = f64::INFINITY;
983 let mut maximum = 0.0_f64;
984 for raw_row_id in 0..graph.node_count() {
985 if raw_row_id.is_multiple_of(256) {
986 check_cancellation(cancellation)?;
987 }
988 let row_id = graph.checked_node_id(raw_row_id)?;
989 let (_, factors) = graph.score_row_checked(row_id)?;
990 let [scale, normalized_norm, _] = factors.persisted_fields();
991 if scale < 0.0 || normalized_norm < 0.0 {
992 return Ok(Self::unbounded());
993 }
994 let center = f64::from(scale) * f64::from(normalized_norm);
995 let normalized_ulp = f64::from(next_up_f32(normalized_norm) - normalized_norm);
996 let uncertainty = f64::from(scale) * normalized_ulp;
997 minimum = minimum.min((center - uncertainty).max(0.0));
998 maximum = maximum.max(center + uncertainty);
999 }
1000 check_cancellation(cancellation)?;
1001 if !minimum.is_finite() || !maximum.is_finite() {
1002 return Ok(Self::unbounded());
1003 }
1004 Ok(Self { minimum, maximum })
1005 }
1006
1007 const fn unbounded() -> Self {
1008 Self {
1009 minimum: 0.0,
1010 maximum: f64::INFINITY,
1011 }
1012 }
1013
1014 #[cfg(test)]
1016 pub(crate) fn endpoint_bits(self) -> (u64, u64) {
1017 (self.minimum.to_bits(), self.maximum.to_bits())
1018 }
1019
1020 pub(crate) fn squared_l2_lower_bound(self, query: &[f32]) -> f64 {
1028 if !self.maximum.is_finite() || query.is_empty() {
1029 return 0.0;
1030 }
1031 let squared_norm = query
1032 .iter()
1033 .map(|value| {
1034 let value = f64::from(*value);
1035 value * value
1036 })
1037 .sum::<f64>();
1038 let relative_error = accumulation_relative_error(query.len());
1039 let query_minimum = next_down_f64((squared_norm / (1.0 + relative_error)).max(0.0).sqrt());
1040 let query_maximum = next_up_f64((squared_norm / (1.0 - relative_error)).max(0.0).sqrt());
1041 let gap = if query_maximum < self.minimum {
1042 self.minimum - query_maximum
1043 } else if query_minimum > self.maximum {
1044 query_minimum - self.maximum
1045 } else {
1046 0.0
1047 };
1048 next_down_f64(gap * gap * (1.0 - relative_error))
1049 }
1050
1051 pub(crate) fn squared_l2_upper_bound(self, query: &[f32]) -> f64 {
1054 if !self.maximum.is_finite() || query.is_empty() {
1055 return f64::INFINITY;
1056 }
1057 let squared_norm = query
1058 .iter()
1059 .map(|value| {
1060 let value = f64::from(*value);
1061 value * value
1062 })
1063 .sum::<f64>();
1064 let relative_error = accumulation_relative_error(query.len());
1065 let query_maximum = next_up_f64((squared_norm / (1.0 - relative_error)).max(0.0).sqrt());
1066 next_up_f64(
1067 (query_maximum + self.maximum)
1068 * (query_maximum + self.maximum)
1069 * (1.0 + relative_error),
1070 )
1071 }
1072}
1073
1074fn accumulation_relative_error(terms: usize) -> f64 {
1075 let scaled = terms as f64 * f64::EPSILON;
1076 (scaled / (1.0 - scaled)) * 2.0
1077}
1078
1079fn next_up_f32(value: f32) -> f32 {
1080 if value == f32::INFINITY {
1081 return value;
1082 }
1083 if value == -0.0 {
1084 return f32::from_bits(1);
1085 }
1086 if value >= 0.0 {
1087 f32::from_bits(value.to_bits().saturating_add(1))
1088 } else {
1089 f32::from_bits(value.to_bits().saturating_sub(1))
1090 }
1091}
1092
1093fn next_up_f64(value: f64) -> f64 {
1094 if value == f64::INFINITY {
1095 return value;
1096 }
1097 if value == -0.0 {
1098 return f64::from_bits(1);
1099 }
1100 if value >= 0.0 {
1101 f64::from_bits(value.to_bits().saturating_add(1))
1102 } else {
1103 f64::from_bits(value.to_bits().saturating_sub(1))
1104 }
1105}
1106
1107fn next_down_f64(value: f64) -> f64 {
1108 if value <= 0.0 {
1109 return 0.0;
1110 }
1111 f64::from_bits(value.to_bits().saturating_sub(1))
1112}
1113
1114#[derive(Debug)]
1116pub struct GraphSearchScratch {
1117 node_count: u32,
1118 max_degree: u8,
1119 visited: Vec<u8>,
1120 epoch: u8,
1121 pool: Vec<ScoredNode>,
1122 frontier: Vec<ScoredNode>,
1123 hop_candidates: Vec<CheckedNodeId>,
1124 rescore_row_ids: Vec<u32>,
1125 rescore_coarse_scores: Vec<f32>,
1126}
1127
1128impl GraphSearchScratch {
1129 pub fn new(node_count: u32, max_degree: u8) -> Result<Self, GraphSearchError> {
1131 Self::with_ef_capacity(node_count, max_degree, 0)
1132 }
1133
1134 pub(crate) fn with_ef_capacity(
1135 node_count: u32,
1136 max_degree: u8,
1137 ef: usize,
1138 ) -> Result<Self, GraphSearchError> {
1139 let node_count_usize = node_count as usize;
1140 let mut visited = Vec::new();
1141 visited
1142 .try_reserve_exact(node_count_usize)
1143 .map_err(|error| {
1144 GraphSearchError::Geometry(format!("visited allocation failed: {error}"))
1145 })?;
1146 visited.resize(node_count_usize, 0);
1147 let mut pool = Vec::new();
1148 pool.try_reserve_exact(ef).map_err(|error| {
1149 GraphSearchError::Geometry(format!("pool allocation failed: {error}"))
1150 })?;
1151 let mut frontier = Vec::new();
1152 frontier.try_reserve_exact(ef).map_err(|error| {
1153 GraphSearchError::Geometry(format!("frontier allocation failed: {error}"))
1154 })?;
1155 let mut hop_candidates = Vec::new();
1156 hop_candidates
1157 .try_reserve_exact(usize::from(max_degree))
1158 .map_err(|error| {
1159 GraphSearchError::Geometry(format!("hop-candidate allocation failed: {error}"))
1160 })?;
1161 let mut rescore_row_ids = Vec::new();
1162 rescore_row_ids.try_reserve_exact(ef).map_err(|error| {
1163 GraphSearchError::Geometry(format!("rescore row-id allocation failed: {error}"))
1164 })?;
1165 let mut rescore_coarse_scores = Vec::new();
1166 rescore_coarse_scores
1167 .try_reserve_exact(ef)
1168 .map_err(|error| {
1169 GraphSearchError::Geometry(format!("rescore score allocation failed: {error}"))
1170 })?;
1171 Ok(Self {
1172 node_count,
1173 max_degree,
1174 visited,
1175 epoch: 0,
1176 pool,
1177 frontier,
1178 hop_candidates,
1179 rescore_row_ids,
1180 rescore_coarse_scores,
1181 })
1182 }
1183
1184 pub fn allocation_bytes(
1186 node_count: u32,
1187 max_degree: u8,
1188 ef: usize,
1189 ) -> Result<usize, GraphSearchError> {
1190 let scored = ef
1191 .checked_mul(std::mem::size_of::<ScoredNode>())
1192 .and_then(|bytes| bytes.checked_mul(2))
1193 .ok_or_else(|| GraphSearchError::Geometry("scratch byte count overflow".to_owned()))?;
1194 let hops = usize::from(max_degree)
1195 .checked_mul(std::mem::size_of::<CheckedNodeId>())
1196 .ok_or_else(|| GraphSearchError::Geometry("scratch byte count overflow".to_owned()))?;
1197 let rescore = ef
1198 .checked_mul(std::mem::size_of::<u32>() + std::mem::size_of::<f32>())
1199 .ok_or_else(|| GraphSearchError::Geometry("scratch byte count overflow".to_owned()))?;
1200 (node_count as usize)
1201 .checked_add(scored)
1202 .and_then(|bytes| bytes.checked_add(hops))
1203 .and_then(|bytes| bytes.checked_add(rescore))
1204 .ok_or_else(|| GraphSearchError::Geometry("scratch byte count overflow".to_owned()))
1205 }
1206
1207 pub fn resident_bytes(&self) -> Result<usize, GraphSearchError> {
1209 let widths = [
1210 (self.visited.capacity(), std::mem::size_of::<u8>()),
1211 (self.pool.capacity(), std::mem::size_of::<ScoredNode>()),
1212 (self.frontier.capacity(), std::mem::size_of::<ScoredNode>()),
1213 (
1214 self.hop_candidates.capacity(),
1215 std::mem::size_of::<CheckedNodeId>(),
1216 ),
1217 (self.rescore_row_ids.capacity(), std::mem::size_of::<u32>()),
1218 (
1219 self.rescore_coarse_scores.capacity(),
1220 std::mem::size_of::<f32>(),
1221 ),
1222 ];
1223 widths.iter().try_fold(0_usize, |total, (capacity, width)| {
1224 capacity
1225 .checked_mul(*width)
1226 .and_then(|bytes| total.checked_add(bytes))
1227 .ok_or_else(|| GraphSearchError::Geometry("scratch byte count overflow".to_owned()))
1228 })
1229 }
1230
1231 pub(crate) fn supports(&self, node_count: u32, max_degree: u8, ef: usize) -> bool {
1232 self.node_count == node_count
1233 && self.max_degree == max_degree
1234 && self.pool.capacity() >= ef
1235 && self.frontier.capacity() >= ef
1236 && self.rescore_row_ids.capacity() >= ef
1237 && self.rescore_coarse_scores.capacity() >= ef
1238 }
1239}
1240
1241pub struct GraphSearcher<'a> {
1243 graph: GraphNodeBlocks<'a>,
1244 rescore: &'a [f32],
1245 rescore_validator: Option<&'a dyn RescoreValidator>,
1246 entries: [CheckedNodeId; 4],
1247 scratch: &'a mut GraphSearchScratch,
1248 #[cfg(any(test, feature = "test-seams"))]
1249 vector_fault: Option<(
1250 &'a crate::scan::vector_fault::VectorFaultController,
1251 crate::scan::vector_fault::VectorRowSource,
1252 crate::scan::vector_fault::VectorSearchTier,
1253 )>,
1254 #[cfg(any(test, feature = "test-seams"))]
1255 hop_cancellation: Option<TestHopCancellation>,
1256}
1257
1258pub(crate) trait RescoreValidator {
1259 fn validate_rows(&self, rows: &[u32]) -> Result<(), String>;
1260}
1261
1262#[cfg(any(test, feature = "test-seams"))]
1263#[derive(Debug)]
1264struct TestHopCancellation {
1265 after_hops: usize,
1266 token: crate::lifecycle::CancelToken,
1267 observed_hops: std::rc::Rc<std::cell::Cell<usize>>,
1268}
1269
1270impl<'a> GraphSearcher<'a> {
1271 pub fn new(
1273 graph: GraphNodeBlocks<'a>,
1274 rescore: &'a [f32],
1275 scratch: &'a mut GraphSearchScratch,
1276 ) -> Result<Self, GraphSearchError> {
1277 let entries = Self::discover_entry_row_ids(graph)?;
1278 Self::with_entry_row_ids(graph, rescore, entries, scratch)
1279 }
1280
1281 pub fn discover_entry_row_ids(
1283 graph: GraphNodeBlocks<'_>,
1284 ) -> Result<[u32; 4], GraphSearchError> {
1285 let mut entries = Vec::with_capacity(4);
1286 for row_id in 0..graph.node_count() {
1287 if graph.block(row_id)?.flags() & 1 != 0 {
1288 entries.push(row_id);
1289 }
1290 }
1291 let entry_count = entries.len();
1292 let expected = usize::try_from(graph.node_count())
1293 .unwrap_or(usize::MAX)
1294 .min(4);
1295 if entry_count != expected {
1296 return Err(GraphSearchError::Geometry(format!(
1297 "graph contains {entry_count} persisted entry seeds, expected {expected}"
1298 )));
1299 }
1300 let Some(first) = entries.first().copied() else {
1301 return Err(GraphSearchError::Geometry(
1302 "graph contains no persisted entry seed".to_owned(),
1303 ));
1304 };
1305 entries.resize(4, first);
1306 entries.try_into().map_err(|_: Vec<u32>| {
1307 GraphSearchError::Geometry(format!(
1308 "graph contains {entry_count} persisted entry seeds, expected at most four"
1309 ))
1310 })
1311 }
1312
1313 pub fn with_entry_row_ids(
1315 graph: GraphNodeBlocks<'a>,
1316 rescore: &'a [f32],
1317 entry_row_ids: [u32; 4],
1318 scratch: &'a mut GraphSearchScratch,
1319 ) -> Result<Self, GraphSearchError> {
1320 let dimensions = graph.layout().dims() as usize;
1321 let expected_rescore = (graph.node_count() as usize)
1322 .checked_mul(dimensions)
1323 .ok_or_else(|| GraphSearchError::Geometry("rescore length overflow".to_owned()))?;
1324 if rescore.len() != expected_rescore {
1325 return Err(GraphSearchError::Geometry(format!(
1326 "rescore has {} values, expected {expected_rescore}",
1327 rescore.len()
1328 )));
1329 }
1330 if scratch.node_count != graph.node_count()
1331 || scratch.max_degree != graph.layout().max_degree()
1332 {
1333 return Err(GraphSearchError::Geometry(format!(
1334 "scratch geometry nodes/degree {}/{}, graph {}/{}",
1335 scratch.node_count,
1336 scratch.max_degree,
1337 graph.node_count(),
1338 graph.layout().max_degree()
1339 )));
1340 }
1341 let mut entries = [None; 4];
1342 for (position, row_id) in entry_row_ids.into_iter().enumerate() {
1343 let block = graph.block(row_id)?;
1344 if block.flags() & 1 == 0 {
1345 return Err(GraphSearchError::Geometry(format!(
1346 "cached entry row {row_id} does not carry the persisted entry flag"
1347 )));
1348 }
1349 let destination = entries.get_mut(position).ok_or_else(|| {
1350 GraphSearchError::Geometry("cached entry position overflow".to_owned())
1351 })?;
1352 *destination = Some(graph.checked_node_id(row_id)?);
1353 }
1354 let entries = entries.map(|entry| {
1355 entry.ok_or_else(|| GraphSearchError::Geometry("cached entry is absent".to_owned()))
1356 });
1357 let [first, second, third, fourth] = entries;
1358 Ok(Self {
1359 graph,
1360 rescore,
1361 rescore_validator: None,
1362 entries: [first?, second?, third?, fourth?],
1363 scratch,
1364 #[cfg(any(test, feature = "test-seams"))]
1365 vector_fault: None,
1366 #[cfg(any(test, feature = "test-seams"))]
1367 hop_cancellation: None,
1368 })
1369 }
1370
1371 pub(crate) fn with_rescore_validator(mut self, validator: &'a dyn RescoreValidator) -> Self {
1372 self.rescore_validator = Some(validator);
1373 self
1374 }
1375
1376 #[cfg(any(test, feature = "test-seams"))]
1377 pub(crate) fn with_vector_fault_controller(
1378 mut self,
1379 controller: &'a crate::scan::vector_fault::VectorFaultController,
1380 source: crate::scan::vector_fault::VectorRowSource,
1381 tier: crate::scan::vector_fault::VectorSearchTier,
1382 ) -> Self {
1383 self.vector_fault = Some((controller, source, tier));
1384 self
1385 }
1386
1387 #[cfg(any(test, feature = "test-seams"))]
1388 pub(crate) fn cancel_after_hops(
1389 &mut self,
1390 after_hops: usize,
1391 token: crate::lifecycle::CancelToken,
1392 ) -> std::rc::Rc<std::cell::Cell<usize>> {
1393 let observed_hops = std::rc::Rc::new(std::cell::Cell::new(0));
1394 self.hop_cancellation = Some(TestHopCancellation {
1395 after_hops,
1396 token,
1397 observed_hops: std::rc::Rc::clone(&observed_hops),
1398 });
1399 observed_hops
1400 }
1401
1402 #[cfg(any(test, feature = "test-seams"))]
1403 fn observe_hop_for_cancellation(&self, hops: usize) {
1404 let Some(source) = &self.hop_cancellation else {
1405 return;
1406 };
1407 source.observed_hops.set(hops);
1408 if hops == source.after_hops {
1409 source.token.cancel();
1410 }
1411 }
1412
1413 pub fn search(
1415 &mut self,
1416 request: GraphSearchRequest<'_>,
1417 cancellation: Option<&QueryCancellation<'_>>,
1418 ) -> Result<GraphSearchResult, GraphSearchError> {
1419 match self.search_inner(request, None, None, cancellation)? {
1420 SearchInnerOutcome::Traversed(result) => Ok(result),
1421 SearchInnerOutcome::VisitedBudgetExceeded { .. } => Err(GraphSearchError::Geometry(
1422 "unfiltered traversal produced a filtered-budget disposition".to_owned(),
1423 )),
1424 }
1425 }
1426
1427 pub fn search_filtered(
1434 &mut self,
1435 request: GraphSearchRequest<'_>,
1436 allow_list: &DocBitmap,
1437 filtered_visited_budget: usize,
1438 cancellation: Option<&QueryCancellation<'_>>,
1439 ) -> Result<FilteredGraphSearchOutcome, GraphSearchError> {
1440 match self.search_inner(
1441 request,
1442 Some(allow_list),
1443 Some(filtered_visited_budget),
1444 cancellation,
1445 )? {
1446 SearchInnerOutcome::Traversed(result) => {
1447 Ok(FilteredGraphSearchOutcome::Traversed(result))
1448 }
1449 SearchInnerOutcome::VisitedBudgetExceeded {
1450 visited,
1451 budget,
1452 counters,
1453 } => Ok(FilteredGraphSearchOutcome::VisitedBudgetExceeded {
1454 visited,
1455 budget,
1456 counters,
1457 }),
1458 }
1459 }
1460
1461 fn search_inner(
1462 &mut self,
1463 request: GraphSearchRequest<'_>,
1464 allow_list: Option<&DocBitmap>,
1465 filtered_visited_budget: Option<usize>,
1466 cancellation: Option<&QueryCancellation<'_>>,
1467 ) -> Result<SearchInnerOutcome, GraphSearchError> {
1468 let entry = self.begin_traversal(request, cancellation)?;
1469 let prepared = PreparedGraphQuery::new(
1470 request.query,
1471 self.graph.layout().padded_dims() as usize,
1472 request.seed,
1473 )?;
1474 match self.traverse_prepared(
1475 request,
1476 allow_list,
1477 filtered_visited_budget,
1478 cancellation,
1479 entry,
1480 &prepared,
1481 )? {
1482 TraversalOutcome::Retained {
1483 counters,
1484 candidate_sequence,
1485 } => self.finalize_traversal(
1486 request,
1487 allow_list,
1488 cancellation,
1489 counters,
1490 candidate_sequence,
1491 ),
1492 TraversalOutcome::Budget(outcome) => Ok(outcome),
1493 }
1494 }
1495
1496 pub(crate) fn search_reusing(
1497 &mut self,
1498 request: GraphSearchRequest<'_>,
1499 cancellation: Option<&QueryCancellation<'_>>,
1500 context: &mut crate::lifecycle::prepared::PreparedVectorQuery<'_>,
1501 segment: &crate::segment::reader::SegmentReader,
1502 source: crate::ingest::RowSource,
1503 ) -> Result<GraphSearchResult, crate::lifecycle::QueryError> {
1504 use crate::lifecycle::map_graph_error;
1505 let entry = self
1506 .begin_traversal(request, cancellation)
1507 .map_err(map_graph_error)?;
1508 let prepared = context.graph(
1509 request.query,
1510 self.graph.layout().padded_dims() as usize,
1511 request.seed,
1512 )?;
1513 let traversal = self
1514 .traverse_prepared(request, None, None, cancellation, entry, prepared)
1515 .map_err(map_graph_error)?;
1516 let TraversalOutcome::Retained {
1517 counters,
1518 candidate_sequence,
1519 } = traversal
1520 else {
1521 return Err(map_graph_error(GraphSearchError::Geometry(
1522 "unfiltered traversal produced a filtered-budget disposition".to_owned(),
1523 )));
1524 };
1525 let rows = self.rescore;
1528 let dimensions = self.graph.layout().dims() as usize;
1529 let pool = self
1530 .prepare_rescore(request, cancellation, counters)
1531 .map_err(map_graph_error)?;
1532 let rescored = context.rescore_graph(
1533 request.query,
1534 rows,
1535 dimensions,
1536 pool,
1537 request.k,
1538 segment,
1539 source,
1540 )?;
1541 match self
1542 .finish_rescore(
1543 request,
1544 None,
1545 cancellation,
1546 counters,
1547 candidate_sequence,
1548 Ok(rescored),
1549 )
1550 .map_err(map_graph_error)?
1551 {
1552 SearchInnerOutcome::Traversed(result) => Ok(result),
1553 SearchInnerOutcome::VisitedBudgetExceeded { .. } => {
1554 Err(map_graph_error(GraphSearchError::Geometry(
1555 "rescore produced a filtered-budget disposition".to_owned(),
1556 )))
1557 }
1558 }
1559 }
1560
1561 fn begin_traversal(
1562 &self,
1563 request: GraphSearchRequest<'_>,
1564 cancellation: Option<&QueryCancellation<'_>>,
1565 ) -> Result<(usize, QueryQosClass, i32), GraphSearchError> {
1566 let ef = self.validate_request(request)?;
1567 check_cancellation(cancellation)?;
1568 #[cfg(any(test, feature = "test-seams"))]
1569 GRAPH_QUERY_TEST_OBSERVATIONS.with(|observations| {
1570 if let Some(observations) = observations.borrow_mut().as_mut() {
1571 observations.traversals += 1;
1572 }
1573 });
1574 let (qos_class, qos_relative_priority) = observed_qos();
1575 Ok((ef, qos_class, qos_relative_priority))
1576 }
1577
1578 #[allow(clippy::too_many_arguments)]
1579 fn traverse_prepared(
1580 &mut self,
1581 request: GraphSearchRequest<'_>,
1582 allow_list: Option<&DocBitmap>,
1583 filtered_visited_budget: Option<usize>,
1584 cancellation: Option<&QueryCancellation<'_>>,
1585 entry: (usize, QueryQosClass, i32),
1586 prepared: &PreparedGraphQuery,
1587 ) -> Result<TraversalOutcome, GraphSearchError> {
1588 let (ef, qos_class, qos_relative_priority) = entry;
1589 let query_norm = prepared.norm;
1590 let prepared = prepared.bit4();
1591 let visited_cap = ef
1592 .checked_mul(usize::from(self.graph.layout().max_degree()))
1593 .and_then(|value| value.checked_mul(4))
1594 .ok_or_else(|| GraphSearchError::Geometry("visited cap overflow".to_owned()))?;
1595 let visited_epoch_cleared = self.next_epoch();
1596 let mut counters = GraphSearchCounters {
1597 effective_ef: ef,
1598 candidates_rescored: 0,
1599 candidate_window_peak: 0,
1600 hops: 0,
1601 candidates_scored: 0,
1602 pushes: 0,
1603 visited: 0,
1604 visited_epoch_cleared,
1605 dims_touched: 0,
1606 bytes_read: 0,
1607 qos_class,
1608 qos_relative_priority,
1609 core_class: request.observed_core_class,
1610 };
1611 self.scratch.pool.clear();
1612 self.scratch.frontier.clear();
1613 reserve_exact_scratch(&mut self.scratch.pool, ef, "pool")?;
1614 reserve_exact_scratch(&mut self.scratch.frontier, ef, "frontier")?;
1615 let mut candidate_sequence = request.trace_candidates.then(|| Vec::with_capacity(ef));
1616 let entries = self.entries;
1617 if request.prefetch.is_enabled() {
1618 for entry in &entries {
1619 self.graph.prefetch_line0(*entry);
1620 }
1621 }
1622 let mut seed_group = [None; 4];
1623 let mut seed_count = 0_usize;
1624 for entry in entries {
1625 if mark_visited(&mut self.scratch.visited, self.scratch.epoch, entry)? {
1626 counters.visited += 1;
1627 ensure_visited_cap(counters.visited, visited_cap)?;
1628 if filtered_visited_budget.is_some_and(|budget| counters.visited > budget) {
1629 return self.filtered_budget_outcome(counters, filtered_visited_budget);
1630 }
1631 let Some(slot) = seed_group.get_mut(seed_count) else {
1632 return Err(GraphSearchError::Geometry(
1633 "entry seed count exceeds fixed seed group".to_owned(),
1634 ));
1635 };
1636 *slot = Some(entry);
1637 seed_count += 1;
1638 }
1639 }
1640 self.score_and_push_group(
1641 prepared,
1642 query_norm,
1643 &seed_group,
1644 seed_count,
1645 ef,
1646 allow_list,
1647 request.prefetch,
1648 &mut counters,
1649 &mut candidate_sequence,
1650 )?;
1651 loop {
1652 check_cancellation(cancellation)?;
1653 let Some(candidate) = min_heap_pop(&mut self.scratch.frontier)? else {
1654 break;
1655 };
1656 if self.scratch.pool.len() >= ef
1657 && self
1658 .scratch
1659 .pool
1660 .first()
1661 .is_some_and(|worst| candidate.distance > worst.distance)
1662 {
1663 break;
1664 }
1665 counters.hops += 1;
1666 #[cfg(any(test, feature = "test-seams"))]
1667 self.observe_hop_for_cancellation(counters.hops);
1668 let (degree, neighbors) = self.graph.adjacency_checked(candidate.row_id)?;
1669 let degree = usize::from(degree);
1670 self.scratch.hop_candidates.clear();
1671 for raw_neighbor in neighbors.take(degree) {
1672 let neighbor = self.graph.checked_node_id(raw_neighbor)?;
1673 if !mark_visited(&mut self.scratch.visited, self.scratch.epoch, neighbor)? {
1674 continue;
1675 }
1676 counters.visited += 1;
1677 ensure_visited_cap(counters.visited, visited_cap)?;
1678 if filtered_visited_budget.is_some_and(|budget| counters.visited > budget) {
1679 return self.filtered_budget_outcome(counters, filtered_visited_budget);
1680 }
1681 if request.prefetch.is_enabled() {
1682 self.graph.prefetch_line0(neighbor);
1683 }
1684 self.scratch.hop_candidates.push(neighbor);
1685 }
1686 let mut group_start = 0_usize;
1687 let hop_candidate_count = self.scratch.hop_candidates.len();
1688 while group_start < hop_candidate_count {
1689 check_cancellation(cancellation)?;
1690 let group_count = (hop_candidate_count - group_start).min(4);
1691 let mut group = [None; 4];
1692 for lane in 0..group_count {
1693 let source = self.scratch.hop_candidates.get(group_start + lane).copied();
1694 let Some(destination) = group.get_mut(lane) else {
1695 return Err(GraphSearchError::Geometry(
1696 "candidate gather lane is unavailable".to_owned(),
1697 ));
1698 };
1699 *destination = source;
1700 }
1701 self.score_and_push_group(
1702 prepared,
1703 query_norm,
1704 &group,
1705 group_count,
1706 ef,
1707 allow_list,
1708 request.prefetch,
1709 &mut counters,
1710 &mut candidate_sequence,
1711 )?;
1712 group_start += group_count;
1713 }
1714 }
1715 Ok(TraversalOutcome::Retained {
1716 counters,
1717 candidate_sequence,
1718 })
1719 }
1720
1721 fn finalize_traversal(
1722 &mut self,
1723 request: GraphSearchRequest<'_>,
1724 allow_list: Option<&DocBitmap>,
1725 cancellation: Option<&QueryCancellation<'_>>,
1726 counters: GraphSearchCounters,
1727 candidate_sequence: Option<Vec<u32>>,
1728 ) -> Result<SearchInnerOutcome, GraphSearchError> {
1729 let rows = self.rescore;
1730 let dimensions = self.graph.layout().dims() as usize;
1731 let pool = self.prepare_rescore(request, cancellation, counters)?;
1732 let rescored = rescore_top_k(request.query, rows, dimensions, pool, request.k);
1733 self.finish_rescore(
1734 request,
1735 allow_list,
1736 cancellation,
1737 counters,
1738 candidate_sequence,
1739 rescored,
1740 )
1741 }
1742
1743 fn prepare_rescore(
1744 &mut self,
1745 request: GraphSearchRequest<'_>,
1746 cancellation: Option<&QueryCancellation<'_>>,
1747 counters: GraphSearchCounters,
1748 ) -> Result<RescorePool<'_>, GraphSearchError> {
1749 check_cancellation(cancellation)?;
1750 self.scratch.rescore_row_ids.clear();
1751 self.scratch.rescore_coarse_scores.clear();
1752 let retained = self.scratch.pool.len();
1753 reserve_exact_scratch(
1754 &mut self.scratch.rescore_row_ids,
1755 retained,
1756 "rescore row-id",
1757 )?;
1758 reserve_exact_scratch(
1759 &mut self.scratch.rescore_coarse_scores,
1760 retained,
1761 "rescore score",
1762 )?;
1763 for candidate in &self.scratch.pool {
1764 self.scratch.rescore_row_ids.push(candidate.row_id.raw());
1765 let score = -(candidate.distance as f32);
1766 if !score.is_finite() {
1767 return Err(GraphSearchError::Geometry(format!(
1768 "coarse distance for row {} is not representable as f32",
1769 candidate.row_id.raw()
1770 )));
1771 }
1772 self.scratch.rescore_coarse_scores.push(score);
1773 }
1774 if let Some(validator) = self.rescore_validator {
1775 validator
1776 .validate_rows(&self.scratch.rescore_row_ids)
1777 .map_err(GraphSearchError::ExactRescoreUnavailable)?;
1778 #[cfg(any(test, feature = "test-seams"))]
1779 GRAPH_QUERY_TEST_OBSERVATIONS.with(|observations| {
1780 if let Some(observations) = observations.borrow_mut().as_mut() {
1781 observations
1782 .validated_pools
1783 .push(self.scratch.rescore_row_ids.clone());
1784 }
1785 });
1786 }
1787 let coarse_bytes_per_row = self
1788 .graph
1789 .layout()
1790 .code_bytes()
1791 .checked_add(12)
1792 .ok_or_else(|| GraphSearchError::Geometry("coarse row bytes overflow".to_owned()))?;
1793 let pool = RescorePool::retained(
1794 &self.scratch.rescore_row_ids,
1795 &self.scratch.rescore_coarse_scores,
1796 RescoreMetric::SquaredL2,
1797 counters.candidates_scored,
1798 coarse_bytes_per_row,
1799 )
1800 .with_prefetch(request.prefetch.is_enabled());
1801 Ok(pool)
1802 }
1803
1804 #[allow(clippy::too_many_arguments)]
1805 fn finish_rescore(
1806 &self,
1807 _request: GraphSearchRequest<'_>,
1808 allow_list: Option<&DocBitmap>,
1809 cancellation: Option<&QueryCancellation<'_>>,
1810 mut counters: GraphSearchCounters,
1811 candidate_sequence: Option<Vec<u32>>,
1812 rescored: Result<crate::quant::RescoreResult, RescoreError>,
1813 ) -> Result<SearchInnerOutcome, GraphSearchError> {
1814 let coarse_bytes_per_row = self
1815 .graph
1816 .layout()
1817 .code_bytes()
1818 .checked_add(12)
1819 .ok_or_else(|| GraphSearchError::Geometry("coarse row bytes overflow".to_owned()))?;
1820 let rescored = match rescored {
1821 Ok(rescored) => rescored,
1822 Err(RescoreError::InsufficientCandidates { k, candidates }) if allow_list.is_some() => {
1823 counters.candidates_rescored = 0;
1824 counters.dims_touched = u64::try_from(counters.candidates_scored)
1825 .ok()
1826 .and_then(|rows| rows.checked_mul(u64::from(self.graph.layout().dims())))
1827 .ok_or_else(|| {
1828 GraphSearchError::Geometry("dimension counter overflow".to_owned())
1829 })?;
1830 counters.bytes_read = counters
1831 .candidates_scored
1832 .checked_mul(coarse_bytes_per_row)
1833 .and_then(|bytes| u64::try_from(bytes).ok())
1834 .ok_or_else(|| {
1835 GraphSearchError::Geometry("byte counter exceeds u64".to_owned())
1836 })?;
1837 return Err(GraphSearchError::FilteredCandidateShortfall {
1838 k,
1839 candidates,
1840 counters,
1841 });
1842 }
1843 Err(error) => return Err(GraphSearchError::Rescore(error)),
1844 };
1845 check_cancellation(cancellation)?;
1846 counters.candidates_rescored = rescored.candidates_rescored;
1847 let touched_rows = counters
1848 .candidates_scored
1849 .checked_add(counters.candidates_rescored)
1850 .ok_or_else(|| GraphSearchError::Geometry("dimension counter overflow".to_owned()))?;
1851 counters.dims_touched = u64::try_from(touched_rows)
1852 .ok()
1853 .and_then(|rows| rows.checked_mul(u64::from(self.graph.layout().dims())))
1854 .ok_or_else(|| GraphSearchError::Geometry("dimension counter overflow".to_owned()))?;
1855 counters.bytes_read = u64::try_from(rescored.bytes.total())
1856 .map_err(|_| GraphSearchError::Geometry("byte counter exceeds u64".to_owned()))?;
1857 let candidates = rescored
1858 .hits
1859 .into_iter()
1860 .map(|hit| {
1861 let row_id = u32::try_from(hit.row_index).map_err(|_| {
1862 GraphSearchError::Geometry("rescored row id exceeds u32".to_owned())
1863 })?;
1864 Ok(GraphSearchCandidate {
1865 row_id,
1866 distance: (-hit.score).max(0.0),
1867 })
1868 })
1869 .collect::<Result<Vec<_>, GraphSearchError>>()?;
1870 Ok(SearchInnerOutcome::Traversed(GraphSearchResult {
1871 candidates,
1872 counters,
1873 candidate_sequence,
1874 }))
1875 }
1876
1877 fn filtered_budget_outcome(
1878 &self,
1879 mut counters: GraphSearchCounters,
1880 filtered_visited_budget: Option<usize>,
1881 ) -> Result<TraversalOutcome, GraphSearchError> {
1882 let budget = filtered_visited_budget.ok_or_else(|| {
1883 GraphSearchError::Geometry(
1884 "filtered budget disposition has no configured budget".to_owned(),
1885 )
1886 })?;
1887 counters.dims_touched = u64::try_from(counters.candidates_scored)
1888 .ok()
1889 .and_then(|rows| rows.checked_mul(u64::from(self.graph.layout().dims())))
1890 .ok_or_else(|| GraphSearchError::Geometry("dimension counter overflow".to_owned()))?;
1891 let coarse_bytes_per_row = self
1892 .graph
1893 .layout()
1894 .code_bytes()
1895 .checked_add(12)
1896 .ok_or_else(|| GraphSearchError::Geometry("coarse row bytes overflow".to_owned()))?;
1897 counters.bytes_read = counters
1898 .candidates_scored
1899 .checked_mul(coarse_bytes_per_row)
1900 .and_then(|bytes| u64::try_from(bytes).ok())
1901 .ok_or_else(|| GraphSearchError::Geometry("byte counter exceeds u64".to_owned()))?;
1902 Ok(TraversalOutcome::Budget(
1903 SearchInnerOutcome::VisitedBudgetExceeded {
1904 visited: counters.visited,
1905 budget,
1906 counters,
1907 },
1908 ))
1909 }
1910
1911 fn validate_request(&self, request: GraphSearchRequest<'_>) -> Result<usize, GraphSearchError> {
1912 let dimensions = self.graph.layout().dims() as usize;
1913 if request.query.len() != dimensions {
1914 return Err(GraphSearchError::Geometry(format!(
1915 "query has {} dimensions, expected {dimensions}",
1916 request.query.len()
1917 )));
1918 }
1919 request
1920 .effective_ef(self.graph.node_count() as usize)
1921 .map_err(Into::into)
1922 }
1923
1924 fn next_epoch(&mut self) -> bool {
1925 self.scratch.epoch = self.scratch.epoch.wrapping_add(1);
1926 if self.scratch.epoch == 0 {
1927 self.scratch.visited.fill(0);
1928 self.scratch.epoch = 1;
1929 return true;
1930 }
1931 false
1932 }
1933
1934 #[allow(clippy::too_many_arguments)]
1935 fn score_and_push_group(
1936 &mut self,
1937 query: &Bit4Query,
1938 query_norm: f64,
1939 row_ids: &[Option<CheckedNodeId>; 4],
1940 row_count: usize,
1941 ef: usize,
1942 allow_list: Option<&DocBitmap>,
1943 prefetch: TraversalPrefetch,
1944 counters: &mut GraphSearchCounters,
1945 candidate_sequence: &mut Option<Vec<u32>>,
1946 ) -> Result<(), GraphSearchError> {
1947 let scored = score_group(&self.graph, query, query_norm, row_ids, row_count)?;
1948 counters.candidates_scored = counters
1949 .candidates_scored
1950 .checked_add(row_count)
1951 .ok_or_else(|| GraphSearchError::Geometry("candidate counter overflow".to_owned()))?;
1952 for candidate in scored.into_iter().flatten() {
1953 let allowed = allow_list.is_none_or(|mask| mask.contains(candidate.row_id.raw()));
1954 #[cfg(any(test, feature = "test-seams"))]
1955 if allowed && let Some((controller, source, tier)) = self.vector_fault {
1956 let local_row = usize::try_from(candidate.row_id.raw()).map_err(|_| {
1957 GraphSearchError::Geometry("graph row id exceeds usize".to_owned())
1958 })?;
1959 if controller.after_eligible_row(source, tier, local_row) {
1960 return Err(GraphSearchError::Cancelled { partial: false });
1961 }
1962 }
1963 if let Some(sequence) = candidate_sequence.as_mut() {
1964 sequence.push(candidate.row_id.raw());
1965 }
1966 let mut pushed = false;
1967 if allowed && max_heap_would_keep(&self.scratch.pool, candidate, ef) {
1968 max_heap_insert_bounded(&mut self.scratch.pool, candidate, ef)?;
1969 min_heap_push(&mut self.scratch.frontier, candidate)?;
1970 counters.pushes = counters.pushes.checked_add(1).ok_or_else(|| {
1971 GraphSearchError::Geometry("push counter overflow".to_owned())
1972 })?;
1973 pushed = true;
1974 } else if !allowed
1975 && (self.scratch.pool.len() < ef
1976 || self
1977 .scratch
1978 .pool
1979 .first()
1980 .is_some_and(|worst| candidate.distance <= worst.distance))
1981 {
1982 min_heap_push(&mut self.scratch.frontier, candidate)?;
1983 counters.pushes = counters.pushes.checked_add(1).ok_or_else(|| {
1984 GraphSearchError::Geometry("push counter overflow".to_owned())
1985 })?;
1986 pushed = true;
1987 }
1988 counters.candidate_window_peak = counters
1989 .candidate_window_peak
1990 .max(self.scratch.pool.len())
1991 .max(self.scratch.frontier.len());
1992 if pushed
1993 && prefetch.is_enabled()
1994 && let Some(head) = self.scratch.frontier.first()
1995 {
1996 self.graph.prefetch_head_block(head.row_id);
1997 }
1998 }
1999 Ok(())
2000 }
2001}
2002
2003fn reserve_exact_scratch<T>(
2004 buffer: &mut Vec<T>,
2005 needed: usize,
2006 what: &'static str,
2007) -> Result<(), GraphSearchError> {
2008 if buffer.capacity() < needed {
2009 buffer.try_reserve_exact(needed).map_err(|error| {
2010 GraphSearchError::Geometry(format!("{what} allocation failed: {error}"))
2011 })?;
2012 }
2013 Ok(())
2014}
2015
2016fn mark_visited(
2017 visited: &mut [u8],
2018 epoch: u8,
2019 row_id: CheckedNodeId,
2020) -> Result<bool, GraphSearchError> {
2021 let marker = visited.get_mut(row_id.raw() as usize).ok_or_else(|| {
2022 GraphSearchError::Geometry(format!("row id {} is unavailable", row_id.raw()))
2023 })?;
2024 if *marker == epoch {
2025 return Ok(false);
2026 }
2027 *marker = epoch;
2028 Ok(true)
2029}
2030
2031fn ensure_visited_cap(visited: usize, cap: usize) -> Result<(), GraphSearchError> {
2032 if visited > cap {
2033 return Err(GraphSearchError::VisitedCapExceeded { visited, cap });
2034 }
2035 Ok(())
2036}
2037
2038fn check_cancellation(
2039 cancellation: Option<&QueryCancellation<'_>>,
2040) -> Result<(), GraphSearchError> {
2041 let Some(cancellation) = cancellation else {
2042 return Ok(());
2043 };
2044 match cancellation.check_graph() {
2045 Ok(()) => Ok(()),
2046 Err(ScanError::Cancelled { partial }) => Err(GraphSearchError::Cancelled { partial }),
2047 Err(ScanError::Timeout { partial }) => Err(GraphSearchError::Timeout { partial }),
2048 Err(ScanError::ReadCancelled { partial }) => {
2049 Err(GraphSearchError::ReadCancelled { partial })
2050 }
2051 Err(error) => Err(GraphSearchError::Scan(error)),
2052 }
2053}
2054
2055fn scored_best_first(left: &ScoredNode, right: &ScoredNode) -> std::cmp::Ordering {
2056 left.distance
2057 .total_cmp(&right.distance)
2058 .then_with(|| left.row_id.raw().cmp(&right.row_id.raw()))
2059}
2060
2061fn score_group(
2062 graph: &GraphNodeBlocks<'_>,
2063 query: &Bit4Query,
2064 query_norm: f64,
2065 row_ids: &[Option<CheckedNodeId>; 4],
2066 row_count: usize,
2067) -> Result<[Option<ScoredNode>; 4], GraphSearchError> {
2068 let first_id = row_ids
2069 .first()
2070 .copied()
2071 .flatten()
2072 .ok_or_else(|| GraphSearchError::Geometry("empty gather group".to_owned()))?;
2073 let (first_row, first_factors) = graph.score_row_checked(first_id)?;
2074 let mut rows = [first_row; 4];
2075 let mut factors = [first_factors; 4];
2076 for lane in 1..row_count {
2077 let row_id =
2078 row_ids.get(lane).copied().flatten().ok_or_else(|| {
2079 GraphSearchError::Geometry(format!("gather lane {lane} is empty"))
2080 })?;
2081 let (row, factors_for_row) = graph.score_row_checked(row_id)?;
2082 let row_slot = rows
2083 .get_mut(lane)
2084 .ok_or_else(|| GraphSearchError::Geometry("gather row lane overflow".to_owned()))?;
2085 *row_slot = row;
2086 let factor_slot = factors
2087 .get_mut(lane)
2088 .ok_or_else(|| GraphSearchError::Geometry("gather factor lane overflow".to_owned()))?;
2089 *factor_slot = factors_for_row;
2090 }
2091 let gathered = Bit4Rows4::from_rows(rows, graph.layout().code_bytes())?;
2092 let (query_codes, query_sum, query_scale_half) = query.kernel_parts();
2093 let mut dots = [0.0_f32; 4];
2094 score_bit4_ptrs(
2095 query_codes,
2096 query_sum,
2097 query_scale_half,
2098 &gathered,
2099 &factors,
2100 &mut dots,
2101 )?;
2102 let mut scored = [None; 4];
2103 for lane in 0..row_count {
2104 let row_id =
2105 row_ids.get(lane).copied().flatten().ok_or_else(|| {
2106 GraphSearchError::Geometry(format!("scored lane {lane} is empty"))
2107 })?;
2108 let dot = f64::from(
2109 *dots
2110 .get(lane)
2111 .ok_or_else(|| GraphSearchError::Geometry("dot lane overflow".to_owned()))?,
2112 );
2113 let norm = factors.get(lane).copied().unwrap_or(first_factors).norm();
2114 let destination = scored
2115 .get_mut(lane)
2116 .ok_or_else(|| GraphSearchError::Geometry("score lane overflow".to_owned()))?;
2117 *destination = Some(ScoredNode {
2118 row_id,
2119 distance: (query_norm + norm * norm - 2.0 * dot).max(0.0),
2120 });
2121 }
2122 Ok(scored)
2123}
2124
2125fn max_heap_would_keep(pool: &[ScoredNode], node: ScoredNode, limit: usize) -> bool {
2126 pool.len() < limit
2127 || pool
2128 .first()
2129 .is_some_and(|worst| scored_best_first(&node, worst).is_lt())
2130}
2131
2132fn max_heap_insert_bounded(
2133 pool: &mut Vec<ScoredNode>,
2134 node: ScoredNode,
2135 limit: usize,
2136) -> Result<(), GraphSearchError> {
2137 if limit == 0 {
2138 return Err(GraphSearchError::Geometry(
2139 "bounded pool capacity is zero".to_owned(),
2140 ));
2141 }
2142 if pool.len() < limit {
2143 pool.push(node);
2144 let child = pool.len().saturating_sub(1);
2145 max_heap_sift_up(pool, child)?;
2146 return Ok(());
2147 }
2148 let root = pool
2149 .first_mut()
2150 .ok_or_else(|| GraphSearchError::Geometry("full pool has no root".to_owned()))?;
2151 *root = node;
2152 max_heap_sift_down(pool, 0)
2153}
2154
2155fn min_heap_push(frontier: &mut Vec<ScoredNode>, node: ScoredNode) -> Result<(), GraphSearchError> {
2156 frontier.push(node);
2157 let child = frontier.len().saturating_sub(1);
2158 min_heap_sift_up(frontier, child)
2159}
2160
2161fn min_heap_pop(frontier: &mut Vec<ScoredNode>) -> Result<Option<ScoredNode>, GraphSearchError> {
2162 let Some(last) = frontier.pop() else {
2163 return Ok(None);
2164 };
2165 let Some(root) = frontier.first_mut() else {
2166 return Ok(Some(last));
2167 };
2168 let result = *root;
2169 *root = last;
2170 min_heap_sift_down(frontier, 0)?;
2171 Ok(Some(result))
2172}
2173
2174fn min_heap_sift_up(heap: &mut [ScoredNode], mut child: usize) -> Result<(), GraphSearchError> {
2175 let start = child;
2176 let moving = heap_node(heap, child)?;
2177 while child != 0 {
2178 let parent = child.saturating_sub(1) / 2;
2179 let parent_node = heap_node(heap, parent)?;
2180 if !scored_best_first(&moving, &parent_node).is_lt() {
2181 break;
2182 }
2183 heap_store(heap, child, parent_node)?;
2184 child = parent;
2185 }
2186 if child != start {
2187 heap_store(heap, child, moving)?;
2188 }
2189 Ok(())
2190}
2191
2192fn max_heap_sift_up(heap: &mut [ScoredNode], mut child: usize) -> Result<(), GraphSearchError> {
2193 let start = child;
2194 let moving = heap_node(heap, child)?;
2195 while child != 0 {
2196 let parent = child.saturating_sub(1) / 2;
2197 let parent_node = heap_node(heap, parent)?;
2198 if !scored_best_first(&moving, &parent_node).is_gt() {
2199 break;
2200 }
2201 heap_store(heap, child, parent_node)?;
2202 child = parent;
2203 }
2204 if child != start {
2205 heap_store(heap, child, moving)?;
2206 }
2207 Ok(())
2208}
2209
2210fn min_heap_sift_down(heap: &mut [ScoredNode], mut parent: usize) -> Result<(), GraphSearchError> {
2211 let start = parent;
2212 let moving = heap_node(heap, parent)?;
2213 loop {
2214 let left = parent
2215 .checked_mul(2)
2216 .and_then(|value| value.checked_add(1))
2217 .ok_or_else(|| GraphSearchError::Geometry("min-heap child overflow".to_owned()))?;
2218 if left >= heap.len() {
2219 break;
2220 }
2221 let right = left.saturating_add(1);
2222 let best_child = if right < heap.len()
2223 && scored_best_first(&heap_node(heap, right)?, &heap_node(heap, left)?).is_lt()
2224 {
2225 right
2226 } else {
2227 left
2228 };
2229 let child_node = heap_node(heap, best_child)?;
2230 if !scored_best_first(&child_node, &moving).is_lt() {
2231 break;
2232 }
2233 heap_store(heap, parent, child_node)?;
2234 parent = best_child;
2235 }
2236 if parent != start {
2237 heap_store(heap, parent, moving)?;
2238 }
2239 Ok(())
2240}
2241
2242fn max_heap_sift_down(heap: &mut [ScoredNode], mut parent: usize) -> Result<(), GraphSearchError> {
2243 let start = parent;
2244 let moving = heap_node(heap, parent)?;
2245 loop {
2246 let left = parent
2247 .checked_mul(2)
2248 .and_then(|value| value.checked_add(1))
2249 .ok_or_else(|| GraphSearchError::Geometry("max-heap child overflow".to_owned()))?;
2250 if left >= heap.len() {
2251 break;
2252 }
2253 let right = left.saturating_add(1);
2254 let worst_child = if right < heap.len()
2255 && scored_best_first(&heap_node(heap, right)?, &heap_node(heap, left)?).is_gt()
2256 {
2257 right
2258 } else {
2259 left
2260 };
2261 let child_node = heap_node(heap, worst_child)?;
2262 if !scored_best_first(&child_node, &moving).is_gt() {
2263 break;
2264 }
2265 heap_store(heap, parent, child_node)?;
2266 parent = worst_child;
2267 }
2268 if parent != start {
2269 heap_store(heap, parent, moving)?;
2270 }
2271 Ok(())
2272}
2273
2274fn heap_node(heap: &[ScoredNode], index: usize) -> Result<ScoredNode, GraphSearchError> {
2275 heap.get(index)
2276 .copied()
2277 .ok_or_else(|| GraphSearchError::Geometry(format!("heap slot {index} is unavailable")))
2278}
2279
2280fn heap_store(
2281 heap: &mut [ScoredNode],
2282 index: usize,
2283 value: ScoredNode,
2284) -> Result<(), GraphSearchError> {
2285 let slot = heap
2286 .get_mut(index)
2287 .ok_or_else(|| GraphSearchError::Geometry(format!("heap slot {index} is unavailable")))?;
2288 *slot = value;
2289 Ok(())
2290}
2291
2292fn squared_norm(values: &[f32]) -> f64 {
2293 values
2294 .iter()
2295 .map(|value| f64::from(*value) * f64::from(*value))
2296 .sum()
2297}
2298
2299#[cfg(target_vendor = "apple")]
2300fn observed_qos() -> (QueryQosClass, i32) {
2301 let mut class = libc::qos_class_t::QOS_CLASS_UNSPECIFIED;
2302 let mut priority = 0_i32;
2303 let status =
2305 unsafe { libc::pthread_get_qos_class_np(libc::pthread_self(), &mut class, &mut priority) };
2306 if status != 0 {
2307 return (QueryQosClass::Unknown, 0);
2308 }
2309 let class = match class {
2310 libc::qos_class_t::QOS_CLASS_USER_INTERACTIVE => QueryQosClass::UserInteractive,
2311 libc::qos_class_t::QOS_CLASS_USER_INITIATED => QueryQosClass::UserInitiated,
2312 libc::qos_class_t::QOS_CLASS_DEFAULT => QueryQosClass::Default,
2313 libc::qos_class_t::QOS_CLASS_UTILITY => QueryQosClass::Utility,
2314 libc::qos_class_t::QOS_CLASS_BACKGROUND => QueryQosClass::Background,
2315 libc::qos_class_t::QOS_CLASS_UNSPECIFIED => QueryQosClass::Unspecified,
2316 };
2317 (class, priority)
2318}
2319
2320#[cfg(not(target_vendor = "apple"))]
2321const fn observed_qos() -> (QueryQosClass, i32) {
2322 (QueryQosClass::Unavailable, 0)
2323}
2324
2325#[cfg(test)]
2326#[allow(
2327 clippy::expect_used,
2328 clippy::indexing_slicing,
2329 clippy::panic,
2330 clippy::unwrap_used
2331)]
2332mod tests {
2333 use std::time::Duration;
2334
2335 use proptest::prelude::*;
2336 use proptest::test_runner::{Config, RngSeed, TestRunner};
2337 use rand::RngCore;
2338
2339 use super::{
2340 AdaptiveEfError, FilteredGraphSearchOutcome, GraphSearchError, GraphSearchProfile,
2341 GraphSearchRequest, GraphSearchScratch, GraphSearcher, GraphSegmentNormRange,
2342 QueryCoreClass, TraversalPrefetch, next_down_f64, next_up_f32, next_up_f64,
2343 };
2344 use crate::graph::block::{
2345 GraphNodeBlockBuild, GraphNodeBlockInput, GraphNodeLayout, decode_node_blocks,
2346 encode_node_blocks,
2347 };
2348 use crate::kernels::GatherShapeError;
2349 use crate::lifecycle::{CancelToken, OpenOptions, QueryCancellation, QueryControl, Store};
2350 use crate::meta::DocBitmap;
2351 use crate::quant::{Bit4Factors, QuantError, RescoreError, quantize_bit4};
2352 use crate::scan::ScanError;
2353
2354 const DIMS: usize = 128;
2355
2356 #[test]
2357 fn adaptive_defaults_follow_the_measured_profile_rules() {
2358 let query = [0.0_f32; DIMS];
2359
2360 let sift = GraphSearchRequest::new(&query, 100, 0x19_05);
2361 assert_eq!(sift.research_floor_ef(), Ok(140));
2362 assert_eq!(sift.effective_ef(1_000_000), Ok(200));
2363 assert_eq!(sift.effective_ef(180), Ok(180));
2364
2365 let angular =
2366 GraphSearchRequest::new(&query, 100, 0x19_05).with_profile(GraphSearchProfile::Angular);
2367 assert_eq!(angular.research_floor_ef(), Ok(400));
2368 assert_eq!(angular.effective_ef(1_000_000), Ok(400));
2369
2370 let explicit = GraphSearchRequest::new(&query, 100, 0x19_05).with_ef(300);
2371 assert_eq!(explicit.effective_ef(1_000_000), Ok(300));
2372 }
2373
2374 #[test]
2375 fn adaptive_ef_overflow_is_typed() {
2376 let request = GraphSearchRequest::new(&[], usize::MAX, 0x19_05);
2377 assert_eq!(
2378 request.effective_ef(usize::MAX),
2379 Err(AdaptiveEfError::ArithmeticOverflow)
2380 );
2381 }
2382
2383 #[test]
2384 fn request_and_search_geometry_rejections_preserve_the_invalid_values() {
2385 let query = vector(0.0);
2386 assert_eq!(
2387 GraphSearchRequest::new(&query, 0, 7).effective_ef(9),
2388 Err(AdaptiveEfError::ZeroK)
2389 );
2390 assert_eq!(
2391 GraphSearchRequest::new(&query, 10, 7).effective_ef(9),
2392 Err(AdaptiveEfError::KExceedsRows { k: 10, rows: 9 })
2393 );
2394 assert_eq!(
2395 GraphSearchRequest::new(&query, 4, 7)
2396 .with_ef(3)
2397 .effective_ef(9),
2398 Err(AdaptiveEfError::ExplicitBelowK { k: 4, ef: 3 })
2399 );
2400 assert_eq!(
2401 GraphSearchRequest::new(&query, 4, 7)
2402 .with_ef(10)
2403 .effective_ef(9),
2404 Err(AdaptiveEfError::ExplicitExceedsRows { ef: 10, rows: 9 })
2405 );
2406
2407 let (encoded, rescore) = best_first_graph_fixture();
2408 let graph = decode_node_blocks(encoded.as_bytes()).expect("fixture graph is valid");
2409 let mut wrong_scratch =
2410 GraphSearchScratch::new(graph.node_count() + 1, graph.layout().max_degree())
2411 .expect("mismatched scratch allocates");
2412 assert!(matches!(
2413 GraphSearcher::new(graph, &rescore, &mut wrong_scratch),
2414 Err(GraphSearchError::Geometry(detail))
2415 if detail.contains("scratch geometry")
2416 ));
2417
2418 let mut scratch = GraphSearchScratch::new(graph.node_count(), graph.layout().max_degree())
2419 .expect("fixture scratch");
2420 assert!(matches!(
2421 GraphSearcher::new(graph, &rescore[..rescore.len() - 1], &mut scratch),
2422 Err(GraphSearchError::Geometry(detail))
2423 if detail.contains("rescore has") && detail.contains("expected")
2424 ));
2425 assert!(matches!(
2426 GraphSearcher::with_entry_row_ids(graph, &rescore, [0, 1, 2, 4], &mut scratch),
2427 Err(GraphSearchError::Geometry(detail))
2428 if detail.contains("cached entry row 4")
2429 ));
2430
2431 let mut searcher =
2432 GraphSearcher::new(graph, &rescore, &mut scratch).expect("valid searcher");
2433 assert!(matches!(
2434 searcher.search(
2435 GraphSearchRequest::new(&query[..DIMS - 1], 1, 7).with_ef(4),
2436 None,
2437 ),
2438 Err(GraphSearchError::Geometry(detail))
2439 if detail.contains("127 dimensions") && detail.contains("expected 128")
2440 ));
2441 let mut non_finite = query.clone();
2442 non_finite[17] = f32::NAN;
2443 assert_eq!(
2444 searcher.search(GraphSearchRequest::new(&non_finite, 1, 7).with_ef(4), None,),
2445 Err(GraphSearchError::Quant(QuantError::NonFinite { index: 17 }))
2446 );
2447 let recovered = searcher
2448 .search(GraphSearchRequest::new(&query, 1, 7).with_ef(4), None)
2449 .expect("a rejected request leaves reusable scratch valid");
2450 assert_eq!(recovered.candidates()[0].row_id(), 8);
2451
2452 let one_entry = factor_graph_fixture(Bit4Factors::from_persisted(1.0, 1.0, 0.0), 1);
2454 let one_entry = decode_node_blocks(one_entry.as_bytes()).expect("one-entry graph is valid");
2455 assert_eq!(
2456 GraphSearcher::discover_entry_row_ids(one_entry)
2457 .expect("a one-row graph carries one seed"),
2458 [0, 0, 0, 0]
2459 );
2460
2461 let no_entry = factor_graph_fixture(Bit4Factors::from_persisted(1.0, 1.0, 0.0), 0);
2463 let no_entry = decode_node_blocks(no_entry.as_bytes()).expect("no-entry graph is valid");
2464 assert!(matches!(
2465 GraphSearcher::discover_entry_row_ids(no_entry),
2466 Err(GraphSearchError::Geometry(detail))
2467 if detail.contains("contains 0 persisted entry seeds")
2468 && detail.contains("expected 1")
2469 ));
2470 }
2471
2472 #[test]
2473 fn segment_norm_bounds_round_outward_or_fail_open() {
2474 let bounded = factor_graph_fixture(Bit4Factors::from_persisted(2.0, 10.0, 0.0), 0);
2475 let bounded = decode_node_blocks(bounded.as_bytes()).expect("bounded graph is valid");
2476 let bounded = GraphSegmentNormRange::from_graph(bounded, None).expect("norm scan succeeds");
2477 let zero_query = vec![0.0_f32; DIMS];
2478 assert!(bounded.squared_l2_lower_bound(&zero_query) <= 400.0);
2479 assert!(bounded.squared_l2_upper_bound(&zero_query) >= 400.0);
2480 assert!(bounded.squared_l2_lower_bound(&zero_query) > 0.0);
2481 assert_eq!(bounded.squared_l2_lower_bound(&[]), 0.0);
2482 assert_eq!(bounded.squared_l2_upper_bound(&[]), f64::INFINITY);
2483
2484 for factors in [
2485 Bit4Factors::from_persisted(-1.0, 1.0, 0.0),
2486 Bit4Factors::from_persisted(f32::MAX, f32::MAX, 0.0),
2487 ] {
2488 let encoded = factor_graph_fixture(factors, 0);
2489 let graph = decode_node_blocks(encoded.as_bytes()).expect("finite factors decode");
2490 let range = GraphSegmentNormRange::from_graph(graph, None).expect("norm scan succeeds");
2491 assert_eq!(range.squared_l2_lower_bound(&zero_query), 0.0);
2492 assert_eq!(range.squared_l2_upper_bound(&zero_query), f64::INFINITY);
2493 }
2494 }
2495
2496 #[test]
2497 fn graph_search_small_persisted_seed_sets() {
2498 fn fixture(rows: usize, flagged: usize) -> crate::graph::block::EncodedNodeBlocks {
2499 let layout =
2500 GraphNodeLayout::new(DIMS as u32, DIMS as u32, 1).expect("small-seed layout");
2501 let codes = vec![0_u8; rows * DIMS.div_ceil(2)];
2502 let factors = Bit4Factors::from_persisted(0.0, 0.0, 0.0);
2503 let nodes = (0..rows)
2504 .map(|row| GraphNodeBlockInput {
2505 codes: &codes[row * DIMS.div_ceil(2)..(row + 1) * DIMS.div_ceil(2)],
2506 factors,
2507 flags: u8::from(row < flagged),
2508 neighbors: &[],
2509 })
2510 .collect::<Vec<_>>();
2511 encode_node_blocks(GraphNodeBlockBuild {
2512 layout,
2513 nodes: &nodes,
2514 })
2515 .expect("small-seed graph")
2516 }
2517
2518 for rows in 1..=3 {
2519 let encoded = fixture(rows, rows);
2520 let graph = decode_node_blocks(encoded.as_bytes()).expect("small graph decodes");
2521 let entries = GraphSearcher::discover_entry_row_ids(graph).expect("real seeds");
2522 assert_eq!(&entries[..rows], &(0..rows as u32).collect::<Vec<_>>());
2523 assert!(entries[rows..].iter().all(|entry| *entry == 0));
2524 }
2525
2526 let encoded = fixture(4, 4);
2527 let graph = decode_node_blocks(encoded.as_bytes()).expect("four-seed graph decodes");
2528 assert_eq!(
2529 GraphSearcher::discover_entry_row_ids(graph).expect("exact four seeds"),
2530 [0, 1, 2, 3]
2531 );
2532 for (rows, flagged) in [(3, 2), (5, 5)] {
2533 let encoded = fixture(rows, flagged);
2534 let graph = decode_node_blocks(encoded.as_bytes()).expect("invalid seed graph decodes");
2535 assert!(matches!(
2536 GraphSearcher::discover_entry_row_ids(graph),
2537 Err(GraphSearchError::Geometry(detail)) if detail.contains("persisted entry seeds")
2538 ));
2539 }
2540
2541 let encoded = fixture(1, 1);
2542 let graph = decode_node_blocks(encoded.as_bytes()).expect("one-row graph decodes");
2543 let rescore = vec![0.0_f32; DIMS];
2544 let mut scratch = GraphSearchScratch::new(1, 1).expect("one-row scratch");
2545 let mut searcher = GraphSearcher::new(graph, &rescore, &mut scratch)
2546 .expect("repeated adapter entries bind");
2547 let result = searcher
2548 .search(GraphSearchRequest::new(&rescore, 1, 7).with_ef(1), None)
2549 .expect("one-row traversal");
2550 assert_eq!(result.counters().visited(), 1);
2551 }
2552
2553 #[test]
2554 fn norm_bound_successors_are_exact_ieee_neighbours() {
2555 assert_eq!(next_up_f32(f32::INFINITY), f32::INFINITY);
2556 assert_eq!(next_up_f32(-0.0).to_bits(), 1);
2557 assert_eq!(next_up_f32(-1.0).to_bits(), (-1.0_f32).to_bits() - 1);
2558 assert_eq!(next_up_f64(f64::INFINITY), f64::INFINITY);
2559 assert_eq!(next_up_f64(-0.0).to_bits(), 1);
2560 assert_eq!(next_up_f64(-1.0).to_bits(), (-1.0_f64).to_bits() - 1);
2561 assert_eq!(next_down_f64(0.0), 0.0);
2562 assert_eq!(next_down_f64(1.0).to_bits(), 1.0_f64.to_bits() - 1);
2563 }
2564
2565 #[test]
2566 fn graph_search_errors_keep_typed_sources_and_actionable_messages() {
2567 let sourced = [
2568 GraphSearchError::from(crate::graph::block::GraphNodeError::InvalidHeader(
2569 "bad graph header".to_owned(),
2570 )),
2571 GraphSearchError::from(QuantError::EmptyVector),
2572 GraphSearchError::from(GatherShapeError::RowBytes {
2573 expected: 64,
2574 actual: 63,
2575 }),
2576 GraphSearchError::from(AdaptiveEfError::ExplicitBelowK { k: 4, ef: 3 }),
2577 GraphSearchError::from(AdaptiveEfError::ZeroK),
2578 GraphSearchError::from(AdaptiveEfError::KExceedsRows { k: 10, rows: 9 }),
2579 GraphSearchError::from(AdaptiveEfError::ExplicitExceedsRows { ef: 10, rows: 9 }),
2580 GraphSearchError::from(AdaptiveEfError::ArithmeticOverflow),
2581 GraphSearchError::from(RescoreError::InsufficientCandidates {
2582 k: 4,
2583 candidates: 3,
2584 }),
2585 GraphSearchError::Scan(ScanError::ZeroDimension),
2586 ];
2587 for error in sourced {
2588 assert!(!error.to_string().is_empty());
2589 assert!(std::error::Error::source(&error).is_some());
2590 }
2591
2592 let unsourced = [
2593 GraphSearchError::Geometry("bad request shape".to_owned()),
2594 GraphSearchError::VisitedCapExceeded {
2595 visited: 17,
2596 cap: 16,
2597 },
2598 GraphSearchError::Cancelled { partial: false },
2599 GraphSearchError::Timeout { partial: false },
2600 GraphSearchError::ReadCancelled { partial: false },
2601 ];
2602 let messages = unsourced
2603 .iter()
2604 .map(ToString::to_string)
2605 .collect::<Vec<_>>();
2606 assert!(messages[0].contains("bad request shape"));
2607 assert!(messages[1].contains("17") && messages[1].contains("16"));
2608 assert!(messages[2].contains("cancelled") && messages[2].contains("partial=false"));
2609 assert!(messages[3].contains("deadline") && messages[3].contains("partial=false"));
2610 assert!(messages[4].contains("store close") && messages[4].contains("partial=false"));
2611 assert!(
2612 unsourced
2613 .iter()
2614 .all(|error| std::error::Error::source(error).is_none())
2615 );
2616 }
2617
2618 #[test]
2619 fn adaptive_ef_stays_between_k_and_rows_under_adversarial_shapes() {
2620 let mut seeded = crate::test_support::seeded_rng(
2621 "graph::search::tests::adaptive_ef_stays_between_k_and_rows_under_adversarial_shapes",
2622 );
2623 let mut runner = TestRunner::new(Config {
2624 cases: 256,
2625 rng_seed: RngSeed::Fixed(seeded.next_u64()),
2626 ..Config::default()
2627 });
2628 let result = runner.run(
2629 &(1_usize..1_000_001, 1_usize..1_000_001, any::<bool>()),
2630 |(rows, raw_k, angular)| {
2631 let k = raw_k.min(rows);
2632 let profile = if angular {
2633 GraphSearchProfile::Angular
2634 } else {
2635 GraphSearchProfile::SiftClass
2636 };
2637 let request = GraphSearchRequest::new(&[], k, 0x19_05).with_profile(profile);
2638 let ef = request.effective_ef(rows).expect("valid shape resolves");
2639 prop_assert!(ef >= k);
2640 prop_assert!(ef <= rows);
2641 Ok(())
2642 },
2643 );
2644
2645 assert!(result.is_ok(), "property result: {result:?}");
2646 }
2647
2648 #[test]
2649 fn traversal_returns_the_same_topk_as_brute_force_at_high_ef() {
2650 let (encoded, rescore) = complete_graph_fixture(12);
2651 let graph = decode_node_blocks(encoded.as_bytes()).expect("fixture graph is valid");
2652 let query = vector(4.25);
2653 let expected = brute_force_topk(&rescore, &query, 5);
2654 let mut scratch = GraphSearchScratch::new(graph.node_count(), graph.layout().max_degree())
2655 .expect("fixture scratch");
2656 let mut searcher =
2657 GraphSearcher::new(graph, &rescore, &mut scratch).expect("fixture geometry is valid");
2658
2659 let result = searcher
2660 .search(GraphSearchRequest::new(&query, 5, 0x19_04), None)
2661 .expect("high-ef traversal succeeds");
2662 let actual = result
2663 .candidates()
2664 .iter()
2665 .map(|candidate| candidate.row_id())
2666 .collect::<Vec<_>>();
2667
2668 assert_eq!(actual, expected);
2669 assert_eq!(result.counters().effective_ef(), 12);
2670 assert_eq!(result.counters().candidates_rescored(), 12);
2671 }
2672
2673 #[test]
2674 fn an_unfiltered_traversal_is_byte_identical_with_the_mask_parameter_absent() {
2675 let (encoded, rescore) = complete_graph_fixture(12);
2676 let graph = decode_node_blocks(encoded.as_bytes()).expect("fixture graph is valid");
2677 let query = vector(4.25);
2678 let request = GraphSearchRequest::new(&query, 5, 0x19_06).with_ef(12);
2679 let mut unfiltered_scratch =
2680 GraphSearchScratch::new(graph.node_count(), graph.layout().max_degree())
2681 .expect("unfiltered scratch");
2682 let mut full_mask_scratch =
2683 GraphSearchScratch::new(graph.node_count(), graph.layout().max_degree())
2684 .expect("full-mask scratch");
2685 let mut unfiltered = GraphSearcher::new(graph, &rescore, &mut unfiltered_scratch)
2686 .expect("unfiltered searcher");
2687 let mut full_mask = GraphSearcher::new(graph, &rescore, &mut full_mask_scratch)
2688 .expect("full-mask searcher");
2689
2690 let expected = unfiltered.search(request, None).expect("unfiltered search");
2691 let actual = full_mask
2692 .search_filtered(
2693 request,
2694 &DocBitmap::full(graph.node_count()),
2695 usize::MAX,
2696 None,
2697 )
2698 .expect("full-mask search");
2699 let FilteredGraphSearchOutcome::Traversed(actual) = actual else {
2700 panic!("an unbounded full mask must traverse");
2701 };
2702
2703 assert_eq!(actual.candidates(), expected.candidates());
2704 assert_eq!(
2705 actual.counters().deterministic_work(),
2706 expected.counters().deterministic_work()
2707 );
2708 assert_eq!(actual.candidate_sequence(), expected.candidate_sequence());
2709 }
2710
2711 #[test]
2712 fn every_filtered_graph_result_satisfies_the_predicate_and_is_alive() {
2713 let (encoded, rescore) = complete_graph_fixture(12);
2714 let graph = decode_node_blocks(encoded.as_bytes()).expect("fixture graph is valid");
2715 let query = vector(4.0);
2716 let allowed = DocBitmap::from_ids([1, 2, 6, 9, 11]);
2717 let mut scratch = GraphSearchScratch::new(graph.node_count(), graph.layout().max_degree())
2718 .expect("filtered scratch");
2719 let mut searcher =
2720 GraphSearcher::new(graph, &rescore, &mut scratch).expect("filtered searcher");
2721
2722 let outcome = searcher
2723 .search_filtered(
2724 GraphSearchRequest::new(&query, 5, 0x19_06).with_ef(12),
2725 &allowed,
2726 usize::MAX,
2727 None,
2728 )
2729 .expect("filtered graph search");
2730 let FilteredGraphSearchOutcome::Traversed(result) = outcome else {
2731 panic!("an unbounded complete graph must traverse");
2732 };
2733
2734 assert_eq!(
2735 result
2736 .candidates()
2737 .iter()
2738 .map(|candidate| candidate.row_id())
2739 .collect::<Vec<_>>(),
2740 vec![2, 6, 1, 9, 11]
2741 );
2742 assert!(
2743 result
2744 .candidates()
2745 .iter()
2746 .all(|candidate| allowed.contains(candidate.row_id()))
2747 );
2748 assert!(
2749 result
2750 .candidates()
2751 .iter()
2752 .all(|candidate| candidate.row_id() != 4),
2753 "the exact nearest row is planted outside the effective alive-and-filter mask"
2754 );
2755 assert_eq!(
2756 result.counters().candidates_rescored(),
2757 allowed.cardinality() as usize
2758 );
2759
2760 let mut seeded = crate::test_support::seeded_rng(
2761 "graph::search::tests::every_filtered_graph_result_satisfies_the_predicate_and_is_alive",
2762 );
2763 let mut runner = TestRunner::new(Config {
2764 cases: 1_024,
2765 rng_seed: RngSeed::Fixed(seeded.next_u64()),
2766 ..Config::default()
2767 });
2768 let property = runner.run(
2769 &(
2770 prop::collection::vec(any::<bool>(), 12),
2771 prop::collection::vec(any::<bool>(), 12),
2772 ),
2773 |(predicate_matches, alive_rows)| {
2774 let effective = DocBitmap::from_ids(
2775 predicate_matches
2776 .iter()
2777 .zip(&alive_rows)
2778 .enumerate()
2779 .filter_map(|(row, (matches, alive))| {
2780 (*matches && *alive).then_some(row as u32)
2781 }),
2782 );
2783 prop_assume!(!effective.is_empty());
2784 let k = effective.cardinality() as usize;
2785 let mut scratch =
2786 GraphSearchScratch::new(graph.node_count(), graph.layout().max_degree())
2787 .expect("property scratch");
2788 let mut searcher =
2789 GraphSearcher::new(graph, &rescore, &mut scratch).expect("property searcher");
2790 let outcome = searcher
2791 .search_filtered(
2792 GraphSearchRequest::new(&query, k, 0x19_0602).with_ef(12),
2793 &effective,
2794 usize::MAX,
2795 None,
2796 )
2797 .expect("property filtered traversal");
2798 let FilteredGraphSearchOutcome::Traversed(result) = outcome else {
2799 return Err(TestCaseError::fail("unbounded complete graph abandoned"));
2800 };
2801 prop_assert_eq!(result.candidates().len(), k);
2802 let all_valid = result.candidates().iter().all(|candidate| {
2803 let row = candidate.row_id() as usize;
2804 effective.contains(candidate.row_id())
2805 && predicate_matches[row]
2806 && alive_rows[row]
2807 });
2808 prop_assert!(all_valid);
2809 Ok(())
2810 },
2811 );
2812 assert!(property.is_ok(), "property result: {property:?}");
2813 }
2814
2815 #[test]
2816 fn vector_row_checkpoint_counts_only_filter_accepted_graph_rows() {
2817 let (encoded, rescore) = complete_graph_fixture(12);
2818 let graph = decode_node_blocks(encoded.as_bytes()).expect("fixture graph is valid");
2819 let query = vector(4.0);
2820 let request = GraphSearchRequest::new(&query, 1, 0x27_24)
2821 .with_ef(12)
2822 .with_candidate_trace();
2823 let mut trace_scratch =
2824 GraphSearchScratch::new(graph.node_count(), graph.layout().max_degree())
2825 .expect("trace scratch");
2826 let mut trace =
2827 GraphSearcher::new(graph, &rescore, &mut trace_scratch).expect("trace searcher");
2828 let trace_result = trace.search(request, None).expect("trace traversal");
2829 let sequence = trace_result.candidate_sequence().expect("candidate trace");
2830 let first = *sequence.first().expect("first scored row");
2831 let accepted = *sequence
2832 .iter()
2833 .find(|&&row| row != first)
2834 .expect("later filter-accepted row");
2835 let allow_list = DocBitmap::from_ids([accepted]);
2836 let controller = crate::scan::vector_fault::VectorFaultController::armed(
2837 crate::scan::vector_fault::VectorFault::CancelAfterRows {
2838 source: crate::scan::vector_fault::VectorRowSource::Sealed([0x27; 16]),
2839 requested_rows: 1,
2840 tier: crate::scan::vector_fault::VectorSearchTier::Graph,
2841 },
2842 27_240,
2843 );
2844 let mut scratch = GraphSearchScratch::new(graph.node_count(), graph.layout().max_degree())
2845 .expect("fault scratch");
2846 let mut searcher = GraphSearcher::new(graph, &rescore, &mut scratch)
2847 .expect("fault searcher")
2848 .with_vector_fault_controller(
2849 &controller,
2850 crate::scan::vector_fault::VectorRowSource::Sealed([0x27; 16]),
2851 crate::scan::vector_fault::VectorSearchTier::Graph,
2852 );
2853 let error = searcher
2854 .search_filtered(request, &allow_list, usize::MAX, None)
2855 .expect_err("first filter-accepted row must cancel");
2856 assert_eq!(error, GraphSearchError::Cancelled { partial: false });
2857 controller.finalize_search(false);
2858 let receipts = controller.take_typed_receipts();
2859 assert_eq!(receipts.len(), 1);
2860 assert!(
2861 matches!(
2862 receipts[0].effect(),
2863 crate::scan::vector_fault::VectorFaultEffect::CancelledAfterRows {
2864 local_row,
2865 requested_rows: 1,
2866 observed_rows: 1,
2867 ..
2868 } if *local_row == accepted
2869 ),
2870 "filtered-out row {first} incorrectly consumed the eligible-row checkpoint: {:?}",
2871 receipts[0].effect()
2872 );
2873 }
2874
2875 #[test]
2876 fn traversal_expands_the_nearest_frontier_node_before_insertion_order() {
2877 let (encoded, rescore) = best_first_graph_fixture();
2878 let graph = decode_node_blocks(encoded.as_bytes()).expect("fixture graph is valid");
2879 let query = vector(0.0);
2880 let mut scratch = GraphSearchScratch::new(graph.node_count(), graph.layout().max_degree())
2881 .expect("fixture scratch");
2882 let mut searcher =
2883 GraphSearcher::new(graph, &rescore, &mut scratch).expect("fixture geometry is valid");
2884
2885 let result = searcher
2886 .search(
2887 GraphSearchRequest::new(&query, 1, 0x19_04_bf)
2888 .with_ef(4)
2889 .with_observed_core_class(QueryCoreClass::Performance)
2890 .with_candidate_trace(),
2891 None,
2892 )
2893 .expect("best-first traversal reaches the hidden nearest row");
2894
2895 assert_eq!(
2896 result
2897 .candidates()
2898 .iter()
2899 .map(|candidate| candidate.row_id())
2900 .collect::<Vec<_>>(),
2901 vec![8]
2902 );
2903 assert_eq!(result.candidate_sequence(), Some(&[0, 1, 2, 3, 8][..]));
2904 assert_eq!(result.counters().hops(), 4);
2905 assert_eq!(result.counters().pushes(), 5);
2906 assert_eq!(result.counters().core_class(), QueryCoreClass::Performance);
2907 assert_eq!(result.counters().candidates_rescored(), 4);
2908 assert!(result.counters().candidates_scored() < 9);
2909 }
2910
2911 #[test]
2912 fn visited_epoch_wraps_correctly_at_255_queries() {
2913 let (encoded, rescore) = complete_graph_fixture(12);
2914 let graph = decode_node_blocks(encoded.as_bytes()).expect("fixture graph is valid");
2915 let query = vector(4.25);
2916 let mut scratch = GraphSearchScratch::new(graph.node_count(), graph.layout().max_degree())
2917 .expect("fixture scratch");
2918 let mut searcher =
2919 GraphSearcher::new(graph, &rescore, &mut scratch).expect("fixture geometry is valid");
2920
2921 for query_index in 0..255 {
2922 let result = searcher
2923 .search(
2924 GraphSearchRequest::new(&query, 5, query_index).with_ef(12),
2925 None,
2926 )
2927 .expect("epoch before wrap succeeds");
2928 assert!(!result.counters().visited_epoch_cleared());
2929 }
2930 let wrapped = searcher
2931 .search(GraphSearchRequest::new(&query, 5, 255).with_ef(12), None)
2932 .expect("first query after 255 epochs succeeds");
2933
2934 assert!(wrapped.counters().visited_epoch_cleared());
2935 assert_eq!(
2936 wrapped
2937 .candidates()
2938 .iter()
2939 .map(|candidate| candidate.row_id())
2940 .collect::<Vec<_>>(),
2941 brute_force_topk(&rescore, &query, 5)
2942 );
2943 }
2944
2945 #[test]
2946 fn no_node_is_scored_twice_within_one_query() {
2947 let (encoded, rescore) = complete_graph_fixture(12);
2948 let graph = decode_node_blocks(encoded.as_bytes()).expect("fixture graph is valid");
2949 let query = vector(4.25);
2950 let mut scratch = GraphSearchScratch::new(graph.node_count(), graph.layout().max_degree())
2951 .expect("fixture scratch");
2952 let mut searcher =
2953 GraphSearcher::new(graph, &rescore, &mut scratch).expect("fixture geometry is valid");
2954
2955 let result = searcher
2956 .search(
2957 GraphSearchRequest::new(&query, 5, 0x19_04)
2958 .with_ef(12)
2959 .with_candidate_trace(),
2960 None,
2961 )
2962 .expect("traced traversal succeeds");
2963 let sequence = result
2964 .candidate_sequence()
2965 .expect("candidate tracing was requested");
2966 let mut seen = [false; 12];
2967 for row_id in sequence {
2968 let slot = &mut seen[*row_id as usize];
2969 assert!(!*slot, "row {row_id} was scored twice");
2970 *slot = true;
2971 }
2972
2973 assert_eq!(sequence.len(), result.counters().candidates_scored());
2974 assert_eq!(sequence.len(), result.counters().visited());
2975 assert_eq!(sequence, &(0_u32..12).collect::<Vec<_>>());
2976 }
2977
2978 #[test]
2979 fn traversal_is_deterministic_under_seed() {
2980 let (encoded, rescore) = complete_graph_fixture(12);
2981 let graph = decode_node_blocks(encoded.as_bytes()).expect("fixture graph is valid");
2982 let query = vector(4.25);
2983 let request = GraphSearchRequest::new(&query, 5, 0x19_0400)
2984 .with_ef(9)
2985 .with_candidate_trace();
2986 let mut first_scratch =
2987 GraphSearchScratch::new(graph.node_count(), graph.layout().max_degree())
2988 .expect("first scratch");
2989 let mut second_scratch =
2990 GraphSearchScratch::new(graph.node_count(), graph.layout().max_degree())
2991 .expect("second scratch");
2992 let mut first = GraphSearcher::new(graph, &rescore, &mut first_scratch)
2993 .expect("fixture geometry is valid");
2994 let mut second = GraphSearcher::new(graph, &rescore, &mut second_scratch)
2995 .expect("fixture geometry is valid");
2996
2997 let first_result = first
2998 .search(request, None)
2999 .expect("first traversal succeeds");
3000 let second_result = second
3001 .search(request, None)
3002 .expect("second traversal succeeds");
3003
3004 assert_eq!(
3005 first_result.candidate_sequence(),
3006 second_result.candidate_sequence()
3007 );
3008 assert_eq!(
3009 first_result.counters().deterministic_work(),
3010 second_result.counters().deterministic_work()
3011 );
3012 assert_eq!(first_result.candidates(), second_result.candidates());
3013 }
3014
3015 #[test]
3016 fn prefetch_switch_preserves_results_work_and_candidate_sequence() {
3017 let (encoded, rescore) = complete_graph_fixture(12);
3018 let graph = decode_node_blocks(encoded.as_bytes()).expect("fixture graph is valid");
3019 let query = vector(4.25);
3020 let mut enabled_scratch =
3021 GraphSearchScratch::new(graph.node_count(), graph.layout().max_degree())
3022 .expect("enabled scratch");
3023 let mut disabled_scratch =
3024 GraphSearchScratch::new(graph.node_count(), graph.layout().max_degree())
3025 .expect("disabled scratch");
3026 let mut enabled =
3027 GraphSearcher::new(graph, &rescore, &mut enabled_scratch).expect("enabled searcher");
3028 let mut disabled =
3029 GraphSearcher::new(graph, &rescore, &mut disabled_scratch).expect("disabled searcher");
3030 let enabled_result = enabled
3031 .search(
3032 GraphSearchRequest::new(&query, 5, 0x19_0400)
3033 .with_ef(9)
3034 .with_candidate_trace(),
3035 None,
3036 )
3037 .expect("prefetch-enabled traversal succeeds");
3038 let disabled_result = disabled
3039 .search(
3040 GraphSearchRequest::new(&query, 5, 0x19_0400)
3041 .with_ef(9)
3042 .with_candidate_trace()
3043 .with_prefetch(TraversalPrefetch::Disabled),
3044 None,
3045 )
3046 .expect("prefetch-disabled traversal succeeds");
3047
3048 assert_eq!(
3049 enabled_result.candidate_sequence(),
3050 disabled_result.candidate_sequence()
3051 );
3052 assert_eq!(
3053 enabled_result.counters().deterministic_work(),
3054 disabled_result.counters().deterministic_work()
3055 );
3056 assert_eq!(enabled_result.candidates(), disabled_result.candidates());
3057 }
3058
3059 #[test]
3060 fn visited_cap_returns_a_typed_error_not_a_wrong_answer() {
3061 let (encoded, rescore) = chain_graph_fixture(10);
3062 let graph = decode_node_blocks(encoded.as_bytes()).expect("fixture graph is valid");
3063 let query = vector(9.0);
3064 let mut scratch = GraphSearchScratch::new(graph.node_count(), graph.layout().max_degree())
3065 .expect("fixture scratch");
3066 let mut searcher =
3067 GraphSearcher::new(graph, &rescore, &mut scratch).expect("fixture geometry is valid");
3068
3069 let error = searcher
3070 .search(
3071 GraphSearchRequest::new(&query, 1, 0x19_0405).with_ef(1),
3072 None,
3073 )
3074 .expect_err("the visited guard must fail closed");
3075
3076 assert_eq!(
3077 error,
3078 super::GraphSearchError::VisitedCapExceeded { visited: 5, cap: 4 }
3079 );
3080 }
3081
3082 #[test]
3083 fn filtered_budget_and_fail_closed_visited_cap_fire_independently() {
3084 let (encoded, rescore) = chain_graph_fixture(10);
3085 let graph = decode_node_blocks(encoded.as_bytes()).expect("fixture graph is valid");
3086 let query = vector(9.0);
3087 let allowed = DocBitmap::full(graph.node_count());
3088 let request = GraphSearchRequest::new(&query, 1, 0x19_0605).with_ef(1);
3089
3090 let mut budget_scratch =
3091 GraphSearchScratch::new(graph.node_count(), graph.layout().max_degree())
3092 .expect("budget scratch");
3093 let mut budget_searcher =
3094 GraphSearcher::new(graph, &rescore, &mut budget_scratch).expect("budget searcher");
3095 let budget_outcome = budget_searcher
3096 .search_filtered(request, &allowed, 2, None)
3097 .expect("selectivity budget is a normal filtered disposition");
3098 assert!(matches!(
3099 budget_outcome,
3100 FilteredGraphSearchOutcome::VisitedBudgetExceeded {
3101 visited: 3,
3102 budget: 2,
3103 ..
3104 }
3105 ));
3106
3107 let mut invariant_scratch =
3108 GraphSearchScratch::new(graph.node_count(), graph.layout().max_degree())
3109 .expect("invariant scratch");
3110 let mut invariant_searcher = GraphSearcher::new(graph, &rescore, &mut invariant_scratch)
3111 .expect("invariant searcher");
3112 let error = invariant_searcher
3113 .search_filtered(request, &allowed, usize::MAX, None)
3114 .expect_err("removing the selectivity budget must expose the invariant guard");
3115 assert_eq!(
3116 error,
3117 GraphSearchError::VisitedCapExceeded { visited: 5, cap: 4 }
3118 );
3119 }
3120
3121 #[test]
3122 fn cancellation_is_visible_across_threads_during_traversal() {
3123 let (encoded, rescore) = long_chain_graph_fixture(40_000);
3124 let graph = decode_node_blocks(encoded.as_bytes()).expect("fixture graph is valid");
3125 let query = vec![0.0_f32; DIMS];
3126 let mut scratch = GraphSearchScratch::new(graph.node_count(), graph.layout().max_degree())
3127 .expect("fixture scratch");
3128 let mut searcher =
3129 GraphSearcher::new(graph, &rescore, &mut scratch).expect("fixture geometry is valid");
3130 let directory = tempfile::tempdir().expect("temporary store directory");
3131 let store = Store::open(directory.path(), OpenOptions::default()).expect("store opens");
3132 let lease = store.snapshot().expect("snapshot lease");
3133 let token = CancelToken::new();
3134 let control = QueryControl::Cancel(token.clone());
3135 let cancellation = QueryCancellation::new(&control, &lease);
3136 let error = std::thread::scope(|scope| {
3137 let canceller = scope.spawn(|| {
3138 std::thread::sleep(Duration::from_millis(2));
3139 token.cancel();
3140 });
3141 let result = searcher.search(
3142 GraphSearchRequest::new(&query, 1, 0x19_0406).with_ef(200),
3143 Some(&cancellation),
3144 );
3145 canceller.join().expect("canceller thread joins");
3146 result
3147 })
3148 .expect_err("in-flight traversal must stop");
3149
3150 assert_eq!(error, super::GraphSearchError::Cancelled { partial: false });
3151 }
3152
3153 #[test]
3154 fn pre_cancelled_token_is_rejected_at_search_entry() {
3155 let (encoded, rescore) = long_chain_graph_fixture(16);
3156 let graph = decode_node_blocks(encoded.as_bytes()).expect("fixture graph is valid");
3157 let query = vec![0.0_f32; DIMS];
3158 let mut scratch = GraphSearchScratch::new(graph.node_count(), graph.layout().max_degree())
3159 .expect("fixture scratch");
3160 let mut searcher =
3161 GraphSearcher::new(graph, &rescore, &mut scratch).expect("fixture geometry is valid");
3162 let directory = tempfile::tempdir().expect("temporary store directory");
3163 let store = Store::open(directory.path(), OpenOptions::default()).expect("store opens");
3164 let lease = store.snapshot().expect("snapshot lease");
3165 let token = CancelToken::new();
3166 token.cancel();
3167 let control = QueryControl::Cancel(token);
3168 let cancellation = QueryCancellation::new(&control, &lease);
3169
3170 let error = searcher
3171 .search(
3172 GraphSearchRequest::new(&query, 1, 0x19_0406).with_ef(16),
3173 Some(&cancellation),
3174 )
3175 .expect_err("pre-cancelled search must stop at entry");
3176
3177 assert_eq!(error, GraphSearchError::Cancelled { partial: false });
3178 }
3179
3180 #[test]
3181 fn counting_token_cancels_after_exact_number_of_hops() {
3182 let (encoded, rescore) = long_chain_graph_fixture(40_000);
3183 let graph = decode_node_blocks(encoded.as_bytes()).expect("fixture graph is valid");
3184 let query = vec![0.0_f32; DIMS];
3185 let mut scratch = GraphSearchScratch::new(graph.node_count(), graph.layout().max_degree())
3186 .expect("fixture scratch");
3187 let mut searcher =
3188 GraphSearcher::new(graph, &rescore, &mut scratch).expect("fixture geometry is valid");
3189 let directory = tempfile::tempdir().expect("temporary store directory");
3190 let store = Store::open(directory.path(), OpenOptions::default()).expect("store opens");
3191 let lease = store.snapshot().expect("snapshot lease");
3192 let token = CancelToken::new();
3193 let observed_hops = searcher.cancel_after_hops(7, token.clone());
3194 let control = QueryControl::Cancel(token);
3195 let cancellation = QueryCancellation::new(&control, &lease);
3196
3197 let error = searcher
3198 .search(
3199 GraphSearchRequest::new(&query, 1, 0x19_0406).with_ef(200),
3200 Some(&cancellation),
3201 )
3202 .expect_err("counting token must cancel during traversal");
3203
3204 assert_eq!(error, GraphSearchError::Cancelled { partial: false });
3205 assert_eq!(observed_hops.get(), 7);
3206 }
3207
3208 #[test]
3209 fn recall_is_monotonic_in_ef() {
3210 let (encoded, rescore) = complete_graph_fixture(12);
3211 let graph = decode_node_blocks(encoded.as_bytes()).expect("fixture graph is valid");
3212 let mut seeded =
3213 crate::test_support::seeded_rng("graph::search::tests::recall_is_monotonic_in_ef");
3214 let mut runner = TestRunner::new(Config {
3215 cases: 64,
3216 rng_seed: RngSeed::Fixed(seeded.next_u64()),
3217 ..Config::default()
3218 });
3219 let result = runner.run(&(5_usize..12, 1_usize..8, -1.0_f32..13.0), |input| {
3220 let (ef1, delta, query_value) = input;
3221 let ef2 = ef1.saturating_add(delta).min(12);
3222 prop_assume!(ef1 < ef2);
3223 let query = vector(query_value);
3224 let truth = brute_force_topk(&rescore, &query, 5);
3225 let mut first_scratch =
3226 GraphSearchScratch::new(graph.node_count(), graph.layout().max_degree())
3227 .expect("first scratch");
3228 let mut second_scratch =
3229 GraphSearchScratch::new(graph.node_count(), graph.layout().max_degree())
3230 .expect("second scratch");
3231 let mut first =
3232 GraphSearcher::new(graph, &rescore, &mut first_scratch).expect("first searcher");
3233 let mut second =
3234 GraphSearcher::new(graph, &rescore, &mut second_scratch).expect("second searcher");
3235 let result1 = first
3236 .search(
3237 GraphSearchRequest::new(&query, 5, 0x19_0408).with_ef(ef1),
3238 None,
3239 )
3240 .expect("ef1 traversal");
3241 let result2 = second
3242 .search(
3243 GraphSearchRequest::new(&query, 5, 0x19_0408).with_ef(ef2),
3244 None,
3245 )
3246 .expect("ef2 traversal");
3247 prop_assert!(result2.recall_against(&truth) >= result1.recall_against(&truth));
3248 Ok(())
3249 });
3250
3251 assert!(result.is_ok(), "property result: {result:?}");
3252 }
3253
3254 #[test]
3255 fn ze76_candidate_window_peak_counts_entries_not_reserved_slots() {
3256 let (encoded, rescore) = complete_graph_fixture(5);
3257 let graph = decode_node_blocks(encoded.as_bytes()).unwrap();
3258 let mut scratch = GraphSearchScratch::with_ef_capacity(
3259 graph.node_count(),
3260 graph.layout().max_degree(),
3261 64,
3262 )
3263 .unwrap();
3264 let query = vec![0.0; DIMS];
3265 let mut searcher = GraphSearcher::new(graph, &rescore, &mut scratch).unwrap();
3266 let result = searcher
3267 .search(GraphSearchRequest::new(&query, 1, 7).with_ef(5), None)
3268 .unwrap();
3269 assert_eq!(result.candidates().len(), 1);
3270 assert_eq!(result.counters().candidate_window_peak(), 5);
3273 }
3274
3275 fn complete_graph_fixture(rows: usize) -> (crate::graph::block::EncodedNodeBlocks, Vec<f32>) {
3276 let layout = GraphNodeLayout::new(DIMS as u32, DIMS as u32, (rows - 1) as u8)
3277 .expect("fixture layout is valid");
3278 let rescore = (0..rows)
3279 .flat_map(|row| vector(row as f32))
3280 .collect::<Vec<_>>();
3281 let mut codes = vec![0_u8; rows * DIMS.div_ceil(2)];
3282 let mut factors = Vec::with_capacity(rows);
3283 for (row, destination) in rescore
3284 .chunks_exact(DIMS)
3285 .zip(codes.chunks_exact_mut(DIMS.div_ceil(2)))
3286 {
3287 factors.push(quantize_bit4(row, destination).expect("fixture row is finite"));
3288 }
3289 let neighbors = (0..rows)
3290 .map(|row| {
3291 (0..rows)
3292 .filter(|candidate| *candidate != row)
3293 .map(|candidate| candidate as u32)
3294 .collect::<Vec<_>>()
3295 })
3296 .collect::<Vec<_>>();
3297 let nodes = (0..rows)
3298 .map(|row| GraphNodeBlockInput {
3299 codes: &codes[row * DIMS.div_ceil(2)..(row + 1) * DIMS.div_ceil(2)],
3300 factors: factors[row],
3301 flags: u8::from(row < 4),
3302 neighbors: &neighbors[row],
3303 })
3304 .collect::<Vec<_>>();
3305 let encoded = encode_node_blocks(GraphNodeBlockBuild {
3306 layout,
3307 nodes: &nodes,
3308 })
3309 .expect("fixture graph encodes");
3310 (encoded, rescore)
3311 }
3312
3313 fn chain_graph_fixture(rows: usize) -> (crate::graph::block::EncodedNodeBlocks, Vec<f32>) {
3314 let layout =
3315 GraphNodeLayout::new(DIMS as u32, DIMS as u32, 1).expect("fixture layout is valid");
3316 let rescore = (0..rows)
3317 .flat_map(|row| vector(row as f32))
3318 .collect::<Vec<_>>();
3319 let mut codes = vec![0_u8; rows * DIMS.div_ceil(2)];
3320 let mut factors = Vec::with_capacity(rows);
3321 for (row, destination) in rescore
3322 .chunks_exact(DIMS)
3323 .zip(codes.chunks_exact_mut(DIMS.div_ceil(2)))
3324 {
3325 factors.push(quantize_bit4(row, destination).expect("fixture row is finite"));
3326 }
3327 let neighbors = (0..rows)
3328 .map(|row| {
3329 (row + 1 < rows)
3330 .then_some((row + 1) as u32)
3331 .into_iter()
3332 .collect::<Vec<_>>()
3333 })
3334 .collect::<Vec<_>>();
3335 let nodes = (0..rows)
3336 .map(|row| GraphNodeBlockInput {
3337 codes: &codes[row * DIMS.div_ceil(2)..(row + 1) * DIMS.div_ceil(2)],
3338 factors: factors[row],
3339 flags: u8::from(row < 4),
3340 neighbors: &neighbors[row],
3341 })
3342 .collect::<Vec<_>>();
3343 let encoded = encode_node_blocks(GraphNodeBlockBuild {
3344 layout,
3345 nodes: &nodes,
3346 })
3347 .expect("fixture graph encodes");
3348 (encoded, rescore)
3349 }
3350
3351 fn best_first_graph_fixture() -> (crate::graph::block::EncodedNodeBlocks, Vec<f32>) {
3352 let values = [10.0_f32, 9.0, 8.0, 1.0, 7.0, 6.0, 5.0, 4.0, 0.0];
3353 let layout =
3354 GraphNodeLayout::new(DIMS as u32, DIMS as u32, 4).expect("best-first layout is valid");
3355 let rescore = values.into_iter().flat_map(vector).collect::<Vec<_>>();
3356 let mut codes = vec![0_u8; values.len() * DIMS.div_ceil(2)];
3357 let factors = rescore
3358 .chunks_exact(DIMS)
3359 .zip(codes.chunks_exact_mut(DIMS.div_ceil(2)))
3360 .map(|(row, destination)| {
3361 quantize_bit4(row, destination).expect("best-first row is finite")
3362 })
3363 .collect::<Vec<_>>();
3364 let neighbors = [
3365 vec![4, 5, 6, 7],
3366 Vec::new(),
3367 Vec::new(),
3368 vec![8],
3369 Vec::new(),
3370 Vec::new(),
3371 Vec::new(),
3372 Vec::new(),
3373 Vec::new(),
3374 ];
3375 let nodes = (0..values.len())
3376 .map(|row| GraphNodeBlockInput {
3377 codes: &codes[row * DIMS.div_ceil(2)..(row + 1) * DIMS.div_ceil(2)],
3378 factors: factors[row],
3379 flags: u8::from(row < 4),
3380 neighbors: &neighbors[row],
3381 })
3382 .collect::<Vec<_>>();
3383 let encoded = encode_node_blocks(GraphNodeBlockBuild {
3384 layout,
3385 nodes: &nodes,
3386 })
3387 .expect("best-first graph encodes");
3388 (encoded, rescore)
3389 }
3390
3391 fn factor_graph_fixture(
3392 factors: Bit4Factors,
3393 flags: u8,
3394 ) -> crate::graph::block::EncodedNodeBlocks {
3395 let layout = GraphNodeLayout::new(DIMS as u32, DIMS as u32, 0)
3396 .expect("factor graph layout is valid");
3397 let codes = vec![0_u8; DIMS.div_ceil(2)];
3398 let node = GraphNodeBlockInput {
3399 codes: &codes,
3400 factors,
3401 flags,
3402 neighbors: &[],
3403 };
3404 encode_node_blocks(GraphNodeBlockBuild {
3405 layout,
3406 nodes: &[node],
3407 })
3408 .expect("factor graph encodes")
3409 }
3410
3411 fn long_chain_graph_fixture(rows: usize) -> (crate::graph::block::EncodedNodeBlocks, Vec<f32>) {
3412 let layout =
3413 GraphNodeLayout::new(DIMS as u32, DIMS as u32, 44).expect("fixture layout is valid");
3414 let codes = vec![0_u8; rows * DIMS.div_ceil(2)];
3415 let factors = (0..rows)
3416 .map(|row| Bit4Factors::from_persisted(1.0, (rows - row) as f32, 0.0))
3417 .collect::<Vec<_>>();
3418 let neighbors = (0..rows)
3419 .map(|row| {
3420 (row + 1 < rows)
3421 .then_some((row + 1) as u32)
3422 .into_iter()
3423 .collect::<Vec<_>>()
3424 })
3425 .collect::<Vec<_>>();
3426 let nodes = (0..rows)
3427 .map(|row| GraphNodeBlockInput {
3428 codes: &codes[row * DIMS.div_ceil(2)..(row + 1) * DIMS.div_ceil(2)],
3429 factors: factors[row],
3430 flags: u8::from(row < 4),
3431 neighbors: &neighbors[row],
3432 })
3433 .collect::<Vec<_>>();
3434 let encoded = encode_node_blocks(GraphNodeBlockBuild {
3435 layout,
3436 nodes: &nodes,
3437 })
3438 .expect("fixture graph encodes");
3439 (encoded, vec![0.0_f32; rows * DIMS])
3440 }
3441
3442 fn vector(value: f32) -> Vec<f32> {
3443 let mut vector = vec![0.0_f32; DIMS];
3444 vector[0] = value;
3445 vector[1] = value * value * 0.1;
3446 vector[17] = value.sin();
3447 vector
3448 }
3449
3450 fn brute_force_topk(base: &[f32], query: &[f32], k: usize) -> Vec<u32> {
3451 let mut scored = base
3452 .chunks_exact(DIMS)
3453 .enumerate()
3454 .map(|(row_id, row)| {
3455 let distance = row
3456 .iter()
3457 .zip(query)
3458 .map(|(left, right)| {
3459 let delta = f64::from(*left) - f64::from(*right);
3460 delta * delta
3461 })
3462 .sum::<f64>();
3463 (distance, row_id as u32)
3464 })
3465 .collect::<Vec<_>>();
3466 scored.sort_unstable_by(|left, right| {
3467 left.0
3468 .total_cmp(&right.0)
3469 .then_with(|| left.1.cmp(&right.1))
3470 });
3471 scored
3472 .into_iter()
3473 .take(k)
3474 .map(|(_, row_id)| row_id)
3475 .collect()
3476 }
3477}