Push dynamic pushdown through CooperativeExec and ProjectionExec (apache#17238)

(cherry picked from commit 4bc0696)

Author: jackkleeman

datafusion/physical-plan/src/coop.rs

@@ -65,17 +65,20 @@
 //! The optimizer rule currently checks the plan for exchange-like operators and leave operators
 //! that report [`SchedulingType::NonCooperative`] in their [plan properties](ExecutionPlan::properties).
 
-#[cfg(any(
-    datafusion_coop = "tokio_fallback",
-    not(any(datafusion_coop = "tokio", datafusion_coop = "per_stream"))
-))]
+use datafusion_common::config::ConfigOptions;
+use datafusion_physical_expr::PhysicalExpr;
+#[cfg(datafusion_coop = "tokio_fallback")]
 use futures::Future;
 use std::any::Any;
 use std::pin::Pin;
 use std::sync::Arc;
 use std::task::{Context, Poll};
 
 use crate::execution_plan::CardinalityEffect::{self, Equal};
+use crate::filter_pushdown::{
+    ChildPushdownResult, FilterDescription, FilterPushdownPhase,
+    FilterPushdownPropagation,
+};
 use crate::{
     DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, RecordBatchStream,
     SendableRecordBatchStream,
@@ -164,6 +167,8 @@ where
                 // after the work has been done and just assume that that succeeded.
                 // The poll result is ignored because we don't want to discard
                 // or buffer the Ready result we got from the inner stream.
+
+                use std::future::Future;
                 let consume = tokio::task::coop::consume_budget();
                 let consume_ref = std::pin::pin!(consume);
                 let _ = consume_ref.poll(cx);
@@ -291,6 +296,24 @@ impl ExecutionPlan for CooperativeExec {
     fn cardinality_effect(&self) -> CardinalityEffect {
         Equal
     }
+
+    fn gather_filters_for_pushdown(
+        &self,
+        _phase: FilterPushdownPhase,
+        parent_filters: Vec<Arc<dyn PhysicalExpr>>,
+        _config: &ConfigOptions,
+    ) -> Result<FilterDescription> {
+        FilterDescription::from_children(parent_filters, &self.children())
+    }
+
+    fn handle_child_pushdown_result(
+        &self,
+        _phase: FilterPushdownPhase,
+        child_pushdown_result: ChildPushdownResult,
+        _config: &ConfigOptions,
+    ) -> Result<FilterPushdownPropagation<Arc<dyn ExecutionPlan>>> {
+        Ok(FilterPushdownPropagation::if_all(child_pushdown_result))
+    }
 }
 
 /// Creates a [`CooperativeStream`] wrapper around the given [`RecordBatchStream`].

datafusion/physical-plan/src/projection.rs

@@ -33,11 +33,16 @@ use super::{
     SendableRecordBatchStream, Statistics,
 };
 use crate::execution_plan::CardinalityEffect;
+use crate::filter_pushdown::{
+    ChildPushdownResult, FilterDescription, FilterPushdownPhase,
+    FilterPushdownPropagation,
+};
 use crate::joins::utils::{ColumnIndex, JoinFilter, JoinOn, JoinOnRef};
 use crate::{ColumnStatistics, DisplayFormatType, ExecutionPlan, PhysicalExpr};
 
 use arrow::datatypes::{Field, Schema, SchemaRef};
 use arrow::record_batch::{RecordBatch, RecordBatchOptions};
+use datafusion_common::config::ConfigOptions;
 use datafusion_common::stats::Precision;
 use datafusion_common::tree_node::{
     Transformed, TransformedResult, TreeNode, TreeNodeRecursion,
@@ -273,6 +278,27 @@ impl ExecutionPlan for ProjectionExec {
             Ok(Some(Arc::new(projection.clone())))
         }
     }
+
+    fn gather_filters_for_pushdown(
+        &self,
+        _phase: FilterPushdownPhase,
+        parent_filters: Vec<Arc<dyn PhysicalExpr>>,
+        _config: &ConfigOptions,
+    ) -> Result<FilterDescription> {
+        // TODO: In future, we can try to handle inverting aliases here.
+        // For the time being, we pass through untransformed filters, so filters on aliases are not handled.
+        // https://github.com/apache/datafusion/issues/17246
+        FilterDescription::from_children(parent_filters, &self.children())
+    }
+
+    fn handle_child_pushdown_result(
+        &self,
+        _phase: FilterPushdownPhase,
+        child_pushdown_result: ChildPushdownResult,
+        _config: &ConfigOptions,
+    ) -> Result<FilterPushdownPropagation<Arc<dyn ExecutionPlan>>> {
+        Ok(FilterPushdownPropagation::if_all(child_pushdown_result))
+    }
 }
 
 fn stats_projection(

datafusion/sqllogictest/test_files/explain_tree.slt

@@ -383,7 +383,7 @@ physical_plan
 44)-----------------------------│    --------------------   ││    --------------------   │
 45)-----------------------------│          files: 1         ││ partition_count(in->out): │
 46)-----------------------------│      format: parquet      ││           1 -> 4          │
-47)-----------------------------│                           ││                           │
+47)-----------------------------│      predicate: true      ││                           │
 48)-----------------------------│                           ││    partitioning_scheme:   │
 49)-----------------------------│                           ││     RoundRobinBatch(4)    │
 50)-----------------------------└───────────────────────────┘└─────────────┬─────────────┘

datafusion/sqllogictest/test_files/topk.slt

@@ -372,7 +372,7 @@ explain select number, letter, age, number as column4, letter as column5 from pa
 physical_plan
 01)SortExec: TopK(fetch=3), expr=[number@0 DESC, letter@1 ASC NULLS LAST, age@2 DESC], preserve_partitioning=[false], sort_prefix=[number@0 DESC, letter@1 ASC NULLS LAST]
 02)--ProjectionExec: expr=[number@0 as number, letter@1 as letter, age@2 as age, number@0 as column4, letter@1 as column5]
-03)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/topk/partial_sorted/1.parquet]]}, projection=[number, letter, age], output_ordering=[number@0 DESC, letter@1 ASC NULLS LAST], file_type=parquet
+03)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/topk/partial_sorted/1.parquet]]}, projection=[number, letter, age], output_ordering=[number@0 DESC, letter@1 ASC NULLS LAST], file_type=parquet, predicate=DynamicFilterPhysicalExpr [ true ]
 
 # Verify that the sort prefix is correctly computed over normalized, order-maintaining projections (number + 1, number, number + 1, age)
 query TT