Sort with limit before topk

Author: Dandandan

datafusion-testing

@@ -18,8 +18,8 @@
 //! TopK: Combination of Sort / LIMIT
 
 use arrow::{
-    array::{Array, AsArray},
-    compute::{interleave_record_batch, prep_null_mask_filter, FilterBuilder},
+    array::{Array, AsArray, BooleanArray, UInt32Array},
+    compute::{FilterBuilder, interleave_record_batch, prep_null_mask_filter, sort_to_indices},
     row::{RowConverter, Rows, SortField},
 };
 use datafusion_expr::{ColumnarValue, Operator};
@@ -172,6 +172,11 @@ fn build_sort_fields(
         .collect::<Result<_>>()
 }
 
+enum TopKSelection {
+    Boolean(BooleanArray),
+    Indices(UInt32Array),
+}
+
 impl TopK {
     /// Create a new [`TopK`] that stores the top `k` values, as
     /// defined by the sort expressions in `expr`.
@@ -251,7 +256,12 @@ impl TopK {
             // nothing to filter, so no need to update
             return Ok(());
         }
-        // only update the keys / rows if the filter does not match all rows
+        // If only a single sort key, and filters out 80% of the rows and existing filter is not very selective,
+        // use sort_to_indices to get the top indices from the input batch
+        if sort_keys.len() == 1 && (self.heap.k as f64) < 0.2 * (num_rows as f64) && (true_count > self.heap.k * 2) {
+            let array = sort_keys[0].as_ref();
+            selected_rows = Some(TopKSelection::Indices(sort_to_indices(array, None, Some(self.heap.k))?));
+        }
         if true_count < num_rows {
             // Indices in `set_indices` should be correct if filter contains nulls
             // So we prepare the filter here. Note this is also done in the `FilterBuilder`
@@ -267,7 +277,7 @@ impl TopK {
             } else {
                 filter_predicate.build()
             };
-            selected_rows = Some(filter);
+            selected_rows = Some(TopKSelection::Boolean(filter));
             sort_keys = sort_keys
                 .iter()
                 .map(|key| filter_predicate.filter(key).map_err(|x| x.into()))
@@ -281,9 +291,13 @@ impl TopK {
         let mut batch_entry = self.heap.register_batch(batch.clone());
 
         let replacements = match selected_rows {
-            Some(filter) => {
+            Some(TopKSelection::Boolean(filter)) => {
                 self.find_new_topk_items(filter.values().set_indices(), &mut batch_entry)
             }
+            Some(TopKSelection::Indices(indices)) => {
+                self.find_new_topk_items(indices.values().iter().map(|i| *i as usize), &mut batch_entry)
+            }
+
             None => self.find_new_topk_items(0..sort_keys[0].len(), &mut batch_entry),
         };