First draft of dynamic filters for TopK

Closes apache#15037

Author: adriangb

datafusion/core/tests/physical_optimizer/mod.rs

@@ -28,3 +28,4 @@ mod projection_pushdown;
 mod replace_with_order_preserving_variants;
 mod sanity_checker;
 mod test_utils;
+mod topk_dynamic_filters;

datafusion/core/tests/physical_optimizer/topk_dynamic_filters.rs

@@ -0,0 +1,7 @@
+#[cfg(test)]
+mod tests {
+    #[tokio::test]
+    async fn test_topk_dynamic_filters() {
+        // TODO: add a test
+    }
+}

datafusion/datasource-avro/src/source.rs

@@ -220,13 +220,6 @@ impl FileSource for AvroSource {
         Arc::new(conf)
     }
 
-    fn with_dynamic_filter(
-        &self,
-        _dynamic_filter: Arc<dyn datafusion_physical_plan::DynamicFilterSource>,
-    ) -> Arc<dyn FileSource> {
-        Arc::new(Self { ..self.clone() })
-    }
-
     fn metrics(&self) -> &ExecutionPlanMetricsSet {
         &self.metrics
     }

datafusion/datasource-csv/src/source.rs

@@ -605,14 +605,6 @@ impl FileSource for CsvSource {
         Arc::new(conf)
     }
 
-    fn with_dynamic_filter(
-        &self,
-        _dynamic_filter: Arc<dyn datafusion_physical_plan::DynamicFilterSource>,
-    ) -> Arc<dyn FileSource> {
-        // AvroSource does not support dynamic filters, so we just return self
-        Arc::new(self.clone())
-    }
-
     fn metrics(&self) -> &ExecutionPlanMetricsSet {
         &self.metrics
     }

datafusion/datasource-json/src/source.rs

@@ -303,13 +303,6 @@ impl FileSource for JsonSource {
         Arc::new(Self { ..self.clone() })
     }
 
-    fn with_dynamic_filter(
-        &self,
-        _dynamic_filter: Arc<dyn datafusion_physical_plan::DynamicFilterSource>,
-    ) -> Arc<dyn FileSource> {
-        Arc::new(Self { ..self.clone() })
-    }
-
     fn metrics(&self) -> &ExecutionPlanMetricsSet {
         &self.metrics
     }

datafusion/datasource-parquet/src/opener.rs

@@ -109,15 +109,9 @@ impl FileOpener for ParquetOpener {
         let dynamic_filters = self
             .dynamic_filters
             .iter()
-            .filter_map(|f| {
-                f.current_filters().ok().and_then(|filters| {
-                    if filters.is_empty() {
-                        None
-                    } else {
-                        Some(filters)
-                    }
-                })
-            })
+            .map(|f| f.current_filters())
+            .collect::<Result<Vec<_>>>()?
+            .into_iter()
             .flatten()
             .collect::<Vec<_>>();
         // Collect dynamic_filters into a single predicate by reducing with AND

datafusion/datasource-parquet/src/source.rs

@@ -531,15 +531,6 @@ impl FileSource for ParquetSource {
         Arc::new(Self { ..self.clone() })
     }
 
-    fn with_dynamic_filter(
-        &self,
-        dynamic_filter: Arc<dyn DynamicFilterSource>,
-    ) -> Arc<dyn FileSource> {
-        let mut conf = self.clone();
-        conf.dynamic_filters.push(dynamic_filter);
-        Arc::new(conf)
-    }
-
     fn metrics(&self) -> &ExecutionPlanMetricsSet {
         &self.metrics
     }

datafusion/datasource/src/file.rs

@@ -52,11 +52,6 @@ pub trait FileSource: Send + Sync {
     fn with_schema(&self, schema: SchemaRef) -> Arc<dyn FileSource>;
     /// Initialize new instance with projection information
     fn with_projection(&self, config: &FileScanConfig) -> Arc<dyn FileSource>;
-    /// Add dynamic filters to the file source
-    fn with_dynamic_filter(
-        &self,
-        dynamic_filter: Arc<dyn DynamicFilterSource>,
-    ) -> Arc<dyn FileSource>;
     /// Initialize new instance with projected statistics
     fn with_statistics(&self, statistics: Statistics) -> Arc<dyn FileSource>;
     /// Return execution plan metrics

datafusion/datasource/src/file_scan_config.rs

@@ -41,7 +41,11 @@ use datafusion_physical_expr::{
     PhysicalSortExpr,
 };
 use datafusion_physical_plan::{
-    display::{display_orderings, ProjectSchemaDisplay}, dynamic_filters::DynamicFilterSource, metrics::ExecutionPlanMetricsSet, projection::{all_alias_free_columns, new_projections_for_columns, ProjectionExec}, DisplayAs, DisplayFormatType, ExecutionPlan
+    display::{display_orderings, ProjectSchemaDisplay},
+    dynamic_filters::DynamicFilterSource,
+    metrics::ExecutionPlanMetricsSet,
+    projection::{all_alias_free_columns, new_projections_for_columns, ProjectionExec},
+    DisplayAs, DisplayFormatType, ExecutionPlan,
 };
 use log::{debug, warn};
 

datafusion/datasource/src/memory.rs

@@ -491,14 +491,6 @@ impl DataSource for MemorySourceConfig {
             })
             .transpose()
     }
-
-    fn with_dynamic_filter(
-        &self,
-        _dynamic_filter: Arc<dyn datafusion_physical_plan::DynamicFilterSource>,
-    ) -> Arc<dyn DataSource> {
-        // MemorySourceConfig does not support dynamic filters
-        Arc::new(self.clone())
-    }
 }
 
 impl MemorySourceConfig {

datafusion/datasource/src/source.rs

@@ -79,10 +79,6 @@ pub trait DataSource: Send + Sync + Debug {
         &self,
         _projection: &ProjectionExec,
     ) -> datafusion_common::Result<Option<Arc<dyn ExecutionPlan>>>;
-    fn with_dynamic_filter(
-        &self,
-        _dynamic_filter: Arc<dyn DynamicFilterSource>,
-    ) -> Arc<dyn DataSource>;
 }
 
 /// [`ExecutionPlan`] handles different file formats like JSON, CSV, AVRO, ARROW, PARQUET
@@ -227,14 +223,6 @@ impl DataSourceExec {
         self
     }
 
-    pub fn with_dynamic_filters(
-        mut self,
-        dynamic_filter: Arc<dyn DynamicFilterSource>,
-    ) -> Self {
-        self.data_source = self.data_source.with_dynamic_filter(dynamic_filter);
-        self
-    }
-
     fn compute_properties(data_source: Arc<dyn DataSource>) -> PlanProperties {
         PlanProperties::new(
             data_source.eq_properties(),

datafusion/datasource/src/test_util.rs

@@ -59,13 +59,6 @@ impl FileSource for MockSource {
         Arc::new(Self { ..self.clone() })
     }
 
-    fn with_dynamic_filter(
-        &self,
-        _dynamic_filter: Arc<dyn datafusion_physical_plan::DynamicFilterSource>,
-    ) -> Arc<dyn FileSource> {
-        Arc::new(Self { ..self.clone() })
-    }
-
     fn with_statistics(&self, statistics: Statistics) -> Arc<dyn FileSource> {
         let mut source = self.clone();
         source.projected_statistics = Some(statistics);

datafusion/physical-optimizer/src/topk_dynamic_filters.rs

@@ -0,0 +1,65 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! An optimizer rule that detects TopK operations that can generate dynamic filters to be pushed down into file scans
+
+use std::sync::Arc;
+
+use crate::PhysicalOptimizerRule;
+use datafusion_common::config::ConfigOptions;
+use datafusion_common::tree_node::{TransformedResult, TreeNode};
+use datafusion_common::Result;
+use datafusion_physical_plan::ExecutionPlan;
+
+/// An optimizer rule that passes a TopK as a DynamicFilterSource to DataSourceExec executors.
+///
+/// This optimizer looks for TopK operators in the plan and connects them to compatible
+/// data sources by registering the TopK as a source of dynamic filters.
+///
+/// When a data source is found that can benefit from a TopK's filter, the source is
+/// modified to include the TopK as a dynamic filter source via `with_dynamic_filter_source`.
+/// During execution, the data source will then consult the TopK's current filter state
+/// to determine which files or partitions can be skipped.
+#[derive(Debug)]
+pub struct TopKDynamicFilters {}
+
+impl PhysicalOptimizerRule for TopKDynamicFilters {
+    fn optimize(
+        &self,
+        plan: Arc<dyn ExecutionPlan>,
+        config: &ConfigOptions,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        if config.optimizer.enable_dynamic_filter_pushdown {
+            plan.transform_down(|_plan| {
+                // TODO: recurse only traversing approved nodes (e.g., FilterExec, ProjectionExec, SortExec, CoalesceBatchesExec, RepartitionExec; namely not aggregations, joins, etc.)
+                // Collect any DataSourceExecs that can use dynamic filters and register the TopK as a source
+                todo!();
+            })
+            .data()
+        } else {
+            Ok(plan)
+        }
+    }
+
+    fn name(&self) -> &str {
+        "TopKDynamicFilters"
+    }
+
+    fn schema_check(&self) -> bool {
+        true
+    }
+}

datafusion/physical-plan/src/dynamic_filters.rs

@@ -0,0 +1,13 @@
+use std::sync::Arc;
+
+use datafusion_common::Result;
+use datafusion_physical_expr::PhysicalExpr;
+
+/// A source of dynamic runtime filters.
+///
+/// During query execution, operators implementing this trait can provide
+/// filter expressions that other operators can use to dynamically prune data.
+pub trait DynamicFilterSource: Send + Sync + std::fmt::Debug + 'static {
+    /// Returns a list of filter expressions that can be used for dynamic pruning.
+    fn current_filters(&self) -> Result<Vec<Arc<dyn PhysicalExpr>>>;
+}

datafusion/physical-plan/src/lib.rs

@@ -40,7 +40,7 @@ pub use datafusion_physical_expr::{
 };
 
 pub use crate::display::{DefaultDisplay, DisplayAs, DisplayFormatType, VerboseDisplay};
-pub use crate::dynamic_filter::DynamicFilterSource;
+pub use crate::dynamic_filters::DynamicFilterSource;
 pub use crate::execution_plan::{
     collect, collect_partitioned, displayable, execute_input_stream, execute_stream,
     execute_stream_partitioned, get_plan_string, with_new_children_if_necessary,

datafusion/physical-plan/src/topk/mod.rs

@@ -28,7 +28,7 @@ use std::{cmp::Ordering, collections::BinaryHeap, sync::Arc};
 use datafusion_expr::Operator;
 
 use super::metrics::{BaselineMetrics, Count, ExecutionPlanMetricsSet, MetricBuilder};
-use crate::dynamic_filter::DynamicFilterSource;
+use crate::dynamic_filters::DynamicFilterSource;
 use crate::spill::get_record_batch_memory_size;
 use crate::{stream::RecordBatchStreamAdapter, SendableRecordBatchStream};
 use arrow::array::{Array, ArrayRef, RecordBatch};