Reduce sorting in TopDocs

benches/agg_bench.rs

@@ -7,9 +7,11 @@ use rand_distr::Distribution;
 use serde_json::json;
 use tantivy::aggregation::agg_req::Aggregations;
 use tantivy::aggregation::AggregationCollector;
+use tantivy::collector::TopDocs;
+use tantivy::fastfield::FastValue;
 use tantivy::query::{AllQuery, TermQuery};
 use tantivy::schema::{IndexRecordOption, Schema, TextFieldIndexing, FAST, STRING};
-use tantivy::{doc, Index, Term};
+use tantivy::{doc, Index, Order, Term};
 
 #[global_allocator]
 pub static GLOBAL: &PeakMemAlloc<std::alloc::System> = &INSTRUMENTED_SYSTEM;
@@ -73,6 +75,12 @@ fn bench_agg(mut group: InputGroup<Index>) {
     register!(group, histogram_with_avg_sub_agg);
     register!(group, avg_and_range_with_avg_sub_agg);
 
+    register!(group, top_docs_small_shallow);
+    register!(group, top_docs_small_deep);
+
+    register!(group, top_docs_large_shallow);
+    register!(group, top_docs_large_deep);
+
     group.run();
 }
 
@@ -359,6 +367,34 @@ fn avg_and_range_with_avg_sub_agg(index: &Index) {
     execute_agg(index, agg_req);
 }
 
+fn execute_top_docs<F: FastValue>(
+    index: &Index,
+    fast_field: &str,
+    order: Order,
+    offset: usize,
+    limit: usize,
+) {
+    let collector = TopDocs::with_limit(limit)
+        .and_offset(offset)
+        .order_by_fast_field::<F>(fast_field, order);
+
+    let reader = index.reader().unwrap();
+    let searcher = reader.searcher();
+    black_box(searcher.search(&AllQuery, &collector).unwrap());
+}
+fn top_docs_small_deep(index: &Index) {
+    execute_top_docs::<u64>(index, "score", Order::Asc, 10000, 10);
+}
+fn top_docs_small_shallow(index: &Index) {
+    execute_top_docs::<u64>(index, "score", Order::Asc, 0, 10);
+}
+fn top_docs_large_deep(index: &Index) {
+    execute_top_docs::<u64>(index, "score", Order::Asc, 10000, 1000);
+}
+fn top_docs_large_shallow(index: &Index) {
+    execute_top_docs::<u64>(index, "score", Order::Asc, 0, 1000);
+}
+
 #[derive(Clone, Copy, Hash, Default, Debug, PartialEq, Eq, PartialOrd, Ord)]
 enum Cardinality {
     /// All documents contain exactly one value.
@@ -402,7 +438,7 @@ fn get_test_index_bench(cardinality: Cardinality) -> tantivy::Result<Index> {
         .collect::<Vec<_>>();
     {
         let mut rng = StdRng::from_seed([1u8; 32]);
-        let mut index_writer = index.writer_with_num_threads(1, 200_000_000)?;
+        let mut index_writer = index.writer_with_num_threads(8, 200_000_000)?;
         // To make the different test cases comparable we just change one doc to force the
         // cardinality
         if cardinality == Cardinality::OptionalDense {

src/collector/top_collector.rs

@@ -114,9 +114,7 @@ where T: PartialOrd + Clone
         }
 
         Ok(top_collector
-            .into_sorted_vec()
-            .into_iter()
-            .skip(self.offset)
+            .into_sorted_after(self.offset)
             .map(|cdoc| (cdoc.feature, cdoc.doc))
             .collect())
     }
@@ -169,7 +167,7 @@ impl<T: PartialOrd + Clone> TopSegmentCollector<T> {
     pub fn harvest(self) -> Vec<(T, DocAddress)> {
         let segment_ord = self.segment_ord;
         self.topn_computer
-            .into_sorted_vec()
+            .into_vec()
             .into_iter()
             .map(|comparable_doc| {
                 (
@@ -206,10 +204,11 @@ mod tests {
         top_collector.collect(5, 0.3);
         assert_eq!(
             top_collector.harvest(),
+            // Note: Individual segments are not sorted.
             vec![
                 (0.8, DocAddress::new(0, 1)),
+                (0.2, DocAddress::new(0, 3)),
                 (0.3, DocAddress::new(0, 5)),
-                (0.2, DocAddress::new(0, 3))
             ]
         );
     }
@@ -224,11 +223,12 @@ mod tests {
         top_collector.collect(9, -0.2);
         assert_eq!(
             top_collector.harvest(),
+            // Note: Individual segments are not sorted.
             vec![
-                (0.9, DocAddress::new(0, 7)),
                 (0.8, DocAddress::new(0, 1)),
+                (0.2, DocAddress::new(0, 3)),
                 (0.3, DocAddress::new(0, 5)),
-                (0.2, DocAddress::new(0, 3))
+                (0.9, DocAddress::new(0, 7)),
             ]
         );
     }

src/collector/top_score_collector.rs

@@ -850,6 +850,24 @@ where
         self.buffer
     }
 
+    /// Returns the elements between `offset` and `top_n` in sorted order.
+    pub fn into_sorted_after(
+        mut self,
+        offset: usize,
+    ) -> impl Iterator<Item = ComparableDoc<Score, D, R>> {
+        if self.buffer.len() > self.top_n {
+            self.truncate_top_n();
+        }
+
+        if offset >= self.buffer.len() {
+            return vec![].into_iter().skip(0);
+        }
+
+        let (_, _, remainder) = self.buffer.select_nth_unstable(offset);
+        remainder.sort_unstable();
+        self.buffer.into_iter().skip(offset)
+    }
+
     /// Returns the top n elements in stored order.
     /// Useful if you do not need the elements in sorted order,
     /// for example when merging the results of multiple segments.