1use std::cmp::Ordering;
10use std::collections::BinaryHeap;
11
12use log::{debug, warn};
13
14use crate::DocId;
15
16const MAX_INITIAL_SCORE_COLLECTOR_CAPACITY: usize = 8 * 1024;
20
21#[derive(Clone, Copy)]
23pub struct HeapEntry {
24 pub doc_id: DocId,
25 pub score: f32,
26 pub ordinal: u16,
27}
28
29impl PartialEq for HeapEntry {
30 fn eq(&self, other: &Self) -> bool {
31 self.score.to_bits() == other.score.to_bits()
32 && self.doc_id == other.doc_id
33 && self.ordinal == other.ordinal
34 }
35}
36
37impl Eq for HeapEntry {}
38
39impl Ord for HeapEntry {
40 fn cmp(&self, other: &Self) -> Ordering {
41 other
45 .score
46 .total_cmp(&self.score)
47 .then_with(|| self.doc_id.cmp(&other.doc_id))
48 .then_with(|| self.ordinal.cmp(&other.ordinal))
49 }
50}
51
52impl PartialOrd for HeapEntry {
53 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
54 Some(self.cmp(other))
55 }
56}
57
58pub struct ScoreCollector {
71 heap: BinaryHeap<HeapEntry>,
73 pub k: usize,
74 cached_threshold: f32,
77 virtual_threshold: Option<f32>,
81}
82
83impl ScoreCollector {
84 pub fn new(k: usize) -> Self {
86 Self {
87 heap: BinaryHeap::with_capacity(k.min(MAX_INITIAL_SCORE_COLLECTOR_CAPACITY)),
88 k,
89 cached_threshold: 0.0,
90 virtual_threshold: None,
91 }
92 }
93
94 #[inline]
96 pub fn threshold(&self) -> f32 {
97 self.cached_threshold
98 }
99
100 #[inline]
102 fn update_threshold(&mut self) {
103 self.cached_threshold = if let Some(threshold) = self.virtual_threshold {
104 threshold
105 } else if self.heap.len() >= self.k {
106 self.heap.peek().map(|e| e.score).unwrap_or(0.0)
107 } else {
108 0.0
109 };
110 }
111
112 #[inline]
115 pub fn insert(&mut self, doc_id: DocId, score: f32) -> bool {
116 self.insert_with_ordinal(doc_id, score, 0)
117 }
118
119 #[inline]
122 pub fn insert_with_ordinal(&mut self, doc_id: DocId, score: f32, ordinal: u16) -> bool {
123 if self.k == 0 {
124 return false;
125 }
126 let entry = HeapEntry {
127 doc_id,
128 score,
129 ordinal,
130 };
131 if self.heap.len() < self.k {
132 if let Some(threshold) = self.virtual_threshold {
133 let sentinel = HeapEntry {
134 doc_id: u32::MAX,
135 score: threshold,
136 ordinal: 0,
137 };
138 if entry >= sentinel {
139 return false;
140 }
141 }
142
143 self.heap.push(entry);
144 if self.heap.len() == self.k {
146 self.virtual_threshold = None;
147 self.update_threshold();
148 }
149 true
150 } else if self.heap.peek().is_some_and(|worst| entry < *worst) {
151 {
152 let mut worst = self.heap.peek_mut().expect("full heap has a root");
153 *worst = entry;
154 }
155 self.update_threshold();
156 true
157 } else {
158 false
159 }
160 }
161
162 #[inline]
164 pub fn would_enter(&self, score: f32) -> bool {
165 self.len() < self.k || score > self.cached_threshold
166 }
167
168 #[inline]
171 pub fn would_enter_candidate(&self, doc_id: DocId, score: f32, ordinal: u16) -> bool {
172 if self.k == 0 {
173 return false;
174 }
175 let entry = HeapEntry {
176 doc_id,
177 score,
178 ordinal,
179 };
180 if let Some(threshold) = self.virtual_threshold {
181 let sentinel = HeapEntry {
182 doc_id: u32::MAX,
183 score: threshold,
184 ordinal: 0,
185 };
186 entry < sentinel
187 } else {
188 self.heap.len() < self.k || self.heap.peek().is_some_and(|worst| entry < *worst)
189 }
190 }
191
192 #[inline]
194 pub fn len(&self) -> usize {
195 if self.virtual_threshold.is_some() {
196 self.k
197 } else {
198 self.heap.len()
199 }
200 }
201
202 #[inline]
204 pub fn real_len(&self) -> usize {
205 self.heap.len()
206 }
207
208 #[inline]
210 pub fn is_empty(&self) -> bool {
211 self.len() == 0
212 }
213
214 pub fn seed_threshold(&mut self, initial_threshold: f32) {
221 if initial_threshold <= 0.0
222 || self.k == 0
223 || (self.len() >= self.k && initial_threshold <= self.cached_threshold)
224 {
225 return;
226 }
227
228 let sentinel = HeapEntry {
229 doc_id: u32::MAX,
230 score: initial_threshold,
231 ordinal: 0,
232 };
233
234 if let Some(current_threshold) = self.virtual_threshold {
238 let current = HeapEntry {
239 doc_id: u32::MAX,
240 score: current_threshold,
241 ordinal: 0,
242 };
243 if sentinel >= current {
244 return;
245 }
246 } else if self.heap.len() >= self.k
247 && !self.heap.peek().is_some_and(|worst| sentinel < *worst)
248 {
249 return;
250 }
251
252 self.virtual_threshold = Some(initial_threshold);
253 while self.heap.peek().is_some_and(|worst| sentinel < *worst) {
254 self.heap.pop();
255 }
256 self.update_threshold();
257 }
258
259 pub fn into_sorted_results(self) -> Vec<(DocId, f32, u16)> {
262 let mut results: Vec<(DocId, f32, u16)> = self
263 .heap
264 .into_vec()
265 .into_iter()
266 .filter(|e| e.doc_id != u32::MAX)
267 .map(|e| (e.doc_id, e.score, e.ordinal))
268 .collect();
269
270 results.sort_unstable_by(|a, b| {
272 b.1.total_cmp(&a.1)
273 .then_with(|| a.0.cmp(&b.0))
274 .then_with(|| a.2.cmp(&b.2))
275 });
276
277 results
278 }
279}
280
281#[derive(Clone, Debug)]
304pub struct SharedThreshold {
305 floor: std::sync::Arc<std::sync::atomic::AtomicU32>,
306 k: usize,
309 deadline: Option<std::time::Instant>,
312 truncated: std::sync::Arc<std::sync::atomic::AtomicBool>,
314}
315
316impl Default for SharedThreshold {
317 fn default() -> Self {
318 Self::new()
319 }
320}
321
322impl SharedThreshold {
323 pub fn new() -> Self {
327 Self::with_depth(usize::MAX)
328 }
329
330 pub fn for_limit(limit: usize) -> Self {
332 Self::with_depth(limit)
333 }
334
335 fn with_depth(k: usize) -> Self {
336 Self {
337 floor: std::sync::Arc::new(std::sync::atomic::AtomicU32::new(0)),
339 k,
340 deadline: None,
341 truncated: std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false)),
342 }
343 }
344
345 pub fn with_deadline(mut self, deadline: Option<std::time::Instant>) -> Self {
347 self.deadline = deadline;
348 self
349 }
350
351 pub fn deadline(&self) -> Option<std::time::Instant> {
353 self.deadline
354 }
355
356 pub(crate) fn budget_only(&self) -> Self {
358 Self {
359 floor: std::sync::Arc::new(std::sync::atomic::AtomicU32::new(0)),
360 k: usize::MAX,
361 deadline: self.deadline,
362 truncated: self.truncated.clone(),
363 }
364 }
365
366 #[inline]
367 pub(crate) fn stop_if_expired(&self) -> bool {
368 if self.expired() {
369 self.mark_truncated();
370 true
371 } else {
372 false
373 }
374 }
375
376 #[inline]
378 pub fn expired(&self) -> bool {
379 self.deadline
380 .is_some_and(|deadline| std::time::Instant::now() >= deadline)
381 }
382
383 pub fn mark_truncated(&self) {
385 self.truncated
386 .store(true, std::sync::atomic::Ordering::Relaxed);
387 }
388
389 pub fn truncated(&self) -> bool {
391 self.truncated.load(std::sync::atomic::Ordering::Relaxed)
392 }
393
394 #[inline]
397 pub(crate) fn covers(&self, heap_depth: usize) -> bool {
398 heap_depth >= self.k
399 }
400
401 #[inline]
403 pub fn get(&self) -> f32 {
404 f32::from_bits(self.floor.load(std::sync::atomic::Ordering::Relaxed))
405 }
406
407 pub fn raise(&self, score: f32) {
412 if score <= 0.0 {
415 return;
416 }
417 use std::sync::atomic::Ordering::Relaxed;
418 let bits = score.to_bits();
419 let mut cur = self.floor.load(Relaxed);
420 while f32::from_bits(cur) < score {
421 match self
422 .floor
423 .compare_exchange_weak(cur, bits, Relaxed, Relaxed)
424 {
425 Ok(_) => break,
426 Err(actual) => cur = actual,
427 }
428 }
429 }
430}
431
432#[derive(Debug, Clone, Copy)]
434pub struct ScoredDoc {
435 pub doc_id: DocId,
436 pub score: f32,
437 pub ordinal: u16,
439}
440
441pub struct MaxScoreExecutor<'a> {
452 metric_index: &'a str,
456 metric_field: &'a str,
457 cursors: Vec<TermCursor<'a>>,
458 prefix_sums: Vec<f32>,
459 collector: ScoreCollector,
460 inv_heap_factor: f32,
461 predicate: Option<super::DocPredicate<'a>>,
462 budget: Option<SharedThreshold>,
465}
466
467#[derive(Clone, Copy)]
471pub enum LengthSource<'a> {
472 Chunks(&'a crate::segment::chunk_map::ChunkMap),
473 Docs(&'a crate::segment::chunk_map::DocLengths),
474}
475
476impl LengthSource<'_> {
477 #[inline]
478 pub fn length(&self, id: u32) -> u32 {
479 match self {
480 LengthSource::Chunks(map) => map.bm25_length(id),
481 LengthSource::Docs(lengths) => lengths.length(id),
482 }
483 }
484}
485
486pub(crate) struct TermCursor<'a> {
495 pub max_score: f32,
496 num_blocks: usize,
497 block_idx: usize,
499 doc_ids: Vec<u32>,
500 scores: Vec<f32>,
501 ordinals: Vec<u16>,
502 pos: usize,
503 block_loaded: bool,
504 exhausted: bool,
505 lazy_ordinals: bool,
509 ordinals_loaded: bool,
511 current_sparse_block: Option<crate::structures::SparseBlock>,
513 variant: CursorVariant<'a>,
515}
516
517#[allow(clippy::large_enum_variant)]
520enum CursorVariant<'a> {
521 Text {
523 list: crate::structures::BlockPostingList,
524 idf: f32,
525 idf_times_k1_plus_1: f32,
527 denom_tf_coeff: f32,
529 denom_const: f32,
531 denom_len_coeff: f32,
534 lengths: Option<LengthSource<'a>>,
537 length_bounds: bool,
541 length_floor: u32,
542 block_bound: CachedScoreBound,
543 group_bound: CachedScoreBound,
544 avg_len: f32,
546 params: super::Bm25Params,
548 tfs: Vec<u32>,
549 deferred_tf: Option<(usize, usize, usize)>,
552 },
553 Sparse {
555 si: &'a crate::segment::SparseIndex,
556 query_weight: f32,
557 skip_start: usize,
558 block_data_offset: u64,
559 },
560}
561
562struct CachedScoreBound(std::sync::atomic::AtomicU64);
565
566impl CachedScoreBound {
567 fn new() -> Self {
568 Self(std::sync::atomic::AtomicU64::new(u64::MAX))
569 }
570
571 fn get_or_compute(&self, key: usize, compute: impl FnOnce() -> f32) -> f32 {
572 use std::sync::atomic::Ordering::Relaxed;
573 let key = key as u64;
575 let entry = self.0.load(Relaxed);
576 if entry >> 32 == key {
577 return f32::from_bits(entry as u32);
578 }
579 let score = compute();
580 self.0
581 .store((key << 32) | u64::from(score.to_bits()), Relaxed);
582 score
583 }
584}
585
586macro_rules! cursor_ensure_block {
594 ($self:ident, $load_block_fn:ident, $($aw:tt)*) => {{
595 if $self.exhausted || $self.block_loaded {
596 return Ok(!$self.exhausted);
597 }
598 match &mut $self.variant {
599 CursorVariant::Text {
600 list,
601 deferred_tf,
602 ..
603 } => {
604 if let Some(state) = list.decode_block_doc_ids_only($self.block_idx, &mut $self.doc_ids) {
605 *deferred_tf = Some(state);
606 $self.scores.clear();
607 $self.pos = 0;
608 $self.block_loaded = true;
609 Ok(true)
610 } else {
611 $self.exhausted = true;
612 Ok(false)
613 }
614 }
615 CursorVariant::Sparse {
616 si,
617 query_weight,
618 skip_start,
619 block_data_offset,
620 ..
621 } => {
622 let block = si
623 .$load_block_fn(*skip_start, *block_data_offset, $self.block_idx)
624 $($aw)* ?;
625 match block {
626 Some(b) => {
627 b.decode_doc_ids_into(&mut $self.doc_ids);
628 b.decode_scored_weights_into(*query_weight, &mut $self.scores);
629 if $self.lazy_ordinals {
630 $self.current_sparse_block = Some(b);
633 $self.ordinals_loaded = false;
634 } else {
635 b.decode_ordinals_into(&mut $self.ordinals);
636 $self.ordinals_loaded = true;
637 $self.current_sparse_block = None;
638 }
639 $self.pos = 0;
640 $self.block_loaded = true;
641 Ok(true)
642 }
643 None => {
644 $self.exhausted = true;
645 Ok(false)
646 }
647 }
648 }
649 }
650 }};
651}
652
653macro_rules! cursor_advance {
654 ($self:ident, $ensure_fn:ident, $($aw:tt)*) => {{
655 if $self.exhausted {
656 return Ok(u32::MAX);
657 }
658 $self.$ensure_fn() $($aw)* ?;
659 if $self.exhausted {
660 return Ok(u32::MAX);
661 }
662 Ok($self.advance_pos())
663 }};
664}
665
666macro_rules! cursor_seek {
667 ($self:ident, $ensure_fn:ident, $target:expr, $($aw:tt)*) => {{
668 if let Some(doc) = $self.seek_prepare($target) {
669 return Ok(doc);
670 }
671 $self.$ensure_fn() $($aw)* ?;
672 if $self.seek_finish($target) {
673 $self.$ensure_fn() $($aw)* ?;
674 }
675 Ok($self.doc())
676 }};
677}
678
679impl<'a> TermCursor<'a> {
680 pub fn text_with_params(
682 posting_list: crate::structures::BlockPostingList,
683 idf: f32,
684 avg_field_len: f32,
685 lengths: Option<LengthSource<'a>>,
686 params: super::Bm25Params,
687 ) -> Self {
688 Self::text_with_lengths(posting_list, idf, avg_field_len, lengths, params)
689 }
690
691 fn text_with_lengths(
692 posting_list: crate::structures::BlockPostingList,
693 idf: f32,
694 avg_field_len: f32,
695 lengths: Option<LengthSource<'a>>,
696 params: super::Bm25Params,
697 ) -> Self {
698 let max_tf = posting_list.max_tf() as f32;
699 let safe_avg = avg_field_len.max(1.0);
700 let length_bounds = lengths.is_some() && posting_list.min_len().is_some();
701 let length_floor = match lengths {
702 Some(LengthSource::Chunks(map)) => map.length_floor(),
703 _ => 0,
704 };
705 let max_score = match posting_list.min_len() {
706 Some(min_len) if length_bounds => params.upper_bound_with_len(
707 max_tf.max(1.0),
708 idf,
709 min_len.max(length_floor) as f32,
710 safe_avg,
711 ),
712 _ => params.upper_bound(max_tf.max(1.0), idf),
713 };
714 let num_blocks = posting_list.num_blocks();
715 Self {
716 max_score,
717 num_blocks,
718 block_idx: 0,
719 doc_ids: Vec::with_capacity(128),
720 scores: Vec::with_capacity(128),
721 ordinals: Vec::new(),
722 pos: 0,
723 block_loaded: false,
724 exhausted: num_blocks == 0,
725 lazy_ordinals: false,
726 ordinals_loaded: true, current_sparse_block: None,
728 variant: CursorVariant::Text {
729 list: posting_list,
730 idf,
731 idf_times_k1_plus_1: idf * (params.k1 + 1.0),
732 denom_tf_coeff: 1.0 + params.k1 * (params.b / safe_avg),
733 denom_const: params.k1 * (1.0 - params.b),
734 denom_len_coeff: params.k1 * params.b / safe_avg,
735 lengths,
736 length_bounds,
737 length_floor,
738 block_bound: CachedScoreBound::new(),
739 group_bound: CachedScoreBound::new(),
740 avg_len: safe_avg,
741 params,
742 tfs: Vec::with_capacity(128),
743 deferred_tf: None,
744 },
745 }
746 }
747
748 pub fn sparse(
751 si: &'a crate::segment::SparseIndex,
752 query_weight: f32,
753 skip_start: usize,
754 skip_count: usize,
755 global_max_weight: f32,
756 block_data_offset: u64,
757 ) -> Self {
758 Self {
759 max_score: query_weight.abs() * global_max_weight,
760 num_blocks: skip_count,
761 block_idx: 0,
762 doc_ids: Vec::with_capacity(256),
763 scores: Vec::with_capacity(256),
764 ordinals: Vec::with_capacity(256),
765 pos: 0,
766 block_loaded: false,
767 exhausted: skip_count == 0,
768 lazy_ordinals: false,
769 ordinals_loaded: true,
770 current_sparse_block: None,
771 variant: CursorVariant::Sparse {
772 si,
773 query_weight,
774 skip_start,
775 block_data_offset,
776 },
777 }
778 }
779
780 #[inline]
783 fn block_first_doc(&self, idx: usize) -> DocId {
784 match &self.variant {
785 CursorVariant::Text { list, .. } => list.block_first_doc(idx).unwrap_or(u32::MAX),
786 CursorVariant::Sparse { si, skip_start, .. } => {
787 si.read_skip_entry(*skip_start + idx).first_doc
788 }
789 }
790 }
791
792 #[inline]
793 fn block_last_doc(&self, idx: usize) -> DocId {
794 match &self.variant {
795 CursorVariant::Text { list, .. } => list.block_last_doc(idx).unwrap_or(0),
796 CursorVariant::Sparse { si, skip_start, .. } => {
797 si.read_skip_entry(*skip_start + idx).last_doc
798 }
799 }
800 }
801
802 #[inline]
805 pub fn doc(&self) -> DocId {
806 if self.exhausted {
807 return u32::MAX;
808 }
809 if self.block_loaded {
810 debug_assert!(self.pos < self.doc_ids.len());
811 unsafe { *self.doc_ids.get_unchecked(self.pos) }
813 } else {
814 self.block_first_doc(self.block_idx)
815 }
816 }
817
818 #[inline]
819 pub fn ordinal(&self) -> u16 {
820 if !self.block_loaded || self.ordinals.is_empty() {
821 return 0;
822 }
823 debug_assert!(self.pos < self.ordinals.len());
824 unsafe { *self.ordinals.get_unchecked(self.pos) }
826 }
827
828 #[inline]
834 pub fn ordinal_mut(&mut self) -> u16 {
835 if !self.block_loaded {
836 return 0;
837 }
838 if !self.ordinals_loaded {
839 if let Some(ref block) = self.current_sparse_block {
840 block.decode_ordinals_into(&mut self.ordinals);
841 }
842 self.ordinals_loaded = true;
843 }
844 if self.ordinals.is_empty() {
845 return 0;
846 }
847 debug_assert!(self.pos < self.ordinals.len());
848 unsafe { *self.ordinals.get_unchecked(self.pos) }
849 }
850
851 #[inline]
852 pub fn score(&self) -> f32 {
853 if !self.block_loaded {
854 return 0.0;
855 }
856 debug_assert!(self.pos < self.scores.len());
857 unsafe { *self.scores.get_unchecked(self.pos) }
859 }
860
861 #[inline]
867 pub fn ensure_scores(&mut self) {
868 if self.block_loaded && self.scores.is_empty() {
869 self.compute_deferred_scores();
870 }
871 }
872
873 #[inline]
874 pub fn current_block_max_score(&self) -> f32 {
875 if self.exhausted {
876 return 0.0;
877 }
878 match &self.variant {
879 CursorVariant::Text { .. } => self.text_block_bound(self.block_idx),
880 CursorVariant::Sparse {
881 si,
882 query_weight,
883 skip_start,
884 ..
885 } => query_weight.abs() * si.read_skip_entry(*skip_start + self.block_idx).max_weight,
886 }
887 }
888
889 #[inline]
893 pub fn current_group_max_score(&self) -> Option<f32> {
894 if self.exhausted {
895 return Some(0.0);
896 }
897 match &self.variant {
898 CursorVariant::Text { .. } => self.text_group_bound(self.block_idx),
899 CursorVariant::Sparse { .. } => None,
900 }
901 }
902
903 #[inline]
906 pub(crate) fn is_text(&self) -> bool {
907 matches!(self.variant, CursorVariant::Text { .. })
908 }
909
910 fn text_block_bound(&self, idx: usize) -> f32 {
912 match &self.variant {
913 CursorVariant::Text {
914 list,
915 idf,
916 length_bounds,
917 length_floor,
918 block_bound,
919 avg_len,
920 params,
921 ..
922 } => block_bound.get_or_compute(idx, || {
923 let (max_tf, min_len) = list.block_bounds(idx).unwrap_or((0, None));
924 match min_len {
925 Some(min_len) if *length_bounds => params.upper_bound_with_len(
926 (max_tf as f32).max(1.0),
927 *idf,
928 min_len.max(*length_floor) as f32,
929 *avg_len,
930 ),
931 _ => params.upper_bound((max_tf as f32).max(1.0), *idf),
932 }
933 }),
934 CursorVariant::Sparse { .. } => self.max_score,
935 }
936 }
937
938 fn text_group_bound(&self, idx: usize) -> Option<f32> {
940 match &self.variant {
941 CursorVariant::Text {
942 list,
943 idf,
944 length_bounds,
945 length_floor,
946 group_bound,
947 avg_len,
948 params,
949 ..
950 } => {
951 let (max_tf, min_len) = list.group_bounds(idx)?;
952 Some(group_bound.get_or_compute(list.next_group_block(idx), || {
953 if *length_bounds {
954 params.upper_bound_with_len(
955 (max_tf as f32).max(1.0),
956 *idf,
957 min_len.max(*length_floor) as f32,
958 *avg_len,
959 )
960 } else {
961 params.upper_bound((max_tf as f32).max(1.0), *idf)
962 }
963 }))
964 }
965 CursorVariant::Sparse { .. } => None,
966 }
967 }
968
969 pub(crate) fn window_upper_bound(&self, from: DocId, to: DocId) -> f32 {
975 if self.exhausted {
976 return 0.0;
977 }
978 let CursorVariant::Text { list, .. } = &self.variant else {
979 return self.max_score;
980 };
981 let start = from.max(self.doc());
984 if start > to {
985 return 0.0;
986 }
987 let Some(mut idx) = list.seek_block(start, self.block_idx) else {
988 return 0.0;
989 };
990 let mut bound = 0.0f32;
991 while idx < self.num_blocks {
992 if list.block_first_doc(idx).unwrap_or(u32::MAX) > to {
993 break;
994 }
995 if list.is_group_start(idx)
996 && list.group_last_doc(idx).is_some_and(|last| last <= to)
997 && let Some(group_bound) = self.text_group_bound(idx)
998 {
999 bound = bound.max(group_bound);
1000 idx = list.next_group_block(idx);
1001 continue;
1002 }
1003 bound = bound.max(self.text_block_bound(idx));
1004 idx += 1;
1005 }
1006 bound
1007 }
1008
1009 pub(crate) fn score_window_sync(
1014 &mut self,
1015 from: DocId,
1016 to: DocId,
1017 scores: &mut [f32],
1018 mask: &mut [u64],
1019 ) -> crate::Result<u32> {
1020 let mut matched = 0u32;
1021 loop {
1022 if self.exhausted {
1023 return Ok(matched);
1024 }
1025 if !self.block_loaded {
1026 if self.block_first_doc(self.block_idx) > to {
1027 return Ok(matched);
1028 }
1029 self.ensure_block_loaded_sync()?;
1030 if self.exhausted {
1031 return Ok(matched);
1032 }
1033 }
1034 if self.doc_ids[self.pos] > to {
1035 return Ok(matched);
1036 }
1037 self.ensure_scores();
1038 let remaining = &self.doc_ids[self.pos..];
1039 let end = if to == u32::MAX {
1040 remaining.len()
1041 } else {
1042 crate::structures::simd::find_first_ge_u32(remaining, to + 1)
1043 };
1044 let block_scores = &self.scores[self.pos..self.pos + end];
1045 for (doc, score) in remaining[..end].iter().zip(block_scores) {
1046 let slot = (doc - from) as usize;
1047 scores[slot] += score;
1048 mask[slot >> 6] |= 1u64 << (slot & 63);
1049 }
1050 matched += end as u32;
1051 self.pos += end;
1052 if self.pos >= self.doc_ids.len() {
1053 self.block_idx += 1;
1054 self.block_loaded = false;
1055 if self.block_idx >= self.num_blocks {
1056 self.exhausted = true;
1057 return Ok(matched);
1058 }
1059 } else {
1060 return Ok(matched);
1061 }
1062 }
1063 }
1064
1065 pub(crate) fn skip_past_sync(&mut self, to: DocId) -> crate::Result<()> {
1068 if to == u32::MAX {
1069 self.exhausted = true;
1070 return Ok(());
1071 }
1072 while !self.exhausted && self.block_last_doc(self.block_idx) <= to {
1073 self.skip_to_next_block();
1074 }
1075 if !self.exhausted && self.doc() <= to {
1076 self.seek_sync(to + 1)?;
1077 }
1078 Ok(())
1079 }
1080
1081 #[inline]
1083 pub fn current_group_last_doc(&self) -> DocId {
1084 match &self.variant {
1085 CursorVariant::Text { list, .. } => list.group_last_doc(self.block_idx).unwrap_or(0),
1086 CursorVariant::Sparse { .. } => self.block_last_doc(self.block_idx),
1087 }
1088 }
1089
1090 pub fn skip_to_next_group(&mut self) -> DocId {
1092 if self.exhausted {
1093 return u32::MAX;
1094 }
1095 let next = match &self.variant {
1096 CursorVariant::Text { list, .. } => list.next_group_block(self.block_idx),
1097 CursorVariant::Sparse { .. } => self.block_idx + 1,
1098 };
1099 self.block_idx = next;
1100 self.block_loaded = false;
1101 if self.block_idx >= self.num_blocks {
1102 self.exhausted = true;
1103 return u32::MAX;
1104 }
1105 self.block_first_doc(self.block_idx)
1106 }
1107
1108 pub fn skip_to_next_block(&mut self) -> DocId {
1111 if self.exhausted {
1112 return u32::MAX;
1113 }
1114 self.block_idx += 1;
1115 self.block_loaded = false;
1116 if self.block_idx >= self.num_blocks {
1117 self.exhausted = true;
1118 return u32::MAX;
1119 }
1120 self.block_first_doc(self.block_idx)
1121 }
1122
1123 #[inline]
1124 fn advance_pos(&mut self) -> DocId {
1125 self.pos += 1;
1126 if self.pos >= self.doc_ids.len() {
1127 self.block_idx += 1;
1128 self.block_loaded = false;
1129 if self.block_idx >= self.num_blocks {
1130 self.exhausted = true;
1131 return u32::MAX;
1132 }
1133 }
1134 self.doc()
1135 }
1136
1137 #[inline(never)]
1139 fn compute_deferred_scores(&mut self) {
1140 if let CursorVariant::Text {
1141 list,
1142 idf_times_k1_plus_1,
1143 denom_tf_coeff,
1144 denom_const,
1145 denom_len_coeff,
1146 lengths,
1147 tfs,
1148 deferred_tf,
1149 ..
1150 } = &mut self.variant
1151 && let Some((block_offset, tf_start, count)) = deferred_tf.take()
1152 {
1153 list.decode_block_tfs_deferred(block_offset, tf_start, count, tfs);
1154 let num_scale = *idf_times_k1_plus_1;
1155 let d_tf = *denom_tf_coeff;
1156 let d_const = *denom_const;
1157 let d_len = *denom_len_coeff;
1158 self.scores.clear();
1159 self.scores.resize(count, 0.0);
1160 match lengths {
1161 Some(source) => {
1163 for i in 0..count {
1164 let tf = unsafe { *tfs.get_unchecked(i) } as f32;
1165 let vid = unsafe { *self.doc_ids.get_unchecked(i) };
1166 let len = source.length(vid) as f32;
1167 let score = (num_scale * tf) / (tf + d_const + d_len * len);
1168 unsafe {
1169 *self.scores.get_unchecked_mut(i) = score;
1170 }
1171 }
1172 }
1173 None => {
1174 for i in 0..count {
1175 let tf = unsafe { *tfs.get_unchecked(i) } as f32;
1176 let score = (num_scale * tf) / (d_tf * tf + d_const);
1177 unsafe {
1178 *self.scores.get_unchecked_mut(i) = score;
1179 }
1180 }
1181 }
1182 }
1183 }
1184 }
1185
1186 pub async fn ensure_block_loaded(&mut self) -> crate::Result<bool> {
1192 cursor_ensure_block!(self, load_block_direct, .await)
1193 }
1194
1195 pub fn ensure_block_loaded_sync(&mut self) -> crate::Result<bool> {
1196 cursor_ensure_block!(self, load_block_direct_sync,)
1197 }
1198
1199 pub async fn advance(&mut self) -> crate::Result<DocId> {
1200 cursor_advance!(self, ensure_block_loaded, .await)
1201 }
1202
1203 pub fn advance_sync(&mut self) -> crate::Result<DocId> {
1204 cursor_advance!(self, ensure_block_loaded_sync,)
1205 }
1206
1207 pub async fn seek(&mut self, target: DocId) -> crate::Result<DocId> {
1208 cursor_seek!(self, ensure_block_loaded, target, .await)
1209 }
1210
1211 pub fn seek_sync(&mut self, target: DocId) -> crate::Result<DocId> {
1212 cursor_seek!(self, ensure_block_loaded_sync, target,)
1213 }
1214
1215 fn seek_prepare(&mut self, target: DocId) -> Option<DocId> {
1216 if self.exhausted {
1217 return Some(u32::MAX);
1218 }
1219
1220 if self.block_loaded
1222 && let Some(&last) = self.doc_ids.last()
1223 {
1224 if last >= target && self.doc_ids[self.pos] < target {
1225 let remaining = &self.doc_ids[self.pos..];
1226 self.pos += crate::structures::simd::find_first_ge_u32(remaining, target);
1227 if self.pos >= self.doc_ids.len() {
1228 self.block_idx += 1;
1229 self.block_loaded = false;
1230 if self.block_idx >= self.num_blocks {
1231 self.exhausted = true;
1232 return Some(u32::MAX);
1233 }
1234 }
1235 return Some(self.doc());
1236 }
1237 if self.doc_ids[self.pos] >= target {
1238 return Some(self.doc());
1239 }
1240 }
1241
1242 let lo = match &self.variant {
1244 CursorVariant::Text { list, .. } => match list.seek_block(target, self.block_idx) {
1246 Some(idx) => idx,
1247 None => {
1248 self.exhausted = true;
1249 return Some(u32::MAX);
1250 }
1251 },
1252 CursorVariant::Sparse { .. } => {
1254 let mut lo = self.block_idx;
1255 let mut hi = self.num_blocks;
1256 while lo < hi {
1257 let mid = lo + (hi - lo) / 2;
1258 if self.block_last_doc(mid) < target {
1259 lo = mid + 1;
1260 } else {
1261 hi = mid;
1262 }
1263 }
1264 lo
1265 }
1266 };
1267 if lo >= self.num_blocks {
1268 self.exhausted = true;
1269 return Some(u32::MAX);
1270 }
1271 if lo != self.block_idx || !self.block_loaded {
1272 self.block_idx = lo;
1273 self.block_loaded = false;
1274 }
1275 None
1276 }
1277
1278 #[inline]
1279 fn seek_finish(&mut self, target: DocId) -> bool {
1280 if self.exhausted {
1281 return false;
1282 }
1283 self.pos = crate::structures::simd::find_first_ge_u32(&self.doc_ids, target);
1284 if self.pos >= self.doc_ids.len() {
1285 self.block_idx += 1;
1286 self.block_loaded = false;
1287 if self.block_idx >= self.num_blocks {
1288 self.exhausted = true;
1289 return false;
1290 }
1291 return true;
1292 }
1293 false
1294 }
1295}
1296
1297macro_rules! bms_execute_loop {
1302 ($self:ident, $ensure:ident, $advance:ident, $seek:ident, $($aw:tt)*) => {{
1303 let n = $self.cursors.len();
1304
1305 for cursor in &mut $self.cursors {
1307 cursor.$ensure() $($aw)* ?;
1308 }
1309
1310 let mut docs_scored = 0u64;
1311 let mut docs_skipped = 0u64;
1312 let mut blocks_skipped = 0u64;
1313 let mut groups_skipped = 0u64;
1314 let mut conjunction_skipped = 0u64;
1315 let mut ordinal_scores: Vec<(u16, f32)> = Vec::with_capacity(n * 2);
1316 let _bms_start = std::time::Instant::now();
1317
1318 let inv_heap_factor = $self.inv_heap_factor;
1319 let mut adjusted_threshold = $self.collector.threshold() * inv_heap_factor - 1e-6;
1320 let mut iterations: u64 = 0;
1321
1322 loop {
1323 iterations += 1;
1327 if iterations & 0xFFF == 0
1328 && let Some(budget) = &$self.budget
1329 && budget.expired()
1330 {
1331 budget.mark_truncated();
1332 log::debug!(
1333 "MaxScoreExecutor: deadline reached after {} iterations, {} scored",
1334 iterations,
1335 docs_scored
1336 );
1337 break;
1338 }
1339 let partition = $self.find_partition();
1340 if partition >= n {
1341 break;
1342 }
1343
1344 let mut min_doc = u32::MAX;
1348 let mut next_other = u32::MAX;
1352 let mut at_min_mask = 0u64; for i in partition..n {
1354 let doc = $self.cursors[i].doc();
1355 match doc.cmp(&min_doc) {
1356 std::cmp::Ordering::Less => {
1357 next_other = min_doc;
1358 min_doc = doc;
1359 at_min_mask = 1u64 << (i as u32);
1360 }
1361 std::cmp::Ordering::Equal => {
1362 at_min_mask |= 1u64 << (i as u32);
1363 }
1364 std::cmp::Ordering::Greater => {
1365 if doc < next_other {
1366 next_other = doc;
1367 }
1368 }
1369 }
1370 }
1371 if min_doc == u32::MAX {
1372 break;
1373 }
1374
1375 let non_essential_upper = if partition > 0 {
1376 $self.prefix_sums[partition - 1]
1377 } else {
1378 0.0
1379 };
1380
1381 if $self.collector.len() >= $self.collector.k {
1383 let mut present_upper: f32 = 0.0;
1384 let mut mask = at_min_mask;
1385 while mask != 0 {
1386 let i = mask.trailing_zeros() as usize;
1387 present_upper += $self.cursors[i].max_score;
1388 mask &= mask - 1;
1389 }
1390
1391 if present_upper + non_essential_upper < adjusted_threshold {
1392 let mut mask = at_min_mask;
1393 while mask != 0 {
1394 let i = mask.trailing_zeros() as usize;
1395 $self.cursors[i].$ensure() $($aw)* ?;
1396 $self.cursors[i].$advance() $($aw)* ?;
1397 mask &= mask - 1;
1398 }
1399 conjunction_skipped += 1;
1400 continue;
1401 }
1402 }
1403
1404 if $self.collector.len() >= $self.collector.k {
1406 let mut block_max_sum: f32 = 0.0;
1407 let mut mask = at_min_mask;
1408 while mask != 0 {
1409 let i = mask.trailing_zeros() as usize;
1410 block_max_sum += $self.cursors[i].current_block_max_score();
1411 mask &= mask - 1;
1412 }
1413
1414 if block_max_sum + non_essential_upper < adjusted_threshold {
1415 let mut group_sum: f32 = 0.0;
1431 let mut mask = at_min_mask;
1432 while mask != 0 {
1433 let i = mask.trailing_zeros() as usize;
1434 group_sum += $self.cursors[i]
1435 .current_group_max_score()
1436 .unwrap_or_else(|| $self.cursors[i].current_block_max_score());
1437 mask &= mask - 1;
1438 }
1439 let group_prunable = group_sum + non_essential_upper < adjusted_threshold;
1440 let mut mask = at_min_mask;
1441 while mask != 0 {
1442 let i = mask.trailing_zeros() as usize;
1443 let by_group =
1444 group_prunable && $self.cursors[i].current_group_max_score().is_some();
1445 let boundary = if by_group {
1446 $self.cursors[i].current_group_last_doc()
1447 } else {
1448 $self.cursors[i].block_last_doc($self.cursors[i].block_idx)
1449 };
1450 if next_other > boundary {
1451 if by_group {
1452 $self.cursors[i].skip_to_next_group();
1453 groups_skipped += 1;
1454 } else {
1455 $self.cursors[i].skip_to_next_block();
1456 }
1457 $self.cursors[i].$ensure() $($aw)* ?;
1458 } else {
1459 $self.cursors[i].$seek(next_other) $($aw)* ?;
1460 }
1461 mask &= mask - 1;
1462 }
1463 blocks_skipped += 1;
1464 continue;
1465 }
1466 }
1467
1468 if let Some(ref pred) = $self.predicate {
1470 if !pred(min_doc) {
1471 let mut mask = at_min_mask;
1472 while mask != 0 {
1473 let i = mask.trailing_zeros() as usize;
1474 $self.cursors[i].$ensure() $($aw)* ?;
1475 $self.cursors[i].$advance() $($aw)* ?;
1476 mask &= mask - 1;
1477 }
1478 continue;
1479 }
1480 }
1481
1482 ordinal_scores.clear();
1484 {
1485 let mut mask = at_min_mask;
1486 while mask != 0 {
1487 let i = mask.trailing_zeros() as usize;
1488 $self.cursors[i].$ensure() $($aw)* ?;
1489 $self.cursors[i].ensure_scores();
1490 while $self.cursors[i].doc() == min_doc {
1491 let ord = $self.cursors[i].ordinal_mut();
1492 let sc = $self.cursors[i].score();
1493 ordinal_scores.push((ord, sc));
1494 $self.cursors[i].$advance() $($aw)* ?;
1495 }
1496 mask &= mask - 1;
1497 }
1498 }
1499
1500 let essential_total: f32 = ordinal_scores.iter().map(|(_, s)| *s).sum();
1501 if $self.collector.len() >= $self.collector.k
1502 && essential_total + non_essential_upper < adjusted_threshold
1503 {
1504 docs_skipped += 1;
1505 continue;
1506 }
1507
1508 let mut running_total = essential_total;
1510 for i in (0..partition).rev() {
1511 if $self.collector.len() >= $self.collector.k
1512 && running_total + $self.prefix_sums[i] < adjusted_threshold
1513 {
1514 break;
1515 }
1516
1517 let doc = $self.cursors[i].$seek(min_doc) $($aw)* ?;
1518 if doc == min_doc {
1519 $self.cursors[i].ensure_scores();
1520 while $self.cursors[i].doc() == min_doc {
1521 let s = $self.cursors[i].score();
1522 running_total += s;
1523 let ord = $self.cursors[i].ordinal_mut();
1524 ordinal_scores.push((ord, s));
1525 $self.cursors[i].$advance() $($aw)* ?;
1526 }
1527 }
1528 }
1529
1530 if ordinal_scores.len() == 1 {
1533 let (ord, score) = ordinal_scores[0];
1534 if $self.collector.insert_with_ordinal(min_doc, score, ord) {
1535 docs_scored += 1;
1536 adjusted_threshold = $self.collector.threshold() * inv_heap_factor - 1e-6;
1537 } else {
1538 docs_skipped += 1;
1539 }
1540 } else if !ordinal_scores.is_empty() {
1541 if ordinal_scores.len() > 2 {
1542 ordinal_scores.sort_unstable_by_key(|(ord, _)| *ord);
1543 } else if ordinal_scores.len() == 2 && ordinal_scores[0].0 > ordinal_scores[1].0 {
1544 ordinal_scores.swap(0, 1);
1545 }
1546 let mut j = 0;
1547 while j < ordinal_scores.len() {
1548 let current_ord = ordinal_scores[j].0;
1549 let mut score = 0.0f32;
1550 while j < ordinal_scores.len() && ordinal_scores[j].0 == current_ord {
1551 score += ordinal_scores[j].1;
1552 j += 1;
1553 }
1554 if $self
1555 .collector
1556 .insert_with_ordinal(min_doc, score, current_ord)
1557 {
1558 docs_scored += 1;
1559 adjusted_threshold = $self.collector.threshold() * inv_heap_factor - 1e-6;
1560 } else {
1561 docs_skipped += 1;
1562 }
1563 }
1564 }
1565 }
1566
1567 let results: Vec<ScoredDoc> = $self
1568 .collector
1569 .into_sorted_results()
1570 .into_iter()
1571 .map(|(doc_id, score, ordinal)| ScoredDoc {
1572 doc_id,
1573 score,
1574 ordinal,
1575 })
1576 .collect();
1577
1578 let _bms_elapsed_ms = _bms_start.elapsed().as_millis() as u64;
1579 if _bms_elapsed_ms > 500 {
1580 warn!(
1581 "slow MaxScore: {}ms, cursors={}, scored={}, skipped={}, blocks_skipped={}, groups_skipped={}, conjunction_skipped={}, returned={}, top_score={:.4}",
1582 _bms_elapsed_ms,
1583 n,
1584 docs_scored,
1585 docs_skipped,
1586 blocks_skipped,
1587 groups_skipped,
1588 conjunction_skipped,
1589 results.len(),
1590 results.first().map(|r| r.score).unwrap_or(0.0)
1591 );
1592 } else {
1593 debug!(
1594 "MaxScoreExecutor: {}ms, scored={}, skipped={}, blocks_skipped={}, groups_skipped={}, conjunction_skipped={}, returned={}, top_score={:.4}",
1595 _bms_elapsed_ms,
1596 docs_scored,
1597 docs_skipped,
1598 blocks_skipped,
1599 groups_skipped,
1600 conjunction_skipped,
1601 results.len(),
1602 results.first().map(|r| r.score).unwrap_or(0.0)
1603 );
1604 }
1605
1606 Ok(results)
1607 }};
1608}
1609
1610impl<'a> MaxScoreExecutor<'a> {
1611 pub(crate) fn new(mut cursors: Vec<TermCursor<'a>>, k: usize, heap_factor: f32) -> Self {
1616 if cursors.len() > super::MAX_QUERY_TERMS {
1620 cursors.sort_unstable_by(|a, b| b.max_score.total_cmp(&a.max_score));
1621 cursors.truncate(super::MAX_QUERY_TERMS);
1622 log::warn!(
1623 "MaxScore cursor count exceeded {}; retaining the strongest cursors",
1624 super::MAX_QUERY_TERMS
1625 );
1626 }
1627
1628 for c in &mut cursors {
1631 c.lazy_ordinals = true;
1632 }
1633
1634 cursors.sort_by(|a, b| {
1636 a.max_score
1637 .partial_cmp(&b.max_score)
1638 .unwrap_or(Ordering::Equal)
1639 });
1640
1641 let mut prefix_sums = Vec::with_capacity(cursors.len());
1642 let mut cumsum = 0.0f32;
1643 for c in &cursors {
1644 cumsum += c.max_score;
1645 prefix_sums.push(cumsum);
1646 }
1647
1648 let clamped_heap_factor = heap_factor.clamp(0.01, 1.0);
1649
1650 debug!(
1651 "Creating MaxScoreExecutor: num_cursors={}, k={}, total_upper={:.4}, heap_factor={:.2}",
1652 cursors.len(),
1653 k,
1654 cumsum,
1655 clamped_heap_factor
1656 );
1657
1658 Self {
1659 cursors,
1660 prefix_sums,
1661 collector: ScoreCollector::new(k),
1662 inv_heap_factor: 1.0 / clamped_heap_factor,
1663 predicate: None,
1664 budget: None,
1665 metric_index: "unknown",
1666 metric_field: "unknown",
1667 }
1668 }
1669
1670 pub fn with_budget(mut self, budget: Option<SharedThreshold>) -> Self {
1672 self.budget = budget.filter(|b| b.deadline().is_some());
1673 self
1674 }
1675
1676 pub fn with_metric_labels(mut self, index: &'a str, field: &'a str) -> Self {
1678 self.metric_index = index;
1679 self.metric_field = field;
1680 self
1681 }
1682
1683 pub fn sparse(
1687 sparse_index: &'a crate::segment::SparseIndex,
1688 query_terms: Vec<(u32, f32)>,
1689 k: usize,
1690 heap_factor: f32,
1691 ) -> Self {
1692 let cursors: Vec<TermCursor<'a>> = query_terms
1693 .iter()
1694 .filter_map(|&(dim_id, qw)| {
1695 let (skip_start, skip_count, global_max, block_data_offset) =
1696 sparse_index.get_skip_range_full(dim_id)?;
1697 Some(TermCursor::sparse(
1698 sparse_index,
1699 qw,
1700 skip_start,
1701 skip_count,
1702 global_max,
1703 block_data_offset,
1704 ))
1705 })
1706 .collect();
1707 Self::new(cursors, k, heap_factor)
1708 }
1709
1710 pub fn text(
1714 posting_lists: Vec<(crate::structures::BlockPostingList, f32)>,
1715 avg_field_len: f32,
1716 k: usize,
1717 lengths: Option<&'a crate::segment::chunk_map::DocLengths>,
1718 params: super::Bm25Params,
1719 heap_factor: f32,
1720 ) -> Self {
1721 let cursors: Vec<TermCursor<'a>> = posting_lists
1722 .into_iter()
1723 .map(|(pl, idf)| {
1724 TermCursor::text_with_params(
1725 pl,
1726 idf,
1727 avg_field_len,
1728 lengths.map(LengthSource::Docs),
1729 params,
1730 )
1731 })
1732 .collect();
1733 Self::new(cursors, k, heap_factor)
1734 }
1735
1736 pub fn text_chunked(
1740 posting_lists: Vec<(crate::structures::BlockPostingList, f32)>,
1741 avg_chunk_len: f32,
1742 k: usize,
1743 lengths: &'a crate::segment::chunk_map::ChunkMap,
1744 params: super::Bm25Params,
1745 heap_factor: f32,
1746 ) -> Self {
1747 let cursors: Vec<TermCursor<'a>> = posting_lists
1748 .into_iter()
1749 .map(|(pl, idf)| {
1750 TermCursor::text_with_params(
1751 pl,
1752 idf,
1753 avg_chunk_len,
1754 Some(LengthSource::Chunks(lengths)),
1755 params,
1756 )
1757 })
1758 .collect();
1759 Self::new(cursors, k, heap_factor)
1760 }
1761
1762 #[inline]
1763 fn find_partition(&self) -> usize {
1764 let threshold = self.collector.threshold() * self.inv_heap_factor;
1768 self.prefix_sums.partition_point(|&sum| sum < threshold)
1771 }
1772
1773 pub fn with_predicate(mut self, predicate: super::DocPredicate<'a>) -> Self {
1779 self.predicate = Some(predicate);
1780 self
1781 }
1782
1783 pub fn seed_threshold(&mut self, initial_threshold: f32) {
1785 self.collector.seed_threshold(initial_threshold);
1786 }
1787
1788 pub async fn execute(mut self) -> crate::Result<Vec<ScoredDoc>> {
1794 if self.cursors.is_empty() {
1795 return Ok(Vec::new());
1796 }
1797 let t = crate::observe::Timer::start();
1798 let results = if self.all_text() {
1799 self.execute_windowed()
1800 } else {
1801 bms_execute_loop!(self, ensure_block_loaded, advance, seek, .await)
1802 };
1803 if let Ok(r) = &results {
1804 crate::observe::maxscore_query(self.metric_index, self.metric_field, t.secs(), r.len());
1805 }
1806 results
1807 }
1808
1809 pub fn execute_sync(mut self) -> crate::Result<Vec<ScoredDoc>> {
1811 if self.cursors.is_empty() {
1812 return Ok(Vec::new());
1813 }
1814 let t = crate::observe::Timer::start();
1815 let results = if self.all_text() {
1816 self.execute_windowed()
1817 } else {
1818 bms_execute_loop!(self, ensure_block_loaded_sync, advance_sync, seek_sync,)
1819 };
1820 if let Ok(r) = &results {
1821 crate::observe::maxscore_query(self.metric_index, self.metric_field, t.secs(), r.len());
1822 }
1823 results
1824 }
1825
1826 #[cfg(test)]
1829 pub(crate) fn execute_doc_at_a_time_sync(mut self) -> crate::Result<Vec<ScoredDoc>> {
1830 if self.cursors.is_empty() {
1831 return Ok(Vec::new());
1832 }
1833 bms_execute_loop!(self, ensure_block_loaded_sync, advance_sync, seek_sync,)
1834 }
1835
1836 fn all_text(&self) -> bool {
1837 self.cursors.iter().all(TermCursor::is_text)
1838 }
1839
1840 pub(crate) fn execute_windowed(&mut self) -> crate::Result<Vec<ScoredDoc>> {
1868 if self
1869 .budget
1870 .as_ref()
1871 .is_some_and(SharedThreshold::stop_if_expired)
1872 {
1873 return Ok(Vec::new());
1874 }
1875 let n = self.cursors.len();
1876 for cursor in &mut self.cursors {
1877 cursor.ensure_block_loaded_sync()?;
1878 }
1879 let inv_heap_factor = self.inv_heap_factor;
1880 let mut window_scores = vec![0.0f32; WINDOW_IDS];
1881 let mut window_mask = vec![0u64; WINDOW_IDS / 64];
1882 let mut cand_docs: Vec<u32> = Vec::with_capacity(WINDOW_IDS);
1883 let mut cand_scores: Vec<f32> = Vec::with_capacity(WINDOW_IDS);
1884 let mut wmax = vec![0.0f32; n];
1885 let mut order: Vec<usize> = (0..n).collect();
1886 let mut wprefix = vec![0.0f32; n];
1887 let mut windows = 0u64;
1888 let mut windows_skipped = 0u64;
1889 let mut candidates = 0u64;
1890 let mut docs_scored = 0u64;
1891 let started = std::time::Instant::now();
1892
1893 loop {
1894 windows += 1;
1895 if windows & 0x3F == 0
1896 && let Some(budget) = &self.budget
1897 && budget.expired()
1898 {
1899 budget.mark_truncated();
1900 log::debug!(
1901 "MaxScoreExecutor(windowed): deadline reached after {} windows, {} scored",
1902 windows,
1903 docs_scored
1904 );
1905 break;
1906 }
1907 let partition = self.find_partition();
1908 if partition >= n {
1909 break;
1910 }
1911 let mut from = u32::MAX;
1914 let mut to = u32::MAX;
1915 for cursor in &self.cursors[partition..] {
1916 if cursor.exhausted {
1917 continue;
1918 }
1919 from = from.min(cursor.doc());
1920 to = to.min(cursor.block_last_doc(cursor.block_idx));
1921 }
1922 if from == u32::MAX {
1923 break;
1924 }
1925 let to = to.max(from).min(from.saturating_add(WINDOW_IDS as u32 - 1));
1926 let width = (to - from) as usize + 1;
1927 let words = width.div_ceil(64);
1928
1929 let heap_full = self.collector.len() >= self.collector.k;
1931 let threshold = if heap_full {
1932 self.collector.threshold() * inv_heap_factor - 1e-6
1933 } else {
1934 0.0
1935 };
1936 for (i, bound) in wmax.iter_mut().enumerate() {
1937 *bound = self.cursors[i].window_upper_bound(from, to);
1938 }
1939 order.sort_unstable_by(|&a, &b| wmax[a].total_cmp(&wmax[b]));
1940 let mut sum = 0.0f32;
1941 for (rank, &i) in order.iter().enumerate() {
1942 sum += wmax[i];
1943 wprefix[rank] = sum;
1944 }
1945 let wpartition = if heap_full {
1946 wprefix.partition_point(|&s| s < threshold)
1947 } else {
1948 0
1949 };
1950 if wpartition >= n {
1951 for cursor in &mut self.cursors {
1953 if !cursor.exhausted && cursor.doc() <= to {
1954 cursor.skip_past_sync(to)?;
1955 }
1956 }
1957 windows_skipped += 1;
1958 continue;
1959 }
1960
1961 window_scores[..width].fill(0.0);
1963 window_mask[..words].fill(0);
1964 for &i in &order[wpartition..] {
1965 let cursor = &mut self.cursors[i];
1966 if cursor.exhausted {
1967 continue;
1968 }
1969 if cursor.doc() < from {
1970 cursor.seek_sync(from)?;
1971 }
1972 if cursor.exhausted || cursor.doc() > to {
1973 continue;
1974 }
1975 cursor.score_window_sync(
1976 from,
1977 to,
1978 &mut window_scores[..width],
1979 &mut window_mask[..words],
1980 )?;
1981 }
1982
1983 cand_docs.clear();
1985 cand_scores.clear();
1986 for (word_idx, word) in window_mask[..words].iter().enumerate() {
1987 let mut bits = *word;
1988 while bits != 0 {
1989 let slot = (word_idx << 6) | bits.trailing_zeros() as usize;
1990 bits &= bits - 1;
1991 cand_docs.push(from + slot as u32);
1992 cand_scores.push(window_scores[slot]);
1993 }
1994 }
1995 if let Some(pred) = &self.predicate {
1996 let mut kept = 0usize;
1997 for j in 0..cand_docs.len() {
1998 let doc = cand_docs[j];
1999 cand_docs[kept] = doc;
2000 cand_scores[kept] = cand_scores[j];
2001 kept += pred(doc) as usize;
2002 }
2003 cand_docs.truncate(kept);
2004 cand_scores.truncate(kept);
2005 }
2006
2007 let mut remaining = if wpartition > 0 {
2009 wprefix[wpartition - 1]
2010 } else {
2011 0.0
2012 };
2013 for rank in (0..wpartition).rev() {
2014 let i = order[rank];
2015 if heap_full {
2016 filter_competitive(&mut cand_docs, &mut cand_scores, remaining, threshold);
2017 }
2018 if cand_docs.is_empty() {
2019 break;
2020 }
2021 if wmax[i] > 0.0 {
2022 let cursor = &mut self.cursors[i];
2023 for (doc, score) in cand_docs.iter().zip(cand_scores.iter_mut()) {
2024 if cursor.seek_sync(*doc)? == *doc {
2025 cursor.ensure_scores();
2026 *score += cursor.score();
2027 }
2028 }
2029 }
2030 remaining -= wmax[i];
2031 }
2032 if heap_full {
2033 filter_competitive(&mut cand_docs, &mut cand_scores, 0.0, threshold);
2034 }
2035 candidates += cand_docs.len() as u64;
2036 for (doc, score) in cand_docs.iter().zip(&cand_scores) {
2037 if self.collector.insert_with_ordinal(*doc, *score, 0) {
2038 docs_scored += 1;
2039 }
2040 }
2041 }
2042
2043 let collector = std::mem::replace(&mut self.collector, ScoreCollector::new(0));
2044 let results: Vec<ScoredDoc> = collector
2045 .into_sorted_results()
2046 .into_iter()
2047 .map(|(doc_id, score, ordinal)| ScoredDoc {
2048 doc_id,
2049 score,
2050 ordinal,
2051 })
2052 .collect();
2053 let elapsed_ms = started.elapsed().as_millis() as u64;
2054 if elapsed_ms > 500 {
2055 warn!(
2056 "slow windowed MaxScore: {}ms, cursors={}, windows={}, windows_skipped={}, candidates={}, scored={}, returned={}, top_score={:.4}",
2057 elapsed_ms,
2058 n,
2059 windows,
2060 windows_skipped,
2061 candidates,
2062 docs_scored,
2063 results.len(),
2064 results.first().map(|r| r.score).unwrap_or(0.0)
2065 );
2066 } else {
2067 debug!(
2068 "MaxScoreExecutor(windowed): {}ms, cursors={}, windows={}, windows_skipped={}, candidates={}, scored={}, returned={}, top_score={:.4}",
2069 elapsed_ms,
2070 n,
2071 windows,
2072 windows_skipped,
2073 candidates,
2074 docs_scored,
2075 results.len(),
2076 results.first().map(|r| r.score).unwrap_or(0.0)
2077 );
2078 }
2079 Ok(results)
2080 }
2081}
2082
2083const WINDOW_IDS: usize = 4096;
2085
2086fn filter_competitive(docs: &mut Vec<u32>, scores: &mut Vec<f32>, remaining: f32, threshold: f32) {
2090 let mut kept = 0usize;
2091 for j in 0..docs.len() {
2092 let doc = docs[j];
2093 let score = scores[j];
2094 docs[kept] = doc;
2095 scores[kept] = score;
2096 kept += (score + remaining >= threshold) as usize;
2097 }
2098 docs.truncate(kept);
2099 scores.truncate(kept);
2100}
2101
2102#[cfg(test)]
2103mod tests {
2104 use super::*;
2105
2106 struct Corpus {
2109 postings: Vec<Vec<(u32, u32)>>,
2111 lengths: Vec<u16>,
2112 n_docs: u32,
2113 }
2114
2115 fn xorshift(state: &mut u64) -> u64 {
2116 *state ^= *state << 13;
2117 *state ^= *state >> 7;
2118 *state ^= *state << 17;
2119 *state
2120 }
2121
2122 fn random_corpus(seed: u64, n_docs: u32, n_terms: usize) -> Corpus {
2125 let mut state = seed.wrapping_mul(0x9E37_79B9_7F4A_7C15) | 1;
2126 let lengths: Vec<u16> = (0..n_docs)
2127 .map(|_| 1 + (xorshift(&mut state) % 400) as u16)
2128 .collect();
2129 let densities = [0.6, 0.25, 0.1, 0.03, 0.005];
2130 let postings = (0..n_terms)
2131 .map(|t| {
2132 let density = densities[t % densities.len()];
2133 let cutoff = (density * u32::MAX as f64) as u64;
2134 let mut postings = Vec::new();
2135 for doc in 0..n_docs {
2136 if (xorshift(&mut state) & 0xFFFF_FFFF) >= cutoff {
2137 continue;
2138 }
2139 let r = xorshift(&mut state) % 100;
2140 let tf = if r < 70 {
2141 1
2142 } else if r < 90 {
2143 2
2144 } else {
2145 3 + (r % 6) as u32
2146 };
2147 postings.push((doc, tf));
2148 }
2149 postings
2150 })
2151 .collect();
2152 Corpus {
2153 postings,
2154 lengths,
2155 n_docs,
2156 }
2157 }
2158
2159 fn build_lists(
2160 corpus: &Corpus,
2161 lengths: Option<&crate::segment::chunk_map::DocLengths>,
2162 ) -> Vec<(crate::structures::BlockPostingList, f32)> {
2163 corpus
2164 .postings
2165 .iter()
2166 .map(|postings| {
2167 let mut list = crate::structures::PostingList::new();
2168 for &(doc, tf) in postings {
2169 list.push(doc, tf);
2170 }
2171 let length_of = lengths.map(|l| move |doc: DocId| l.length(doc));
2172 let block_list = crate::structures::BlockPostingList::from_posting_list_with(
2173 &list,
2174 false,
2175 length_of.as_ref().map(|f| f as &dyn Fn(DocId) -> u32),
2176 )
2177 .unwrap();
2178 let idf = super::super::bm25_idf(postings.len() as f32, corpus.n_docs as f32);
2179 (block_list, idf)
2180 })
2181 .collect()
2182 }
2183
2184 fn exhaustive(
2186 corpus: &Corpus,
2187 lists: &[(crate::structures::BlockPostingList, f32)],
2188 real_lengths: bool,
2189 avg: f32,
2190 params: super::super::Bm25Params,
2191 ) -> std::collections::HashMap<u32, f32> {
2192 let mut scores: std::collections::HashMap<u32, f32> = std::collections::HashMap::new();
2193 for (postings, (_, idf)) in corpus.postings.iter().zip(lists) {
2194 for &(doc, tf) in postings {
2195 let len = if real_lengths {
2196 corpus.lengths[doc as usize] as f32
2197 } else {
2198 tf as f32
2199 };
2200 *scores.entry(doc).or_insert(0.0) += params.score(tf as f32, *idf, len, avg);
2201 }
2202 }
2203 scores
2204 }
2205
2206 fn check_top_k(
2207 label: &str,
2208 results: &[ScoredDoc],
2209 exhaustive: &std::collections::HashMap<u32, f32>,
2210 k: usize,
2211 predicate: Option<&dyn Fn(u32) -> bool>,
2212 ) {
2213 let mut expected: Vec<(u32, f32)> = exhaustive
2214 .iter()
2215 .filter(|(doc, _)| predicate.is_none_or(|p| p(**doc)))
2216 .map(|(doc, score)| (*doc, *score))
2217 .collect();
2218 expected.sort_by(|a, b| b.1.total_cmp(&a.1).then(a.0.cmp(&b.0)));
2219 let want = k.min(expected.len());
2220 assert_eq!(results.len(), want, "{label}: result count");
2221 for (rank, (got, exp)) in results.iter().zip(&expected).enumerate() {
2222 let tolerance = 1e-4 * exp.1.abs().max(1.0);
2223 assert!(
2224 (got.score - exp.1).abs() <= tolerance,
2225 "{label}: rank {rank} score {} vs exhaustive {} (doc {} vs {})",
2226 got.score,
2227 exp.1,
2228 got.doc_id,
2229 exp.0
2230 );
2231 let own = exhaustive[&got.doc_id];
2232 assert!(
2233 (got.score - own).abs() <= tolerance,
2234 "{label}: doc {} scored {} but exhaustive says {}",
2235 got.doc_id,
2236 got.score,
2237 own
2238 );
2239 if let Some(p) = predicate {
2240 assert!(
2241 p(got.doc_id),
2242 "{label}: doc {} fails the predicate",
2243 got.doc_id
2244 );
2245 }
2246 }
2247 for pair in results.windows(2) {
2248 assert!(
2249 pair[0].score >= pair[1].score,
2250 "{label}: results not sorted"
2251 );
2252 }
2253 }
2254
2255 #[test]
2260 fn windowed_text_maxscore_matches_exhaustive_and_doc_at_a_time() {
2261 let params = super::super::Bm25Params::default();
2262 let predicate_fn = |doc: u32| !doc.is_multiple_of(3);
2263 let mut cases = 0usize;
2264 for seed in 1..=6u64 {
2265 for &n_docs in &[300u32, 2_500, 12_000] {
2266 for &n_terms in &[1usize, 2, 4, 9] {
2267 let corpus = random_corpus(seed, n_docs, n_terms);
2268 let doc_lengths =
2269 crate::segment::chunk_map::DocLengths::from_lengths(&corpus.lengths);
2270 for real_lengths in [true, false] {
2271 let lengths = real_lengths.then_some(&doc_lengths);
2272 let lists = build_lists(&corpus, lengths);
2273 let avg = if real_lengths {
2274 doc_lengths.avg_len()
2275 } else {
2276 1.0
2277 };
2278 let truth = exhaustive(&corpus, &lists, real_lengths, avg, params);
2279 for &k in &[1usize, 10, 100] {
2280 for with_predicate in [false, true] {
2281 let label = format!(
2282 "seed={seed} docs={n_docs} terms={n_terms} lengths={real_lengths} k={k} pred={with_predicate}"
2283 );
2284 let pred: Option<&dyn Fn(u32) -> bool> =
2285 with_predicate.then_some(&predicate_fn);
2286 let make = |seeded: f32| {
2287 let mut executor = MaxScoreExecutor::text(
2288 lists.clone(),
2289 avg,
2290 k,
2291 lengths,
2292 params,
2293 1.0,
2294 );
2295 if with_predicate {
2296 executor = executor.with_predicate(Box::new(predicate_fn));
2297 }
2298 if seeded > 0.0 {
2299 executor.seed_threshold(seeded);
2300 }
2301 executor
2302 };
2303 let windowed = make(0.0).execute_windowed().unwrap();
2304 check_top_k(
2305 &format!("windowed {label}"),
2306 &windowed,
2307 &truth,
2308 k,
2309 pred,
2310 );
2311 let reference = make(0.0).execute_doc_at_a_time_sync().unwrap();
2312 check_top_k(
2313 &format!("reference {label}"),
2314 &reference,
2315 &truth,
2316 k,
2317 pred,
2318 );
2319 if let Some(kth) = windowed.last().map(|r| r.score)
2321 && windowed.len() == k
2322 {
2323 let seeded = make(kth * 0.9).execute_windowed().unwrap();
2324 check_top_k(
2325 &format!("seeded {label}"),
2326 &seeded,
2327 &truth,
2328 k,
2329 pred,
2330 );
2331 let above =
2333 make(windowed[0].score * 1.5).execute_windowed().unwrap();
2334 assert!(above.is_empty(), "{label}: floor above all scores");
2335 }
2336 cases += 1;
2337 }
2338 }
2339 }
2340 }
2341 }
2342 }
2343 assert!(cases > 400);
2344 }
2345
2346 #[test]
2349 fn windowed_text_maxscore_heap_factor_is_a_subset_with_exact_scores() {
2350 let params = super::super::Bm25Params::default();
2351 let corpus = random_corpus(7, 20_000, 6);
2352 let doc_lengths = crate::segment::chunk_map::DocLengths::from_lengths(&corpus.lengths);
2353 let lists = build_lists(&corpus, Some(&doc_lengths));
2354 let avg = doc_lengths.avg_len();
2355 let truth = exhaustive(&corpus, &lists, true, avg, params);
2356 let exact = MaxScoreExecutor::text(lists.clone(), avg, 50, Some(&doc_lengths), params, 1.0)
2357 .execute_windowed()
2358 .unwrap();
2359 check_top_k("exact", &exact, &truth, 50, None);
2360 let approx = MaxScoreExecutor::text(lists, avg, 50, Some(&doc_lengths), params, 0.6)
2361 .execute_windowed()
2362 .unwrap();
2363 assert_eq!(approx.len(), 50);
2364 for hit in &approx {
2365 let own = truth[&hit.doc_id];
2366 assert!((hit.score - own).abs() <= 1e-4 * own.max(1.0));
2367 }
2368 let exact_kth = exact.last().unwrap().score;
2371 assert!(approx.iter().all(|hit| hit.score >= exact_kth * 0.6 - 1e-4));
2372 assert_eq!(approx[0].doc_id, exact[0].doc_id);
2373 let overlap = approx
2374 .iter()
2375 .filter(|hit| exact.iter().any(|e| e.doc_id == hit.doc_id))
2376 .count();
2377 assert!(overlap >= 25, "overlap {overlap} of 50");
2378 }
2379
2380 #[test]
2381 fn text_heap_factor_below_one_actually_changes_pruning() {
2382 let params = super::super::Bm25Params::default();
2383 let corpus = random_corpus(7, 20_000, 6);
2384 let lengths = crate::segment::chunk_map::DocLengths::from_lengths(&corpus.lengths);
2385 let lists = build_lists(&corpus, Some(&lengths));
2386 let mut exact = MaxScoreExecutor::text(
2387 lists.clone(),
2388 lengths.avg_len(),
2389 50,
2390 Some(&lengths),
2391 params,
2392 1.0,
2393 );
2394 let mut approximate =
2395 MaxScoreExecutor::text(lists, lengths.avg_len(), 50, Some(&lengths), params, 0.01);
2396 let exact = exact.execute_windowed().unwrap();
2397 let approximate = approximate.execute_windowed().unwrap();
2398 assert_ne!(
2399 exact.iter().map(|hit| hit.doc_id).collect::<Vec<_>>(),
2400 approximate.iter().map(|hit| hit.doc_id).collect::<Vec<_>>()
2401 );
2402 }
2403
2404 #[test]
2405 fn text_and_sparse_factors_have_identical_conventions() {
2406 for factor in [0.0, 0.01, 0.5, 1.0] {
2407 let text = MaxScoreExecutor::text(
2408 Vec::new(),
2409 1.0,
2410 1,
2411 None,
2412 super::super::Bm25Params::default(),
2413 factor,
2414 );
2415 let sparse = MaxScoreExecutor::new(Vec::new(), 1, factor);
2416 assert_eq!(text.inv_heap_factor, sparse.inv_heap_factor);
2417 assert_eq!(text.inv_heap_factor, factor.clamp(0.01, 1.0).recip());
2418 }
2419 }
2420
2421 #[test]
2422 fn chunk_bounds_use_the_same_length_floor_as_scoring() {
2423 use crate::segment::chunk_map::{ChunkMapBuilder, read_chunk_maps, write_chunk_maps};
2424 use crate::structures::{BlockPostingList, PostingList};
2425 let mut chunks = ChunkMapBuilder::default();
2426 let mut postings = PostingList::new();
2427 for id in 0..2048 {
2428 chunks
2429 .push(id / 4, (id % 4) as u16, if id % 4 == 0 { 1 } else { 200 })
2430 .unwrap();
2431 postings.push(id, 1);
2432 }
2433 let mut bytes = Vec::new();
2434 write_chunk_maps(&mut bytes, &[(0, &chunks)], &[]).unwrap();
2435 let maps = read_chunk_maps(crate::directories::OwnedBytes::new(bytes)).unwrap();
2436 let map = &maps.chunk_maps[&0];
2437 let list =
2438 BlockPostingList::from_posting_list_with(&postings, false, Some(&|id| map.length(id)))
2439 .unwrap();
2440 assert_eq!(list.min_len(), Some(1));
2441 let params = super::super::Bm25Params::default();
2442 let cursor = TermCursor::text_with_params(
2443 list,
2444 2.0,
2445 map.avg_len(),
2446 Some(LengthSource::Chunks(map)),
2447 params,
2448 );
2449 let expected = params.upper_bound_with_len(1.0, 2.0, 200.0, map.avg_len());
2450 assert_eq!(cursor.max_score.to_bits(), expected.to_bits());
2451 for block in 0..cursor.num_blocks {
2452 assert_eq!(cursor.text_block_bound(block).to_bits(), expected.to_bits());
2453 assert_eq!(
2454 cursor.text_group_bound(block).unwrap().to_bits(),
2455 expected.to_bits()
2456 );
2457 }
2458 }
2459
2460 #[test]
2461 fn nested_scorers_keep_deadline_but_not_outer_score_floor() {
2462 let shared = SharedThreshold::for_limit(10).with_deadline(Some(std::time::Instant::now()));
2463 shared.raise(100.0);
2464 let options = super::super::ScorerOptions {
2465 shared_threshold: Some(shared.clone()),
2466 initial_threshold: 100.0,
2467 ..Default::default()
2468 }
2469 .without_threshold();
2470 assert_eq!(options.initial_threshold, 0.0);
2471 let nested = options.shared_threshold.unwrap();
2472 assert_eq!(nested.get(), 0.0);
2473 assert!(!nested.covers(10));
2474 assert!(nested.stop_if_expired());
2475 assert!(shared.truncated());
2476 nested.raise(200.0);
2477 assert_eq!(shared.get(), 100.0);
2478 }
2479
2480 #[test]
2481 fn score_bound_cache_reuses_only_its_last_key() {
2482 let cache = CachedScoreBound::new();
2483 assert_eq!(cache.get_or_compute(0, || 1.5), 1.5);
2484 assert_eq!(
2485 cache.get_or_compute(0, || panic!("recomputed active bound")),
2486 1.5
2487 );
2488 assert_eq!(cache.get_or_compute(1, || 2.5), 2.5);
2489 assert_eq!(cache.get_or_compute(0, || 3.5), 3.5);
2490 }
2491
2492 #[test]
2493 fn test_shared_threshold_monotonic_raise() {
2494 let shared = SharedThreshold::new();
2495 assert_eq!(shared.get(), 0.0);
2496
2497 shared.raise(2.5);
2498 assert_eq!(shared.get(), 2.5);
2499
2500 shared.raise(1.0);
2502 assert_eq!(shared.get(), 2.5);
2503
2504 shared.raise(4.0);
2506 assert_eq!(shared.get(), 4.0);
2507
2508 shared.raise(0.0);
2510 shared.raise(-3.0);
2511 shared.raise(f32::NAN);
2512 assert_eq!(shared.get(), 4.0);
2513
2514 let clone = shared.clone();
2516 clone.raise(9.0);
2517 assert_eq!(shared.get(), 9.0);
2518 }
2519
2520 #[test]
2521 fn test_shared_threshold_seed_matches_manual() {
2522 let mut seeded = ScoreCollector::new(2);
2525 seeded.seed_threshold(3.0);
2526 assert_eq!(seeded.threshold(), 3.0);
2527 assert!(!seeded.would_enter(3.0));
2529 assert!(seeded.would_enter(3.5));
2530 seeded.insert(1, 5.0);
2533 seeded.insert(2, 4.0);
2534 let results = seeded.into_sorted_results();
2535 assert_eq!(results.len(), 2);
2536 assert_eq!(results[0].0, 1);
2537 assert_eq!(results[1].0, 2);
2538 }
2539
2540 #[test]
2541 fn test_shared_threshold_can_raise_after_real_inserts() {
2542 let mut collector = ScoreCollector::new(3);
2543 collector.insert(1, 10.0);
2544 collector.insert(2, 4.0);
2545 assert_eq!(collector.real_len(), 2);
2546
2547 collector.seed_threshold(6.0);
2550 assert_eq!(collector.threshold(), 6.0);
2551 assert_eq!(collector.real_len(), 1);
2552
2553 assert!(collector.would_enter_candidate(3, 6.0, 0));
2556 assert!(collector.insert(3, 6.0));
2557 assert_eq!(collector.real_len(), 2);
2558 let results = collector.into_sorted_results();
2559 assert_eq!(results, vec![(1, 10.0, 0), (3, 6.0, 0)]);
2560 }
2561
2562 #[test]
2563 fn test_large_seed_uses_virtual_sentinels() {
2564 let k = 1_000_000_000;
2565 let mut collector = ScoreCollector::new(k);
2566 assert!(collector.heap.capacity() <= MAX_INITIAL_SCORE_COLLECTOR_CAPACITY);
2567
2568 collector.seed_threshold(42.0);
2569
2570 assert_eq!(collector.heap.len(), 0);
2573 assert!(collector.heap.capacity() <= MAX_INITIAL_SCORE_COLLECTOR_CAPACITY);
2574 assert_eq!(collector.len(), k);
2575 assert_eq!(collector.real_len(), 0);
2576 assert_eq!(collector.threshold(), 42.0);
2577 assert!(!collector.is_empty());
2578
2579 assert!(collector.insert_with_ordinal(9, 42.0, 7));
2582 assert!(!collector.insert(10, 41.0));
2583 assert!(collector.insert(11, 43.0));
2584 assert_eq!(collector.len(), k);
2585 assert_eq!(collector.real_len(), 2);
2586 assert_eq!(
2587 collector.into_sorted_results(),
2588 vec![(11, 43.0, 0), (9, 42.0, 7)]
2589 );
2590 }
2591
2592 #[test]
2593 fn test_virtual_sentinels_preserve_tie_order_when_filled() {
2594 let mut collector = ScoreCollector::new(3);
2595 collector.seed_threshold(5.0);
2596
2597 assert!(collector.insert_with_ordinal(3, 5.0, 2));
2598 assert!(collector.insert_with_ordinal(2, 5.0, 8));
2599 assert!(collector.insert_with_ordinal(1, 5.0, 4));
2600 assert_eq!(collector.real_len(), 3);
2601 assert!(collector.virtual_threshold.is_none());
2602
2603 assert!(collector.insert_with_ordinal(2, 5.0, 1));
2606 assert!(!collector.insert_with_ordinal(4, 5.0, 0));
2607 assert_eq!(
2608 collector.into_sorted_results(),
2609 vec![(1, 5.0, 4), (2, 5.0, 1), (2, 5.0, 8)]
2610 );
2611 }
2612
2613 #[test]
2614 fn test_score_collector_basic() {
2615 let mut collector = ScoreCollector::new(3);
2616
2617 collector.insert(1, 1.0);
2618 collector.insert(2, 2.0);
2619 collector.insert(3, 3.0);
2620 assert_eq!(collector.threshold(), 1.0);
2621
2622 collector.insert(4, 4.0);
2623 assert_eq!(collector.threshold(), 2.0);
2624
2625 let results = collector.into_sorted_results();
2626 assert_eq!(results.len(), 3);
2627 assert_eq!(results[0].0, 4); assert_eq!(results[1].0, 3);
2629 assert_eq!(results[2].0, 2);
2630 }
2631
2632 #[test]
2633 fn test_score_collector_threshold() {
2634 let mut collector = ScoreCollector::new(2);
2635
2636 collector.insert(1, 5.0);
2637 collector.insert(2, 3.0);
2638 assert_eq!(collector.threshold(), 3.0);
2639
2640 assert!(!collector.would_enter(2.0));
2642 assert!(!collector.insert(3, 2.0));
2643
2644 assert!(collector.would_enter(4.0));
2646 assert!(collector.insert(4, 4.0));
2647 assert_eq!(collector.threshold(), 4.0);
2648 }
2649
2650 #[test]
2651 fn test_heap_entry_ordering() {
2652 let mut heap = BinaryHeap::new();
2653 heap.push(HeapEntry {
2654 doc_id: 1,
2655 score: 3.0,
2656 ordinal: 0,
2657 });
2658 heap.push(HeapEntry {
2659 doc_id: 2,
2660 score: 1.0,
2661 ordinal: 0,
2662 });
2663 heap.push(HeapEntry {
2664 doc_id: 3,
2665 score: 2.0,
2666 ordinal: 0,
2667 });
2668
2669 assert_eq!(heap.pop().unwrap().score, 1.0);
2671 assert_eq!(heap.pop().unwrap().score, 2.0);
2672 assert_eq!(heap.pop().unwrap().score, 3.0);
2673 }
2674}