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
44 .score
45 .total_cmp(&self.score)
46 .then_with(|| self.doc_id.cmp(&other.doc_id))
47 .then_with(|| self.ordinal.cmp(&other.ordinal))
48 }
49}
50
51impl PartialOrd for HeapEntry {
52 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
53 Some(self.cmp(other))
54 }
55}
56
57pub struct ScoreCollector {
70 heap: BinaryHeap<HeapEntry>,
72 pub k: usize,
73 cached_threshold: f32,
76 virtual_threshold: Option<f32>,
80}
81
82impl ScoreCollector {
83 pub fn new(k: usize) -> Self {
85 Self {
86 heap: BinaryHeap::with_capacity(k.min(MAX_INITIAL_SCORE_COLLECTOR_CAPACITY)),
87 k,
88 cached_threshold: 0.0,
89 virtual_threshold: None,
90 }
91 }
92
93 #[inline]
95 pub fn threshold(&self) -> f32 {
96 self.cached_threshold
97 }
98
99 #[inline]
101 fn update_threshold(&mut self) {
102 self.cached_threshold = if let Some(threshold) = self.virtual_threshold {
103 threshold
104 } else if self.heap.len() >= self.k {
105 self.heap.peek().map(|e| e.score).unwrap_or(0.0)
106 } else {
107 0.0
108 };
109 }
110
111 #[inline]
114 pub fn insert(&mut self, doc_id: DocId, score: f32) -> bool {
115 self.insert_with_ordinal(doc_id, score, 0)
116 }
117
118 #[inline]
121 pub fn insert_with_ordinal(&mut self, doc_id: DocId, score: f32, ordinal: u16) -> bool {
122 if self.k == 0 {
123 return false;
124 }
125 let entry = HeapEntry {
126 doc_id,
127 score,
128 ordinal,
129 };
130 if self.heap.len() < self.k {
131 if let Some(threshold) = self.virtual_threshold {
132 let sentinel = HeapEntry {
133 doc_id: u32::MAX,
134 score: threshold,
135 ordinal: 0,
136 };
137 if entry >= sentinel {
138 return false;
139 }
140 }
141
142 self.heap.push(entry);
143 if self.heap.len() == self.k {
145 self.virtual_threshold = None;
146 self.update_threshold();
147 }
148 true
149 } else if self.heap.peek().is_some_and(|worst| entry < *worst) {
150 {
151 let mut worst = self.heap.peek_mut().expect("full heap has a root");
152 *worst = entry;
153 }
154 self.update_threshold();
155 true
156 } else {
157 false
158 }
159 }
160
161 #[inline]
163 pub fn would_enter(&self, score: f32) -> bool {
164 self.len() < self.k || score > self.cached_threshold
165 }
166
167 #[inline]
170 pub fn would_enter_candidate(&self, doc_id: DocId, score: f32, ordinal: u16) -> bool {
171 if self.k == 0 {
172 return false;
173 }
174 let entry = HeapEntry {
175 doc_id,
176 score,
177 ordinal,
178 };
179 if let Some(threshold) = self.virtual_threshold {
180 let sentinel = HeapEntry {
181 doc_id: u32::MAX,
182 score: threshold,
183 ordinal: 0,
184 };
185 entry < sentinel
186 } else {
187 self.heap.len() < self.k || self.heap.peek().is_some_and(|worst| entry < *worst)
188 }
189 }
190
191 #[inline]
193 pub fn len(&self) -> usize {
194 if self.virtual_threshold.is_some() {
195 self.k
196 } else {
197 self.heap.len()
198 }
199 }
200
201 #[inline]
203 pub fn real_len(&self) -> usize {
204 self.heap.len()
205 }
206
207 #[inline]
209 pub fn is_empty(&self) -> bool {
210 self.len() == 0
211 }
212
213 pub fn seed_threshold(&mut self, initial_threshold: f32) {
220 if initial_threshold <= 0.0
221 || self.k == 0
222 || (self.len() >= self.k && initial_threshold <= self.cached_threshold)
223 {
224 return;
225 }
226
227 let sentinel = HeapEntry {
228 doc_id: u32::MAX,
229 score: initial_threshold,
230 ordinal: 0,
231 };
232
233 if let Some(current_threshold) = self.virtual_threshold {
237 let current = HeapEntry {
238 doc_id: u32::MAX,
239 score: current_threshold,
240 ordinal: 0,
241 };
242 if sentinel >= current {
243 return;
244 }
245 } else if self.heap.len() >= self.k
246 && !self.heap.peek().is_some_and(|worst| sentinel < *worst)
247 {
248 return;
249 }
250
251 self.virtual_threshold = Some(initial_threshold);
252 while self.heap.peek().is_some_and(|worst| sentinel < *worst) {
253 self.heap.pop();
254 }
255 self.update_threshold();
256 }
257
258 pub fn into_sorted_results(self) -> Vec<(DocId, f32, u16)> {
261 let mut results: Vec<(DocId, f32, u16)> = self
262 .heap
263 .into_vec()
264 .into_iter()
265 .filter(|e| e.doc_id != u32::MAX)
266 .map(|e| (e.doc_id, e.score, e.ordinal))
267 .collect();
268
269 results.sort_unstable_by(|a, b| {
271 b.1.total_cmp(&a.1)
272 .then_with(|| a.0.cmp(&b.0))
273 .then_with(|| a.2.cmp(&b.2))
274 });
275
276 results
277 }
278}
279
280#[derive(Clone, Debug)]
303pub struct SharedThreshold {
304 floor: std::sync::Arc<std::sync::atomic::AtomicU32>,
305 k: usize,
308}
309
310impl Default for SharedThreshold {
311 fn default() -> Self {
312 Self::new()
313 }
314}
315
316impl SharedThreshold {
317 pub fn new() -> Self {
321 Self::with_depth(usize::MAX)
322 }
323
324 pub fn for_limit(limit: usize) -> Self {
326 Self::with_depth(limit)
327 }
328
329 fn with_depth(k: usize) -> Self {
330 Self {
331 floor: std::sync::Arc::new(std::sync::atomic::AtomicU32::new(0)),
333 k,
334 }
335 }
336
337 #[inline]
340 pub(crate) fn covers(&self, heap_depth: usize) -> bool {
341 heap_depth >= self.k
342 }
343
344 #[inline]
346 pub fn get(&self) -> f32 {
347 f32::from_bits(self.floor.load(std::sync::atomic::Ordering::Relaxed))
348 }
349
350 pub fn raise(&self, score: f32) {
355 if score <= 0.0 {
358 return;
359 }
360 use std::sync::atomic::Ordering::Relaxed;
361 let bits = score.to_bits();
362 let mut cur = self.floor.load(Relaxed);
363 while f32::from_bits(cur) < score {
364 match self
365 .floor
366 .compare_exchange_weak(cur, bits, Relaxed, Relaxed)
367 {
368 Ok(_) => break,
369 Err(actual) => cur = actual,
370 }
371 }
372 }
373}
374
375#[derive(Debug, Clone, Copy)]
377pub struct ScoredDoc {
378 pub doc_id: DocId,
379 pub score: f32,
380 pub ordinal: u16,
382}
383
384pub struct MaxScoreExecutor<'a> {
395 metric_index: &'a str,
399 metric_field: &'a str,
400 cursors: Vec<TermCursor<'a>>,
401 prefix_sums: Vec<f32>,
402 collector: ScoreCollector,
403 inv_heap_factor: f32,
404 predicate: Option<super::DocPredicate<'a>>,
405}
406
407pub(crate) struct TermCursor<'a> {
416 pub max_score: f32,
417 num_blocks: usize,
418 block_idx: usize,
420 doc_ids: Vec<u32>,
421 scores: Vec<f32>,
422 ordinals: Vec<u16>,
423 pos: usize,
424 block_loaded: bool,
425 exhausted: bool,
426 lazy_ordinals: bool,
430 ordinals_loaded: bool,
432 current_sparse_block: Option<crate::structures::SparseBlock>,
434 variant: CursorVariant<'a>,
436}
437
438enum CursorVariant<'a> {
439 Text {
441 list: crate::structures::BlockPostingList,
442 idf: f32,
443 idf_times_k1_plus_1: f32,
445 denom_tf_coeff: f32,
447 denom_const: f32,
449 tfs: Vec<u32>,
450 deferred_tf: Option<(usize, usize, usize)>,
453 },
454 Sparse {
456 si: &'a crate::segment::SparseIndex,
457 query_weight: f32,
458 skip_start: usize,
459 block_data_offset: u64,
460 },
461}
462
463macro_rules! cursor_ensure_block {
471 ($self:ident, $load_block_fn:ident, $($aw:tt)*) => {{
472 if $self.exhausted || $self.block_loaded {
473 return Ok(!$self.exhausted);
474 }
475 match &mut $self.variant {
476 CursorVariant::Text {
477 list,
478 deferred_tf,
479 ..
480 } => {
481 if let Some(state) = list.decode_block_doc_ids_only($self.block_idx, &mut $self.doc_ids) {
482 *deferred_tf = Some(state);
483 $self.scores.clear();
484 $self.pos = 0;
485 $self.block_loaded = true;
486 Ok(true)
487 } else {
488 $self.exhausted = true;
489 Ok(false)
490 }
491 }
492 CursorVariant::Sparse {
493 si,
494 query_weight,
495 skip_start,
496 block_data_offset,
497 ..
498 } => {
499 let block = si
500 .$load_block_fn(*skip_start, *block_data_offset, $self.block_idx)
501 $($aw)* ?;
502 match block {
503 Some(b) => {
504 b.decode_doc_ids_into(&mut $self.doc_ids);
505 b.decode_scored_weights_into(*query_weight, &mut $self.scores);
506 if $self.lazy_ordinals {
507 $self.current_sparse_block = Some(b);
510 $self.ordinals_loaded = false;
511 } else {
512 b.decode_ordinals_into(&mut $self.ordinals);
513 $self.ordinals_loaded = true;
514 $self.current_sparse_block = None;
515 }
516 $self.pos = 0;
517 $self.block_loaded = true;
518 Ok(true)
519 }
520 None => {
521 $self.exhausted = true;
522 Ok(false)
523 }
524 }
525 }
526 }
527 }};
528}
529
530macro_rules! cursor_advance {
531 ($self:ident, $ensure_fn:ident, $($aw:tt)*) => {{
532 if $self.exhausted {
533 return Ok(u32::MAX);
534 }
535 $self.$ensure_fn() $($aw)* ?;
536 if $self.exhausted {
537 return Ok(u32::MAX);
538 }
539 Ok($self.advance_pos())
540 }};
541}
542
543macro_rules! cursor_seek {
544 ($self:ident, $ensure_fn:ident, $target:expr, $($aw:tt)*) => {{
545 if let Some(doc) = $self.seek_prepare($target) {
546 return Ok(doc);
547 }
548 $self.$ensure_fn() $($aw)* ?;
549 if $self.seek_finish($target) {
550 $self.$ensure_fn() $($aw)* ?;
551 }
552 Ok($self.doc())
553 }};
554}
555
556impl<'a> TermCursor<'a> {
557 pub fn text(
559 posting_list: crate::structures::BlockPostingList,
560 idf: f32,
561 avg_field_len: f32,
562 ) -> Self {
563 let max_tf = posting_list.max_tf() as f32;
564 let max_score = super::bm25_upper_bound(max_tf.max(1.0), idf);
565 let num_blocks = posting_list.num_blocks();
566 let safe_avg = avg_field_len.max(1.0);
567 Self {
568 max_score,
569 num_blocks,
570 block_idx: 0,
571 doc_ids: Vec::with_capacity(128),
572 scores: Vec::with_capacity(128),
573 ordinals: Vec::new(),
574 pos: 0,
575 block_loaded: false,
576 exhausted: num_blocks == 0,
577 lazy_ordinals: false,
578 ordinals_loaded: true, current_sparse_block: None,
580 variant: CursorVariant::Text {
581 list: posting_list,
582 idf,
583 idf_times_k1_plus_1: idf * (super::BM25_K1 + 1.0),
584 denom_tf_coeff: 1.0 + super::BM25_K1 * (super::BM25_B / safe_avg),
585 denom_const: super::BM25_K1 * (1.0 - super::BM25_B),
586 tfs: Vec::with_capacity(128),
587 deferred_tf: None,
588 },
589 }
590 }
591
592 pub fn sparse(
595 si: &'a crate::segment::SparseIndex,
596 query_weight: f32,
597 skip_start: usize,
598 skip_count: usize,
599 global_max_weight: f32,
600 block_data_offset: u64,
601 ) -> Self {
602 Self {
603 max_score: query_weight.abs() * global_max_weight,
604 num_blocks: skip_count,
605 block_idx: 0,
606 doc_ids: Vec::with_capacity(256),
607 scores: Vec::with_capacity(256),
608 ordinals: Vec::with_capacity(256),
609 pos: 0,
610 block_loaded: false,
611 exhausted: skip_count == 0,
612 lazy_ordinals: false,
613 ordinals_loaded: true,
614 current_sparse_block: None,
615 variant: CursorVariant::Sparse {
616 si,
617 query_weight,
618 skip_start,
619 block_data_offset,
620 },
621 }
622 }
623
624 #[inline]
627 fn block_first_doc(&self, idx: usize) -> DocId {
628 match &self.variant {
629 CursorVariant::Text { list, .. } => list.block_first_doc(idx).unwrap_or(u32::MAX),
630 CursorVariant::Sparse { si, skip_start, .. } => {
631 si.read_skip_entry(*skip_start + idx).first_doc
632 }
633 }
634 }
635
636 #[inline]
637 fn block_last_doc(&self, idx: usize) -> DocId {
638 match &self.variant {
639 CursorVariant::Text { list, .. } => list.block_last_doc(idx).unwrap_or(0),
640 CursorVariant::Sparse { si, skip_start, .. } => {
641 si.read_skip_entry(*skip_start + idx).last_doc
642 }
643 }
644 }
645
646 #[inline]
649 pub fn doc(&self) -> DocId {
650 if self.exhausted {
651 return u32::MAX;
652 }
653 if self.block_loaded {
654 debug_assert!(self.pos < self.doc_ids.len());
655 unsafe { *self.doc_ids.get_unchecked(self.pos) }
657 } else {
658 self.block_first_doc(self.block_idx)
659 }
660 }
661
662 #[inline]
663 pub fn ordinal(&self) -> u16 {
664 if !self.block_loaded || self.ordinals.is_empty() {
665 return 0;
666 }
667 debug_assert!(self.pos < self.ordinals.len());
668 unsafe { *self.ordinals.get_unchecked(self.pos) }
670 }
671
672 #[inline]
678 pub fn ordinal_mut(&mut self) -> u16 {
679 if !self.block_loaded {
680 return 0;
681 }
682 if !self.ordinals_loaded {
683 if let Some(ref block) = self.current_sparse_block {
684 block.decode_ordinals_into(&mut self.ordinals);
685 }
686 self.ordinals_loaded = true;
687 }
688 if self.ordinals.is_empty() {
689 return 0;
690 }
691 debug_assert!(self.pos < self.ordinals.len());
692 unsafe { *self.ordinals.get_unchecked(self.pos) }
693 }
694
695 #[inline]
696 pub fn score(&self) -> f32 {
697 if !self.block_loaded {
698 return 0.0;
699 }
700 debug_assert!(self.pos < self.scores.len());
701 unsafe { *self.scores.get_unchecked(self.pos) }
703 }
704
705 #[inline]
711 pub fn ensure_scores(&mut self) {
712 if self.block_loaded && self.scores.is_empty() {
713 self.compute_deferred_scores();
714 }
715 }
716
717 #[inline]
718 pub fn current_block_max_score(&self) -> f32 {
719 if self.exhausted {
720 return 0.0;
721 }
722 match &self.variant {
723 CursorVariant::Text { list, idf, .. } => {
724 let block_max_tf = list.block_max_tf(self.block_idx).unwrap_or(0) as f32;
725 super::bm25_upper_bound(block_max_tf.max(1.0), *idf)
726 }
727 CursorVariant::Sparse {
728 si,
729 query_weight,
730 skip_start,
731 ..
732 } => query_weight.abs() * si.read_skip_entry(*skip_start + self.block_idx).max_weight,
733 }
734 }
735
736 pub fn skip_to_next_block(&mut self) -> DocId {
739 if self.exhausted {
740 return u32::MAX;
741 }
742 self.block_idx += 1;
743 self.block_loaded = false;
744 if self.block_idx >= self.num_blocks {
745 self.exhausted = true;
746 return u32::MAX;
747 }
748 self.block_first_doc(self.block_idx)
749 }
750
751 #[inline]
752 fn advance_pos(&mut self) -> DocId {
753 self.pos += 1;
754 if self.pos >= self.doc_ids.len() {
755 self.block_idx += 1;
756 self.block_loaded = false;
757 if self.block_idx >= self.num_blocks {
758 self.exhausted = true;
759 return u32::MAX;
760 }
761 }
762 self.doc()
763 }
764
765 #[inline(never)]
767 fn compute_deferred_scores(&mut self) {
768 if let CursorVariant::Text {
769 list,
770 idf_times_k1_plus_1,
771 denom_tf_coeff,
772 denom_const,
773 tfs,
774 deferred_tf,
775 ..
776 } = &mut self.variant
777 && let Some((block_offset, tf_start, count)) = deferred_tf.take()
778 {
779 list.decode_block_tfs_deferred(block_offset, tf_start, count, tfs);
780 let num_scale = *idf_times_k1_plus_1;
781 let d_tf = *denom_tf_coeff;
782 let d_const = *denom_const;
783 self.scores.clear();
784 self.scores.resize(count, 0.0);
785 for i in 0..count {
786 let tf = unsafe { *tfs.get_unchecked(i) } as f32;
787 let score = (num_scale * tf) / (d_tf * tf + d_const);
788 unsafe {
789 *self.scores.get_unchecked_mut(i) = score;
790 }
791 }
792 }
793 }
794
795 pub async fn ensure_block_loaded(&mut self) -> crate::Result<bool> {
801 cursor_ensure_block!(self, load_block_direct, .await)
802 }
803
804 pub fn ensure_block_loaded_sync(&mut self) -> crate::Result<bool> {
805 cursor_ensure_block!(self, load_block_direct_sync,)
806 }
807
808 pub async fn advance(&mut self) -> crate::Result<DocId> {
809 cursor_advance!(self, ensure_block_loaded, .await)
810 }
811
812 pub fn advance_sync(&mut self) -> crate::Result<DocId> {
813 cursor_advance!(self, ensure_block_loaded_sync,)
814 }
815
816 pub async fn seek(&mut self, target: DocId) -> crate::Result<DocId> {
817 cursor_seek!(self, ensure_block_loaded, target, .await)
818 }
819
820 pub fn seek_sync(&mut self, target: DocId) -> crate::Result<DocId> {
821 cursor_seek!(self, ensure_block_loaded_sync, target,)
822 }
823
824 fn seek_prepare(&mut self, target: DocId) -> Option<DocId> {
825 if self.exhausted {
826 return Some(u32::MAX);
827 }
828
829 if self.block_loaded
831 && let Some(&last) = self.doc_ids.last()
832 {
833 if last >= target && self.doc_ids[self.pos] < target {
834 let remaining = &self.doc_ids[self.pos..];
835 self.pos += crate::structures::simd::find_first_ge_u32(remaining, target);
836 if self.pos >= self.doc_ids.len() {
837 self.block_idx += 1;
838 self.block_loaded = false;
839 if self.block_idx >= self.num_blocks {
840 self.exhausted = true;
841 return Some(u32::MAX);
842 }
843 }
844 return Some(self.doc());
845 }
846 if self.doc_ids[self.pos] >= target {
847 return Some(self.doc());
848 }
849 }
850
851 let lo = match &self.variant {
853 CursorVariant::Text { list, .. } => match list.seek_block(target, self.block_idx) {
855 Some(idx) => idx,
856 None => {
857 self.exhausted = true;
858 return Some(u32::MAX);
859 }
860 },
861 CursorVariant::Sparse { .. } => {
863 let mut lo = self.block_idx;
864 let mut hi = self.num_blocks;
865 while lo < hi {
866 let mid = lo + (hi - lo) / 2;
867 if self.block_last_doc(mid) < target {
868 lo = mid + 1;
869 } else {
870 hi = mid;
871 }
872 }
873 lo
874 }
875 };
876 if lo >= self.num_blocks {
877 self.exhausted = true;
878 return Some(u32::MAX);
879 }
880 if lo != self.block_idx || !self.block_loaded {
881 self.block_idx = lo;
882 self.block_loaded = false;
883 }
884 None
885 }
886
887 #[inline]
888 fn seek_finish(&mut self, target: DocId) -> bool {
889 if self.exhausted {
890 return false;
891 }
892 self.pos = crate::structures::simd::find_first_ge_u32(&self.doc_ids, target);
893 if self.pos >= self.doc_ids.len() {
894 self.block_idx += 1;
895 self.block_loaded = false;
896 if self.block_idx >= self.num_blocks {
897 self.exhausted = true;
898 return false;
899 }
900 return true;
901 }
902 false
903 }
904}
905
906macro_rules! bms_execute_loop {
911 ($self:ident, $ensure:ident, $advance:ident, $seek:ident, $($aw:tt)*) => {{
912 let n = $self.cursors.len();
913
914 for cursor in &mut $self.cursors {
916 cursor.$ensure() $($aw)* ?;
917 }
918
919 let mut docs_scored = 0u64;
920 let mut docs_skipped = 0u64;
921 let mut blocks_skipped = 0u64;
922 let mut conjunction_skipped = 0u64;
923 let mut ordinal_scores: Vec<(u16, f32)> = Vec::with_capacity(n * 2);
924 let _bms_start = std::time::Instant::now();
925
926 let inv_heap_factor = $self.inv_heap_factor;
927 let mut adjusted_threshold = $self.collector.threshold() * inv_heap_factor - 1e-6;
928
929 loop {
930 let partition = $self.find_partition();
931 if partition >= n {
932 break;
933 }
934
935 let mut min_doc = u32::MAX;
939 let mut at_min_mask = 0u64; for i in partition..n {
941 let doc = $self.cursors[i].doc();
942 match doc.cmp(&min_doc) {
943 std::cmp::Ordering::Less => {
944 min_doc = doc;
945 at_min_mask = 1u64 << (i as u32);
946 }
947 std::cmp::Ordering::Equal => {
948 at_min_mask |= 1u64 << (i as u32);
949 }
950 _ => {}
951 }
952 }
953 if min_doc == u32::MAX {
954 break;
955 }
956
957 let non_essential_upper = if partition > 0 {
958 $self.prefix_sums[partition - 1]
959 } else {
960 0.0
961 };
962
963 if $self.collector.len() >= $self.collector.k {
965 let mut present_upper: f32 = 0.0;
966 let mut mask = at_min_mask;
967 while mask != 0 {
968 let i = mask.trailing_zeros() as usize;
969 present_upper += $self.cursors[i].max_score;
970 mask &= mask - 1;
971 }
972
973 if present_upper + non_essential_upper < adjusted_threshold {
974 let mut mask = at_min_mask;
975 while mask != 0 {
976 let i = mask.trailing_zeros() as usize;
977 $self.cursors[i].$ensure() $($aw)* ?;
978 $self.cursors[i].$advance() $($aw)* ?;
979 mask &= mask - 1;
980 }
981 conjunction_skipped += 1;
982 continue;
983 }
984 }
985
986 if $self.collector.len() >= $self.collector.k {
988 let mut block_max_sum: f32 = 0.0;
989 let mut mask = at_min_mask;
990 while mask != 0 {
991 let i = mask.trailing_zeros() as usize;
992 block_max_sum += $self.cursors[i].current_block_max_score();
993 mask &= mask - 1;
994 }
995
996 if block_max_sum + non_essential_upper < adjusted_threshold {
997 let mut mask = at_min_mask;
998 while mask != 0 {
999 let i = mask.trailing_zeros() as usize;
1000 $self.cursors[i].skip_to_next_block();
1001 $self.cursors[i].$ensure() $($aw)* ?;
1002 mask &= mask - 1;
1003 }
1004 blocks_skipped += 1;
1005 continue;
1006 }
1007 }
1008
1009 if let Some(ref pred) = $self.predicate {
1011 if !pred(min_doc) {
1012 let mut mask = at_min_mask;
1013 while mask != 0 {
1014 let i = mask.trailing_zeros() as usize;
1015 $self.cursors[i].$ensure() $($aw)* ?;
1016 $self.cursors[i].$advance() $($aw)* ?;
1017 mask &= mask - 1;
1018 }
1019 continue;
1020 }
1021 }
1022
1023 ordinal_scores.clear();
1025 {
1026 let mut mask = at_min_mask;
1027 while mask != 0 {
1028 let i = mask.trailing_zeros() as usize;
1029 $self.cursors[i].$ensure() $($aw)* ?;
1030 $self.cursors[i].ensure_scores();
1031 while $self.cursors[i].doc() == min_doc {
1032 let ord = $self.cursors[i].ordinal_mut();
1033 let sc = $self.cursors[i].score();
1034 ordinal_scores.push((ord, sc));
1035 $self.cursors[i].$advance() $($aw)* ?;
1036 }
1037 mask &= mask - 1;
1038 }
1039 }
1040
1041 let essential_total: f32 = ordinal_scores.iter().map(|(_, s)| *s).sum();
1042 if $self.collector.len() >= $self.collector.k
1043 && essential_total + non_essential_upper < adjusted_threshold
1044 {
1045 docs_skipped += 1;
1046 continue;
1047 }
1048
1049 let mut running_total = essential_total;
1051 for i in (0..partition).rev() {
1052 if $self.collector.len() >= $self.collector.k
1053 && running_total + $self.prefix_sums[i] < adjusted_threshold
1054 {
1055 break;
1056 }
1057
1058 let doc = $self.cursors[i].$seek(min_doc) $($aw)* ?;
1059 if doc == min_doc {
1060 $self.cursors[i].ensure_scores();
1061 while $self.cursors[i].doc() == min_doc {
1062 let s = $self.cursors[i].score();
1063 running_total += s;
1064 let ord = $self.cursors[i].ordinal_mut();
1065 ordinal_scores.push((ord, s));
1066 $self.cursors[i].$advance() $($aw)* ?;
1067 }
1068 }
1069 }
1070
1071 if ordinal_scores.len() == 1 {
1074 let (ord, score) = ordinal_scores[0];
1075 if $self.collector.insert_with_ordinal(min_doc, score, ord) {
1076 docs_scored += 1;
1077 adjusted_threshold = $self.collector.threshold() * inv_heap_factor - 1e-6;
1078 } else {
1079 docs_skipped += 1;
1080 }
1081 } else if !ordinal_scores.is_empty() {
1082 if ordinal_scores.len() > 2 {
1083 ordinal_scores.sort_unstable_by_key(|(ord, _)| *ord);
1084 } else if ordinal_scores.len() == 2 && ordinal_scores[0].0 > ordinal_scores[1].0 {
1085 ordinal_scores.swap(0, 1);
1086 }
1087 let mut j = 0;
1088 while j < ordinal_scores.len() {
1089 let current_ord = ordinal_scores[j].0;
1090 let mut score = 0.0f32;
1091 while j < ordinal_scores.len() && ordinal_scores[j].0 == current_ord {
1092 score += ordinal_scores[j].1;
1093 j += 1;
1094 }
1095 if $self
1096 .collector
1097 .insert_with_ordinal(min_doc, score, current_ord)
1098 {
1099 docs_scored += 1;
1100 adjusted_threshold = $self.collector.threshold() * inv_heap_factor - 1e-6;
1101 } else {
1102 docs_skipped += 1;
1103 }
1104 }
1105 }
1106 }
1107
1108 let results: Vec<ScoredDoc> = $self
1109 .collector
1110 .into_sorted_results()
1111 .into_iter()
1112 .map(|(doc_id, score, ordinal)| ScoredDoc {
1113 doc_id,
1114 score,
1115 ordinal,
1116 })
1117 .collect();
1118
1119 let _bms_elapsed_ms = _bms_start.elapsed().as_millis() as u64;
1120 if _bms_elapsed_ms > 500 {
1121 warn!(
1122 "slow MaxScore: {}ms, cursors={}, scored={}, skipped={}, blocks_skipped={}, conjunction_skipped={}, returned={}, top_score={:.4}",
1123 _bms_elapsed_ms,
1124 n,
1125 docs_scored,
1126 docs_skipped,
1127 blocks_skipped,
1128 conjunction_skipped,
1129 results.len(),
1130 results.first().map(|r| r.score).unwrap_or(0.0)
1131 );
1132 } else {
1133 debug!(
1134 "MaxScoreExecutor: {}ms, scored={}, skipped={}, blocks_skipped={}, conjunction_skipped={}, returned={}, top_score={:.4}",
1135 _bms_elapsed_ms,
1136 docs_scored,
1137 docs_skipped,
1138 blocks_skipped,
1139 conjunction_skipped,
1140 results.len(),
1141 results.first().map(|r| r.score).unwrap_or(0.0)
1142 );
1143 }
1144
1145 Ok(results)
1146 }};
1147}
1148
1149impl<'a> MaxScoreExecutor<'a> {
1150 pub(crate) fn new(mut cursors: Vec<TermCursor<'a>>, k: usize, heap_factor: f32) -> Self {
1155 if cursors.len() > super::MAX_QUERY_TERMS {
1159 cursors.sort_unstable_by(|a, b| b.max_score.total_cmp(&a.max_score));
1160 cursors.truncate(super::MAX_QUERY_TERMS);
1161 log::warn!(
1162 "MaxScore cursor count exceeded {}; retaining the strongest cursors",
1163 super::MAX_QUERY_TERMS
1164 );
1165 }
1166
1167 for c in &mut cursors {
1170 c.lazy_ordinals = true;
1171 }
1172
1173 cursors.sort_by(|a, b| {
1175 a.max_score
1176 .partial_cmp(&b.max_score)
1177 .unwrap_or(Ordering::Equal)
1178 });
1179
1180 let mut prefix_sums = Vec::with_capacity(cursors.len());
1181 let mut cumsum = 0.0f32;
1182 for c in &cursors {
1183 cumsum += c.max_score;
1184 prefix_sums.push(cumsum);
1185 }
1186
1187 let clamped_heap_factor = heap_factor.clamp(0.01, 1.0);
1188
1189 debug!(
1190 "Creating MaxScoreExecutor: num_cursors={}, k={}, total_upper={:.4}, heap_factor={:.2}",
1191 cursors.len(),
1192 k,
1193 cumsum,
1194 clamped_heap_factor
1195 );
1196
1197 Self {
1198 cursors,
1199 prefix_sums,
1200 collector: ScoreCollector::new(k),
1201 inv_heap_factor: 1.0 / clamped_heap_factor,
1202 predicate: None,
1203 metric_index: "unknown",
1204 metric_field: "unknown",
1205 }
1206 }
1207
1208 pub fn with_metric_labels(mut self, index: &'a str, field: &'a str) -> Self {
1210 self.metric_index = index;
1211 self.metric_field = field;
1212 self
1213 }
1214
1215 pub fn sparse(
1219 sparse_index: &'a crate::segment::SparseIndex,
1220 query_terms: Vec<(u32, f32)>,
1221 k: usize,
1222 heap_factor: f32,
1223 ) -> Self {
1224 let cursors: Vec<TermCursor<'a>> = query_terms
1225 .iter()
1226 .filter_map(|&(dim_id, qw)| {
1227 let (skip_start, skip_count, global_max, block_data_offset) =
1228 sparse_index.get_skip_range_full(dim_id)?;
1229 Some(TermCursor::sparse(
1230 sparse_index,
1231 qw,
1232 skip_start,
1233 skip_count,
1234 global_max,
1235 block_data_offset,
1236 ))
1237 })
1238 .collect();
1239 Self::new(cursors, k, heap_factor)
1240 }
1241
1242 pub fn text(
1246 posting_lists: Vec<(crate::structures::BlockPostingList, f32)>,
1247 avg_field_len: f32,
1248 k: usize,
1249 ) -> Self {
1250 let cursors: Vec<TermCursor<'a>> = posting_lists
1251 .into_iter()
1252 .map(|(pl, idf)| TermCursor::text(pl, idf, avg_field_len))
1253 .collect();
1254 Self::new(cursors, k, 1.0)
1255 }
1256
1257 #[inline]
1258 fn find_partition(&self) -> usize {
1259 let threshold = self.collector.threshold() * self.inv_heap_factor;
1263 self.prefix_sums.partition_point(|&sum| sum < threshold)
1266 }
1267
1268 pub fn with_predicate(mut self, predicate: super::DocPredicate<'a>) -> Self {
1274 self.predicate = Some(predicate);
1275 self
1276 }
1277
1278 pub fn seed_threshold(&mut self, initial_threshold: f32) {
1280 self.collector.seed_threshold(initial_threshold);
1281 }
1282
1283 pub async fn execute(mut self) -> crate::Result<Vec<ScoredDoc>> {
1285 if self.cursors.is_empty() {
1286 return Ok(Vec::new());
1287 }
1288 let t = crate::observe::Timer::start();
1289 let results = bms_execute_loop!(self, ensure_block_loaded, advance, seek, .await);
1290 if let Ok(r) = &results {
1291 crate::observe::maxscore_query(self.metric_index, self.metric_field, t.secs(), r.len());
1292 }
1293 results
1294 }
1295
1296 pub fn execute_sync(mut self) -> crate::Result<Vec<ScoredDoc>> {
1298 if self.cursors.is_empty() {
1299 return Ok(Vec::new());
1300 }
1301 let t = crate::observe::Timer::start();
1302 let results = bms_execute_loop!(self, ensure_block_loaded_sync, advance_sync, seek_sync,);
1303 if let Ok(r) = &results {
1304 crate::observe::maxscore_query(self.metric_index, self.metric_field, t.secs(), r.len());
1305 }
1306 results
1307 }
1308}
1309
1310#[cfg(test)]
1311mod tests {
1312 use super::*;
1313
1314 #[test]
1315 fn test_shared_threshold_monotonic_raise() {
1316 let shared = SharedThreshold::new();
1317 assert_eq!(shared.get(), 0.0);
1318
1319 shared.raise(2.5);
1320 assert_eq!(shared.get(), 2.5);
1321
1322 shared.raise(1.0);
1324 assert_eq!(shared.get(), 2.5);
1325
1326 shared.raise(4.0);
1328 assert_eq!(shared.get(), 4.0);
1329
1330 shared.raise(0.0);
1332 shared.raise(-3.0);
1333 shared.raise(f32::NAN);
1334 assert_eq!(shared.get(), 4.0);
1335
1336 let clone = shared.clone();
1338 clone.raise(9.0);
1339 assert_eq!(shared.get(), 9.0);
1340 }
1341
1342 #[test]
1343 fn test_shared_threshold_seed_matches_manual() {
1344 let mut seeded = ScoreCollector::new(2);
1347 seeded.seed_threshold(3.0);
1348 assert_eq!(seeded.threshold(), 3.0);
1349 assert!(!seeded.would_enter(3.0));
1351 assert!(seeded.would_enter(3.5));
1352 seeded.insert(1, 5.0);
1355 seeded.insert(2, 4.0);
1356 let results = seeded.into_sorted_results();
1357 assert_eq!(results.len(), 2);
1358 assert_eq!(results[0].0, 1);
1359 assert_eq!(results[1].0, 2);
1360 }
1361
1362 #[test]
1363 fn test_shared_threshold_can_raise_after_real_inserts() {
1364 let mut collector = ScoreCollector::new(3);
1365 collector.insert(1, 10.0);
1366 collector.insert(2, 4.0);
1367 assert_eq!(collector.real_len(), 2);
1368
1369 collector.seed_threshold(6.0);
1372 assert_eq!(collector.threshold(), 6.0);
1373 assert_eq!(collector.real_len(), 1);
1374
1375 assert!(collector.would_enter_candidate(3, 6.0, 0));
1378 assert!(collector.insert(3, 6.0));
1379 assert_eq!(collector.real_len(), 2);
1380 let results = collector.into_sorted_results();
1381 assert_eq!(results, vec![(1, 10.0, 0), (3, 6.0, 0)]);
1382 }
1383
1384 #[test]
1385 fn test_large_seed_uses_virtual_sentinels() {
1386 let k = 1_000_000_000;
1387 let mut collector = ScoreCollector::new(k);
1388 assert!(collector.heap.capacity() <= MAX_INITIAL_SCORE_COLLECTOR_CAPACITY);
1389
1390 collector.seed_threshold(42.0);
1391
1392 assert_eq!(collector.heap.len(), 0);
1395 assert!(collector.heap.capacity() <= MAX_INITIAL_SCORE_COLLECTOR_CAPACITY);
1396 assert_eq!(collector.len(), k);
1397 assert_eq!(collector.real_len(), 0);
1398 assert_eq!(collector.threshold(), 42.0);
1399 assert!(!collector.is_empty());
1400
1401 assert!(collector.insert_with_ordinal(9, 42.0, 7));
1404 assert!(!collector.insert(10, 41.0));
1405 assert!(collector.insert(11, 43.0));
1406 assert_eq!(collector.len(), k);
1407 assert_eq!(collector.real_len(), 2);
1408 assert_eq!(
1409 collector.into_sorted_results(),
1410 vec![(11, 43.0, 0), (9, 42.0, 7)]
1411 );
1412 }
1413
1414 #[test]
1415 fn test_virtual_sentinels_preserve_tie_order_when_filled() {
1416 let mut collector = ScoreCollector::new(3);
1417 collector.seed_threshold(5.0);
1418
1419 assert!(collector.insert_with_ordinal(3, 5.0, 2));
1420 assert!(collector.insert_with_ordinal(2, 5.0, 8));
1421 assert!(collector.insert_with_ordinal(1, 5.0, 4));
1422 assert_eq!(collector.real_len(), 3);
1423 assert!(collector.virtual_threshold.is_none());
1424
1425 assert!(collector.insert_with_ordinal(2, 5.0, 1));
1428 assert!(!collector.insert_with_ordinal(4, 5.0, 0));
1429 assert_eq!(
1430 collector.into_sorted_results(),
1431 vec![(1, 5.0, 4), (2, 5.0, 1), (2, 5.0, 8)]
1432 );
1433 }
1434
1435 #[test]
1436 fn test_score_collector_basic() {
1437 let mut collector = ScoreCollector::new(3);
1438
1439 collector.insert(1, 1.0);
1440 collector.insert(2, 2.0);
1441 collector.insert(3, 3.0);
1442 assert_eq!(collector.threshold(), 1.0);
1443
1444 collector.insert(4, 4.0);
1445 assert_eq!(collector.threshold(), 2.0);
1446
1447 let results = collector.into_sorted_results();
1448 assert_eq!(results.len(), 3);
1449 assert_eq!(results[0].0, 4); assert_eq!(results[1].0, 3);
1451 assert_eq!(results[2].0, 2);
1452 }
1453
1454 #[test]
1455 fn test_score_collector_threshold() {
1456 let mut collector = ScoreCollector::new(2);
1457
1458 collector.insert(1, 5.0);
1459 collector.insert(2, 3.0);
1460 assert_eq!(collector.threshold(), 3.0);
1461
1462 assert!(!collector.would_enter(2.0));
1464 assert!(!collector.insert(3, 2.0));
1465
1466 assert!(collector.would_enter(4.0));
1468 assert!(collector.insert(4, 4.0));
1469 assert_eq!(collector.threshold(), 4.0);
1470 }
1471
1472 #[test]
1473 fn test_heap_entry_ordering() {
1474 let mut heap = BinaryHeap::new();
1475 heap.push(HeapEntry {
1476 doc_id: 1,
1477 score: 3.0,
1478 ordinal: 0,
1479 });
1480 heap.push(HeapEntry {
1481 doc_id: 2,
1482 score: 1.0,
1483 ordinal: 0,
1484 });
1485 heap.push(HeapEntry {
1486 doc_id: 3,
1487 score: 2.0,
1488 ordinal: 0,
1489 });
1490
1491 assert_eq!(heap.pop().unwrap().score, 1.0);
1493 assert_eq!(heap.pop().unwrap().score, 2.0);
1494 assert_eq!(heap.pop().unwrap().score, 3.0);
1495 }
1496}