fix: Allow BMW to be used for compound ORDER BY with a score prefix (#167)

Remove the override of `collect_segment_top_k` for
`SortBySimilarityScore`, and instead bake score awareness into the
default implementation of that method.

To do so, introduces the `supports_bm25_pruning` and
`bm25_pruning_threshold` methods to `SortKeyComputer`, which can pass
through to wrapped implementations if BMW pruning is safe (for example,
if a score is a prefix).

This additionally allows for simplifying `SharedThreshold`, because
equal values and tiebreakers are now handled directly in
`collect_segment_top_k`.

Benchmarks in
paradedb/paradedb#5496 (review)
demonstrate a speedup for `top_k-score-desc-tiebreaker`.

Author: stuhood

src/collector/sort_key/mod.rs

@@ -457,5 +457,61 @@ pub(crate) mod tests {
                 top_n_results
             );
         }
+
+        #[test]
+        fn test_order_by_score_then_fast_field_prop(
+            limit in 1..64_usize,
+            offset in 0..64_usize,
+            segments_docs in
+              proptest::collection::vec(
+                  proptest::collection::vec(
+                      // (doc_terms, fast_field_val)
+                      (proptest::collection::vec(0..5_u8, 1..10_usize), 0..100_u64),
+                      1..32_usize
+                  ),
+                  0..8_usize,
+              )
+        ) {
+            use crate::query::TermQuery;
+            use crate::schema::IndexRecordOption;
+            use crate::Term;
+
+            let mut schema_builder = Schema::builder();
+            let text = schema_builder.add_text_field("text", TEXT);
+            let value = schema_builder.add_u64_field("value", FAST);
+            let schema = schema_builder.build();
+            let index = Index::create_in_ram(schema);
+            let mut index_writer = index.writer_for_tests().unwrap();
+
+            for segment_docs in segments_docs.into_iter() {
+                for (term_ids, val) in segment_docs.into_iter() {
+                    let term_str = term_ids.into_iter().map(|id| format!("term{id}")).collect::<Vec<_>>().join(" ");
+                    index_writer.add_document(doc!(
+                        text => term_str,
+                        value => val,
+                    )).unwrap();
+                }
+                index_writer.commit().unwrap();
+            }
+
+            let searcher = index.reader().unwrap().searcher();
+            let query = TermQuery::new(Term::from_field_text(text, "term0"), IndexRecordOption::WithFreqs);
+
+            // 1. TopDocs with BMW (Score, u64)
+            let top_collector = TopDocs::with_limit(limit).and_offset(offset).order_by::<(Score, Option<u64>)>((
+                (SortBySimilarityScore::new(), Order::Desc),
+                (SortByStaticFastValue::for_field("value"), Order::Desc),
+            ));
+            let top_n_results: Vec<((Score, Option<u64>), DocAddress)> = searcher.search(&query, &top_collector).unwrap();
+
+            // 2. Golden baseline: collect all with scoring
+            let all_collector = TopDocs::with_limit(1000).order_by::<(Score, Option<u64>)>((
+                (SortBySimilarityScore::new(), Order::Desc),
+                (SortByStaticFastValue::for_field("value"), Order::Desc),
+            ));
+            let expected_docs = searcher.search(&query, &all_collector).unwrap().into_iter().skip(offset).take(limit).collect::<Vec<_>>();
+
+            prop_assert_eq!(expected_docs, top_n_results);
+        }
     }
 }

src/collector/sort_key/order.rs

@@ -524,9 +524,14 @@ where
         segment_reader: &crate::SegmentReader,
     ) -> crate::Result<Self::Child> {
         let child = self.0.segment_sort_key_computer(segment_reader)?;
+        let bmw_supported = matches!(
+            self.1,
+            ComparatorEnum::Natural | ComparatorEnum::NaturalNoneHigher
+        );
         Ok(SegmentSortKeyComputerWithComparator {
             segment_sort_key_computer: child,
             comparator: self.comparator(),
+            bmw_supported,
         })
     }
 }
@@ -570,9 +575,11 @@ where
         segment_reader: &crate::SegmentReader,
     ) -> crate::Result<Self::Child> {
         let child = self.0.segment_sort_key_computer(segment_reader)?;
+        let bmw_supported = self.1 == Order::Desc;
         Ok(SegmentSortKeyComputerWithComparator {
             segment_sort_key_computer: child,
             comparator: self.comparator(),
+            bmw_supported,
         })
     }
 }
@@ -581,6 +588,7 @@ where
 pub struct SegmentSortKeyComputerWithComparator<TSegmentSortKeyComputer, TComparator> {
     segment_sort_key_computer: TSegmentSortKeyComputer,
     comparator: TComparator,
+    bmw_supported: bool,
 }
 
 impl<TSegmentSortKeyComputer, TSegmentSortKey, TComparator> SegmentSortKeyComputer
@@ -611,6 +619,20 @@ where
         self.segment_sort_key_computer
             .convert_segment_sort_key(sort_key)
     }
+
+    fn supports_bm25_pruning(&self) -> bool {
+        self.bmw_supported && self.segment_sort_key_computer.supports_bm25_pruning()
+    }
+
+    fn bm25_pruning_threshold(
+        &self,
+        threshold: &Self::SegmentSortKey,
+        segment_ord: crate::SegmentOrdinal,
+        threshold_ord: crate::SegmentOrdinal,
+    ) -> Option<Score> {
+        self.segment_sort_key_computer
+            .bm25_pruning_threshold(threshold, segment_ord, threshold_ord)
+    }
 }
 
 #[cfg(test)]

src/collector/sort_key/shared_threshold.rs

@@ -34,35 +34,6 @@ pub trait SharedThreshold<T>: Send + Sync {
         expected_threshold: &Option<(T, SegmentOrdinal)>,
         new_threshold: (T, SegmentOrdinal),
     ) -> Result<(), Option<(T, SegmentOrdinal)>>;
-
-    /// Returns a threshold value `T` that is strictly better than the current shared threshold
-    /// if the given `segment_ord` would lose a tie-break against the `threshold_ord`.
-    ///
-    /// This is used for Block-Max WAND pruning to ensure determinism when multiple segments
-    /// have documents with the same score. Currently, this is only used for score-based
-    /// pushdown; for other types, it can safely return the value unchanged.
-    ///
-    /// # Tie-breaking Logic
-    /// Tantivy breaks ties on sort keys by favoring documents with a lower [`DocAddress`].
-    /// A [`DocAddress`] is composed of a `segment_ord` and a `doc_id`. Since documents within
-    /// a segment are processed in ascending `doc_id` order, any new document in a segment
-    /// with `segment_ord >= threshold_ord` will have a strictly higher [`DocAddress`] than the
-    /// document that set the threshold.
-    ///
-    /// Therefore, if `segment_ord >= threshold_ord`, this method should return a threshold that
-    /// is "one step better" than `value`. For floating-point scores, this is typically
-    /// `value.next_up()`. If `segment_ord < threshold_ord`, the document could still win the
-    /// tie-break, so the `value` should be returned unchanged.
-    fn competitive_threshold(
-        &self,
-        value: T,
-        threshold_ord: SegmentOrdinal,
-        segment_ord: SegmentOrdinal,
-    ) -> T {
-        let _ = threshold_ord;
-        let _ = segment_ord;
-        value
-    }
 }
 
 #[inline]
@@ -151,25 +122,6 @@ impl SharedThreshold<Score> for AtomicSharedThreshold {
             }
         }
     }
-
-    fn competitive_threshold(
-        &self,
-        value: Score,
-        threshold_ord: SegmentOrdinal,
-        segment_ord: SegmentOrdinal,
-    ) -> Score {
-        if segment_ord < threshold_ord {
-            // If our segment wins tie-breaks against the threshold setter, we want to accept
-            // documents with a score GREATER THAN OR EQUAL to the threshold.
-            // Since Tantivy's pruning loop uses a strict `score > threshold` check, we
-            // return `next_down()` to effectively relax the check to `>=`.
-            value.next_down()
-        } else {
-            // If our segment loses tie-breaks, we demand a score STRICTLY GREATER than the
-            // threshold. The pruning loop's `score > threshold` check already achieves this.
-            value
-        }
-    }
 }
 
 /// A shared threshold for `Option<u64>` values.
@@ -215,15 +167,6 @@ impl SharedThreshold<Option<u64>> for RwLockSharedThresholdOptionU64 {
             Err(*guard)
         }
     }
-
-    fn competitive_threshold(
-        &self,
-        value: Option<u64>,
-        _threshold_ord: SegmentOrdinal,
-        _segment_ord: SegmentOrdinal,
-    ) -> Option<u64> {
-        value
-    }
 }
 
 #[cfg(test)]
@@ -305,34 +248,6 @@ mod tests {
         assert_eq!(t.load(), Some((0.9, 10)));
     }
 
-    #[test]
-    fn test_competitive_threshold() {
-        let t = AtomicSharedThreshold::default();
-        let is_better = |a: &Score, a_ord: SegmentOrdinal, b: &Score, b_ord: SegmentOrdinal| {
-            if a > b {
-                true
-            } else if a == b {
-                a_ord < b_ord
-            } else {
-                false
-            }
-        };
-        update_helper(&t, 0.5, 5, &(), &is_better);
-
-        // Segment 5 itself: should require strictly greater score for new docs.
-        // The pruning loop uses `score > threshold`, so returning 0.5 unchanged
-        // means we only accept scores > 0.5.
-        assert_eq!(t.competitive_threshold(0.5, 5, 5), 0.5);
-
-        // Segment 6 (later): should require strictly greater score.
-        assert_eq!(t.competitive_threshold(0.5, 5, 6), 0.5);
-
-        // Segment 4 (earlier): can accept equal score (>= 0.5).
-        // Since pruning loop is `score > threshold`, we must return `next_down(0.5)`
-        // so that `0.5 > next_down(0.5)` is true.
-        assert_eq!(t.competitive_threshold(0.5, 5, 4), 0.5f32.next_down());
-    }
-
     #[test]
     fn test_rwlock_shared_threshold_option_u64_asc() {
         let t = RwLockSharedThresholdOptionU64::new();

src/collector/sort_key/sort_by_bytes.rs

@@ -88,6 +88,18 @@ impl SegmentSortKeyComputer for ByBytesColumnSegmentSortKeyComputer {
             .ok()?;
         Some(bytes)
     }
+
+    fn supports_bm25_pruning(&self) -> bool {
+        false
+    }
+    fn bm25_pruning_threshold(
+        &self,
+        _threshold: &Self::SegmentSortKey,
+        _segment_ord: crate::SegmentOrdinal,
+        _threshold_ord: crate::SegmentOrdinal,
+    ) -> Option<crate::Score> {
+        None
+    }
 }
 
 #[cfg(test)]

src/collector/sort_key/sort_by_erased_type.rs

@@ -246,6 +246,18 @@ impl SegmentSortKeyComputer for ErasedColumnSegmentSortKeyComputer {
     fn convert_segment_sort_key(&self, segment_sort_key: Self::SegmentSortKey) -> OwnedValue {
         self.inner.convert_segment_sort_key(segment_sort_key)
     }
+
+    fn supports_bm25_pruning(&self) -> bool {
+        false
+    }
+    fn bm25_pruning_threshold(
+        &self,
+        _threshold: &Self::SegmentSortKey,
+        _segment_ord: crate::SegmentOrdinal,
+        _threshold_ord: crate::SegmentOrdinal,
+    ) -> Option<crate::Score> {
+        None
+    }
 }
 
 #[cfg(test)]

src/collector/sort_key/sort_by_score.rs

@@ -2,7 +2,6 @@ use std::sync::Arc;
 
 use super::shared_threshold::{AtomicSharedThreshold, SharedThresholdArcOpt};
 use crate::collector::sort_key::NaturalComparator;
-use crate::collector::sort_key_top_collector::TopBySortKeySegmentCollector;
 use crate::collector::{SegmentSortKeyComputer, SortKeyComputer};
 use crate::{DocId, Score, SegmentOrdinal};
 
@@ -65,42 +64,6 @@ impl SortKeyComputer for SortBySimilarityScore {
     ) -> crate::Result<Self::Child> {
         Ok(self.clone())
     }
-
-    fn collect_segment_top_k(
-        &self,
-        weight: &dyn crate::query::Weight,
-        reader: &crate::SegmentReader,
-        segment_collector: &mut TopBySortKeySegmentCollector<Self::Child, Self::Comparator>,
-    ) -> crate::Result<()> {
-        let top_n = &mut segment_collector.topn_computer;
-
-        let (initial_score, initial_ord) = top_n
-            .shared_threshold
-            .as_ref()
-            .and_then(|s| s.load())
-            .unwrap_or((Score::MIN, SegmentOrdinal::MAX));
-
-        top_n.set_threshold((initial_score, initial_ord));
-
-        let pruning_threshold = top_n.pruning_threshold.unwrap_or(Score::MIN);
-
-        if let Some(alive_bitset) = reader.alive_bitset() {
-            weight.for_each_pruning(pruning_threshold, reader, &mut |doc, score| {
-                if alive_bitset.is_deleted(doc) {
-                    return top_n.pruning_threshold.unwrap_or(Score::MIN);
-                }
-                top_n.push(score, doc);
-                top_n.pruning_threshold.unwrap_or(Score::MIN)
-            })?;
-        } else {
-            weight.for_each_pruning(pruning_threshold, reader, &mut |doc, score| {
-                top_n.push(score, doc);
-                top_n.pruning_threshold.unwrap_or(Score::MIN)
-            })?;
-        }
-
-        Ok(())
-    }
 }
 
 impl SegmentSortKeyComputer for SortBySimilarityScore {
@@ -116,4 +79,21 @@ impl SegmentSortKeyComputer for SortBySimilarityScore {
     fn convert_segment_sort_key(&self, score: Score) -> Score {
         score
     }
+
+    fn supports_bm25_pruning(&self) -> bool {
+        true
+    }
+
+    fn bm25_pruning_threshold(
+        &self,
+        threshold: &Score,
+        segment_ord: SegmentOrdinal,
+        threshold_ord: SegmentOrdinal,
+    ) -> Option<Score> {
+        if segment_ord < threshold_ord {
+            Some(threshold.next_down())
+        } else {
+            Some(*threshold)
+        }
+    }
 }

src/collector/sort_key/sort_by_static_fast_value.rs

@@ -132,4 +132,16 @@ impl<T: FastValue> SegmentSortKeyComputer for SortByFastValueSegmentSortKeyCompu
     fn convert_segment_sort_key(&self, sort_key: Self::SegmentSortKey) -> Self::SortKey {
         sort_key.map(T::from_u64)
     }
+
+    fn supports_bm25_pruning(&self) -> bool {
+        false
+    }
+    fn bm25_pruning_threshold(
+        &self,
+        _threshold: &Self::SegmentSortKey,
+        _segment_ord: crate::SegmentOrdinal,
+        _threshold_ord: crate::SegmentOrdinal,
+    ) -> Option<crate::Score> {
+        None
+    }
 }

src/collector/sort_key/sort_by_string.rs

@@ -82,4 +82,16 @@ impl SegmentSortKeyComputer for ByStringColumnSegmentSortKeyComputer {
             .ok()?;
         String::try_from(bytes).ok()
     }
+
+    fn supports_bm25_pruning(&self) -> bool {
+        false
+    }
+    fn bm25_pruning_threshold(
+        &self,
+        _threshold: &Self::SegmentSortKey,
+        _segment_ord: crate::SegmentOrdinal,
+        _threshold_ord: crate::SegmentOrdinal,
+    ) -> Option<crate::Score> {
+        None
+    }
 }