Add test for top k

Author: Shiv Bhatia

datafusion/physical-plan/src/sorts/sort.rs

@@ -1454,14 +1454,14 @@ mod tests {
     use arrow::compute::SortOptions;
     use arrow::datatypes::*;
     use datafusion_common::cast::as_primitive_array;
-    use datafusion_common::test_util::batches_to_string;
     use datafusion_common::config::ConfigOptions;
+    use datafusion_common::test_util::batches_to_string;
     use datafusion_common::{DataFusionError, Result, ScalarValue};
     use datafusion_execution::RecordBatchStream;
+    use datafusion_execution::config::SessionConfig;
     use datafusion_execution::memory_pool::{
         GreedyMemoryPool, MemoryConsumer, MemoryPool,
     };
-    use datafusion_execution::config::SessionConfig;
     use datafusion_execution::runtime_env::RuntimeEnvBuilder;
     use datafusion_physical_expr::EquivalenceProperties;
     use datafusion_physical_expr::expressions::{Column, Literal};
@@ -2938,4 +2938,48 @@ mod tests {
 
         Ok(())
     }
+
+    #[test]
+    fn test_sort_with_fetch_allows_topk_self_filter_in_post_phase() -> Result<()> {
+        let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)]));
+        let input = Arc::new(EmptyExec::new(Arc::clone(&schema)));
+        let sort = SortExec::new(
+            [PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0)))].into(),
+            input,
+        )
+        .with_fetch(Some(10));
+
+        // with_fetch(Some(_)) creates the TopK dynamic filter automatically.
+        assert!(sort.filter.is_some(), "TopK filter should be created");
+
+        let parent_filter: Arc<dyn PhysicalExpr> = Arc::new(Column::new("a", 0));
+        let mut config = ConfigOptions::new();
+        config.optimizer.enable_topk_dynamic_filter_pushdown = true;
+
+        let desc = sort.gather_filters_for_pushdown(
+            FilterPushdownPhase::Post,
+            vec![parent_filter],
+            &config,
+        )?;
+
+        // Parent filters should be blocked in Post phase when fetch is set.
+        let parent_filters = desc.parent_filters();
+        assert_eq!(parent_filters.len(), 1);
+        assert_eq!(parent_filters[0].len(), 1);
+        assert!(
+            matches!(parent_filters[0][0].discriminant, PushedDown::No),
+            "Parent filter should be unsupported in Post phase when sort has fetch"
+        );
+
+        // The TopK self-filter should still be allowed through.
+        let self_filters = desc.self_filters();
+        assert_eq!(self_filters.len(), 1);
+        assert_eq!(
+            self_filters[0].len(),
+            1,
+            "TopK dynamic self-filter should be pushed down"
+        );
+
+        Ok(())
+    }
 }