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 denom_len_coeff: f32,
452 lengths: Option<&'a crate::segment::chunk_map::ChunkMap>,
455 tfs: Vec<u32>,
456 deferred_tf: Option<(usize, usize, usize)>,
459 },
460 Sparse {
462 si: &'a crate::segment::SparseIndex,
463 query_weight: f32,
464 skip_start: usize,
465 block_data_offset: u64,
466 },
467}
468
469macro_rules! cursor_ensure_block {
477 ($self:ident, $load_block_fn:ident, $($aw:tt)*) => {{
478 if $self.exhausted || $self.block_loaded {
479 return Ok(!$self.exhausted);
480 }
481 match &mut $self.variant {
482 CursorVariant::Text {
483 list,
484 deferred_tf,
485 ..
486 } => {
487 if let Some(state) = list.decode_block_doc_ids_only($self.block_idx, &mut $self.doc_ids) {
488 *deferred_tf = Some(state);
489 $self.scores.clear();
490 $self.pos = 0;
491 $self.block_loaded = true;
492 Ok(true)
493 } else {
494 $self.exhausted = true;
495 Ok(false)
496 }
497 }
498 CursorVariant::Sparse {
499 si,
500 query_weight,
501 skip_start,
502 block_data_offset,
503 ..
504 } => {
505 let block = si
506 .$load_block_fn(*skip_start, *block_data_offset, $self.block_idx)
507 $($aw)* ?;
508 match block {
509 Some(b) => {
510 b.decode_doc_ids_into(&mut $self.doc_ids);
511 b.decode_scored_weights_into(*query_weight, &mut $self.scores);
512 if $self.lazy_ordinals {
513 $self.current_sparse_block = Some(b);
516 $self.ordinals_loaded = false;
517 } else {
518 b.decode_ordinals_into(&mut $self.ordinals);
519 $self.ordinals_loaded = true;
520 $self.current_sparse_block = None;
521 }
522 $self.pos = 0;
523 $self.block_loaded = true;
524 Ok(true)
525 }
526 None => {
527 $self.exhausted = true;
528 Ok(false)
529 }
530 }
531 }
532 }
533 }};
534}
535
536macro_rules! cursor_advance {
537 ($self:ident, $ensure_fn:ident, $($aw:tt)*) => {{
538 if $self.exhausted {
539 return Ok(u32::MAX);
540 }
541 $self.$ensure_fn() $($aw)* ?;
542 if $self.exhausted {
543 return Ok(u32::MAX);
544 }
545 Ok($self.advance_pos())
546 }};
547}
548
549macro_rules! cursor_seek {
550 ($self:ident, $ensure_fn:ident, $target:expr, $($aw:tt)*) => {{
551 if let Some(doc) = $self.seek_prepare($target) {
552 return Ok(doc);
553 }
554 $self.$ensure_fn() $($aw)* ?;
555 if $self.seek_finish($target) {
556 $self.$ensure_fn() $($aw)* ?;
557 }
558 Ok($self.doc())
559 }};
560}
561
562impl<'a> TermCursor<'a> {
563 pub fn text(
565 posting_list: crate::structures::BlockPostingList,
566 idf: f32,
567 avg_field_len: f32,
568 ) -> Self {
569 Self::text_with_lengths(posting_list, idf, avg_field_len, None)
570 }
571
572 pub fn text_chunked(
575 posting_list: crate::structures::BlockPostingList,
576 idf: f32,
577 avg_chunk_len: f32,
578 lengths: &'a crate::segment::chunk_map::ChunkMap,
579 ) -> Self {
580 Self::text_with_lengths(posting_list, idf, avg_chunk_len, Some(lengths))
581 }
582
583 fn text_with_lengths(
584 posting_list: crate::structures::BlockPostingList,
585 idf: f32,
586 avg_field_len: f32,
587 lengths: Option<&'a crate::segment::chunk_map::ChunkMap>,
588 ) -> Self {
589 let max_tf = posting_list.max_tf() as f32;
590 let max_score = super::bm25_upper_bound(max_tf.max(1.0), idf);
591 let num_blocks = posting_list.num_blocks();
592 let safe_avg = avg_field_len.max(1.0);
593 Self {
594 max_score,
595 num_blocks,
596 block_idx: 0,
597 doc_ids: Vec::with_capacity(128),
598 scores: Vec::with_capacity(128),
599 ordinals: Vec::new(),
600 pos: 0,
601 block_loaded: false,
602 exhausted: num_blocks == 0,
603 lazy_ordinals: false,
604 ordinals_loaded: true, current_sparse_block: None,
606 variant: CursorVariant::Text {
607 list: posting_list,
608 idf,
609 idf_times_k1_plus_1: idf * (super::BM25_K1 + 1.0),
610 denom_tf_coeff: 1.0 + super::BM25_K1 * (super::BM25_B / safe_avg),
611 denom_const: super::BM25_K1 * (1.0 - super::BM25_B),
612 denom_len_coeff: super::BM25_K1 * super::BM25_B / safe_avg,
613 lengths,
614 tfs: Vec::with_capacity(128),
615 deferred_tf: None,
616 },
617 }
618 }
619
620 pub fn sparse(
623 si: &'a crate::segment::SparseIndex,
624 query_weight: f32,
625 skip_start: usize,
626 skip_count: usize,
627 global_max_weight: f32,
628 block_data_offset: u64,
629 ) -> Self {
630 Self {
631 max_score: query_weight.abs() * global_max_weight,
632 num_blocks: skip_count,
633 block_idx: 0,
634 doc_ids: Vec::with_capacity(256),
635 scores: Vec::with_capacity(256),
636 ordinals: Vec::with_capacity(256),
637 pos: 0,
638 block_loaded: false,
639 exhausted: skip_count == 0,
640 lazy_ordinals: false,
641 ordinals_loaded: true,
642 current_sparse_block: None,
643 variant: CursorVariant::Sparse {
644 si,
645 query_weight,
646 skip_start,
647 block_data_offset,
648 },
649 }
650 }
651
652 #[inline]
655 fn block_first_doc(&self, idx: usize) -> DocId {
656 match &self.variant {
657 CursorVariant::Text { list, .. } => list.block_first_doc(idx).unwrap_or(u32::MAX),
658 CursorVariant::Sparse { si, skip_start, .. } => {
659 si.read_skip_entry(*skip_start + idx).first_doc
660 }
661 }
662 }
663
664 #[inline]
665 fn block_last_doc(&self, idx: usize) -> DocId {
666 match &self.variant {
667 CursorVariant::Text { list, .. } => list.block_last_doc(idx).unwrap_or(0),
668 CursorVariant::Sparse { si, skip_start, .. } => {
669 si.read_skip_entry(*skip_start + idx).last_doc
670 }
671 }
672 }
673
674 #[inline]
677 pub fn doc(&self) -> DocId {
678 if self.exhausted {
679 return u32::MAX;
680 }
681 if self.block_loaded {
682 debug_assert!(self.pos < self.doc_ids.len());
683 unsafe { *self.doc_ids.get_unchecked(self.pos) }
685 } else {
686 self.block_first_doc(self.block_idx)
687 }
688 }
689
690 #[inline]
691 pub fn ordinal(&self) -> u16 {
692 if !self.block_loaded || self.ordinals.is_empty() {
693 return 0;
694 }
695 debug_assert!(self.pos < self.ordinals.len());
696 unsafe { *self.ordinals.get_unchecked(self.pos) }
698 }
699
700 #[inline]
706 pub fn ordinal_mut(&mut self) -> u16 {
707 if !self.block_loaded {
708 return 0;
709 }
710 if !self.ordinals_loaded {
711 if let Some(ref block) = self.current_sparse_block {
712 block.decode_ordinals_into(&mut self.ordinals);
713 }
714 self.ordinals_loaded = true;
715 }
716 if self.ordinals.is_empty() {
717 return 0;
718 }
719 debug_assert!(self.pos < self.ordinals.len());
720 unsafe { *self.ordinals.get_unchecked(self.pos) }
721 }
722
723 #[inline]
724 pub fn score(&self) -> f32 {
725 if !self.block_loaded {
726 return 0.0;
727 }
728 debug_assert!(self.pos < self.scores.len());
729 unsafe { *self.scores.get_unchecked(self.pos) }
731 }
732
733 #[inline]
739 pub fn ensure_scores(&mut self) {
740 if self.block_loaded && self.scores.is_empty() {
741 self.compute_deferred_scores();
742 }
743 }
744
745 #[inline]
746 pub fn current_block_max_score(&self) -> f32 {
747 if self.exhausted {
748 return 0.0;
749 }
750 match &self.variant {
751 CursorVariant::Text { list, idf, .. } => {
752 let block_max_tf = list.block_max_tf(self.block_idx).unwrap_or(0) as f32;
753 super::bm25_upper_bound(block_max_tf.max(1.0), *idf)
754 }
755 CursorVariant::Sparse {
756 si,
757 query_weight,
758 skip_start,
759 ..
760 } => query_weight.abs() * si.read_skip_entry(*skip_start + self.block_idx).max_weight,
761 }
762 }
763
764 pub fn skip_to_next_block(&mut self) -> DocId {
767 if self.exhausted {
768 return u32::MAX;
769 }
770 self.block_idx += 1;
771 self.block_loaded = false;
772 if self.block_idx >= self.num_blocks {
773 self.exhausted = true;
774 return u32::MAX;
775 }
776 self.block_first_doc(self.block_idx)
777 }
778
779 #[inline]
780 fn advance_pos(&mut self) -> DocId {
781 self.pos += 1;
782 if self.pos >= self.doc_ids.len() {
783 self.block_idx += 1;
784 self.block_loaded = false;
785 if self.block_idx >= self.num_blocks {
786 self.exhausted = true;
787 return u32::MAX;
788 }
789 }
790 self.doc()
791 }
792
793 #[inline(never)]
795 fn compute_deferred_scores(&mut self) {
796 if let CursorVariant::Text {
797 list,
798 idf_times_k1_plus_1,
799 denom_tf_coeff,
800 denom_const,
801 denom_len_coeff,
802 lengths,
803 tfs,
804 deferred_tf,
805 ..
806 } = &mut self.variant
807 && let Some((block_offset, tf_start, count)) = deferred_tf.take()
808 {
809 list.decode_block_tfs_deferred(block_offset, tf_start, count, tfs);
810 let num_scale = *idf_times_k1_plus_1;
811 let d_tf = *denom_tf_coeff;
812 let d_const = *denom_const;
813 let d_len = *denom_len_coeff;
814 self.scores.clear();
815 self.scores.resize(count, 0.0);
816 match lengths {
817 Some(map) => {
819 for i in 0..count {
820 let tf = unsafe { *tfs.get_unchecked(i) } as f32;
821 let vid = unsafe { *self.doc_ids.get_unchecked(i) };
822 let len = map.length(vid) as f32;
823 let score = (num_scale * tf) / (tf + d_const + d_len * len);
824 unsafe {
825 *self.scores.get_unchecked_mut(i) = score;
826 }
827 }
828 }
829 None => {
830 for i in 0..count {
831 let tf = unsafe { *tfs.get_unchecked(i) } as f32;
832 let score = (num_scale * tf) / (d_tf * tf + d_const);
833 unsafe {
834 *self.scores.get_unchecked_mut(i) = score;
835 }
836 }
837 }
838 }
839 }
840 }
841
842 pub async fn ensure_block_loaded(&mut self) -> crate::Result<bool> {
848 cursor_ensure_block!(self, load_block_direct, .await)
849 }
850
851 pub fn ensure_block_loaded_sync(&mut self) -> crate::Result<bool> {
852 cursor_ensure_block!(self, load_block_direct_sync,)
853 }
854
855 pub async fn advance(&mut self) -> crate::Result<DocId> {
856 cursor_advance!(self, ensure_block_loaded, .await)
857 }
858
859 pub fn advance_sync(&mut self) -> crate::Result<DocId> {
860 cursor_advance!(self, ensure_block_loaded_sync,)
861 }
862
863 pub async fn seek(&mut self, target: DocId) -> crate::Result<DocId> {
864 cursor_seek!(self, ensure_block_loaded, target, .await)
865 }
866
867 pub fn seek_sync(&mut self, target: DocId) -> crate::Result<DocId> {
868 cursor_seek!(self, ensure_block_loaded_sync, target,)
869 }
870
871 fn seek_prepare(&mut self, target: DocId) -> Option<DocId> {
872 if self.exhausted {
873 return Some(u32::MAX);
874 }
875
876 if self.block_loaded
878 && let Some(&last) = self.doc_ids.last()
879 {
880 if last >= target && self.doc_ids[self.pos] < target {
881 let remaining = &self.doc_ids[self.pos..];
882 self.pos += crate::structures::simd::find_first_ge_u32(remaining, target);
883 if self.pos >= self.doc_ids.len() {
884 self.block_idx += 1;
885 self.block_loaded = false;
886 if self.block_idx >= self.num_blocks {
887 self.exhausted = true;
888 return Some(u32::MAX);
889 }
890 }
891 return Some(self.doc());
892 }
893 if self.doc_ids[self.pos] >= target {
894 return Some(self.doc());
895 }
896 }
897
898 let lo = match &self.variant {
900 CursorVariant::Text { list, .. } => match list.seek_block(target, self.block_idx) {
902 Some(idx) => idx,
903 None => {
904 self.exhausted = true;
905 return Some(u32::MAX);
906 }
907 },
908 CursorVariant::Sparse { .. } => {
910 let mut lo = self.block_idx;
911 let mut hi = self.num_blocks;
912 while lo < hi {
913 let mid = lo + (hi - lo) / 2;
914 if self.block_last_doc(mid) < target {
915 lo = mid + 1;
916 } else {
917 hi = mid;
918 }
919 }
920 lo
921 }
922 };
923 if lo >= self.num_blocks {
924 self.exhausted = true;
925 return Some(u32::MAX);
926 }
927 if lo != self.block_idx || !self.block_loaded {
928 self.block_idx = lo;
929 self.block_loaded = false;
930 }
931 None
932 }
933
934 #[inline]
935 fn seek_finish(&mut self, target: DocId) -> bool {
936 if self.exhausted {
937 return false;
938 }
939 self.pos = crate::structures::simd::find_first_ge_u32(&self.doc_ids, target);
940 if self.pos >= self.doc_ids.len() {
941 self.block_idx += 1;
942 self.block_loaded = false;
943 if self.block_idx >= self.num_blocks {
944 self.exhausted = true;
945 return false;
946 }
947 return true;
948 }
949 false
950 }
951}
952
953macro_rules! bms_execute_loop {
958 ($self:ident, $ensure:ident, $advance:ident, $seek:ident, $($aw:tt)*) => {{
959 let n = $self.cursors.len();
960
961 for cursor in &mut $self.cursors {
963 cursor.$ensure() $($aw)* ?;
964 }
965
966 let mut docs_scored = 0u64;
967 let mut docs_skipped = 0u64;
968 let mut blocks_skipped = 0u64;
969 let mut conjunction_skipped = 0u64;
970 let mut ordinal_scores: Vec<(u16, f32)> = Vec::with_capacity(n * 2);
971 let _bms_start = std::time::Instant::now();
972
973 let inv_heap_factor = $self.inv_heap_factor;
974 let mut adjusted_threshold = $self.collector.threshold() * inv_heap_factor - 1e-6;
975
976 loop {
977 let partition = $self.find_partition();
978 if partition >= n {
979 break;
980 }
981
982 let mut min_doc = u32::MAX;
986 let mut at_min_mask = 0u64; for i in partition..n {
988 let doc = $self.cursors[i].doc();
989 match doc.cmp(&min_doc) {
990 std::cmp::Ordering::Less => {
991 min_doc = doc;
992 at_min_mask = 1u64 << (i as u32);
993 }
994 std::cmp::Ordering::Equal => {
995 at_min_mask |= 1u64 << (i as u32);
996 }
997 _ => {}
998 }
999 }
1000 if min_doc == u32::MAX {
1001 break;
1002 }
1003
1004 let non_essential_upper = if partition > 0 {
1005 $self.prefix_sums[partition - 1]
1006 } else {
1007 0.0
1008 };
1009
1010 if $self.collector.len() >= $self.collector.k {
1012 let mut present_upper: f32 = 0.0;
1013 let mut mask = at_min_mask;
1014 while mask != 0 {
1015 let i = mask.trailing_zeros() as usize;
1016 present_upper += $self.cursors[i].max_score;
1017 mask &= mask - 1;
1018 }
1019
1020 if present_upper + non_essential_upper < adjusted_threshold {
1021 let mut mask = at_min_mask;
1022 while mask != 0 {
1023 let i = mask.trailing_zeros() as usize;
1024 $self.cursors[i].$ensure() $($aw)* ?;
1025 $self.cursors[i].$advance() $($aw)* ?;
1026 mask &= mask - 1;
1027 }
1028 conjunction_skipped += 1;
1029 continue;
1030 }
1031 }
1032
1033 if $self.collector.len() >= $self.collector.k {
1035 let mut block_max_sum: f32 = 0.0;
1036 let mut mask = at_min_mask;
1037 while mask != 0 {
1038 let i = mask.trailing_zeros() as usize;
1039 block_max_sum += $self.cursors[i].current_block_max_score();
1040 mask &= mask - 1;
1041 }
1042
1043 if block_max_sum + non_essential_upper < adjusted_threshold {
1044 let mut mask = at_min_mask;
1045 while mask != 0 {
1046 let i = mask.trailing_zeros() as usize;
1047 $self.cursors[i].skip_to_next_block();
1048 $self.cursors[i].$ensure() $($aw)* ?;
1049 mask &= mask - 1;
1050 }
1051 blocks_skipped += 1;
1052 continue;
1053 }
1054 }
1055
1056 if let Some(ref pred) = $self.predicate {
1058 if !pred(min_doc) {
1059 let mut mask = at_min_mask;
1060 while mask != 0 {
1061 let i = mask.trailing_zeros() as usize;
1062 $self.cursors[i].$ensure() $($aw)* ?;
1063 $self.cursors[i].$advance() $($aw)* ?;
1064 mask &= mask - 1;
1065 }
1066 continue;
1067 }
1068 }
1069
1070 ordinal_scores.clear();
1072 {
1073 let mut mask = at_min_mask;
1074 while mask != 0 {
1075 let i = mask.trailing_zeros() as usize;
1076 $self.cursors[i].$ensure() $($aw)* ?;
1077 $self.cursors[i].ensure_scores();
1078 while $self.cursors[i].doc() == min_doc {
1079 let ord = $self.cursors[i].ordinal_mut();
1080 let sc = $self.cursors[i].score();
1081 ordinal_scores.push((ord, sc));
1082 $self.cursors[i].$advance() $($aw)* ?;
1083 }
1084 mask &= mask - 1;
1085 }
1086 }
1087
1088 let essential_total: f32 = ordinal_scores.iter().map(|(_, s)| *s).sum();
1089 if $self.collector.len() >= $self.collector.k
1090 && essential_total + non_essential_upper < adjusted_threshold
1091 {
1092 docs_skipped += 1;
1093 continue;
1094 }
1095
1096 let mut running_total = essential_total;
1098 for i in (0..partition).rev() {
1099 if $self.collector.len() >= $self.collector.k
1100 && running_total + $self.prefix_sums[i] < adjusted_threshold
1101 {
1102 break;
1103 }
1104
1105 let doc = $self.cursors[i].$seek(min_doc) $($aw)* ?;
1106 if doc == min_doc {
1107 $self.cursors[i].ensure_scores();
1108 while $self.cursors[i].doc() == min_doc {
1109 let s = $self.cursors[i].score();
1110 running_total += s;
1111 let ord = $self.cursors[i].ordinal_mut();
1112 ordinal_scores.push((ord, s));
1113 $self.cursors[i].$advance() $($aw)* ?;
1114 }
1115 }
1116 }
1117
1118 if ordinal_scores.len() == 1 {
1121 let (ord, score) = ordinal_scores[0];
1122 if $self.collector.insert_with_ordinal(min_doc, score, ord) {
1123 docs_scored += 1;
1124 adjusted_threshold = $self.collector.threshold() * inv_heap_factor - 1e-6;
1125 } else {
1126 docs_skipped += 1;
1127 }
1128 } else if !ordinal_scores.is_empty() {
1129 if ordinal_scores.len() > 2 {
1130 ordinal_scores.sort_unstable_by_key(|(ord, _)| *ord);
1131 } else if ordinal_scores.len() == 2 && ordinal_scores[0].0 > ordinal_scores[1].0 {
1132 ordinal_scores.swap(0, 1);
1133 }
1134 let mut j = 0;
1135 while j < ordinal_scores.len() {
1136 let current_ord = ordinal_scores[j].0;
1137 let mut score = 0.0f32;
1138 while j < ordinal_scores.len() && ordinal_scores[j].0 == current_ord {
1139 score += ordinal_scores[j].1;
1140 j += 1;
1141 }
1142 if $self
1143 .collector
1144 .insert_with_ordinal(min_doc, score, current_ord)
1145 {
1146 docs_scored += 1;
1147 adjusted_threshold = $self.collector.threshold() * inv_heap_factor - 1e-6;
1148 } else {
1149 docs_skipped += 1;
1150 }
1151 }
1152 }
1153 }
1154
1155 let results: Vec<ScoredDoc> = $self
1156 .collector
1157 .into_sorted_results()
1158 .into_iter()
1159 .map(|(doc_id, score, ordinal)| ScoredDoc {
1160 doc_id,
1161 score,
1162 ordinal,
1163 })
1164 .collect();
1165
1166 let _bms_elapsed_ms = _bms_start.elapsed().as_millis() as u64;
1167 if _bms_elapsed_ms > 500 {
1168 warn!(
1169 "slow MaxScore: {}ms, cursors={}, scored={}, skipped={}, blocks_skipped={}, conjunction_skipped={}, returned={}, top_score={:.4}",
1170 _bms_elapsed_ms,
1171 n,
1172 docs_scored,
1173 docs_skipped,
1174 blocks_skipped,
1175 conjunction_skipped,
1176 results.len(),
1177 results.first().map(|r| r.score).unwrap_or(0.0)
1178 );
1179 } else {
1180 debug!(
1181 "MaxScoreExecutor: {}ms, scored={}, skipped={}, blocks_skipped={}, conjunction_skipped={}, returned={}, top_score={:.4}",
1182 _bms_elapsed_ms,
1183 docs_scored,
1184 docs_skipped,
1185 blocks_skipped,
1186 conjunction_skipped,
1187 results.len(),
1188 results.first().map(|r| r.score).unwrap_or(0.0)
1189 );
1190 }
1191
1192 Ok(results)
1193 }};
1194}
1195
1196impl<'a> MaxScoreExecutor<'a> {
1197 pub(crate) fn new(mut cursors: Vec<TermCursor<'a>>, k: usize, heap_factor: f32) -> Self {
1202 if cursors.len() > super::MAX_QUERY_TERMS {
1206 cursors.sort_unstable_by(|a, b| b.max_score.total_cmp(&a.max_score));
1207 cursors.truncate(super::MAX_QUERY_TERMS);
1208 log::warn!(
1209 "MaxScore cursor count exceeded {}; retaining the strongest cursors",
1210 super::MAX_QUERY_TERMS
1211 );
1212 }
1213
1214 for c in &mut cursors {
1217 c.lazy_ordinals = true;
1218 }
1219
1220 cursors.sort_by(|a, b| {
1222 a.max_score
1223 .partial_cmp(&b.max_score)
1224 .unwrap_or(Ordering::Equal)
1225 });
1226
1227 let mut prefix_sums = Vec::with_capacity(cursors.len());
1228 let mut cumsum = 0.0f32;
1229 for c in &cursors {
1230 cumsum += c.max_score;
1231 prefix_sums.push(cumsum);
1232 }
1233
1234 let clamped_heap_factor = heap_factor.clamp(0.01, 1.0);
1235
1236 debug!(
1237 "Creating MaxScoreExecutor: num_cursors={}, k={}, total_upper={:.4}, heap_factor={:.2}",
1238 cursors.len(),
1239 k,
1240 cumsum,
1241 clamped_heap_factor
1242 );
1243
1244 Self {
1245 cursors,
1246 prefix_sums,
1247 collector: ScoreCollector::new(k),
1248 inv_heap_factor: 1.0 / clamped_heap_factor,
1249 predicate: None,
1250 metric_index: "unknown",
1251 metric_field: "unknown",
1252 }
1253 }
1254
1255 pub fn with_metric_labels(mut self, index: &'a str, field: &'a str) -> Self {
1257 self.metric_index = index;
1258 self.metric_field = field;
1259 self
1260 }
1261
1262 pub fn sparse(
1266 sparse_index: &'a crate::segment::SparseIndex,
1267 query_terms: Vec<(u32, f32)>,
1268 k: usize,
1269 heap_factor: f32,
1270 ) -> Self {
1271 let cursors: Vec<TermCursor<'a>> = query_terms
1272 .iter()
1273 .filter_map(|&(dim_id, qw)| {
1274 let (skip_start, skip_count, global_max, block_data_offset) =
1275 sparse_index.get_skip_range_full(dim_id)?;
1276 Some(TermCursor::sparse(
1277 sparse_index,
1278 qw,
1279 skip_start,
1280 skip_count,
1281 global_max,
1282 block_data_offset,
1283 ))
1284 })
1285 .collect();
1286 Self::new(cursors, k, heap_factor)
1287 }
1288
1289 pub fn text(
1293 posting_lists: Vec<(crate::structures::BlockPostingList, f32)>,
1294 avg_field_len: f32,
1295 k: usize,
1296 ) -> Self {
1297 let cursors: Vec<TermCursor<'a>> = posting_lists
1298 .into_iter()
1299 .map(|(pl, idf)| TermCursor::text(pl, idf, avg_field_len))
1300 .collect();
1301 Self::new(cursors, k, 1.0)
1302 }
1303
1304 pub fn text_chunked(
1308 posting_lists: Vec<(crate::structures::BlockPostingList, f32)>,
1309 avg_chunk_len: f32,
1310 k: usize,
1311 lengths: &'a crate::segment::chunk_map::ChunkMap,
1312 ) -> Self {
1313 let cursors: Vec<TermCursor<'a>> = posting_lists
1314 .into_iter()
1315 .map(|(pl, idf)| TermCursor::text_chunked(pl, idf, avg_chunk_len, lengths))
1316 .collect();
1317 Self::new(cursors, k, 1.0)
1318 }
1319
1320 #[inline]
1321 fn find_partition(&self) -> usize {
1322 let threshold = self.collector.threshold() * self.inv_heap_factor;
1326 self.prefix_sums.partition_point(|&sum| sum < threshold)
1329 }
1330
1331 pub fn with_predicate(mut self, predicate: super::DocPredicate<'a>) -> Self {
1337 self.predicate = Some(predicate);
1338 self
1339 }
1340
1341 pub fn seed_threshold(&mut self, initial_threshold: f32) {
1343 self.collector.seed_threshold(initial_threshold);
1344 }
1345
1346 pub async fn execute(mut self) -> crate::Result<Vec<ScoredDoc>> {
1348 if self.cursors.is_empty() {
1349 return Ok(Vec::new());
1350 }
1351 let t = crate::observe::Timer::start();
1352 let results = bms_execute_loop!(self, ensure_block_loaded, advance, seek, .await);
1353 if let Ok(r) = &results {
1354 crate::observe::maxscore_query(self.metric_index, self.metric_field, t.secs(), r.len());
1355 }
1356 results
1357 }
1358
1359 pub fn execute_sync(mut self) -> crate::Result<Vec<ScoredDoc>> {
1361 if self.cursors.is_empty() {
1362 return Ok(Vec::new());
1363 }
1364 let t = crate::observe::Timer::start();
1365 let results = bms_execute_loop!(self, ensure_block_loaded_sync, advance_sync, seek_sync,);
1366 if let Ok(r) = &results {
1367 crate::observe::maxscore_query(self.metric_index, self.metric_field, t.secs(), r.len());
1368 }
1369 results
1370 }
1371}
1372
1373#[cfg(test)]
1374mod tests {
1375 use super::*;
1376
1377 #[test]
1378 fn test_shared_threshold_monotonic_raise() {
1379 let shared = SharedThreshold::new();
1380 assert_eq!(shared.get(), 0.0);
1381
1382 shared.raise(2.5);
1383 assert_eq!(shared.get(), 2.5);
1384
1385 shared.raise(1.0);
1387 assert_eq!(shared.get(), 2.5);
1388
1389 shared.raise(4.0);
1391 assert_eq!(shared.get(), 4.0);
1392
1393 shared.raise(0.0);
1395 shared.raise(-3.0);
1396 shared.raise(f32::NAN);
1397 assert_eq!(shared.get(), 4.0);
1398
1399 let clone = shared.clone();
1401 clone.raise(9.0);
1402 assert_eq!(shared.get(), 9.0);
1403 }
1404
1405 #[test]
1406 fn test_shared_threshold_seed_matches_manual() {
1407 let mut seeded = ScoreCollector::new(2);
1410 seeded.seed_threshold(3.0);
1411 assert_eq!(seeded.threshold(), 3.0);
1412 assert!(!seeded.would_enter(3.0));
1414 assert!(seeded.would_enter(3.5));
1415 seeded.insert(1, 5.0);
1418 seeded.insert(2, 4.0);
1419 let results = seeded.into_sorted_results();
1420 assert_eq!(results.len(), 2);
1421 assert_eq!(results[0].0, 1);
1422 assert_eq!(results[1].0, 2);
1423 }
1424
1425 #[test]
1426 fn test_shared_threshold_can_raise_after_real_inserts() {
1427 let mut collector = ScoreCollector::new(3);
1428 collector.insert(1, 10.0);
1429 collector.insert(2, 4.0);
1430 assert_eq!(collector.real_len(), 2);
1431
1432 collector.seed_threshold(6.0);
1435 assert_eq!(collector.threshold(), 6.0);
1436 assert_eq!(collector.real_len(), 1);
1437
1438 assert!(collector.would_enter_candidate(3, 6.0, 0));
1441 assert!(collector.insert(3, 6.0));
1442 assert_eq!(collector.real_len(), 2);
1443 let results = collector.into_sorted_results();
1444 assert_eq!(results, vec![(1, 10.0, 0), (3, 6.0, 0)]);
1445 }
1446
1447 #[test]
1448 fn test_large_seed_uses_virtual_sentinels() {
1449 let k = 1_000_000_000;
1450 let mut collector = ScoreCollector::new(k);
1451 assert!(collector.heap.capacity() <= MAX_INITIAL_SCORE_COLLECTOR_CAPACITY);
1452
1453 collector.seed_threshold(42.0);
1454
1455 assert_eq!(collector.heap.len(), 0);
1458 assert!(collector.heap.capacity() <= MAX_INITIAL_SCORE_COLLECTOR_CAPACITY);
1459 assert_eq!(collector.len(), k);
1460 assert_eq!(collector.real_len(), 0);
1461 assert_eq!(collector.threshold(), 42.0);
1462 assert!(!collector.is_empty());
1463
1464 assert!(collector.insert_with_ordinal(9, 42.0, 7));
1467 assert!(!collector.insert(10, 41.0));
1468 assert!(collector.insert(11, 43.0));
1469 assert_eq!(collector.len(), k);
1470 assert_eq!(collector.real_len(), 2);
1471 assert_eq!(
1472 collector.into_sorted_results(),
1473 vec![(11, 43.0, 0), (9, 42.0, 7)]
1474 );
1475 }
1476
1477 #[test]
1478 fn test_virtual_sentinels_preserve_tie_order_when_filled() {
1479 let mut collector = ScoreCollector::new(3);
1480 collector.seed_threshold(5.0);
1481
1482 assert!(collector.insert_with_ordinal(3, 5.0, 2));
1483 assert!(collector.insert_with_ordinal(2, 5.0, 8));
1484 assert!(collector.insert_with_ordinal(1, 5.0, 4));
1485 assert_eq!(collector.real_len(), 3);
1486 assert!(collector.virtual_threshold.is_none());
1487
1488 assert!(collector.insert_with_ordinal(2, 5.0, 1));
1491 assert!(!collector.insert_with_ordinal(4, 5.0, 0));
1492 assert_eq!(
1493 collector.into_sorted_results(),
1494 vec![(1, 5.0, 4), (2, 5.0, 1), (2, 5.0, 8)]
1495 );
1496 }
1497
1498 #[test]
1499 fn test_score_collector_basic() {
1500 let mut collector = ScoreCollector::new(3);
1501
1502 collector.insert(1, 1.0);
1503 collector.insert(2, 2.0);
1504 collector.insert(3, 3.0);
1505 assert_eq!(collector.threshold(), 1.0);
1506
1507 collector.insert(4, 4.0);
1508 assert_eq!(collector.threshold(), 2.0);
1509
1510 let results = collector.into_sorted_results();
1511 assert_eq!(results.len(), 3);
1512 assert_eq!(results[0].0, 4); assert_eq!(results[1].0, 3);
1514 assert_eq!(results[2].0, 2);
1515 }
1516
1517 #[test]
1518 fn test_score_collector_threshold() {
1519 let mut collector = ScoreCollector::new(2);
1520
1521 collector.insert(1, 5.0);
1522 collector.insert(2, 3.0);
1523 assert_eq!(collector.threshold(), 3.0);
1524
1525 assert!(!collector.would_enter(2.0));
1527 assert!(!collector.insert(3, 2.0));
1528
1529 assert!(collector.would_enter(4.0));
1531 assert!(collector.insert(4, 4.0));
1532 assert_eq!(collector.threshold(), 4.0);
1533 }
1534
1535 #[test]
1536 fn test_heap_entry_ordering() {
1537 let mut heap = BinaryHeap::new();
1538 heap.push(HeapEntry {
1539 doc_id: 1,
1540 score: 3.0,
1541 ordinal: 0,
1542 });
1543 heap.push(HeapEntry {
1544 doc_id: 2,
1545 score: 1.0,
1546 ordinal: 0,
1547 });
1548 heap.push(HeapEntry {
1549 doc_id: 3,
1550 score: 2.0,
1551 ordinal: 0,
1552 });
1553
1554 assert_eq!(heap.pop().unwrap().score, 1.0);
1556 assert_eq!(heap.pop().unwrap().score, 2.0);
1557 assert_eq!(heap.pop().unwrap().score, 3.0);
1558 }
1559}