diff --git a/datafusion-examples/examples/custom_data_source/custom_datasource.rs b/datafusion-examples/examples/custom_data_source/custom_datasource.rs index a2d7d7699927f..a5a38edf0b6f5 100644 --- a/datafusion-examples/examples/custom_data_source/custom_datasource.rs +++ b/datafusion-examples/examples/custom_data_source/custom_datasource.rs @@ -27,6 +27,7 @@ use datafusion::arrow::array::{UInt8Builder, UInt64Builder}; use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use datafusion::arrow::record_batch::RecordBatch; use datafusion::common::assert_batches_eq; +use datafusion::common::tree_node::TreeNodeRecursion; use datafusion::datasource::{TableProvider, TableType, provider_as_source}; use datafusion::error::Result; use datafusion::execution::context::TaskContext; @@ -314,4 +315,13 @@ impl ExecutionPlan for CustomExec { None, )?)) } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut( + &Arc, + ) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } } diff --git a/datafusion-examples/examples/execution_monitoring/memory_pool_execution_plan.rs b/datafusion-examples/examples/execution_monitoring/memory_pool_execution_plan.rs index ca765774d141f..6decb84b55be1 100644 --- a/datafusion-examples/examples/execution_monitoring/memory_pool_execution_plan.rs +++ b/datafusion-examples/examples/execution_monitoring/memory_pool_execution_plan.rs @@ -29,6 +29,7 @@ use arrow::array::record_batch; use arrow::record_batch::RecordBatch; use arrow_schema::SchemaRef; +use datafusion::common::tree_node::TreeNodeRecursion; use datafusion::common::{exec_datafusion_err, internal_err}; use datafusion::datasource::{DefaultTableSource, memory::MemTable}; use datafusion::error::Result; @@ -291,4 +292,13 @@ impl ExecutionPlan for BufferingExecutionPlan { }), ))) } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut( + &Arc, + ) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } } diff --git a/datafusion-examples/examples/proto/composed_extension_codec.rs b/datafusion-examples/examples/proto/composed_extension_codec.rs index 6077a982c320d..d5197fe61bea7 100644 --- a/datafusion-examples/examples/proto/composed_extension_codec.rs +++ b/datafusion-examples/examples/proto/composed_extension_codec.rs @@ -37,6 +37,7 @@ use std::sync::Arc; use datafusion::common::Result; use datafusion::common::internal_err; +use datafusion::common::tree_node::TreeNodeRecursion; use datafusion::execution::TaskContext; use datafusion::physical_plan::{DisplayAs, ExecutionPlan}; use datafusion::prelude::SessionContext; @@ -124,6 +125,15 @@ impl ExecutionPlan for ParentExec { ) -> Result { unreachable!() } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut( + &Arc, + ) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } } /// A PhysicalExtensionCodec that can serialize and deserialize ParentExec @@ -202,6 +212,15 @@ impl ExecutionPlan for ChildExec { ) -> Result { unreachable!() } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut( + &Arc, + ) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } } /// A PhysicalExtensionCodec that can serialize and deserialize ChildExec diff --git a/datafusion-examples/examples/relation_planner/table_sample.rs b/datafusion-examples/examples/relation_planner/table_sample.rs index b0ccff8d10d8c..4a4dfff1da3e8 100644 --- a/datafusion-examples/examples/relation_planner/table_sample.rs +++ b/datafusion-examples/examples/relation_planner/table_sample.rs @@ -117,7 +117,7 @@ use datafusion::{ }; use datafusion_common::{ DFSchemaRef, DataFusionError, Result, Statistics, internal_err, not_impl_err, - plan_datafusion_err, plan_err, + plan_datafusion_err, plan_err, tree_node::TreeNodeRecursion, }; use datafusion_expr::physical_planning_context::PhysicalPlanningContext; use datafusion_expr::{ @@ -749,6 +749,22 @@ impl ExecutionPlan for SampleExec { Ok(Arc::new(stats)) } + + fn apply_expressions( + &self, + f: &mut dyn FnMut( + &Arc, + ) -> Result, + ) -> Result { + datafusion::physical_plan::apply_expression_roots( + self.cache + .output_ordering() + .into_iter() + .flatten() + .map(|sort_expr| &sort_expr.expr), + f, + ) + } } /// Bernoulli sampler: includes each row with probability `(upper - lower)`. diff --git a/datafusion/catalog/src/memory/table.rs b/datafusion/catalog/src/memory/table.rs index 5d07133799ffc..ef5669a3a13f0 100644 --- a/datafusion/catalog/src/memory/table.rs +++ b/datafusion/catalog/src/memory/table.rs @@ -31,6 +31,7 @@ use arrow::compute::{and, filter_record_batch}; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use arrow::record_batch::RecordBatch; use datafusion_common::error::Result; +use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::{Constraints, DFSchema, SchemaExt, not_impl_err, plan_err}; use datafusion_datasource::memory::{MemSink, MemorySourceConfig}; use datafusion_datasource::sink::DataSinkExec; @@ -39,13 +40,13 @@ use datafusion_expr::dml::InsertOp; use datafusion_expr::physical_planning_context::PhysicalPlanningContext; use datafusion_expr::{Expr, SortExpr, TableType}; use datafusion_physical_expr::{ - LexOrdering, PhysicalExpr, create_physical_expr, create_physical_sort_exprs, + LexOrdering, create_physical_expr, create_physical_sort_exprs, }; use datafusion_physical_plan::repartition::RepartitionExec; use datafusion_physical_plan::stream::RecordBatchStreamAdapter; use datafusion_physical_plan::{ - DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, - collect_partitioned, + DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PhysicalExpr, + PlanProperties, collect_partitioned, }; use datafusion_session::Session; @@ -597,4 +598,11 @@ impl ExecutionPlan for DmlResultExec { stream, ))) } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } } diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index 4e914556b4cc0..723617006b422 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -4711,6 +4711,13 @@ mod tests { ) -> Result { unimplemented!("NoOpExecutionPlan::execute"); } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } } struct ExpressionExtensionPlanner; @@ -4878,6 +4885,12 @@ digraph { ) -> Result { unimplemented!() } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } } impl DisplayAs for OkExtensionNode { fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { @@ -4924,6 +4937,12 @@ digraph { ) -> Result { unimplemented!() } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } } impl DisplayAs for InvariantFailsExtensionNode { fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { @@ -5048,6 +5067,12 @@ digraph { ) -> Result { unimplemented!() } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } } impl DisplayAs for ExecutableInvariantFails { fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { diff --git a/datafusion/core/tests/custom_sources_cases/mod.rs b/datafusion/core/tests/custom_sources_cases/mod.rs index c70722cb2f2ff..9582c13669358 100644 --- a/datafusion/core/tests/custom_sources_cases/mod.rs +++ b/datafusion/core/tests/custom_sources_cases/mod.rs @@ -38,6 +38,7 @@ use datafusion_catalog::Session; use datafusion_common::cast::as_primitive_array; use datafusion_common::project_schema; use datafusion_common::stats::Precision; +use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_physical_expr::EquivalenceProperties; use datafusion_physical_plan::PlanProperties; use datafusion_physical_plan::StatisticsArgs; @@ -209,6 +210,22 @@ impl ExecutionPlan for CustomExecutionPlan { .collect(), })) } + + fn apply_expressions( + &self, + f: &mut dyn FnMut( + &Arc, + ) -> Result, + ) -> Result { + datafusion::physical_plan::apply_expression_roots( + self.cache + .output_ordering() + .into_iter() + .flatten() + .map(|sort_expr| &sort_expr.expr), + f, + ) + } } #[async_trait] diff --git a/datafusion/core/tests/custom_sources_cases/provider_filter_pushdown.rs b/datafusion/core/tests/custom_sources_cases/provider_filter_pushdown.rs index 18695accd0f2e..45fcc1b423de9 100644 --- a/datafusion/core/tests/custom_sources_cases/provider_filter_pushdown.rs +++ b/datafusion/core/tests/custom_sources_cases/provider_filter_pushdown.rs @@ -35,6 +35,7 @@ use datafusion::prelude::*; use datafusion::scalar::ScalarValue; use datafusion_catalog::Session; use datafusion_common::cast::as_primitive_array; +use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::{DataFusionError, internal_err, not_impl_err}; use datafusion_expr::expr::{BinaryExpr, Cast}; use datafusion_functions_aggregate::expr_fn::count; @@ -148,6 +149,22 @@ impl ExecutionPlan for CustomPlan { })), ))) } + + fn apply_expressions( + &self, + f: &mut dyn FnMut( + &Arc, + ) -> Result, + ) -> Result { + datafusion::physical_plan::apply_expression_roots( + self.cache + .output_ordering() + .into_iter() + .flatten() + .map(|sort_expr| &sort_expr.expr), + f, + ) + } } #[derive(Clone, Debug)] diff --git a/datafusion/core/tests/custom_sources_cases/statistics.rs b/datafusion/core/tests/custom_sources_cases/statistics.rs index d289b5c348b3c..5665b4b193aed 100644 --- a/datafusion/core/tests/custom_sources_cases/statistics.rs +++ b/datafusion/core/tests/custom_sources_cases/statistics.rs @@ -33,6 +33,7 @@ use datafusion::{ scalar::ScalarValue, }; use datafusion_catalog::Session; +use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::{project_schema, stats::Precision}; use datafusion_physical_expr::EquivalenceProperties; use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType}; @@ -185,6 +186,22 @@ impl ExecutionPlan for StatisticsValidation { Ok(Arc::new(self.stats.clone())) } } + + fn apply_expressions( + &self, + f: &mut dyn FnMut( + &Arc, + ) -> Result, + ) -> Result { + datafusion::physical_plan::apply_expression_roots( + self.cache + .output_ordering() + .into_iter() + .flatten() + .map(|sort_expr| &sort_expr.expr), + f, + ) + } } fn init_ctx(stats: Statistics, schema: Schema) -> Result { diff --git a/datafusion/core/tests/fuzz_cases/once_exec.rs b/datafusion/core/tests/fuzz_cases/once_exec.rs index 9b57141061518..e3ee52a4872e5 100644 --- a/datafusion/core/tests/fuzz_cases/once_exec.rs +++ b/datafusion/core/tests/fuzz_cases/once_exec.rs @@ -17,6 +17,7 @@ use arrow_schema::SchemaRef; use datafusion_common::internal_datafusion_err; +use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_execution::{SendableRecordBatchStream, TaskContext}; use datafusion_physical_expr::{EquivalenceProperties, Partitioning}; use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType}; @@ -105,4 +106,20 @@ impl ExecutionPlan for OnceExec { stream.ok_or_else(|| internal_datafusion_err!("Stream already consumed")) } + + fn apply_expressions( + &self, + f: &mut dyn FnMut( + &Arc, + ) -> datafusion_common::Result, + ) -> datafusion_common::Result { + datafusion_physical_plan::apply_expression_roots( + self.cache + .output_ordering() + .into_iter() + .flatten() + .map(|sort_expr| &sort_expr.expr), + f, + ) + } } diff --git a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs index 189650fe4afca..ccfc8fb055eb1 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs +++ b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs @@ -40,7 +40,9 @@ use datafusion::prelude::{SessionConfig, SessionContext}; use datafusion_common::ScalarValue; use datafusion_common::config::CsvOptions; use datafusion_common::error::Result; -use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode}; +use datafusion_common::tree_node::{ + Transformed, TransformedResult, TreeNode, TreeNodeRecursion, +}; use datafusion_datasource::file_groups::FileGroup; use datafusion_datasource::file_scan_config::FileScanConfigBuilder; use datafusion_expr::{JoinType, Operator}; @@ -203,6 +205,20 @@ impl ExecutionPlan for SortRequiredExec { ))) } + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + datafusion_physical_plan::apply_expression_roots( + self.cache + .output_ordering() + .into_iter() + .flatten() + .map(|sort_expr| &sort_expr.expr), + f, + ) + } + fn execute( &self, _partition: usize, @@ -290,6 +306,13 @@ impl ExecutionPlan for SinglePartitionMaintainsOrderExec { Ok(Arc::new(Self::new(child))) } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + fn execute( &self, _partition: usize, diff --git a/datafusion/core/tests/physical_optimizer/ensure_requirements.rs b/datafusion/core/tests/physical_optimizer/ensure_requirements.rs index 2c6c46c82985a..4baa0380c7208 100644 --- a/datafusion/core/tests/physical_optimizer/ensure_requirements.rs +++ b/datafusion/core/tests/physical_optimizer/ensure_requirements.rs @@ -22,6 +22,7 @@ //! suite and can use real `ExecutionPlan`s where convenient. use datafusion_common::config::ConfigOptions; +use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_physical_optimizer::PhysicalOptimizerRule; use datafusion_physical_optimizer::ensure_requirements::EnsureRequirements; @@ -111,6 +112,12 @@ impl ExecutionPlan for MockMultiPartitionExec { fn children(&self) -> Vec<&Arc> { vec![] } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } fn with_new_children( self: Arc, _children: Vec>, @@ -974,6 +981,27 @@ impl ExecutionPlan for MockReqExec { fn children(&self) -> Vec<&Arc> { vec![&self.input] } + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + #[expect(deprecated)] + let distribution = if let Distribution::HashPartitioned(exprs) + | Distribution::KeyPartitioned(exprs) = &self.dist + { + exprs.as_slice() + } else { + &[] + }; + datafusion_physical_plan::apply_expression_roots( + self.ord + .iter() + .flatten() + .map(|sort_expr| &sort_expr.expr) + .chain(distribution), + f, + ) + } fn input_distribution_requirements( &self, ) -> datafusion_physical_plan::InputDistributionRequirements { diff --git a/datafusion/core/tests/physical_optimizer/filter_pushdown.rs b/datafusion/core/tests/physical_optimizer/filter_pushdown.rs index 7593fe351548e..a099c10747939 100644 --- a/datafusion/core/tests/physical_optimizer/filter_pushdown.rs +++ b/datafusion/core/tests/physical_optimizer/filter_pushdown.rs @@ -34,7 +34,7 @@ use datafusion::{ scalar::ScalarValue, }; use datafusion_catalog::memory::DataSourceExec; -use datafusion_common::config::ConfigOptions; +use datafusion_common::{JoinType, config::ConfigOptions, tree_node::TreeNodeRecursion}; use datafusion_datasource::{ PartitionedFile, file_groups::FileGroup, file_scan_config::FileScanConfigBuilder, }; @@ -46,7 +46,9 @@ use datafusion_functions_aggregate::{ min_max::{max_udaf, min_udaf}, }; use datafusion_physical_expr::{ - LexOrdering, PhysicalSortExpr, expressions::col, utils::conjunction, + LexOrdering, PhysicalSortExpr, + expressions::{DynamicFilterPhysicalExpr, col}, + utils::conjunction, }; use datafusion_physical_expr::{ Partitioning, ScalarFunctionExpr, aggregate::AggregateExprBuilder, @@ -60,6 +62,7 @@ use datafusion_physical_plan::{ coalesce_partitions::CoalescePartitionsExec, collect, filter::{FilterExec, FilterExecBuilder}, + joins::{HashJoinExec, PartitionMode}, projection::ProjectionExec, repartition::RepartitionExec, sorts::sort::SortExec, @@ -3085,6 +3088,106 @@ async fn test_filter_with_projection_pushdown() { assert_batches_eq!(expected, &result); } +/// Test that ExecutionPlan::apply_expressions() can discover dynamic filters across the plan tree. +/// +/// Not portable to sqllogictest: asserts by walking the plan tree with +/// `apply_expressions` + `downcast_ref::` and +/// counting nodes. Neither API is observable from SQL. +#[tokio::test] +async fn test_discover_dynamic_filters_via_expressions_api() { + fn count_dynamic_filters(plan: &Arc) -> usize { + let mut count = 0; + + // Check expressions from this node using apply_expressions + let _ = plan.apply_expressions(&mut |expr| { + if let Some(_df) = expr.downcast_ref::() { + count += 1; + } + Ok(TreeNodeRecursion::Continue) + }); + + // Recursively visit children + for child in plan.children() { + count += count_dynamic_filters(child); + } + + count + } + + // Create build side (left) + let build_batches = + vec![record_batch!(("a", Utf8, ["foo", "bar"]), ("b", Int32, [1, 2])).unwrap()]; + let build_schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Utf8, false), + Field::new("b", DataType::Int32, false), + ])); + let build_scan = TestScanBuilder::new(build_schema.clone()) + .with_support(true) + .with_batches(build_batches) + .build(); + + // Create probe side (right) + let probe_batches = vec![ + record_batch!( + ("a", Utf8, ["foo", "bar", "baz", "qux"]), + ("c", Float64, [1.0, 2.0, 3.0, 4.0]) + ) + .unwrap(), + ]; + let probe_schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Utf8, false), + Field::new("c", DataType::Float64, false), + ])); + let probe_scan = TestScanBuilder::new(probe_schema.clone()) + .with_support(true) + .with_batches(probe_batches) + .build(); + + // Create HashJoinExec + let plan = Arc::new( + HashJoinExec::try_new( + build_scan, + probe_scan, + vec![( + col("a", &build_schema).unwrap(), + col("a", &probe_schema).unwrap(), + )], + None, + &JoinType::Inner, + None, + PartitionMode::CollectLeft, + datafusion_common::NullEquality::NullEqualsNothing, + false, + ) + .unwrap(), + ) as Arc; + + // Before optimization: no dynamic filters + let count_before = count_dynamic_filters(&plan); + assert_eq!( + count_before, 0, + "Before optimization, should have no dynamic filters" + ); + + // Apply filter pushdown optimization (this creates dynamic filters) + let mut config = ConfigOptions::default(); + config.optimizer.enable_dynamic_filter_pushdown = true; + config.execution.parquet.pushdown_filters = true; + let optimized_plan = FilterPushdown::new_post_optimization() + .optimize(plan, &config) + .unwrap(); + + // After optimization: should discover dynamic filters + // We expect 2 dynamic filters: + // 1. In the HashJoinExec (producer) + // 2. In the DataSourceExec (consumer, pushed down to the probe side) + let count_after = count_dynamic_filters(&optimized_plan); + assert_eq!( + count_after, 2, + "After optimization, should discover exactly 2 dynamic filters (1 in HashJoinExec, 1 in DataSourceExec), found {count_after}" + ); +} + // ==== Filter pushdown through SortExec tests ==== /// FilterExec above a plain SortExec (no fetch) should be pushed below it. diff --git a/datafusion/core/tests/physical_optimizer/join_selection.rs b/datafusion/core/tests/physical_optimizer/join_selection.rs index 3827e6e98b5e6..41f8960e8c21b 100644 --- a/datafusion/core/tests/physical_optimizer/join_selection.rs +++ b/datafusion/core/tests/physical_optimizer/join_selection.rs @@ -25,6 +25,7 @@ use std::{ use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use arrow::record_batch::RecordBatch; use datafusion_common::config::ConfigOptions; +use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::{ColumnStatistics, JoinType, ScalarValue, stats::Precision}; use datafusion_common::{JoinSide, NullEquality}; use datafusion_common::{Result, Statistics}; @@ -1125,6 +1126,20 @@ impl ExecutionPlan for UnboundedExec { batch: self.batch.clone(), })) } + + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + datafusion_physical_plan::apply_expression_roots( + self.cache + .output_ordering() + .into_iter() + .flatten() + .map(|sort_expr| &sort_expr.expr), + f, + ) + } } #[derive(Eq, PartialEq, Debug)] @@ -1230,6 +1245,20 @@ impl ExecutionPlan for StatisticsExec { self.stats.clone() })) } + + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + datafusion_physical_plan::apply_expression_roots( + self.cache + .output_ordering() + .into_iter() + .flatten() + .map(|sort_expr| &sort_expr.expr), + f, + ) + } } #[test] diff --git a/datafusion/core/tests/physical_optimizer/pushdown_utils.rs b/datafusion/core/tests/physical_optimizer/pushdown_utils.rs index 2ffd1899b3c1d..4f8b9ad42b6c8 100644 --- a/datafusion/core/tests/physical_optimizer/pushdown_utils.rs +++ b/datafusion/core/tests/physical_optimizer/pushdown_utils.rs @@ -18,6 +18,7 @@ use arrow::datatypes::SchemaRef; use arrow::{array::RecordBatch, compute::concat_batches}; use datafusion::{datasource::object_store::ObjectStoreUrl, physical_plan::PhysicalExpr}; +use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::{Result, config::ConfigOptions, internal_err}; use datafusion_datasource::{ PartitionedFile, file::FileSource, file_scan_config::FileScanConfig, @@ -234,6 +235,21 @@ impl FileSource for TestSource { fn table_schema(&self) -> &datafusion_datasource::TableSchema { &self.table_schema } + + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + datafusion_physical_plan::apply_expression_roots( + self.predicate.iter().chain( + self.projection + .iter() + .flatten() + .map(|proj_expr| &proj_expr.expr), + ), + f, + ) + } } #[derive(Debug, Clone)] @@ -549,4 +565,11 @@ impl ExecutionPlan for TestNode { Ok(res) } } + + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + datafusion_physical_plan::apply_expression_roots([&self.predicate], f) + } } diff --git a/datafusion/core/tests/physical_optimizer/test_utils.rs b/datafusion/core/tests/physical_optimizer/test_utils.rs index 74230b24e2ab5..a570089b94f01 100644 --- a/datafusion/core/tests/physical_optimizer/test_utils.rs +++ b/datafusion/core/tests/physical_optimizer/test_utils.rs @@ -30,7 +30,9 @@ use datafusion::datasource::physical_plan::ParquetSource; use datafusion::datasource::source::{DataSource, DataSourceExec}; use datafusion_common::config::ConfigOptions; use datafusion_common::stats::Precision; -use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode}; +use datafusion_common::tree_node::{ + Transformed, TransformedResult, TreeNode, TreeNodeRecursion, +}; use datafusion_common::utils::expr::COUNT_STAR_EXPANSION; use datafusion_common::{ ColumnStatistics, JoinType, NullEquality, Result, Statistics, internal_err, @@ -526,6 +528,19 @@ impl ExecutionPlan for RequirementsTestExec { ) -> Result { unimplemented!("Test exec does not support execution") } + + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + datafusion_physical_plan::apply_expression_roots( + self.required_input_ordering + .iter() + .flatten() + .map(|sort_expr| &sort_expr.expr), + f, + ) + } } /// A [`PlanContext`] object is susceptible to being left in an inconsistent state after @@ -1074,6 +1089,20 @@ impl ExecutionPlan for TestScan { }) } } + + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + datafusion_physical_plan::apply_expression_roots( + self.output_ordering + .iter() + .flatten() + .chain(self.requested_ordering.iter().flatten()) + .map(|sort_expr| &sort_expr.expr), + f, + ) + } } /// Helper function to create a TestScan with ordering @@ -1090,6 +1119,13 @@ struct InexactMemorySource { } impl DataSource for InexactMemorySource { + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + self.inner.apply_expressions(f) + } + fn open( &self, partition: usize, diff --git a/datafusion/core/tests/user_defined/insert_operation.rs b/datafusion/core/tests/user_defined/insert_operation.rs index f3d3f70bdf925..9efe1ec9eba15 100644 --- a/datafusion/core/tests/user_defined/insert_operation.rs +++ b/datafusion/core/tests/user_defined/insert_operation.rs @@ -25,6 +25,7 @@ use datafusion::{ }; use datafusion_catalog::{Session, TableProvider}; use datafusion_common::config::Dialect; +use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_expr::{Expr, TableType, dml::InsertOp}; use datafusion_physical_expr::{EquivalenceProperties, Partitioning}; use datafusion_physical_plan::execution_plan::SchedulingType; @@ -176,6 +177,22 @@ impl ExecutionPlan for TestInsertExec { ) -> Result { unimplemented!("TestInsertExec is a stub for testing.") } + + fn apply_expressions( + &self, + f: &mut dyn FnMut( + &Arc, + ) -> Result, + ) -> Result { + datafusion_physical_plan::apply_expression_roots( + self.plan_properties + .output_ordering() + .into_iter() + .flatten() + .map(|sort_expr| &sort_expr.expr), + f, + ) + } } fn make_count_schema() -> SchemaRef { diff --git a/datafusion/core/tests/user_defined/user_defined_plan.rs b/datafusion/core/tests/user_defined/user_defined_plan.rs index 354a1b3110250..07d44c8de39fc 100644 --- a/datafusion/core/tests/user_defined/user_defined_plan.rs +++ b/datafusion/core/tests/user_defined/user_defined_plan.rs @@ -90,7 +90,9 @@ use datafusion::{ prelude::{SessionConfig, SessionContext}, }; use datafusion_common::config::ConfigOptions; -use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode}; +use datafusion_common::tree_node::{ + Transformed, TransformedResult, TreeNode, TreeNodeRecursion, +}; use datafusion_common::{ScalarValue, assert_eq_or_internal_err, assert_or_internal_err}; use datafusion_expr::{FetchType, InvariantLevel, Projection, SortExpr}; use datafusion_optimizer::AnalyzerRule; @@ -748,6 +750,22 @@ impl ExecutionPlan for TopKExec { state: BTreeMap::new(), })) } + + fn apply_expressions( + &self, + f: &mut dyn FnMut( + &Arc, + ) -> Result, + ) -> Result { + datafusion::physical_plan::apply_expression_roots( + self.cache + .output_ordering() + .into_iter() + .flatten() + .map(|sort_expr| &sort_expr.expr), + f, + ) + } } // A very specialized TopK implementation diff --git a/datafusion/datasource-arrow/src/source.rs b/datafusion/datasource-arrow/src/source.rs index 27533052ce03f..3da60697fb4de 100644 --- a/datafusion/datasource-arrow/src/source.rs +++ b/datafusion/datasource-arrow/src/source.rs @@ -40,6 +40,7 @@ use arrow::buffer::Buffer; use arrow::ipc::reader::{FileDecoder, FileReader, StreamReader}; use datafusion_common::error::Result; use datafusion_common::exec_datafusion_err; +use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_datasource::PartitionedFile; use datafusion_datasource::file::FileSource; use datafusion_datasource::file_scan_config::FileScanConfig; @@ -392,6 +393,21 @@ impl FileSource for ArrowSource { fn projection(&self) -> Option<&ProjectionExprs> { Some(&self.projection.source) } + + fn apply_expressions( + &self, + f: &mut dyn FnMut( + &Arc, + ) -> Result, + ) -> Result { + datafusion_physical_plan::apply_expression_roots( + self.projection + .source + .iter() + .map(|proj_expr| &proj_expr.expr), + f, + ) + } } /// `FileOpener` wrapper for both Arrow IPC file and stream formats diff --git a/datafusion/datasource-avro/src/source.rs b/datafusion/datasource-avro/src/source.rs index e3be9d8a401d0..a069c1b4287e6 100644 --- a/datafusion/datasource-avro/src/source.rs +++ b/datafusion/datasource-avro/src/source.rs @@ -22,6 +22,7 @@ use std::sync::Arc; use arrow::datatypes::{Schema, SchemaRef}; use arrow_avro::reader::{Reader, ReaderBuilder}; use datafusion_common::error::Result; +use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_datasource::TableSchema; use datafusion_datasource::file::FileSource; use datafusion_datasource::file_scan_config::FileScanConfig; @@ -168,6 +169,21 @@ impl FileSource for AvroSource { // Avro OCF does not support safe byte-range splitting in this reader path. false } + + fn apply_expressions( + &self, + f: &mut dyn FnMut( + &Arc, + ) -> Result, + ) -> Result { + datafusion_physical_plan::apply_expression_roots( + self.projection + .source + .iter() + .map(|proj_expr| &proj_expr.expr), + f, + ) + } } mod private { diff --git a/datafusion/datasource-csv/src/source.rs b/datafusion/datasource-csv/src/source.rs index 25ec311880405..725e36e33cf32 100644 --- a/datafusion/datasource-csv/src/source.rs +++ b/datafusion/datasource-csv/src/source.rs @@ -33,6 +33,7 @@ use datafusion_datasource::{ use arrow::csv; use datafusion_common::config::CsvOptions; +use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::{DataFusionError, Result, exec_datafusion_err}; use datafusion_common_runtime::JoinSet; use datafusion_datasource::file::FileSource; @@ -308,6 +309,21 @@ impl FileSource for CsvSource { DisplayFormatType::TreeRender => Ok(()), } } + + fn apply_expressions( + &self, + f: &mut dyn FnMut( + &Arc, + ) -> Result, + ) -> Result { + datafusion_physical_plan::apply_expression_roots( + self.projection + .source + .iter() + .map(|proj_expr| &proj_expr.expr), + f, + ) + } } impl FileOpener for CsvOpener { diff --git a/datafusion/datasource-json/src/source.rs b/datafusion/datasource-json/src/source.rs index 8632d6b942bc1..f19f6a71d8f67 100644 --- a/datafusion/datasource-json/src/source.rs +++ b/datafusion/datasource-json/src/source.rs @@ -27,6 +27,7 @@ use crate::utils::{ChannelReader, JsonArrayToNdjsonReader}; use datafusion_common::error::{DataFusionError, Result}; use datafusion_common::exec_datafusion_err; +use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common_runtime::{JoinSet, SpawnedTask}; use datafusion_datasource::boundary_stream::AlignedBoundaryStream; use datafusion_datasource::decoder::{DecoderDeserializer, deserialize_stream}; @@ -231,6 +232,21 @@ impl FileSource for JsonSource { fn file_type(&self) -> &str { "json" } + + fn apply_expressions( + &self, + f: &mut dyn FnMut( + &Arc, + ) -> Result, + ) -> Result { + datafusion_physical_plan::apply_expression_roots( + self.projection + .source + .iter() + .map(|proj_expr| &proj_expr.expr), + f, + ) + } } impl FileOpener for JsonOpener { diff --git a/datafusion/datasource-parquet/src/source.rs b/datafusion/datasource-parquet/src/source.rs index 3443b08475e0d..b2674efd404c3 100644 --- a/datafusion/datasource-parquet/src/source.rs +++ b/datafusion/datasource-parquet/src/source.rs @@ -40,6 +40,7 @@ use arrow::array::timezone::Tz; use arrow::datatypes::TimeUnit; use datafusion_common::DataFusionError; use datafusion_common::config::TableParquetOptions; +use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_datasource::TableSchema; use datafusion_datasource::file::FileSource; use datafusion_datasource::file_scan_config::FileScanConfig; @@ -1047,6 +1048,20 @@ impl FileSource for ParquetSource { inner: Arc::new(new_source) as Arc, }) } + + fn apply_expressions( + &self, + f: &mut dyn FnMut( + &Arc, + ) -> datafusion_common::Result, + ) -> datafusion_common::Result { + datafusion_physical_plan::apply_expression_roots( + self.predicate + .iter() + .chain(self.projection.iter().map(|proj_expr| &proj_expr.expr)), + f, + ) + } } /// Returns the a [`TableSchema`] containing a [`RowNumber`] virtual column and a [`Column`] expression referencing its row index column. diff --git a/datafusion/datasource/src/file.rs b/datafusion/datasource/src/file.rs index 07460b23694b7..155daaa501c4b 100644 --- a/datafusion/datasource/src/file.rs +++ b/datafusion/datasource/src/file.rs @@ -29,6 +29,7 @@ use crate::morsel::{FileOpenerMorselizer, Morselizer}; #[expect(deprecated)] use crate::schema_adapter::SchemaAdapterFactory; use datafusion_common::config::ConfigOptions; +use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::{Result, not_impl_err}; use datafusion_physical_expr::projection::ProjectionExprs; use datafusion_physical_expr::{EquivalenceProperties, LexOrdering, PhysicalExpr}; @@ -351,6 +352,27 @@ pub trait FileSource: Any + Send + Sync { fn schema_adapter_factory(&self) -> Option> { None } + + /// Apply a function to all physical expressions used by this file source. + /// + /// This includes: + /// - Filter predicates (which may contain dynamic filters) + /// - Projection expressions + /// + /// The function `f` is called once for each expression. The function should + /// return `TreeNodeRecursion::Continue` to continue visiting other expressions, + /// or `TreeNodeRecursion::Stop` to stop visiting expressions early. + /// + /// Implementations must explicitly visit all expressions. There is no default + /// implementation to ensure that all FileSource implementations handle this correctly. + /// + /// See [`ExecutionPlan::apply_expressions`] for more details and examples. + /// + /// [`ExecutionPlan::apply_expressions`]: datafusion_physical_plan::ExecutionPlan::apply_expressions + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result; } impl dyn FileSource { diff --git a/datafusion/datasource/src/file_scan_config/mod.rs b/datafusion/datasource/src/file_scan_config/mod.rs index d1dd3c11fca7d..f38a95b1ef879 100644 --- a/datafusion/datasource/src/file_scan_config/mod.rs +++ b/datafusion/datasource/src/file_scan_config/mod.rs @@ -30,6 +30,7 @@ use crate::{ use arrow::datatypes::Fields; use arrow::datatypes::{DataType, Schema, SchemaRef}; use datafusion_common::config::ConfigOptions; +use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::{ Constraints, Result, ScalarValue, Statistics, internal_datafusion_err, internal_err, }; @@ -82,7 +83,9 @@ use std::{fmt::Debug, fmt::Formatter, fmt::Result as FmtResult, sync::Arc}; /// # use arrow::datatypes::{Field, Fields, DataType, Schema, SchemaRef}; /// # use object_store::ObjectStore; /// # use datafusion_common::Result; +/// # use datafusion_common::tree_node::TreeNodeRecursion; /// # use datafusion_datasource::file::FileSource; +/// # use datafusion_physical_plan::PhysicalExpr; /// # use datafusion_datasource::file_groups::FileGroup; /// # use datafusion_datasource::PartitionedFile; /// # use datafusion_datasource::file_scan_config::{FileScanConfig, FileScanConfigBuilder}; @@ -112,6 +115,7 @@ use std::{fmt::Debug, fmt::Formatter, fmt::Result as FmtResult, sync::Arc}; /// # fn file_type(&self) -> &str { "parquet" } /// # // Note that this implementation drops the projection on the floor, it is not complete! /// # fn try_pushdown_projection(&self, projection: &ProjectionExprs) -> Result>> { Ok(Some(Arc::new(self.clone()) as Arc)) } +/// # fn apply_expressions(&self, _f: &mut dyn FnMut(&Arc) -> Result) -> Result { Ok(TreeNodeRecursion::Continue) } /// # } /// # impl ParquetSource { /// # fn new(table_schema: impl Into) -> Self { Self {table_schema: table_schema.into()} } @@ -1155,6 +1159,14 @@ impl DataSource for FileScanConfig { Some(Arc::new(new_config)) } + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + // Delegate to the file source + self.file_source.apply_expressions(f) + } + /// Create any shared state that should be passed between sibling streams /// during one execution. /// @@ -1566,9 +1578,11 @@ mod tests { use arrow::datatypes::Field; use datafusion_common::ColumnStatistics; use datafusion_common::stats::Precision; + use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::{Result, assert_batches_eq, internal_err}; use datafusion_execution::TaskContext; use datafusion_expr::SortExpr; + use datafusion_physical_expr::PhysicalExpr; use datafusion_physical_expr::create_physical_sort_expr; use datafusion_physical_expr::expressions::Literal; use datafusion_physical_expr::projection::ProjectionExpr; @@ -1631,6 +1645,13 @@ mod tests { inner: Arc::new(self.clone()) as Arc, }) } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } } #[test] @@ -2896,6 +2917,13 @@ mod tests { inner: Arc::new(self.clone()) as Arc, }) } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } } #[test] diff --git a/datafusion/datasource/src/memory.rs b/datafusion/datasource/src/memory.rs index 255dd76cbd6b4..874457e249db6 100644 --- a/datafusion/datasource/src/memory.rs +++ b/datafusion/datasource/src/memory.rs @@ -28,6 +28,7 @@ use crate::source::{DataSource, DataSourceExec}; use arrow::array::{RecordBatch, RecordBatchOptions}; use arrow::datatypes::{Schema, SchemaRef}; +use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::{ Result, ScalarValue, assert_or_internal_err, plan_err, project_schema, }; @@ -256,6 +257,19 @@ impl DataSource for MemorySourceConfig { }) .transpose() } + + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + datafusion_physical_plan::apply_expression_roots( + self.sort_information + .iter() + .flatten() + .map(|sort_expr| &sort_expr.expr), + f, + ) + } } impl MemorySourceConfig { diff --git a/datafusion/datasource/src/sink.rs b/datafusion/datasource/src/sink.rs index 18ebe80773e8a..881c5258ca388 100644 --- a/datafusion/datasource/src/sink.rs +++ b/datafusion/datasource/src/sink.rs @@ -24,9 +24,10 @@ use std::sync::Arc; use arrow::array::{ArrayRef, RecordBatch, UInt64Array}; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::{Result, assert_eq_or_internal_err}; use datafusion_execution::TaskContext; -use datafusion_physical_expr::{Distribution, EquivalenceProperties}; +use datafusion_physical_expr::{Distribution, EquivalenceProperties, PhysicalExpr}; use datafusion_physical_expr_common::sort_expr::{LexRequirement, OrderingRequirements}; use datafusion_physical_plan::metrics::MetricsSet; use datafusion_physical_plan::stream::RecordBatchStreamAdapter; @@ -231,6 +232,19 @@ impl ExecutionPlan for DataSinkExec { ))) } + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + datafusion_physical_plan::apply_expression_roots( + self.sort_order + .iter() + .flatten() + .map(|requirement| &requirement.expr), + f, + ) + } + /// Execute the plan and return a stream of `RecordBatch`es for /// the specified partition. fn execute( diff --git a/datafusion/datasource/src/source.rs b/datafusion/datasource/src/source.rs index c280470bb0d0b..a4a167137c46d 100644 --- a/datafusion/datasource/src/source.rs +++ b/datafusion/datasource/src/source.rs @@ -40,6 +40,7 @@ use itertools::Itertools; use crate::file::FileSource; use crate::file_scan_config::FileScanConfig; use datafusion_common::config::ConfigOptions; +use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::{Constraints, Result, Statistics}; use datafusion_execution::{SendableRecordBatchStream, TaskContext}; use datafusion_physical_expr::{EquivalenceProperties, Partitioning, PhysicalExpr}; @@ -225,6 +226,22 @@ pub trait DataSource: Any + Send + Sync + Debug { None } + /// Apply a closure to each expression used by this data source. + /// + /// This includes filter predicates (which may contain dynamic filters) and any + /// other expressions used during data scanning. + /// + /// Implementations must override this method. If the data source has no expressions, + /// return `Ok(TreeNodeRecursion::Continue)` immediately. + /// + /// See [`ExecutionPlan::apply_expressions`] for more details and implementation examples. + /// + /// [`ExecutionPlan::apply_expressions`]: datafusion_physical_plan::ExecutionPlan::apply_expressions + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result; + /// Injects arbitrary run-time state into this DataSource, returning a new instance /// that incorporates that state *if* it is relevant to the concrete DataSource implementation. /// @@ -359,6 +376,14 @@ impl ExecutionPlan for DataSourceExec { Vec::new() } + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + // Delegate to the underlying data source + self.data_source.apply_expressions(f) + } + fn with_new_children( self: Arc, _: Vec>, diff --git a/datafusion/datasource/src/test_util.rs b/datafusion/datasource/src/test_util.rs index 5bef6d44e1408..20dfae5b3ac79 100644 --- a/datafusion/datasource/src/test_util.rs +++ b/datafusion/datasource/src/test_util.rs @@ -22,7 +22,7 @@ use crate::{ use std::sync::Arc; use arrow::datatypes::Schema; -use datafusion_common::Result; +use datafusion_common::{Result, tree_node::TreeNodeRecursion}; use datafusion_physical_expr::{PhysicalExpr, expressions::Column}; use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet; use object_store::ObjectStore; @@ -125,6 +125,13 @@ impl FileSource for MockSource { ) -> Option<&datafusion_physical_plan::projection::ProjectionExprs> { Some(&self.projection.source) } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } } /// Create a column expression diff --git a/datafusion/ffi/src/execution_plan.rs b/datafusion/ffi/src/execution_plan.rs index 087a351b697cc..2ba1d0c406c6e 100644 --- a/datafusion/ffi/src/execution_plan.rs +++ b/datafusion/ffi/src/execution_plan.rs @@ -20,6 +20,7 @@ use std::pin::Pin; use std::sync::Arc; use datafusion_common::config::ConfigOptions; +use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::{DataFusionError, Result, Statistics}; use datafusion_execution::{SendableRecordBatchStream, TaskContext}; use datafusion_physical_expr_common::metrics::MetricsSet; @@ -33,6 +34,7 @@ use tokio::runtime::Handle; use crate::config::FFI_ConfigOptions; use crate::execution::FFI_TaskContext; +use crate::physical_expr::FFI_PhysicalExpr; use crate::physical_expr::metrics::FFI_MetricsSet; use crate::plan_properties::FFI_PlanProperties; use crate::record_batch_stream::FFI_RecordBatchStream; @@ -50,6 +52,10 @@ pub struct FFI_ExecutionPlan { /// Return a vector of children plans pub children: unsafe extern "C" fn(plan: &Self) -> SVec, + /// Return the physical expression roots owned by this plan node. + pub apply_expressions: + unsafe extern "C" fn(plan: &Self) -> FFI_Result>, + pub with_new_children: unsafe extern "C" fn(plan: &Self, children: SVec) -> FFI_Result, @@ -91,6 +97,9 @@ pub struct FFI_ExecutionPlan { /// Release the memory of the private data when it is no longer being used. pub release: unsafe extern "C" fn(arg: &mut Self), + /// Return the major DataFusion version number of this provider. + pub version: unsafe extern "C" fn() -> u64, + /// Internal data. This is only to be accessed by the provider of the plan. /// A [`ForeignExecutionPlan`] should never attempt to access this data. pub private_data: *mut c_void, @@ -138,6 +147,17 @@ unsafe extern "C" fn children_fn_wrapper( .collect() } +unsafe extern "C" fn apply_expressions_fn_wrapper( + plan: &FFI_ExecutionPlan, +) -> FFI_Result> { + let mut expressions = SVec::new(); + let result = plan.inner().apply_expressions(&mut |expr| { + expressions.push(FFI_PhysicalExpr::from(Arc::clone(expr))); + Ok(TreeNodeRecursion::Continue) + }); + sresult!(result.map(|_| expressions)) +} + unsafe extern "C" fn with_new_children_fn_wrapper( plan: &FFI_ExecutionPlan, children: SVec, @@ -306,6 +326,7 @@ impl FFI_ExecutionPlan { Self { properties: properties_fn_wrapper, children: children_fn_wrapper, + apply_expressions: apply_expressions_fn_wrapper, with_new_children: with_new_children_fn_wrapper, name: name_fn_wrapper, execute: execute_fn_wrapper, @@ -314,6 +335,7 @@ impl FFI_ExecutionPlan { partition_statistics: partition_statistics_fn_wrapper, clone: clone_fn_wrapper, release: release_fn_wrapper, + version: crate::version, private_data: Box::into_raw(private_data) as *mut c_void, library_marker_id: crate::get_library_marker_id, } @@ -442,6 +464,24 @@ impl ExecutionPlan for ForeignExecutionPlan { } } + fn apply_expressions( + &self, + f: &mut dyn FnMut( + &Arc, + ) -> Result, + ) -> Result { + let expressions = + df_result!(unsafe { (self.plan.apply_expressions)(&self.plan) })?; + datafusion_physical_plan::apply_expression_roots( + expressions.iter().map(|expression| { + let expression: Arc = + expression.into(); + expression + }), + f, + ) + } + fn repartitioned( &self, target_partitions: usize, @@ -474,8 +514,8 @@ impl ExecutionPlan for ForeignExecutionPlan { #[cfg(any(test, feature = "integration-tests"))] pub mod tests { - use datafusion_physical_plan::Partitioning; use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType}; + use datafusion_physical_plan::{Partitioning, PhysicalExpr}; use super::*; @@ -483,6 +523,7 @@ pub mod tests { pub struct EmptyExec { props: Arc, children: Vec>, + expressions: Vec>, metrics: Option, statistics: Option, } @@ -497,6 +538,7 @@ pub mod tests { Boundedness::Bounded, )), children: Vec::default(), + expressions: Vec::default(), metrics: None, statistics: None, } @@ -511,6 +553,14 @@ pub mod tests { self.statistics = Some(statistics); self } + + pub fn with_expressions( + mut self, + expressions: Vec>, + ) -> Self { + self.expressions = expressions; + self + } } impl DisplayAs for EmptyExec { @@ -543,6 +593,7 @@ pub mod tests { Ok(Arc::new(EmptyExec { props: Arc::clone(&self.props), children, + expressions: self.expressions.clone(), metrics: self.metrics.clone(), statistics: self.statistics.clone(), })) @@ -569,6 +620,13 @@ pub mod tests { Statistics::new_unknown(self.props.eq_properties.schema()) }))) } + + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + datafusion_physical_plan::apply_expression_roots(&self.expressions, f) + } } #[test] @@ -600,6 +658,38 @@ pub mod tests { Ok(()) } + #[test] + fn test_ffi_execution_plan_apply_expressions() -> Result<()> { + use datafusion_physical_expr::expressions::{DynamicFilterPhysicalExpr, lit}; + + let schema = Arc::new(arrow::datatypes::Schema::empty()); + let dynamic_filter = Arc::new(DynamicFilterPhysicalExpr::new(vec![], lit(true))); + let expected_id = dynamic_filter + .expression_id() + .expect("dynamic filters always have an expression ID"); + let expression: Arc = + Arc::::clone(&dynamic_filter); + let original_plan = + Arc::new(EmptyExec::new(schema).with_expressions(vec![expression])); + + let mut ffi_plan = FFI_ExecutionPlan::new(original_plan, None); + ffi_plan.library_marker_id = crate::mock_foreign_marker_id; + let foreign_plan: Arc = (&ffi_plan).try_into()?; + + let mut retained = None; + foreign_plan.apply_expressions(&mut |expr| { + retained = Some(Arc::clone(expr)); + Ok(TreeNodeRecursion::Continue) + })?; + drop(foreign_plan); + + assert_eq!( + retained.and_then(|expr| expr.expression_id()), + Some(expected_id) + ); + Ok(()) + } + #[test] fn test_ffi_execution_plan_children() -> Result<()> { let schema = Arc::new(arrow::datatypes::Schema::new(vec![ diff --git a/datafusion/ffi/src/physical_expr/mod.rs b/datafusion/ffi/src/physical_expr/mod.rs index 9a3ee273936c3..3544dcc6605f8 100644 --- a/datafusion/ffi/src/physical_expr/mod.rs +++ b/datafusion/ffi/src/physical_expr/mod.rs @@ -120,6 +120,8 @@ pub struct FFI_PhysicalExpr { pub is_volatile_node: unsafe extern "C" fn(&Self) -> bool, + pub expression_id: unsafe extern "C" fn(&Self) -> FFI_Option, + // Display trait pub display: unsafe extern "C" fn(&Self) -> SString, @@ -434,6 +436,7 @@ unsafe extern "C" fn clone_fn_wrapper(expr: &FFI_PhysicalExpr) -> FFI_PhysicalEx snapshot: snapshot_fn_wrapper, snapshot_generation: snapshot_generation_fn_wrapper, is_volatile_node: is_volatile_node_fn_wrapper, + expression_id: expression_id_fn_wrapper, display: display_fn_wrapper, hash: hash_fn_wrapper, clone: clone_fn_wrapper, @@ -445,6 +448,12 @@ unsafe extern "C" fn clone_fn_wrapper(expr: &FFI_PhysicalExpr) -> FFI_PhysicalEx } } +unsafe extern "C" fn expression_id_fn_wrapper( + expr: &FFI_PhysicalExpr, +) -> FFI_Option { + expr.inner().expression_id().into() +} + impl Drop for FFI_PhysicalExpr { fn drop(&mut self) { unsafe { (self.release)(self) } @@ -477,6 +486,7 @@ impl From> for FFI_PhysicalExpr { snapshot: snapshot_fn_wrapper, snapshot_generation: snapshot_generation_fn_wrapper, is_volatile_node: is_volatile_node_fn_wrapper, + expression_id: expression_id_fn_wrapper, display: display_fn_wrapper, hash: hash_fn_wrapper, clone: clone_fn_wrapper, @@ -713,6 +723,10 @@ impl PhysicalExpr for ForeignPhysicalExpr { fn is_volatile_node(&self) -> bool { unsafe { (self.expr.is_volatile_node)(&self.expr) } } + + fn expression_id(&self) -> Option { + unsafe { (self.expr.expression_id)(&self.expr) }.into() + } } impl Eq for ForeignPhysicalExpr {} @@ -747,7 +761,9 @@ mod tests { use datafusion_expr::interval_arithmetic::Interval; #[expect(deprecated)] use datafusion_expr::statistics::Distribution; - use datafusion_physical_expr::expressions::{Column, NegativeExpr, NotExpr}; + use datafusion_physical_expr::expressions::{ + Column, DynamicFilterPhysicalExpr, NegativeExpr, NotExpr, lit, + }; use datafusion_physical_expr_common::physical_expr::{PhysicalExpr, fmt_sql}; use crate::physical_expr::FFI_PhysicalExpr; @@ -762,6 +778,21 @@ mod tests { (original, foreign_expr) } + #[test] + fn ffi_physical_expr_expression_id() { + let dynamic_filter = Arc::new(DynamicFilterPhysicalExpr::new(vec![], lit(true))); + let expected_id = dynamic_filter + .expression_id() + .expect("dynamic filters always have an expression ID"); + let expression: Arc = + Arc::::clone(&dynamic_filter); + let mut ffi_expr = FFI_PhysicalExpr::from(expression); + ffi_expr.library_marker_id = crate::mock_foreign_marker_id; + + let foreign_expr: Arc = (&ffi_expr).into(); + assert_eq!(foreign_expr.expression_id(), Some(expected_id)); + } + fn test_record_batch() -> RecordBatch { record_batch!(("a", Int32, [1, 2, 3])).unwrap() } diff --git a/datafusion/ffi/src/tests/async_provider.rs b/datafusion/ffi/src/tests/async_provider.rs index 9821c3e501f67..f14ee2f3efd60 100644 --- a/datafusion/ffi/src/tests/async_provider.rs +++ b/datafusion/ffi/src/tests/async_provider.rs @@ -32,6 +32,7 @@ use arrow::array::RecordBatch; use arrow::datatypes::Schema; use async_trait::async_trait; use datafusion_catalog::{MemoryCatalogProvider, TableProvider}; +use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::{Result, exec_err}; use datafusion_execution::RecordBatchStream; use datafusion_expr::Expr; @@ -227,6 +228,22 @@ impl ExecutionPlan for AsyncTestExecutionPlan { batch_receiver: self.batch_receiver.resubscribe(), })) } + + fn apply_expressions( + &self, + f: &mut dyn FnMut( + &Arc, + ) -> Result, + ) -> Result { + datafusion_physical_plan::apply_expression_roots( + self.properties + .output_ordering() + .into_iter() + .flatten() + .map(|sort_expr| &sort_expr.expr), + f, + ) + } } impl datafusion_physical_plan::DisplayAs for AsyncTestExecutionPlan { diff --git a/datafusion/ffi/src/tests/mod.rs b/datafusion/ffi/src/tests/mod.rs index d372dcf9177e6..7b2511753d15f 100644 --- a/datafusion/ffi/src/tests/mod.rs +++ b/datafusion/ffi/src/tests/mod.rs @@ -28,6 +28,8 @@ use datafusion_common::stats::Precision; use datafusion_common::{ColumnStatistics, Statistics}; use datafusion_common::{Result, ScalarValue}; use datafusion_expr::{Expr, TableType}; +use datafusion_physical_expr::PhysicalExpr; +use datafusion_physical_expr::expressions::{DynamicFilterPhysicalExpr, lit}; use datafusion_physical_plan::ExecutionPlan; use sync_provider::create_sync_table_provider; use udf_udaf_udwf::{ @@ -107,6 +109,8 @@ pub struct ForeignLibraryModule { pub create_empty_exec: extern "C" fn() -> FFI_ExecutionPlan, + pub create_exec_with_expressions: extern "C" fn() -> FFI_ExecutionPlan, + pub create_exec_with_statistics: extern "C" fn() -> FFI_ExecutionPlan, pub create_table_with_statistics: @@ -161,6 +165,14 @@ pub(crate) extern "C" fn create_empty_exec() -> FFI_ExecutionPlan { FFI_ExecutionPlan::new(plan, None) } +pub(crate) extern "C" fn create_exec_with_expressions() -> FFI_ExecutionPlan { + let schema = Arc::new(Schema::empty()); + let expression: Arc = + Arc::new(DynamicFilterPhysicalExpr::new(vec![], lit(true))); + let plan = Arc::new(EmptyExec::new(schema).with_expressions(vec![expression])); + FFI_ExecutionPlan::new(plan, None) +} + /// Returns canonical statistics used by both the producer and consumer sides of /// the integration tests so round-trips can be asserted without hard-coding /// the values in two places. @@ -259,6 +271,7 @@ pub extern "C" fn datafusion_ffi_get_module() -> ForeignLibraryModule { create_rank_udwf: create_ffi_rank_func, create_extension_options: config::create_extension_options, create_empty_exec, + create_exec_with_expressions, create_exec_with_statistics, create_table_with_statistics, create_physical_optimizer_rule: diff --git a/datafusion/ffi/tests/ffi_execution_plan.rs b/datafusion/ffi/tests/ffi_execution_plan.rs index 7d04e828bd4a5..96d3af8ab6b51 100644 --- a/datafusion/ffi/tests/ffi_execution_plan.rs +++ b/datafusion/ffi/tests/ffi_execution_plan.rs @@ -21,6 +21,7 @@ mod tests { use arrow::datatypes::Schema; use arrow_schema::DataType; use datafusion_common::DataFusionError; + use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_ffi::execution_plan::FFI_ExecutionPlan; use datafusion_ffi::execution_plan::ForeignExecutionPlan; use datafusion_ffi::execution_plan::{ExecutionPlanPrivateData, tests::EmptyExec}; @@ -63,6 +64,30 @@ mod tests { Ok(()) } + #[test] + fn test_ffi_execution_plan_expressions_cross_library() -> Result<(), DataFusionError> + { + let module = get_module()?; + let plan = (module.create_exec_with_expressions)(); + let plan: Arc = (&plan).try_into()?; + assert!(plan.is::()); + + let mut retained = None; + plan.apply_expressions(&mut |expr| { + retained = Some(Arc::clone(expr)); + Ok(TreeNodeRecursion::Continue) + })?; + drop(plan); + + assert!( + retained + .as_ref() + .and_then(|expr| expr.expression_id()) + .is_some() + ); + Ok(()) + } + #[test] fn test_ffi_execution_plan_new_sets_runtimes_on_children() -> Result<(), DataFusionError> { diff --git a/datafusion/physical-optimizer/src/ensure_coop.rs b/datafusion/physical-optimizer/src/ensure_coop.rs index e7aacb2321b67..10da9e4e75174 100644 --- a/datafusion/physical-optimizer/src/ensure_coop.rs +++ b/datafusion/physical-optimizer/src/ensure_coop.rs @@ -264,10 +264,11 @@ mod tests { // Test that cooperative context is reset when encountering an eager evaluation boundary. use arrow::datatypes::Schema; use datafusion_common::internal_err; + use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_execution::TaskContext; use datafusion_physical_expr::EquivalenceProperties; use datafusion_physical_plan::{ - DisplayAs, DisplayFormatType, Partitioning, PlanProperties, + DisplayAs, DisplayFormatType, Partitioning, PhysicalExpr, PlanProperties, SendableRecordBatchStream, execution_plan::{Boundedness, EmissionType}, }; @@ -345,6 +346,13 @@ mod tests { ) -> Result { internal_err!("DummyExec does not support execution") } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } } // Build a plan similar to the original test: diff --git a/datafusion/physical-optimizer/src/output_requirements.rs b/datafusion/physical-optimizer/src/output_requirements.rs index b9d0d06da1dda..59a92f971974e 100644 --- a/datafusion/physical-optimizer/src/output_requirements.rs +++ b/datafusion/physical-optimizer/src/output_requirements.rs @@ -27,7 +27,9 @@ use std::sync::Arc; use crate::PhysicalOptimizerRule; use datafusion_common::config::ConfigOptions; -use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode}; +use datafusion_common::tree_node::{ + Transformed, TransformedResult, TreeNode, TreeNodeRecursion, +}; use datafusion_common::{Result, Statistics, internal_err}; use datafusion_execution::TaskContext; use datafusion_physical_expr::Distribution; @@ -324,6 +326,33 @@ impl ExecutionPlan for OutputRequirementExec { fn fetch(&self) -> Option { self.fetch } + + fn apply_expressions( + &self, + f: &mut dyn FnMut( + &Arc, + ) -> Result, + ) -> Result { + #[expect(deprecated)] + let distribution = if let Distribution::HashPartitioned(exprs) + | Distribution::KeyPartitioned(exprs) = + &self.dist_requirement + { + exprs.as_slice() + } else { + &[] + }; + let ordering = self + .order_requirement + .iter() + .flat_map(|requirements| match requirements { + OrderingRequirements::Hard(alternatives) + | OrderingRequirements::Soft(alternatives) => alternatives, + }) + .flatten() + .map(|sort_expr| &sort_expr.expr); + datafusion_physical_plan::apply_expression_roots(ordering.chain(distribution), f) + } } impl PhysicalOptimizerRule for OutputRequirements { diff --git a/datafusion/physical-plan/benches/compute_statistics.rs b/datafusion/physical-plan/benches/compute_statistics.rs index 56a518c95292e..93c95ea4ba099 100644 --- a/datafusion/physical-plan/benches/compute_statistics.rs +++ b/datafusion/physical-plan/benches/compute_statistics.rs @@ -33,6 +33,7 @@ use std::sync::Arc; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; use datafusion_common::ScalarValue; +use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::{Result, Statistics}; use datafusion_execution::TaskContext; use datafusion_physical_expr::EquivalenceProperties; @@ -97,6 +98,13 @@ impl ExecutionPlan for BenchLeaf { vec![] } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + fn with_new_children( self: Arc, _children: Vec>, diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index db1eb951d6fbc..27dbd62884e02 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -176,6 +176,7 @@ use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use arrow::record_batch::RecordBatch; use arrow_schema::FieldRef; use datafusion_common::stats::Precision; +use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion}; use datafusion_common::{ Constraint, Constraints, Result, ScalarValue, assert_eq_or_internal_err, internal_err, not_impl_err, @@ -763,6 +764,32 @@ struct AggrDynFilter { supported_accumulators_info: Vec, } +fn plan_contains_expression_id( + plan: &Arc, + expression_id: u64, +) -> Result { + let mut found = false; + plan.apply(|node| { + node.apply_expressions(&mut |root| { + root.apply(|expr| { + if expr.expression_id() == Some(expression_id) { + found = true; + Ok(TreeNodeRecursion::Stop) + } else { + Ok(TreeNodeRecursion::Continue) + } + }) + })?; + + Ok(if found { + TreeNodeRecursion::Stop + } else { + TreeNodeRecursion::Continue + }) + })?; + Ok(found) +} + // ---- Aggregate Dynamic Filter Utility Structs ---- /// Aggregate expressions that support the dynamic filter pushdown in aggregation. @@ -1960,6 +1987,33 @@ impl ExecutionPlan for AggregateExec { vec![&self.input] } + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + let group_by = self.group_by.input_exprs(); + let aggregates = self.aggr_expr.iter().flat_map(|aggr| { + let expressions = aggr.all_expressions(); + expressions + .args + .into_iter() + .chain(expressions.order_by_exprs) + }); + let filters = self.filter_expr.iter().flatten().cloned(); + let dynamic_filter = self.dynamic_filter.iter().map(|dynamic_filter| { + Arc::::clone(&dynamic_filter.filter) + as Arc + }); + crate::apply_expression_roots( + group_by + .into_iter() + .chain(aggregates) + .chain(filters) + .chain(dynamic_filter), + f, + ) + } + fn with_new_children( self: Arc, children: Vec>, @@ -2097,27 +2151,12 @@ impl ExecutionPlan for AggregateExec { if phase == FilterPushdownPhase::Post && let Some(dyn_filter) = &self.dynamic_filter { - // let child_accepts_dyn_filter = child_pushdown_result - // .self_filters - // .first() - // .map(|filters| { - // assert_eq_or_internal_err!( - // filters.len(), - // 1, - // "Aggregate only pushdown one self dynamic filter" - // ); - // let filter = filters.get(0).unwrap(); // Asserted above - // Ok(matches!(filter.discriminant, PushedDown::Yes)) - // }) - // .unwrap_or_else(|| internal_err!("The length of self filters equals to the number of child of this ExecutionPlan, so it must be 1"))?; - - // HACK: The above snippet should be used, however, now the child reply - // `PushDown::No` can indicate they're not able to push down row-level - // filter, but still keep the filter for statistics pruning. - // So here, we try to use ref count to determine if the dynamic filter - // has actually be pushed down. - // Issue: - let child_accepts_dyn_filter = Arc::strong_count(dyn_filter) > 1; + let child_accepts_dyn_filter = dyn_filter + .filter + .expression_id() + .map(|id| plan_contains_expression_id(&self.input, id)) + .transpose()? + .unwrap_or(false); if !child_accepts_dyn_filter { // Child can't consume the self dynamic filter, so disable it by setting @@ -2506,6 +2545,8 @@ impl AggregateExec { })?; aggregate.with_dynamic_filter_expr(dynamic_filter)? } else { + let mut aggregate = aggregate; + aggregate.dynamic_filter = None; aggregate }; @@ -3009,6 +3050,7 @@ mod tests { use crate::empty::EmptyExec; use crate::execution_plan::Boundedness; use crate::expressions::col; + use crate::filter::FilterExecBuilder; use crate::metrics::MetricValue; use crate::statistics::{StatisticsArgs, StatisticsContext}; use crate::test::TestMemoryExec; @@ -3044,7 +3086,7 @@ mod tests { use datafusion_physical_expr::Partitioning; use datafusion_physical_expr::PhysicalSortExpr; use datafusion_physical_expr::aggregate::AggregateExprBuilder; - use datafusion_physical_expr::expressions::Literal; + use datafusion_physical_expr::expressions::{Literal, NotExpr}; use crate::projection::ProjectionExec; use datafusion_physical_expr::projection::ProjectionExpr; @@ -3564,6 +3606,13 @@ mod tests { vec![] } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + fn with_new_children( self: Arc, _: Vec>, @@ -7579,6 +7628,38 @@ mod tests { Ok(()) } + #[test] + fn test_plan_contains_expression_id_recurses_plans_and_expressions() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + let empty: Arc = Arc::new(EmptyExec::new(Arc::clone(&schema))); + let dynamic_filter = Arc::new(DynamicFilterPhysicalExpr::new( + vec![col("a", &schema)?], + lit(true), + )); + let expression_id = dynamic_filter + .expression_id() + .expect("dynamic filters always have an expression ID"); + + assert!(!plan_contains_expression_id(&empty, expression_id)?); + + let dynamic_filter_expr: Arc = + Arc::::clone(&dynamic_filter); + let predicate: Arc = + Arc::new(NotExpr::new(dynamic_filter_expr)); + let filter: Arc = + Arc::new(FilterExecBuilder::new(predicate, empty).build()?); + let projection: Arc = Arc::new(ProjectionExec::try_new( + [ProjectionExpr::new_from_expression( + col("a", &schema)?, + &schema, + )?], + filter, + )?); + + assert!(plan_contains_expression_id(&projection, expression_id)?); + Ok(()) + } + /// Test that [`AggregateExec::with_dynamic_filter_expr`] errors when the aggregate does not support dynamic filtering #[test] fn test_with_dynamic_filter_error_unsupported() -> Result<()> { diff --git a/datafusion/physical-plan/src/analyze.rs b/datafusion/physical-plan/src/analyze.rs index 31e0a27410ff9..0ab28f56e169c 100644 --- a/datafusion/physical-plan/src/analyze.rs +++ b/datafusion/physical-plan/src/analyze.rs @@ -32,11 +32,13 @@ use crate::{DisplayFormatType, ExecutionPlan, Partitioning}; use arrow::{array::StringBuilder, datatypes::SchemaRef, record_batch::RecordBatch}; use datafusion_common::format::ExplainFormat; use datafusion_common::instant::Instant; +use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::{ DataFusionError, Result, assert_eq_or_internal_err, internal_err, }; use datafusion_execution::TaskContext; use datafusion_physical_expr::EquivalenceProperties; +use datafusion_physical_expr::PhysicalExpr; use futures::StreamExt; @@ -219,6 +221,13 @@ impl ExecutionPlan for AnalyzeExec { ]) } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + fn with_new_children( self: Arc, mut children: Vec>, diff --git a/datafusion/physical-plan/src/async_func.rs b/datafusion/physical-plan/src/async_func.rs index e13a5b986aa2c..1696e2a677ccf 100644 --- a/datafusion/physical-plan/src/async_func.rs +++ b/datafusion/physical-plan/src/async_func.rs @@ -153,6 +153,13 @@ impl ExecutionPlan for AsyncFuncExec { vec![&self.input] } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + fn with_new_children( self: Arc, mut children: Vec>, diff --git a/datafusion/physical-plan/src/buffer.rs b/datafusion/physical-plan/src/buffer.rs index 3be331a1ee1ba..a1c3c7ea01658 100644 --- a/datafusion/physical-plan/src/buffer.rs +++ b/datafusion/physical-plan/src/buffer.rs @@ -32,6 +32,7 @@ use crate::{ }; use arrow::array::RecordBatch; use datafusion_common::config::ConfigOptions; +use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::{Result, Statistics, internal_err, plan_err}; use datafusion_common_runtime::SpawnedTask; use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; @@ -159,6 +160,13 @@ impl ExecutionPlan for BufferExec { vec![&self.input] } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + fn with_new_children( self: Arc, mut children: Vec>, diff --git a/datafusion/physical-plan/src/coalesce_batches.rs b/datafusion/physical-plan/src/coalesce_batches.rs index c5b91767777f2..511e5d793b873 100644 --- a/datafusion/physical-plan/src/coalesce_batches.rs +++ b/datafusion/physical-plan/src/coalesce_batches.rs @@ -34,6 +34,7 @@ use crate::{ use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; use datafusion_common::Result; +use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_execution::TaskContext; use datafusion_physical_expr::PhysicalExpr; @@ -173,6 +174,13 @@ impl ExecutionPlan for CoalesceBatchesExec { vec![false] } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + fn with_new_children( self: Arc, mut children: Vec>, diff --git a/datafusion/physical-plan/src/coalesce_partitions.rs b/datafusion/physical-plan/src/coalesce_partitions.rs index f9694e0d16817..5cd7a707c23b3 100644 --- a/datafusion/physical-plan/src/coalesce_partitions.rs +++ b/datafusion/physical-plan/src/coalesce_partitions.rs @@ -35,6 +35,7 @@ use crate::{DisplayFormatType, ExecutionPlan, Partitioning, check_if_same_proper use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; use datafusion_common::config::ConfigOptions; +use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::{Result, assert_eq_or_internal_err, internal_err}; use datafusion_execution::TaskContext; use datafusion_physical_expr::PhysicalExpr; @@ -143,6 +144,13 @@ impl ExecutionPlan for CoalescePartitionsExec { vec![false] } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + fn with_new_children( self: Arc, mut children: Vec>, diff --git a/datafusion/physical-plan/src/coop.rs b/datafusion/physical-plan/src/coop.rs index a5b57f546bbfa..dc7d98891e114 100644 --- a/datafusion/physical-plan/src/coop.rs +++ b/datafusion/physical-plan/src/coop.rs @@ -71,6 +71,7 @@ //! that report [`SchedulingType::NonCooperative`] in their [plan properties](ExecutionPlan::properties). use datafusion_common::config::ConfigOptions; +use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_physical_expr::PhysicalExpr; #[cfg(datafusion_coop = "tokio_fallback")] use futures::Future; @@ -267,6 +268,13 @@ impl ExecutionPlan for CooperativeExec { vec![&self.input] } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + fn with_new_children( self: Arc, mut children: Vec>, diff --git a/datafusion/physical-plan/src/display.rs b/datafusion/physical-plan/src/display.rs index 34493a5f51742..89029aea87546 100644 --- a/datafusion/physical-plan/src/display.rs +++ b/datafusion/physical-plan/src/display.rs @@ -1504,8 +1504,11 @@ mod tests { use std::fmt::Write; use std::sync::Arc; - use datafusion_common::{Result, Statistics, internal_datafusion_err}; + use datafusion_common::{ + Result, Statistics, internal_datafusion_err, tree_node::TreeNodeRecursion, + }; use datafusion_execution::{SendableRecordBatchStream, TaskContext}; + use datafusion_physical_expr::PhysicalExpr; use crate::statistics::StatisticsArgs; use crate::{DisplayAs, ExecutionPlan, PlanProperties}; @@ -1549,6 +1552,13 @@ mod tests { unimplemented!() } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + fn execute( &self, _: usize, @@ -1706,6 +1716,17 @@ mod tests { fn children(&self) -> Vec<&Arc> { vec![&self.inner] } + fn apply_expressions( + &self, + f: &mut dyn FnMut( + &Arc, + ) -> Result< + datafusion_common::tree_node::TreeNodeRecursion, + >, + ) -> Result + { + self.inner.apply_expressions(f) + } fn with_new_children( self: Arc, _: Vec>, diff --git a/datafusion/physical-plan/src/empty.rs b/datafusion/physical-plan/src/empty.rs index 3bd38bf238dc1..bd91ec742d48c 100644 --- a/datafusion/physical-plan/src/empty.rs +++ b/datafusion/physical-plan/src/empty.rs @@ -29,9 +29,10 @@ use crate::{ use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; use datafusion_common::stats::Precision; +use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::{ColumnStatistics, Result, ScalarValue, assert_or_internal_err}; use datafusion_execution::TaskContext; -use datafusion_physical_expr::EquivalenceProperties; +use datafusion_physical_expr::{EquivalenceProperties, PhysicalExpr}; use crate::execution_plan::SchedulingType; use crate::statistics::StatisticsArgs; @@ -119,6 +120,13 @@ impl ExecutionPlan for EmptyExec { vec![] } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + fn with_new_children( self: Arc, _: Vec>, diff --git a/datafusion/physical-plan/src/execution_plan.rs b/datafusion/physical-plan/src/execution_plan.rs index 11a8d69a37669..0b85370f9d88c 100644 --- a/datafusion/physical-plan/src/execution_plan.rs +++ b/datafusion/physical-plan/src/execution_plan.rs @@ -28,7 +28,9 @@ pub use crate::stream::EmptyRecordBatchStream; use arrow_schema::Schema; pub use datafusion_common::hash_utils; -use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode}; +use datafusion_common::tree_node::{ + Transformed, TransformedResult, TreeNode, TreeNodeRecursion, +}; pub use datafusion_common::utils::project_schema; pub use datafusion_common::{ColumnStatistics, Statistics, internal_err}; pub use datafusion_execution::{RecordBatchStream, SendableRecordBatchStream}; @@ -39,6 +41,7 @@ pub use datafusion_physical_expr::{ }; use std::any::Any; +use std::borrow::Borrow; use std::fmt::Debug; use std::sync::{Arc, LazyLock}; @@ -67,6 +70,28 @@ use datafusion_physical_expr_common::sort_expr::{ use futures::stream::{StreamExt, TryStreamExt}; +/// Applies `f` to a shallow sequence of physical expression roots. +/// +/// [`TreeNodeRecursion::Stop`] stops iteration and is returned immediately. +/// [`TreeNodeRecursion::Jump`] is normalized to [`TreeNodeRecursion::Continue`] +/// because this function does not visit expression children. +pub fn apply_expression_roots( + roots: I, + f: &mut dyn FnMut(&Arc) -> Result, +) -> Result +where + I: IntoIterator, + I::Item: Borrow>, +{ + for root in roots { + match f(root.borrow())? { + TreeNodeRecursion::Stop => return Ok(TreeNodeRecursion::Stop), + TreeNodeRecursion::Continue | TreeNodeRecursion::Jump => {} + } + } + Ok(TreeNodeRecursion::Continue) +} + /// Represent nodes in the DataFusion Physical Plan. /// /// Calling [`execute`] produces an `async` [`SendableRecordBatchStream`] of @@ -247,6 +272,67 @@ pub trait ExecutionPlan: Any + Debug + DisplayAs + Send + Sync { /// joins). fn children(&self) -> Vec<&Arc>; + /// Apply a closure `f` to each expression (non-recursively) in the current + /// physical plan node. This does not include expressions in any children. + /// + /// The closure `f` is applied to expressions in the order they appear in the plan. + /// The closure can return `TreeNodeRecursion::Continue` to continue visiting, + /// `TreeNodeRecursion::Stop` to stop visiting immediately, or `TreeNodeRecursion::Jump` + /// to continue with the next expression. `Jump` does not skip anything here because + /// this method does not recurse into expression children. + /// + /// The expressions visited do not necessarily represent or even contribute + /// to the output schema of this node. For example, `FilterExec` visits the + /// filter predicate even though the output of a Filter has the same columns + /// as the input. + /// + /// # Example Usage + /// ``` + /// # use std::sync::Arc; + /// # use datafusion_physical_plan::ExecutionPlan; + /// # use datafusion_common::tree_node::TreeNodeRecursion; + /// # fn example(plan: Arc) -> datafusion_common::Result<()> { + /// // Count the number of expressions + /// let mut count = 0; + /// plan.apply_expressions(&mut |_expr| { + /// count += 1; + /// Ok(TreeNodeRecursion::Continue) + /// })?; + /// # Ok(()) + /// # } + /// ``` + /// + /// # Implementation Examples + /// + /// ## Node with no expressions (e.g., EmptyExec, MemoryExec) + /// ```ignore + /// fn apply_expressions( + /// &self, + /// _f: &mut dyn FnMut(&Arc) -> Result, + /// ) -> Result { + /// Ok(TreeNodeRecursion::Continue) + /// } + /// ``` + /// + /// ## Node with expressions (e.g., FilterExec, ProjectionExec) + /// + /// Use [`apply_expression_roots`] to iterate over expressions. It propagates + /// [`TreeNodeRecursion::Stop`] and normalizes [`TreeNodeRecursion::Jump`] to + /// [`TreeNodeRecursion::Continue`] because this method does not recurse into + /// expression children. + /// ```ignore + /// fn apply_expressions( + /// &self, + /// f: &mut dyn FnMut(&Arc) -> Result, + /// ) -> Result { + /// apply_expression_roots([&self.predicate], f) + /// } + /// ``` + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result; + /// Returns a new `ExecutionPlan` where all existing children were replaced /// by the `children`, in order fn with_new_children( @@ -1772,6 +1858,13 @@ mod tests { unimplemented!() } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + fn execute( &self, _partition: usize, @@ -1828,6 +1921,13 @@ mod tests { vec![] } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + fn with_new_children( self: Arc, _: Vec>, @@ -1878,6 +1978,13 @@ mod tests { vec![] } + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + self.0.apply_expressions(f) + } + fn with_new_children( self: Arc, _: Vec>, @@ -1937,6 +2044,12 @@ mod tests { fn children(&self) -> Vec<&Arc> { vec![] } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } fn with_new_children( self: Arc, _: Vec>, @@ -2000,6 +2113,12 @@ mod tests { fn children(&self) -> Vec<&Arc> { vec![&self.input] } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } fn with_new_children( self: Arc, mut children: Vec>, @@ -2091,6 +2210,12 @@ mod tests { fn children(&self) -> Vec<&Arc> { vec![&self.input] } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } fn with_new_children( self: Arc, mut children: Vec>, @@ -2232,6 +2357,171 @@ mod tests { Ok(()) } + /// A test node that holds a fixed list of expressions, used to test + /// `apply_expressions` behavior. + #[derive(Debug)] + struct MultiExprExec { + exprs: Vec>, + children: Vec>, + } + + impl DisplayAs for MultiExprExec { + fn fmt_as( + &self, + _t: DisplayFormatType, + _f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + unimplemented!() + } + } + + impl ExecutionPlan for MultiExprExec { + fn name(&self) -> &'static str { + "MultiExprExec" + } + + fn properties(&self) -> &Arc { + unimplemented!() + } + + fn children(&self) -> Vec<&Arc> { + self.children.iter().collect() + } + + fn with_new_children( + self: Arc, + _: Vec>, + ) -> Result> { + unimplemented!() + } + + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + apply_expression_roots(&self.exprs, f) + } + + fn execute( + &self, + _partition: usize, + _context: Arc, + ) -> Result { + unimplemented!() + } + + fn partition_statistics( + &self, + _partition: Option, + ) -> Result> { + unimplemented!() + } + } + + /// Returns a simple literal `Arc` for use in tests. + fn lit_expr(val: i64) -> Arc { + use datafusion_physical_expr::expressions::Literal; + Arc::new(Literal::new(datafusion_common::ScalarValue::Int64(Some( + val, + )))) + } + + /// `apply_expressions` visits all expressions when `f` always returns `Continue`. + #[test] + fn test_apply_expressions_continue_visits_all() -> Result<()> { + let plan = MultiExprExec { + exprs: vec![lit_expr(1), lit_expr(2), lit_expr(3)], + children: vec![], + }; + let mut visited = 0usize; + plan.apply_expressions(&mut |_expr| { + visited += 1; + Ok(TreeNodeRecursion::Continue) + })?; + assert_eq!(visited, 3); + Ok(()) + } + + #[test] + fn test_apply_expressions_stop_halts_early() -> Result<()> { + let plan = MultiExprExec { + exprs: vec![lit_expr(1), lit_expr(2), lit_expr(3)], + children: vec![], + }; + let mut visited = 0usize; + let tnr = plan.apply_expressions(&mut |_expr| { + visited += 1; + Ok(TreeNodeRecursion::Stop) + })?; + // Only the first expression is visited; the rest are skipped. + assert_eq!(visited, 1); + assert_eq!(tnr, TreeNodeRecursion::Stop); + Ok(()) + } + + #[test] + fn test_apply_expressions_jump_visits_next_root() -> Result<()> { + let plan = MultiExprExec { + exprs: vec![lit_expr(1), lit_expr(2), lit_expr(3)], + children: vec![], + }; + let mut visited = 0usize; + let tnr = plan.apply_expressions(&mut |_expr| { + visited += 1; + Ok(TreeNodeRecursion::Jump) + })?; + assert_eq!(visited, 3); + assert_eq!(tnr, TreeNodeRecursion::Continue); + Ok(()) + } + + #[test] + fn test_apply_expressions_does_not_recurse() -> Result<()> { + use datafusion_physical_expr::expressions::NegativeExpr; + + let child: Arc = Arc::new(MultiExprExec { + exprs: vec![lit_expr(2)], + children: vec![], + }); + let nested: Arc = Arc::new(NegativeExpr::new(lit_expr(1))); + let plan = MultiExprExec { + exprs: vec![nested], + children: vec![child], + }; + + let mut visited = 0; + plan.apply_expressions(&mut |expr| { + visited += 1; + assert!(expr.is::()); + Ok(TreeNodeRecursion::Continue) + })?; + assert_eq!(visited, 1); + Ok(()) + } + + #[test] + fn test_apply_expressions_callback_can_retain_arc() -> Result<()> { + let expected = lit_expr(1); + let plan = MultiExprExec { + exprs: vec![Arc::clone(&expected)], + children: vec![], + }; + let mut retained = None; + plan.apply_expressions(&mut |expr| { + retained = Some(Arc::clone(expr)); + Ok(TreeNodeRecursion::Continue) + })?; + drop(plan); + + assert!(Arc::ptr_eq( + &expected, + retained + .as_ref() + .expect("callback should retain expression") + )); + Ok(()) + } + #[test] fn test_execution_plan_name() { let schema1 = Arc::new(Schema::empty()); diff --git a/datafusion/physical-plan/src/explain.rs b/datafusion/physical-plan/src/explain.rs index a270a003eba17..72ceae81f3724 100644 --- a/datafusion/physical-plan/src/explain.rs +++ b/datafusion/physical-plan/src/explain.rs @@ -26,9 +26,10 @@ use crate::{DisplayFormatType, ExecutionPlan, Partitioning}; use arrow::{array::StringBuilder, datatypes::SchemaRef, record_batch::RecordBatch}; use datafusion_common::display::StringifiedPlan; +use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::{Result, assert_eq_or_internal_err}; use datafusion_execution::TaskContext; -use datafusion_physical_expr::EquivalenceProperties; +use datafusion_physical_expr::{EquivalenceProperties, PhysicalExpr}; use log::trace; @@ -116,6 +117,13 @@ impl ExecutionPlan for ExplainExec { vec![] } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + fn with_new_children( self: Arc, _: Vec>, diff --git a/datafusion/physical-plan/src/filter.rs b/datafusion/physical-plan/src/filter.rs index 511be0bbdd9e2..c6800283c5c23 100644 --- a/datafusion/physical-plan/src/filter.rs +++ b/datafusion/physical-plan/src/filter.rs @@ -55,6 +55,7 @@ use arrow::record_batch::RecordBatch; use datafusion_common::cast::as_boolean_array; use datafusion_common::config::ConfigOptions; use datafusion_common::stats::Precision; +use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::{ DataFusionError, Result, ScalarValue, internal_err, plan_err, project_schema, }; @@ -535,6 +536,13 @@ impl ExecutionPlan for FilterExec { vec![&self.input] } + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + crate::apply_expression_roots([&self.predicate], f) + } + fn maintains_input_order(&self) -> Vec { // Tell optimizer this operator doesn't reorder its input vec![true] diff --git a/datafusion/physical-plan/src/joins/cross_join.rs b/datafusion/physical-plan/src/joins/cross_join.rs index 1a631aac980ab..d56ef7d0881dc 100644 --- a/datafusion/physical-plan/src/joins/cross_join.rs +++ b/datafusion/physical-plan/src/joins/cross_join.rs @@ -43,11 +43,13 @@ use arrow::array::{RecordBatch, RecordBatchOptions}; use arrow::compute::concat_batches; use arrow::datatypes::{Fields, Schema, SchemaRef}; use datafusion_common::stats::Precision; +use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::{ JoinType, Result, ScalarValue, assert_eq_or_internal_err, internal_err, }; use datafusion_execution::TaskContext; use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; +use datafusion_physical_expr::PhysicalExpr; use datafusion_physical_expr::equivalence::join_equivalence_properties; use async_trait::async_trait; @@ -290,6 +292,14 @@ impl ExecutionPlan for CrossJoinExec { Some(self.metrics.clone_inner()) } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + // CrossJoin has no join conditions or expressions + Ok(TreeNodeRecursion::Continue) + } + fn with_new_children( self: Arc, children: Vec>, diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index 9d9c867c2724b..190311979f133 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -73,6 +73,7 @@ use arrow::record_batch::RecordBatch; use arrow::util::bit_util; use arrow_schema::{DataType, Schema}; use datafusion_common::config::ConfigOptions; +use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::utils::memory::{RecordBatchMemoryCounter, estimate_memory_size}; use datafusion_common::{ JoinSide, JoinType, NullEquality, Result, assert_or_internal_err, internal_err, @@ -1330,6 +1331,25 @@ impl ExecutionPlan for HashJoinExec { vec![&self.left, &self.right] } + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + let join_keys = self + .on + .iter() + .flat_map(|(left, right)| [Arc::clone(left), Arc::clone(right)]); + let filter = self + .filter + .iter() + .map(|filter| Arc::clone(filter.expression())); + let dynamic_filter = self.dynamic_filter.iter().map(|dynamic_filter| { + Arc::::clone(&dynamic_filter.filter) + as Arc + }); + crate::apply_expression_roots(join_keys.chain(filter).chain(dynamic_filter), f) + } + /// Creates a new HashJoinExec with different children while preserving configuration. /// /// This method is called during query optimization when the optimizer creates new @@ -2461,6 +2481,13 @@ mod tests { vec![] } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + fn with_new_children( self: Arc, _: Vec>, diff --git a/datafusion/physical-plan/src/joins/nested_loop_join.rs b/datafusion/physical-plan/src/joins/nested_loop_join.rs index 515dcc2931c05..7069a8b44805c 100644 --- a/datafusion/physical-plan/src/joins/nested_loop_join.rs +++ b/datafusion/physical-plan/src/joins/nested_loop_join.rs @@ -61,6 +61,7 @@ use arrow::datatypes::{Schema, SchemaRef}; use arrow::record_batch::RecordBatch; use arrow_schema::DataType; use datafusion_common::cast::as_boolean_array; +use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::{ JoinSide, NullEquality, Result, ScalarValue, Statistics, arrow_err, assert_eq_or_internal_err, internal_datafusion_err, internal_err, project_schema, @@ -562,6 +563,17 @@ impl ExecutionPlan for NestedLoopJoinExec { vec![&self.left, &self.right] } + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + // Apply to join filter expressions if present + crate::apply_expression_roots( + self.filter.iter().map(|filter| filter.expression()), + f, + ) + } + fn with_new_children( self: Arc, children: Vec>, diff --git a/datafusion/physical-plan/src/joins/piecewise_merge_join/exec.rs b/datafusion/physical-plan/src/joins/piecewise_merge_join/exec.rs index 5ec564295ece1..b60ec1c784de5 100644 --- a/datafusion/physical-plan/src/joins/piecewise_merge_join/exec.rs +++ b/datafusion/physical-plan/src/joins/piecewise_merge_join/exec.rs @@ -23,6 +23,7 @@ use arrow::{ }; use arrow_schema::{SchemaRef, SortOptions}; use datafusion_common::not_impl_err; +use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::{JoinSide, Result, internal_err}; use datafusion_execution::{ SendableRecordBatchStream, @@ -482,6 +483,14 @@ impl ExecutionPlan for PiecewiseMergeJoinExec { vec![&self.buffered, &self.streamed] } + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + // Apply to the two expressions being compared in the range predicate + crate::apply_expression_roots([&self.on.0, &self.on.1], f) + } + fn required_input_distribution(&self) -> Vec { self.input_distribution_requirements().into_per_child() } diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs index 3b597323b2e7b..4d82eed547a46 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs @@ -47,6 +47,7 @@ use crate::{ use arrow::compute::SortOptions; use arrow::datatypes::SchemaRef; +use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::{ JoinSide, JoinType, NullEquality, Result, assert_eq_or_internal_err, internal_err, plan_err, @@ -441,6 +442,15 @@ impl ExecutionPlan for SortMergeJoinExec { vec![&self.left, &self.right] } + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + let join_keys = self.on.iter().flat_map(|(left, right)| [left, right]); + let filter = self.filter.iter().map(|filter| filter.expression()); + crate::apply_expression_roots(join_keys.chain(filter), f) + } + fn with_new_children( self: Arc, children: Vec>, diff --git a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs index 95f4c35871431..4b161a7b47473 100644 --- a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs +++ b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs @@ -66,6 +66,7 @@ use arrow::compute::concat_batches; use arrow::datatypes::{ArrowNativeType, Schema, SchemaRef}; use arrow::record_batch::RecordBatch; use datafusion_common::hash_utils::create_hashes; +use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::utils::bisect; use datafusion_common::{ HashSet, JoinSide, JoinType, NullEquality, Result, assert_eq_or_internal_err, @@ -452,6 +453,15 @@ impl ExecutionPlan for SymmetricHashJoinExec { vec![&self.left, &self.right] } + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + let join_keys = self.on.iter().flat_map(|(left, right)| [left, right]); + let filter = self.filter.iter().map(|filter| filter.expression()); + crate::apply_expression_roots(join_keys.chain(filter), f) + } + fn with_new_children( self: Arc, children: Vec>, diff --git a/datafusion/physical-plan/src/lib.rs b/datafusion/physical-plan/src/lib.rs index 8cba650b79770..6e1df1f840af0 100644 --- a/datafusion/physical-plan/src/lib.rs +++ b/datafusion/physical-plan/src/lib.rs @@ -45,9 +45,9 @@ pub use crate::distribution_requirements::{ ChildSatisfactionOptions, InputDistributionRequirements, }; pub use crate::execution_plan::{ - ExecutionPlan, ExecutionPlanProperties, PlanProperties, collect, collect_partitioned, - displayable, execute_input_stream, execute_stream, execute_stream_partitioned, - get_plan_string, with_new_children_if_necessary, + ExecutionPlan, ExecutionPlanProperties, PlanProperties, apply_expression_roots, + collect, collect_partitioned, displayable, execute_input_stream, execute_stream, + execute_stream_partitioned, get_plan_string, with_new_children_if_necessary, }; pub use crate::metrics::Metric; pub use crate::ordering::InputOrderMode; diff --git a/datafusion/physical-plan/src/limit.rs b/datafusion/physical-plan/src/limit.rs index ddce680fc18ad..0e588ba1aacc3 100644 --- a/datafusion/physical-plan/src/limit.rs +++ b/datafusion/physical-plan/src/limit.rs @@ -35,10 +35,11 @@ use crate::{ use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; +use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::{Result, assert_eq_or_internal_err, internal_err}; use datafusion_execution::TaskContext; -use datafusion_physical_expr::LexOrdering; +use datafusion_physical_expr::{LexOrdering, PhysicalExpr}; use futures::stream::{Stream, StreamExt}; use log::trace; @@ -167,6 +168,19 @@ impl ExecutionPlan for GlobalLimitExec { vec![false] } + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + crate::apply_expression_roots( + self.required_ordering + .iter() + .flatten() + .map(|sort_expr| &sort_expr.expr), + f, + ) + } + fn with_new_children( self: Arc, mut children: Vec>, @@ -398,6 +412,19 @@ impl ExecutionPlan for LocalLimitExec { vec![true] } + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + crate::apply_expression_roots( + self.required_ordering + .iter() + .flatten() + .map(|sort_expr| &sort_expr.expr), + f, + ) + } + fn with_new_children( self: Arc, children: Vec>, @@ -648,7 +675,6 @@ mod tests { use arrow::array::RecordBatchOptions; use arrow::datatypes::Schema; use datafusion_common::stats::Precision; - use datafusion_physical_expr::PhysicalExpr; use datafusion_physical_expr::expressions::col; #[tokio::test] diff --git a/datafusion/physical-plan/src/memory.rs b/datafusion/physical-plan/src/memory.rs index ad54905f474aa..eb141b8c70d5e 100644 --- a/datafusion/physical-plan/src/memory.rs +++ b/datafusion/physical-plan/src/memory.rs @@ -32,10 +32,11 @@ use crate::{ use arrow::array::RecordBatch; use arrow::datatypes::SchemaRef; +use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::{Result, assert_eq_or_internal_err, assert_or_internal_err}; use datafusion_execution::TaskContext; use datafusion_execution::memory_pool::MemoryReservation; -use datafusion_physical_expr::EquivalenceProperties; +use datafusion_physical_expr::{EquivalenceProperties, PhysicalExpr}; use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; use futures::Stream; @@ -311,6 +312,13 @@ impl ExecutionPlan for LazyMemoryExec { vec![] } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + fn with_new_children( self: Arc, children: Vec>, diff --git a/datafusion/physical-plan/src/operator_statistics/mod.rs b/datafusion/physical-plan/src/operator_statistics/mod.rs index 142768fcf49d2..ec54201e7b3d5 100644 --- a/datafusion/physical-plan/src/operator_statistics/mod.rs +++ b/datafusion/physical-plan/src/operator_statistics/mod.rs @@ -1038,6 +1038,7 @@ mod tests { use std::fmt; use crate::execution_plan::{Boundedness, EmissionType}; + use datafusion_common::tree_node::TreeNodeRecursion; fn make_schema() -> Arc { Arc::new(Schema::new(vec![ @@ -1117,6 +1118,13 @@ mod tests { &self.cache } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + fn execute( &self, _partition: usize, @@ -1219,6 +1227,13 @@ mod tests { self.input.properties() } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + fn execute( &self, _partition: usize, diff --git a/datafusion/physical-plan/src/placeholder_row.rs b/datafusion/physical-plan/src/placeholder_row.rs index 5d71058269f49..de07529bab70c 100644 --- a/datafusion/physical-plan/src/placeholder_row.rs +++ b/datafusion/physical-plan/src/placeholder_row.rs @@ -29,9 +29,11 @@ use crate::{ use arrow::array::{ArrayRef, NullArray, RecordBatch, RecordBatchOptions}; use arrow::datatypes::{DataType, Field, Fields, Schema, SchemaRef}; +use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::{Result, assert_or_internal_err}; use datafusion_execution::TaskContext; use datafusion_physical_expr::EquivalenceProperties; +use datafusion_physical_expr::PhysicalExpr; use crate::statistics::StatisticsArgs; use log::trace; @@ -136,6 +138,13 @@ impl ExecutionPlan for PlaceholderRowExec { vec![] } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + fn with_new_children( self: Arc, _: Vec>, diff --git a/datafusion/physical-plan/src/projection.rs b/datafusion/physical-plan/src/projection.rs index 42501f22395b4..cf1db9328ada2 100644 --- a/datafusion/physical-plan/src/projection.rs +++ b/datafusion/physical-plan/src/projection.rs @@ -302,6 +302,20 @@ impl ExecutionPlan for ProjectionExec { vec![&self.input] } + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + crate::apply_expression_roots( + self.projector + .projection() + .as_ref() + .iter() + .map(|proj_expr| &proj_expr.expr), + f, + ) + } + fn with_new_children( self: Arc, mut children: Vec>, diff --git a/datafusion/physical-plan/src/recursive_query.rs b/datafusion/physical-plan/src/recursive_query.rs index 00df227cb87db..4c6f0493adf40 100644 --- a/datafusion/physical-plan/src/recursive_query.rs +++ b/datafusion/physical-plan/src/recursive_query.rs @@ -37,12 +37,14 @@ use arrow::array::{BooleanArray, BooleanBuilder}; use arrow::compute::filter_record_batch; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; +use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode}; use datafusion_common::{ Result, exec_datafusion_err, internal_datafusion_err, not_impl_err, }; use datafusion_execution::TaskContext; use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; +use datafusion_physical_expr::PhysicalExpr; use datafusion_physical_expr::{EquivalenceProperties, Partitioning}; use futures::{Stream, StreamExt, ready}; @@ -153,6 +155,13 @@ impl ExecutionPlan for RecursiveQueryExec { vec![&self.static_term, &self.recursive_term] } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + // TODO: control these hints and see whether we can // infer some from the child plans (static/recursive terms). fn maintains_input_order(&self) -> Vec { diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs index 3473aad9b3fc0..dce5805e16bd4 100644 --- a/datafusion/physical-plan/src/repartition/mod.rs +++ b/datafusion/physical-plan/src/repartition/mod.rs @@ -53,6 +53,7 @@ use arrow::datatypes::{SchemaRef, UInt32Type}; use arrow_schema::SortOptions; use datafusion_common::config::ConfigOptions; use datafusion_common::stats::Precision; +use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::utils::{compare_rows, extract_row_at_idx_to_buf, transpose}; use datafusion_common::{ ColumnStatistics, DataFusionError, HashMap, ScalarValue, SplitPoint, @@ -1337,6 +1338,20 @@ impl ExecutionPlan for RepartitionExec { vec![&self.input] } + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + match self.partitioning() { + Partitioning::Hash(exprs, _) => crate::apply_expression_roots(exprs, f), + Partitioning::Range(range) => crate::apply_expression_roots( + range.ordering().iter().map(|sort_expr| &sort_expr.expr), + f, + ), + _ => Ok(TreeNodeRecursion::Continue), + } + } + fn with_new_children( self: Arc, mut children: Vec>, @@ -3045,6 +3060,13 @@ mod tests { vec![] } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + fn with_new_children( self: Arc, _: Vec>, @@ -4279,6 +4301,13 @@ mod test { )?); let exec = RepartitionExec::try_new(source, partitioning)?; + let mut expressions = vec![]; + exec.apply_expressions(&mut |expr| { + expressions.push(expr.to_string()); + Ok(TreeNodeRecursion::Continue) + })?; + assert_eq!(expressions, ["c0@0"]); + // Range partition count is fixed by split points, so repartitioned() // cannot change it to an arbitrary target. let result = exec.repartitioned(10, &Default::default())?; diff --git a/datafusion/physical-plan/src/scalar_subquery.rs b/datafusion/physical-plan/src/scalar_subquery.rs index 73acb2ab13480..ee3c2e5d077f5 100644 --- a/datafusion/physical-plan/src/scalar_subquery.rs +++ b/datafusion/physical-plan/src/scalar_subquery.rs @@ -27,9 +27,11 @@ use std::fmt; use std::sync::Arc; +use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::{Result, ScalarValue, Statistics, exec_err, internal_err}; use datafusion_execution::TaskContext; use datafusion_expr::physical_planning_context::{ScalarSubqueryResults, SubqueryIndex}; +use datafusion_physical_expr::PhysicalExpr; use crate::execution_plan::{CardinalityEffect, ExecutionPlan, PlanProperties}; use crate::joins::utils::{OnceAsync, OnceFut}; @@ -224,6 +226,13 @@ impl ExecutionPlan for ScalarSubqueryExec { ))) } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + fn maintains_input_order(&self) -> Vec { // Only the main input (first child); subquery children don't contribute // to ordering. @@ -453,6 +462,13 @@ mod tests { ))) } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + fn execute( &self, partition: usize, diff --git a/datafusion/physical-plan/src/sorts/partial_sort.rs b/datafusion/physical-plan/src/sorts/partial_sort.rs index 3a4ddd2fbeaf7..5a7c2b2e4557e 100644 --- a/datafusion/physical-plan/src/sorts/partial_sort.rs +++ b/datafusion/physical-plan/src/sorts/partial_sort.rs @@ -70,9 +70,10 @@ use arrow::compute::concat_batches; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; use datafusion_common::Result; +use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::utils::evaluate_partition_ranges; use datafusion_execution::{RecordBatchStream, TaskContext}; -use datafusion_physical_expr::LexOrdering; +use datafusion_physical_expr::{LexOrdering, PhysicalExpr}; use futures::{Stream, StreamExt, ready}; use log::trace; @@ -307,6 +308,16 @@ impl ExecutionPlan for PartialSortExec { vec![&self.input] } + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + crate::apply_expression_roots( + self.expr.iter().map(|sort_expr| &sort_expr.expr), + f, + ) + } + fn with_new_children( self: Arc, children: Vec>, diff --git a/datafusion/physical-plan/src/sorts/partitioned_topk.rs b/datafusion/physical-plan/src/sorts/partitioned_topk.rs index 730440a429c68..78dd9b9696d45 100644 --- a/datafusion/physical-plan/src/sorts/partitioned_topk.rs +++ b/datafusion/physical-plan/src/sorts/partitioned_topk.rs @@ -37,6 +37,7 @@ use std::sync::Arc; use arrow::datatypes::SchemaRef; use arrow::row::SortField; use datafusion_common::Result; +use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_execution::TaskContext; use datafusion_execution::runtime_env::RuntimeEnv; use datafusion_physical_expr::PhysicalExpr; @@ -379,6 +380,16 @@ impl ExecutionPlan for PartitionedTopKExec { )?)) } + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + crate::apply_expression_roots( + self.expr.iter().map(|sort_expr| &sort_expr.expr), + f, + ) + } + fn execute( &self, partition: usize, diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs index 4b30aede7d02a..b11db4633e979 100644 --- a/datafusion/physical-plan/src/sorts/sort.rs +++ b/datafusion/physical-plan/src/sorts/sort.rs @@ -60,6 +60,7 @@ use arrow::array::{RecordBatch, RecordBatchOptions}; use arrow::compute::{concat_batches, lexsort_to_indices, take_arrays}; use arrow::datatypes::SchemaRef; use datafusion_common::config::SpillCompression; +use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::{ DataFusionError, Result, assert_or_internal_err, internal_datafusion_err, unwrap_or_internal_err, @@ -1274,6 +1275,23 @@ impl ExecutionPlan for SortExec { vec![&self.input] } + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + let dynamic_filter = self + .filter + .as_ref() + .map(|filter| filter.read().expr() as Arc); + crate::apply_expression_roots( + self.expr + .iter() + .map(|sort_expr| &sort_expr.expr) + .chain(dynamic_filter.iter()), + f, + ) + } + fn benefits_from_input_partitioning(&self) -> Vec { vec![false] } @@ -1760,6 +1778,13 @@ mod tests { Ok(self) } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + fn execute( &self, _partition: usize, diff --git a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs index 2add3e1eb82f0..ac6f5d18cd2ff 100644 --- a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs +++ b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs @@ -31,9 +31,11 @@ use crate::{ check_if_same_properties, }; +use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::{Result, assert_eq_or_internal_err, internal_err}; use datafusion_execution::TaskContext; use datafusion_execution::memory_pool::MemoryConsumer; +use datafusion_physical_expr::PhysicalExpr; use datafusion_physical_expr_common::sort_expr::{LexOrdering, OrderingRequirements}; use crate::execution_plan::{CardinalityEffect, EvaluationType, SchedulingType}; @@ -281,6 +283,16 @@ impl ExecutionPlan for SortPreservingMergeExec { vec![&self.input] } + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + crate::apply_expression_roots( + self.expr.iter().map(|sort_expr| &sort_expr.expr), + f, + ) + } + fn with_new_children( self: Arc, mut children: Vec>, @@ -1573,6 +1585,12 @@ mod tests { fn children(&self) -> Vec<&Arc> { vec![] } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } fn with_new_children( self: Arc, _: Vec>, diff --git a/datafusion/physical-plan/src/streaming.rs b/datafusion/physical-plan/src/streaming.rs index 61a9b9cc6d0de..a82f8d9441e95 100644 --- a/datafusion/physical-plan/src/streaming.rs +++ b/datafusion/physical-plan/src/streaming.rs @@ -33,8 +33,10 @@ use crate::stream::RecordBatchStreamAdapter; use crate::{ExecutionPlan, Partitioning, SendableRecordBatchStream}; use arrow::datatypes::{Schema, SchemaRef}; +use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::{Result, internal_err, plan_err}; use datafusion_execution::TaskContext; +use datafusion_physical_expr::PhysicalExpr; use datafusion_physical_expr::projection::ProjectionMapping; use datafusion_physical_expr::{EquivalenceProperties, LexOrdering}; @@ -271,6 +273,13 @@ impl ExecutionPlan for StreamingTableExec { vec![] } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + fn with_new_children( self: Arc, children: Vec>, diff --git a/datafusion/physical-plan/src/test.rs b/datafusion/physical-plan/src/test.rs index e8c775a786578..fd66cc9e05b2f 100644 --- a/datafusion/physical-plan/src/test.rs +++ b/datafusion/physical-plan/src/test.rs @@ -36,6 +36,7 @@ use crate::{DisplayAs, DisplayFormatType, PlanProperties}; use arrow::array::{Array, ArrayRef, Int32Array, RecordBatch}; use arrow_schema::{DataType, Field, Schema, SchemaRef}; +use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::{ Result, Statistics, assert_or_internal_err, config::ConfigOptions, project_schema, }; @@ -45,7 +46,9 @@ use datafusion_physical_expr::equivalence::{ }; use datafusion_physical_expr::expressions::Column; use datafusion_physical_expr::utils::collect_columns; -use datafusion_physical_expr::{EquivalenceProperties, LexOrdering, Partitioning}; +use datafusion_physical_expr::{ + EquivalenceProperties, LexOrdering, Partitioning, PhysicalExpr, +}; use futures::{Future, FutureExt}; @@ -138,6 +141,19 @@ impl ExecutionPlan for TestMemoryExec { Vec::new() } + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + crate::apply_expression_roots( + self.sort_information + .iter() + .flatten() + .map(|sort_expr| &sort_expr.expr), + f, + ) + } + fn with_new_children( self: Arc, _: Vec>, diff --git a/datafusion/physical-plan/src/test/exec.rs b/datafusion/physical-plan/src/test/exec.rs index b92008c6b219b..22e6859e5ee06 100644 --- a/datafusion/physical-plan/src/test/exec.rs +++ b/datafusion/physical-plan/src/test/exec.rs @@ -35,9 +35,10 @@ use std::{ use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use arrow::record_batch::RecordBatch; +use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::{DataFusionError, Result, internal_err}; use datafusion_execution::TaskContext; -use datafusion_physical_expr::EquivalenceProperties; +use datafusion_physical_expr::{EquivalenceProperties, PhysicalExpr}; use futures::Stream; use tokio::sync::Barrier; @@ -195,6 +196,13 @@ impl ExecutionPlan for MockExec { vec![] } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + fn with_new_children( self: Arc, _: Vec>, @@ -432,6 +440,13 @@ impl ExecutionPlan for BarrierExec { unimplemented!() } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + /// Returns a stream which yields data fn execute( &self, @@ -568,6 +583,13 @@ impl ExecutionPlan for ErrorExec { unimplemented!() } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + /// Returns a stream which yields data fn execute( &self, @@ -647,6 +669,13 @@ impl ExecutionPlan for StatisticsExec { vec![] } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + fn with_new_children( self: Arc, _: Vec>, @@ -758,6 +787,13 @@ impl ExecutionPlan for BlockingExec { internal_err!("Children cannot be replaced in {self:?}") } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + fn execute( &self, _partition: usize, @@ -893,6 +929,13 @@ impl ExecutionPlan for PanicExec { vec![] } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + fn with_new_children( self: Arc, _: Vec>, diff --git a/datafusion/physical-plan/src/union.rs b/datafusion/physical-plan/src/union.rs index 4722329ea55a4..9be852e6cc649 100644 --- a/datafusion/physical-plan/src/union.rs +++ b/datafusion/physical-plan/src/union.rs @@ -50,6 +50,7 @@ use arrow::datatypes::{Field, Schema, SchemaRef}; use arrow::record_batch::RecordBatch; use datafusion_common::config::ConfigOptions; use datafusion_common::stats::NdvFallback; +use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::{ Result, assert_or_internal_err, exec_err, internal_datafusion_err, }; @@ -252,6 +253,13 @@ impl ExecutionPlan for UnionExec { self.inputs.iter().collect() } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + fn with_new_children( self: Arc, children: Vec>, @@ -624,6 +632,13 @@ impl ExecutionPlan for InterleaveExec { vec![false; self.inputs().len()] } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + fn with_new_children( self: Arc, children: Vec>, diff --git a/datafusion/physical-plan/src/unnest.rs b/datafusion/physical-plan/src/unnest.rs index fbe849229941a..f3d75b7e7d263 100644 --- a/datafusion/physical-plan/src/unnest.rs +++ b/datafusion/physical-plan/src/unnest.rs @@ -44,6 +44,7 @@ use arrow::datatypes::{DataType, Int64Type, Schema, SchemaRef}; use arrow::record_batch::RecordBatch; use arrow_ord::cmp::lt; use async_trait::async_trait; +use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::{ Constraints, HashMap, HashSet, Result, UnnestOptions, exec_datafusion_err, exec_err, internal_err, @@ -227,6 +228,13 @@ impl ExecutionPlan for UnnestExec { vec![&self.input] } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + fn with_new_children( self: Arc, mut children: Vec>, diff --git a/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs b/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs index d5863080895f6..938f6138cd760 100644 --- a/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs +++ b/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs @@ -50,6 +50,7 @@ use arrow::{ }; use datafusion_common::hash_utils::create_hashes; use datafusion_common::stats::Precision; +use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::utils::{ evaluate_partition_ranges, get_at_indices, get_row_at_idx, }; @@ -311,6 +312,21 @@ impl ExecutionPlan for BoundedWindowAggExec { vec![&self.input] } + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + let expressions = self.window_expr.iter().flat_map(|window_expr| { + let expressions = window_expr.all_expressions(); + expressions + .args + .into_iter() + .chain(expressions.partition_by_exprs) + .chain(expressions.order_by_exprs) + }); + crate::apply_expression_roots(expressions, f) + } + fn required_input_ordering(&self) -> Vec> { let partition_bys = self.window_expr()[0].partition_by(); let order_keys = self.window_expr()[0].order_by(); diff --git a/datafusion/physical-plan/src/windows/window_agg_exec.rs b/datafusion/physical-plan/src/windows/window_agg_exec.rs index 81838300cf5c7..31ff5189253cf 100644 --- a/datafusion/physical-plan/src/windows/window_agg_exec.rs +++ b/datafusion/physical-plan/src/windows/window_agg_exec.rs @@ -45,6 +45,7 @@ use arrow::datatypes::SchemaRef; use arrow::error::ArrowError; use arrow::record_batch::RecordBatch; use datafusion_common::stats::Precision; +use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::utils::{evaluate_partition_ranges, transpose}; use datafusion_common::{Result, assert_eq_or_internal_err}; use datafusion_execution::TaskContext; @@ -214,6 +215,21 @@ impl ExecutionPlan for WindowAggExec { vec![&self.input] } + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + let expressions = self.window_expr.iter().flat_map(|window_expr| { + let expressions = window_expr.all_expressions(); + expressions + .args + .into_iter() + .chain(expressions.partition_by_exprs) + .chain(expressions.order_by_exprs) + }); + crate::apply_expression_roots(expressions, f) + } + fn maintains_input_order(&self) -> Vec { vec![true] } diff --git a/datafusion/physical-plan/src/work_table.rs b/datafusion/physical-plan/src/work_table.rs index c92face1e5404..83cd0a15a6d26 100644 --- a/datafusion/physical-plan/src/work_table.rs +++ b/datafusion/physical-plan/src/work_table.rs @@ -32,10 +32,11 @@ use crate::{ use crate::statistics::StatisticsArgs; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; +use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::{Result, assert_eq_or_internal_err, internal_datafusion_err}; use datafusion_execution::TaskContext; use datafusion_execution::memory_pool::MemoryReservation; -use datafusion_physical_expr::{EquivalenceProperties, Partitioning}; +use datafusion_physical_expr::{EquivalenceProperties, Partitioning, PhysicalExpr}; /// A vector of record batches with a memory reservation. #[derive(Debug)] @@ -186,6 +187,13 @@ impl ExecutionPlan for WorkTableExec { vec![] } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + fn with_new_children( self: Arc, _: Vec>, diff --git a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs index 864e6d68676ee..c218a8c1810bb 100644 --- a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs @@ -106,6 +106,7 @@ use datafusion_common::file_options::json_writer::JsonWriterOptions; use datafusion_common::format::ExplainFormat; use datafusion_common::parsers::CompressionTypeVariant; use datafusion_common::stats::Precision; +use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::{ DataFusionError, JoinSide, NullEquality, Result, UnnestOptions, exec_datafusion_err, internal_datafusion_err, internal_err, not_impl_err, @@ -321,6 +322,13 @@ impl ExecutionPlan for DowncastDelegatingExec { self.inner.children() } + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + self.inner.apply_expressions(f) + } + fn with_new_children( self: Arc, children: Vec>, @@ -4477,6 +4485,61 @@ fn test_aggregate_with_dynamic_filter_roundtrip() -> Result<()> { Ok(()) } +#[test] +fn test_aggregate_without_dynamic_filter_roundtrip() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + let col_a: Arc = Arc::new(Column::new("a", 0)); + let child = Arc::new(EmptyExec::new(Arc::clone(&schema))); + let aggregate = Arc::new(AggregateExec::try_new( + AggregateMode::Partial, + PhysicalGroupBy::new_single(vec![]), + vec![ + AggregateExprBuilder::new( + datafusion::functions_aggregate::min_max::min_udaf(), + vec![col_a], + ) + .schema(Arc::clone(&schema)) + .alias("min_a") + .build() + .map(Arc::new)?, + ], + vec![None], + child, + Arc::clone(&schema), + )?) as Arc; + + let mut config = ConfigOptions::default(); + config.optimizer.enable_aggregate_dynamic_filter_pushdown = true; + let optimizer = FilterPushdown::new_post_optimization(); + let plan = optimizer.optimize(aggregate, &config)?; + assert!( + plan.downcast_ref::() + .expect("Should be AggregateExec") + .dynamic_filter_expr() + .is_none() + ); + + let ctx = SessionContext::new(); + let codec = DefaultPhysicalExtensionCodec {}; + let converter = DefaultPhysicalProtoConverter {}; + let bytes = physical_plan_to_bytes_with_proto_converter(plan, &codec, &converter)?; + let deserialized = physical_plan_from_bytes_with_proto_converter( + bytes.as_ref(), + ctx.task_ctx().as_ref(), + &codec, + &converter, + )?; + + assert!( + deserialized + .downcast_ref::() + .expect("Should be AggregateExec") + .dynamic_filter_expr() + .is_none() + ); + Ok(()) +} + /// Test that plan containing a SortExec with dynamic filter pushdown /// can be serialized and deserialized while preserving references to the dynamic filter. #[test] @@ -4595,6 +4658,13 @@ impl ExecutionPlan for CustomExecWithExprs { vec![&self.child] } + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + datafusion_physical_plan::apply_expression_roots(&self.exprs, f) + } + fn with_new_children( self: Arc, _children: Vec>, diff --git a/docs/source/library-user-guide/custom-table-providers.md b/docs/source/library-user-guide/custom-table-providers.md index 540782e3e8bf7..c6a316aa74b94 100644 --- a/docs/source/library-user-guide/custom-table-providers.md +++ b/docs/source/library-user-guide/custom-table-providers.md @@ -766,6 +766,7 @@ impl DatePartitionedTable { # fn children(&self) -> Vec<&Arc> { vec![] } # fn with_new_children(self: Arc, _: Vec>) -> Result> { Ok(self) } # fn execute(&self, _: usize, _: Arc) -> Result { todo!() } +# fn apply_expressions(&self, _f: &mut dyn FnMut(&Arc) -> Result) -> Result { Ok(TreeNodeRecursion::Continue) } # } ``` @@ -909,6 +910,13 @@ impl ExecutionPlan for CountingExec { batch_stream, ))) } + +# fn apply_expressions( +# &self, +# _f: &mut dyn FnMut(&Arc) -> Result, +# ) -> Result { +# Ok(TreeNodeRecursion::Continue) +# } } ``` diff --git a/docs/source/library-user-guide/upgrading/54.0.0.md b/docs/source/library-user-guide/upgrading/54.0.0.md index f8e7ac93c08d8..b067f6e782fd3 100644 --- a/docs/source/library-user-guide/upgrading/54.0.0.md +++ b/docs/source/library-user-guide/upgrading/54.0.0.md @@ -165,6 +165,55 @@ where string types are preferred (`UNION`, `CASE THEN/ELSE`, `NVL2`). string-preferring behavior - Crates that call `get_coerce_type_for_case_expression` +### `ExecutionPlan::apply_expressions` is now a required method + +`apply_expressions` has been added as a **required** method on the `ExecutionPlan` trait (no default implementation). The same applies to the `FileSource` and `DataSource` traits. Any custom implementation of these traits must now implement `apply_expressions`. + +**Who is affected:** + +- Users who implement custom `ExecutionPlan` nodes +- Users who implement custom `FileSource` or `DataSource` sources + +**Migration guide:** + +Add `apply_expressions` to your implementation. Call `f` on each top-level +`PhysicalExpr` your node owns. Do not recurse into child plans or expression +children. Use `apply_expression_roots` for nodes with expressions so `Stop` +short-circuits the iteration and `Jump` proceeds to the next root expression. + +**Node with no expressions:** + +```rust,ignore +fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, +) -> Result { + Ok(TreeNodeRecursion::Continue) +} +``` + +**Node with a single expression:** + +```rust,ignore +fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> Result, +) -> Result { + apply_expression_roots([&self.predicate], f) +} +``` + +**Node with multiple expressions:** + +```rust,ignore +fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> Result, +) -> Result { + apply_expression_roots(&self.expressions, f) +} +``` + ### `ExecutionPlan::partition_statistics` now returns `Arc` `ExecutionPlan::partition_statistics` now returns `Result>` instead of `Result`. This avoids cloning `Statistics` when it is shared across multiple consumers.