always require a filter in topk

adriangb · adriangb · commit 8de60633a427 · 2025-06-18T18:37:50.000-05:00
diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs
@@ -38,7 +38,7 @@ use crate::spill::get_record_batch_memory_size;
 use crate::spill::in_progress_spill_file::InProgressSpillFile;
 use crate::spill::spill_manager::SpillManager;
 use crate::stream::RecordBatchStreamAdapter;
-use crate::topk::TopK;
+use crate::topk::{TopK, TopKDynamicFilters};
 use crate::{
     DisplayAs, DisplayFormatType, Distribution, EmptyRecordBatchStream, ExecutionPlan,
     ExecutionPlanProperties, Partitioning, PlanProperties, SendableRecordBatchStream,
@@ -846,8 +846,10 @@ pub struct SortExec {
     common_sort_prefix: Vec<PhysicalSortExpr>,
     /// Cache holding plan properties like equivalences, output partitioning etc.
     cache: PlanProperties,
-    /// Filter matching the state of the sort for dynamic filter pushdown
-    filter: Option<Arc<DynamicFilterPhysicalExpr>>,
+    /// Filter matching the state of the sort for dynamic filter pushdown.
+    /// If `fetch` is `Some`, this will also be set and a TopK operator may be used.
+    /// If `fetch` is `None`, this will be `None`.
+    filter: Option<TopKDynamicFilters>,
 }
 
 impl SortExec {
@@ -920,7 +922,10 @@ impl SortExec {
                     .iter()
                     .map(|sort_expr| Arc::clone(&sort_expr.expr))
                     .collect::<Vec<_>>();
-                Arc::new(DynamicFilterPhysicalExpr::new(children, lit(true)))
+                TopKDynamicFilters::new(Arc::new(DynamicFilterPhysicalExpr::new(
+                    children,
+                    lit(true),
+                )))
             })
         });
         SortExec {
@@ -935,11 +940,6 @@ impl SortExec {
         }
     }
 
-    pub fn with_filter(mut self, filter: Arc<DynamicFilterPhysicalExpr>) -> Self {
-        self.filter = Some(filter);
-        self
-    }
-
     /// Input schema
     pub fn input(&self) -> &Arc<dyn ExecutionPlan> {
         &self.input
@@ -955,11 +955,6 @@ impl SortExec {
         self.fetch
     }
 
-    /// If `Some(filter)`, returns the filter expression that matches the state of the sort.
-    pub fn filter(&self) -> Option<Arc<DynamicFilterPhysicalExpr>> {
-        self.filter.clone()
-    }
-
     fn output_partitioning_helper(
         input: &Arc<dyn ExecutionPlan>,
         preserve_partitioning: bool,
@@ -1038,7 +1033,7 @@ impl DisplayAs for SortExec {
                     Some(fetch) => {
                         write!(f, "SortExec: TopK(fetch={fetch}), expr=[{}], preserve_partitioning=[{preserve_partitioning}]", self.expr)?;
                         if let Some(filter) = &self.filter {
-                            if let Ok(current) = filter.current() {
+                            if let Ok(current) = filter.expr().current() {
                                 if !current.eq(&lit(true)) {
                                     write!(f, ", filter=[{current}]")?;
                                 }
@@ -1158,7 +1153,10 @@ impl ExecutionPlan for SortExec {
                     context.session_config().batch_size(),
                     context.runtime_env(),
                     &self.metrics_set,
-                    self.filter.clone(),
+                    self.filter
+                        .as_ref()
+                        .expect("Filter should be set when fetch is Some")
+                        .clone(),
                 )?;
                 Ok(Box::pin(RecordBatchStreamAdapter::new(
                     self.schema(),
@@ -1278,10 +1276,9 @@ impl ExecutionPlan for SortExec {
         }
         if let Some(filter) = &self.filter {
             if config.optimizer.enable_dynamic_filter_pushdown {
-                let filter = Arc::clone(filter) as Arc<dyn PhysicalExpr>;
                 return Ok(FilterDescription::new_with_child_count(1)
                     .all_parent_filters_supported(parent_filters)
-                    .with_self_filter(filter));
+                    .with_self_filter(filter.expr()));
             }
         }
         Ok(FilterDescription::new_with_child_count(1)
diff --git a/datafusion/physical-plan/src/topk/mod.rs b/datafusion/physical-plan/src/topk/mod.rs
@@ -23,6 +23,7 @@ use arrow::{
     row::{RowConverter, Rows, SortField},
 };
 use datafusion_expr::{ColumnarValue, Operator};
+use parking_lot::RwLock;
 use std::mem::size_of;
 use std::{cmp::Ordering, collections::BinaryHeap, sync::Arc};
 
@@ -121,13 +122,36 @@ pub struct TopK {
     /// Common sort prefix between the input and the sort expressions to allow early exit optimization
     common_sort_prefix: Arc<[PhysicalSortExpr]>,
     /// Filter matching the state of the `TopK` heap used for dynamic filter pushdown
-    filter: Option<Arc<DynamicFilterPhysicalExpr>>,
+    filter: TopKDynamicFilters,
     /// If true, indicates that all rows of subsequent batches are guaranteed
     /// to be greater (by byte order, after row conversion) than the top K,
     /// which means the top K won't change and the computation can be finished early.
     pub(crate) finished: bool,
 }
 
+#[derive(Debug, Clone)]
+pub struct TopKDynamicFilters {
+    /// The current *global* threshold for the dynamic filter.
+    /// This is shared across all partitions and is updated by any of them.
+    thresholds: Arc<RwLock<Option<Vec<ScalarValue>>>>,
+    /// The expression used to evaluate the dynamic filter
+    expr: Arc<DynamicFilterPhysicalExpr>,
+}
+
+impl TopKDynamicFilters {
+    /// Create a new `TopKDynamicFilters` with the given expression
+    pub fn new(expr: Arc<DynamicFilterPhysicalExpr>) -> Self {
+        Self {
+            thresholds: Arc::new(RwLock::new(None)),
+            expr,
+        }
+    }
+
+    pub fn expr(&self) -> Arc<DynamicFilterPhysicalExpr> {
+        Arc::clone(&self.expr)
+    }
+}
+
 // Guesstimate for memory allocation: estimated number of bytes used per row in the RowConverter
 const ESTIMATED_BYTES_PER_ROW: usize = 20;
 
@@ -160,7 +184,7 @@ impl TopK {
         batch_size: usize,
         runtime: Arc<RuntimeEnv>,
         metrics: &ExecutionPlanMetricsSet,
-        filter: Option<Arc<DynamicFilterPhysicalExpr>>,
+        filter: TopKDynamicFilters,
     ) -> Result<Self> {
         let reservation = MemoryConsumer::new(format!("TopK[{partition_id}]"))
             .register(&runtime.memory_pool);
@@ -214,41 +238,39 @@ impl TopK {
 
         let mut selected_rows = None;
 
-        if let Some(filter) = self.filter.as_ref() {
-            // If a filter is provided, update it with the new rows
-            let filter = filter.current()?;
-            let filtered = filter.evaluate(&batch)?;
-            let num_rows = batch.num_rows();
-            let array = filtered.into_array(num_rows)?;
-            let mut filter = array.as_boolean().clone();
-            let true_count = filter.true_count();
-            if true_count == 0 {
-                // nothing to filter, so no need to update
-                return Ok(());
+        // If a filter is provided, update it with the new rows
+        let filter = self.filter.expr.current()?;
+        let filtered = filter.evaluate(&batch)?;
+        let num_rows = batch.num_rows();
+        let array = filtered.into_array(num_rows)?;
+        let mut filter = array.as_boolean().clone();
+        let true_count = filter.true_count();
+        if true_count == 0 {
+            // nothing to filter, so no need to update
+            return Ok(());
+        }
+        // only update the keys / rows if the filter does not match all rows
+        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`
+            // so there is no overhead to do this here.
+            if filter.nulls().is_some() {
+                filter = prep_null_mask_filter(&filter);
             }
-            // only update the keys / rows if the filter does not match all rows
-            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`
-                // so there is no overhead to do this here.
-                if filter.nulls().is_some() {
-                    filter = prep_null_mask_filter(&filter);
-                }
 
-                let filter_predicate = FilterBuilder::new(&filter);
-                let filter_predicate = if sort_keys.len() > 1 {
-                    // Optimize filter when it has multiple sort keys
-                    filter_predicate.optimize().build()
-                } else {
-                    filter_predicate.build()
-                };
-                selected_rows = Some(filter);
-                sort_keys = sort_keys
-                    .iter()
-                    .map(|key| filter_predicate.filter(key).map_err(|x| x.into()))
-                    .collect::<Result<Vec<_>>>()?;
-            }
-        };
+            let filter_predicate = FilterBuilder::new(&filter);
+            let filter_predicate = if sort_keys.len() > 1 {
+                // Optimize filter when it has multiple sort keys
+                filter_predicate.optimize().build()
+            } else {
+                filter_predicate.build()
+            };
+            selected_rows = Some(filter);
+            sort_keys = sort_keys
+                .iter()
+                .map(|key| filter_predicate.filter(key).map_err(|x| x.into()))
+                .collect::<Result<Vec<_>>>()?;
+        }
         // reuse existing `Rows` to avoid reallocations
         let rows = &mut self.scratch_rows;
         rows.clear();
@@ -319,13 +341,88 @@ impl TopK {
     /// (a > 2 OR (a = 2 AND b < 3))
     /// ```
     fn update_filter(&mut self) -> Result<()> {
-        let Some(filter) = &self.filter else {
-            return Ok(());
-        };
         let Some(thresholds) = self.heap.get_threshold_values(&self.expr)? else {
             return Ok(());
         };
 
+        // Are the new thresholds more selective than our existing ones?
+        let should_update = {
+            if let Some(current) = self.filter.thresholds.write().as_mut() {
+                assert!(current.len() == thresholds.len());
+                // Check if new thresholds are more selective than current ones
+                let mut more_selective = false;
+                for ((current_value, new_value), sort_expr) in
+                    current.iter().zip(thresholds.iter()).zip(self.expr.iter())
+                {
+                    // Handle null cases
+                    let (current_is_null, new_is_null) =
+                        (current_value.is_null(), new_value.is_null());
+
+                    match (current_is_null, new_is_null) {
+                        (true, true) => {
+                            // Both null, continue checking next values
+                        }
+                        (true, false) => {
+                            // Current is null, new is not null
+                            // For nulls_first: null < non-null, so new value is less selective
+                            // For nulls_last: null > non-null, so new value is more selective
+                            more_selective = !sort_expr.options.nulls_first;
+                            break;
+                        }
+                        (false, true) => {
+                            // Current is not null, new is null
+                            // For nulls_first: non-null > null, so new value is more selective
+                            // For nulls_last: non-null < null, so new value is less selective
+                            more_selective = sort_expr.options.nulls_first;
+                            break;
+                        }
+                        (false, false) => {
+                            // Neither is null, compare values
+                            match current_value.partial_cmp(new_value) {
+                                Some(ordering) => {
+                                    match ordering {
+                                        Ordering::Equal => {
+                                            // Continue checking next values
+                                        }
+                                        Ordering::Less => {
+                                            // For descending sort: new > current means more selective
+                                            // For ascending sort: new > current means less selective
+                                            more_selective = sort_expr.options.descending;
+                                            break;
+                                        }
+                                        Ordering::Greater => {
+                                            // For descending sort: new < current means less selective
+                                            // For ascending sort: new < current means more selective
+                                            more_selective =
+                                                !sort_expr.options.descending;
+                                            break;
+                                        }
+                                    }
+                                }
+                                None => {
+                                    // If values can't be compared, don't update
+                                    more_selective = false;
+                                    break;
+                                }
+                            }
+                        }
+                    }
+                }
+                // If the new thresholds are more selective, update the current ones
+                if more_selective {
+                    *current = thresholds.clone();
+                }
+                more_selective
+            } else {
+                // No current thresholds, so update with the new ones
+                true
+            }
+        };
+
+        if !should_update {
+            return Ok(());
+        }
+
         // Create filter expressions for each threshold
         let mut filters: Vec<Arc<dyn PhysicalExpr>> =
             Vec::with_capacity(thresholds.len());
@@ -405,7 +502,7 @@ impl TopK {
 
         if let Some(predicate) = dynamic_predicate {
             if !predicate.eq(&lit(true)) {
-                filter.update(predicate)?;
+                self.filter.expr.update(predicate)?;
             }
         }
 
@@ -1053,7 +1150,10 @@ mod tests {
             2,
             runtime,
             &metrics,
-            None,
+            TopKDynamicFilters::new(Arc::new(DynamicFilterPhysicalExpr::new(
+                vec![],
+                lit(true),
+            ))),
         )?;
 
         // Create the first batch with two columns:
