From 58f0e1a65118a252ea5b0bfad1d43e6701efce9b Mon Sep 17 00:00:00 2001 From: Jayant Shrivastava Date: Thu, 30 Jul 2026 21:23:42 +0000 Subject: [PATCH 1/3] Reapply "Add `ExecutionPlan::apply_expressions()` (#20337)" (#22437) --- .../custom_data_source/custom_datasource.rs | 17 ++ .../memory_pool_execution_plan.rs | 17 ++ .../proto/composed_extension_codec.rs | 19 ++ .../examples/relation_planner/table_sample.rs | 18 +- datafusion/catalog/src/memory/table.rs | 14 +- datafusion/core/src/physical_planner.rs | 32 +++ .../core/tests/custom_sources_cases/mod.rs | 17 ++ .../provider_filter_pushdown.rs | 17 ++ .../tests/custom_sources_cases/statistics.rs | 17 ++ datafusion/core/tests/fuzz_cases/once_exec.rs | 17 ++ .../enforce_distribution.rs | 25 +- .../physical_optimizer/ensure_requirements.rs | 27 +++ .../physical_optimizer/filter_pushdown.rs | 105 +++++++++ .../physical_optimizer/join_selection.rs | 29 +++ .../physical_optimizer/pushdown_utils.rs | 29 +++ .../tests/physical_optimizer/test_utils.rs | 47 +++- .../tests/user_defined/insert_operation.rs | 17 ++ .../tests/user_defined/user_defined_plan.rs | 20 +- datafusion/datasource-arrow/src/source.rs | 15 ++ datafusion/datasource-avro/src/source.rs | 15 ++ datafusion/datasource-csv/src/source.rs | 15 ++ datafusion/datasource-json/src/source.rs | 15 ++ datafusion/datasource-parquet/src/source.rs | 21 ++ datafusion/datasource/src/file.rs | 22 ++ .../datasource/src/file_scan_config/mod.rs | 28 +++ datafusion/datasource/src/memory.rs | 15 ++ datafusion/datasource/src/sink.rs | 16 +- datafusion/datasource/src/source.rs | 25 ++ datafusion/datasource/src/test_util.rs | 9 +- datafusion/ffi/src/execution_plan.rs | 33 +++ datafusion/ffi/src/tests/async_provider.rs | 17 ++ .../physical-optimizer/src/ensure_coop.rs | 10 +- .../src/output_requirements.rs | 37 ++- .../benches/compute_statistics.rs | 8 + .../physical-plan/src/aggregates/mod.rs | 38 +++ datafusion/physical-plan/src/analyze.rs | 9 + datafusion/physical-plan/src/async_func.rs | 7 + datafusion/physical-plan/src/buffer.rs | 8 + .../physical-plan/src/coalesce_batches.rs | 8 + .../physical-plan/src/coalesce_partitions.rs | 8 + datafusion/physical-plan/src/coop.rs | 8 + datafusion/physical-plan/src/display.rs | 20 +- datafusion/physical-plan/src/empty.rs | 10 +- .../physical-plan/src/execution_plan.rs | 220 +++++++++++++++++- datafusion/physical-plan/src/explain.rs | 10 +- datafusion/physical-plan/src/filter.rs | 8 + .../physical-plan/src/joins/cross_join.rs | 10 + .../physical-plan/src/joins/hash_join/exec.rs | 32 +++ .../src/joins/nested_loop_join.rs | 12 + .../src/joins/piecewise_merge_join/exec.rs | 9 + .../src/joins/sort_merge_join/exec.rs | 18 ++ .../src/joins/symmetric_hash_join.rs | 18 ++ datafusion/physical-plan/src/limit.rs | 32 ++- datafusion/physical-plan/src/memory.rs | 10 +- .../src/operator_statistics/mod.rs | 15 ++ .../physical-plan/src/placeholder_row.rs | 9 + datafusion/physical-plan/src/projection.rs | 11 + .../physical-plan/src/recursive_query.rs | 9 + .../physical-plan/src/repartition/mod.rs | 28 +++ .../physical-plan/src/scalar_subquery.rs | 16 ++ .../physical-plan/src/sorts/partial_sort.rs | 14 +- .../src/sorts/partitioned_topk.rs | 12 + datafusion/physical-plan/src/sorts/sort.rs | 27 +++ .../src/sorts/sort_preserving_merge.rs | 19 ++ datafusion/physical-plan/src/streaming.rs | 9 + datafusion/physical-plan/src/test.rs | 19 +- datafusion/physical-plan/src/test/exec.rs | 45 +++- datafusion/physical-plan/src/union.rs | 15 ++ datafusion/physical-plan/src/unnest.rs | 8 + .../src/windows/bounded_window_agg_exec.rs | 14 ++ .../src/windows/window_agg_exec.rs | 14 ++ datafusion/physical-plan/src/work_table.rs | 10 +- .../tests/cases/roundtrip_physical_plan.rs | 19 ++ .../custom-table-providers.md | 8 + .../library-user-guide/upgrading/54.0.0.md | 67 ++++++ 75 files changed, 1646 insertions(+), 22 deletions(-) diff --git a/datafusion-examples/examples/custom_data_source/custom_datasource.rs b/datafusion-examples/examples/custom_data_source/custom_datasource.rs index a2d7d7699927f..e5a8838d797ac 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,20 @@ impl ExecutionPlan for CustomExec { None, )?)) } + + fn apply_expressions( + &self, + f: &mut dyn FnMut( + &dyn datafusion::physical_plan::PhysicalExpr, + ) -> Result, + ) -> Result { + // Visit expressions in the output ordering from equivalence properties + let mut tnr = TreeNodeRecursion::Continue; + if let Some(ordering) = self.cache.output_ordering() { + for sort_expr in ordering { + tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; + } + } + Ok(tnr) + } } 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..f761e9b6e6d50 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,20 @@ impl ExecutionPlan for BufferingExecutionPlan { }), ))) } + + fn apply_expressions( + &self, + f: &mut dyn FnMut( + &dyn datafusion::physical_plan::PhysicalExpr, + ) -> Result, + ) -> Result { + // Visit expressions in the output ordering from equivalence properties + let mut tnr = TreeNodeRecursion::Continue; + if let Some(ordering) = self.properties.output_ordering() { + for sort_expr in ordering { + tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; + } + } + Ok(tnr) + } } diff --git a/datafusion-examples/examples/proto/composed_extension_codec.rs b/datafusion-examples/examples/proto/composed_extension_codec.rs index 6077a982c320d..3405cc6a63aca 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( + &dyn datafusion::physical_plan::PhysicalExpr, + ) -> 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( + &dyn datafusion::physical_plan::PhysicalExpr, + ) -> 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..ba5b474b2e240 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( + &dyn datafusion::physical_plan::PhysicalExpr, + ) -> Result, + ) -> Result { + // Visit expressions in the output ordering from equivalence properties + let mut tnr = TreeNodeRecursion::Continue; + if let Some(ordering) = self.cache.output_ordering() { + for sort_expr in ordering { + tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; + } + } + Ok(tnr) + } } /// 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..bc71e8c010353 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(&dyn PhysicalExpr) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } } diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index 4e914556b4cc0..14c8d2045cf5e 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -4711,6 +4711,20 @@ mod tests { ) -> Result { unimplemented!("NoOpExecutionPlan::execute"); } + + fn apply_expressions( + &self, + f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + ) -> Result { + // Visit expressions in the output ordering from equivalence properties + let mut tnr = TreeNodeRecursion::Continue; + if let Some(ordering) = self.cache.output_ordering() { + for sort_expr in ordering { + tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; + } + } + Ok(tnr) + } } struct ExpressionExtensionPlanner; @@ -4878,6 +4892,12 @@ digraph { ) -> Result { unimplemented!() } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } } impl DisplayAs for OkExtensionNode { fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { @@ -4924,6 +4944,12 @@ digraph { ) -> Result { unimplemented!() } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } } impl DisplayAs for InvariantFailsExtensionNode { fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { @@ -5048,6 +5074,12 @@ digraph { ) -> Result { unimplemented!() } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&dyn PhysicalExpr) -> 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..3b8548c49e4f0 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( + &dyn datafusion::physical_plan::PhysicalExpr, + ) -> Result, + ) -> Result { + // Visit expressions in the output ordering from equivalence properties + let mut tnr = TreeNodeRecursion::Continue; + if let Some(ordering) = self.cache.output_ordering() { + for sort_expr in ordering { + tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; + } + } + Ok(tnr) + } } #[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..e52c559ec79ef 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( + &dyn datafusion::physical_plan::PhysicalExpr, + ) -> Result, + ) -> Result { + // Visit expressions in the output ordering from equivalence properties + let mut tnr = TreeNodeRecursion::Continue; + if let Some(ordering) = self.cache.output_ordering() { + for sort_expr in ordering { + tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; + } + } + Ok(tnr) + } } #[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..da5ff444f15a8 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( + &dyn datafusion::physical_plan::PhysicalExpr, + ) -> Result, + ) -> Result { + // Visit expressions in the output ordering from equivalence properties + let mut tnr = TreeNodeRecursion::Continue; + if let Some(ordering) = self.cache.output_ordering() { + for sort_expr in ordering { + tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; + } + } + Ok(tnr) + } } 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..403e377a690e2 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( + &dyn datafusion_physical_plan::PhysicalExpr, + ) -> datafusion_common::Result, + ) -> datafusion_common::Result { + // Visit expressions in the output ordering from equivalence properties + let mut tnr = TreeNodeRecursion::Continue; + if let Some(ordering) = self.cache.output_ordering() { + for sort_expr in ordering { + tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; + } + } + Ok(tnr) + } } diff --git a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs index 189650fe4afca..5e1182c09c1f7 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(&dyn PhysicalExpr) -> Result, + ) -> Result { + // Visit expressions in the output ordering from equivalence properties + let mut tnr = TreeNodeRecursion::Continue; + if let Some(ordering) = self.cache.output_ordering() { + for sort_expr in ordering { + tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; + } + } + Ok(tnr) + } + 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(&dyn PhysicalExpr) -> 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..5bf361c8ada0c 100644 --- a/datafusion/core/tests/physical_optimizer/ensure_requirements.rs +++ b/datafusion/core/tests/physical_optimizer/ensure_requirements.rs @@ -30,6 +30,7 @@ use std::sync::Arc; use arrow::compute::SortOptions; use arrow::datatypes::{DataType, Field, Schema}; use datafusion_common::Result; +use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_physical_expr::expressions::Column; use datafusion_physical_expr::{ EquivalenceProperties, LexOrdering, PhysicalExpr, PhysicalSortExpr, @@ -111,6 +112,12 @@ impl ExecutionPlan for MockMultiPartitionExec { fn children(&self) -> Vec<&Arc> { vec![] } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } fn with_new_children( self: Arc, _children: Vec>, @@ -974,6 +981,26 @@ impl ExecutionPlan for MockReqExec { fn children(&self) -> Vec<&Arc> { vec![&self.input] } + fn apply_expressions( + &self, + f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + ) -> Result { + let mut tnr = TreeNodeRecursion::Continue; + if let Some(ordering) = &self.ord { + for sort_expr in ordering { + tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; + } + } + #[expect(deprecated)] + if let Distribution::HashPartitioned(exprs) + | Distribution::KeyPartitioned(exprs) = &self.dist + { + for expr in exprs { + tnr = tnr.visit_sibling(|| f(expr.as_ref()))?; + } + } + Ok(tnr) + } 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..ad6ca6fba0998 100644 --- a/datafusion/core/tests/physical_optimizer/filter_pushdown.rs +++ b/datafusion/core/tests/physical_optimizer/filter_pushdown.rs @@ -3085,6 +3085,111 @@ 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() { + use datafusion_common::JoinType; + use datafusion_common::tree_node::TreeNodeRecursion; + use datafusion_physical_expr::expressions::DynamicFilterPhysicalExpr; + use datafusion_physical_plan::joins::{HashJoinExec, PartitionMode}; + + 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..f8d2004bf1989 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(&dyn PhysicalExpr) -> Result, + ) -> Result { + // Visit expressions in the output ordering from equivalence properties + let mut tnr = TreeNodeRecursion::Continue; + if let Some(ordering) = self.cache.output_ordering() { + for sort_expr in ordering { + tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; + } + } + Ok(tnr) + } } #[derive(Eq, PartialEq, Debug)] @@ -1230,6 +1245,20 @@ impl ExecutionPlan for StatisticsExec { self.stats.clone() })) } + + fn apply_expressions( + &self, + f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + ) -> Result { + // Visit expressions in the output ordering from equivalence properties + let mut tnr = TreeNodeRecursion::Continue; + if let Some(ordering) = self.cache.output_ordering() { + for sort_expr in ordering { + tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; + } + } + Ok(tnr) + } } #[test] diff --git a/datafusion/core/tests/physical_optimizer/pushdown_utils.rs b/datafusion/core/tests/physical_optimizer/pushdown_utils.rs index 2ffd1899b3c1d..9333dac290948 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,25 @@ impl FileSource for TestSource { fn table_schema(&self) -> &datafusion_datasource::TableSchema { &self.table_schema } + + fn apply_expressions( + &self, + f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + ) -> Result { + // Visit predicate (filter) expression if present + if let Some(predicate) = &self.predicate { + f(predicate.as_ref())?; + } + + // Visit projection expressions if present + if let Some(projection) = &self.projection { + for proj_expr in projection { + f(proj_expr.expr.as_ref())?; + } + } + + Ok(TreeNodeRecursion::Continue) + } } #[derive(Debug, Clone)] @@ -549,4 +569,13 @@ impl ExecutionPlan for TestNode { Ok(res) } } + + fn apply_expressions( + &self, + f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + ) -> Result { + // Visit the predicate expression + f(self.predicate.as_ref())?; + Ok(TreeNodeRecursion::Continue) + } } diff --git a/datafusion/core/tests/physical_optimizer/test_utils.rs b/datafusion/core/tests/physical_optimizer/test_utils.rs index 74230b24e2ab5..09a322b47e43c 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,20 @@ impl ExecutionPlan for RequirementsTestExec { ) -> Result { unimplemented!("Test exec does not support execution") } + + fn apply_expressions( + &self, + f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + ) -> Result { + // Visit expressions in required_input_ordering if present + let mut tnr = TreeNodeRecursion::Continue; + if let Some(ordering) = &self.required_input_ordering { + for sort_expr in ordering { + tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; + } + } + Ok(tnr) + } } /// A [`PlanContext`] object is susceptible to being left in an inconsistent state after @@ -1074,6 +1090,28 @@ impl ExecutionPlan for TestScan { }) } } + + fn apply_expressions( + &self, + f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + ) -> Result { + // Visit expressions in output_ordering + let mut tnr = TreeNodeRecursion::Continue; + for ordering in &self.output_ordering { + for sort_expr in ordering { + tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; + } + } + + // Visit expressions in requested_ordering if present + if let Some(ordering) = &self.requested_ordering { + for sort_expr in ordering { + tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; + } + } + + Ok(tnr) + } } /// Helper function to create a TestScan with ordering @@ -1090,6 +1128,13 @@ struct InexactMemorySource { } impl DataSource for InexactMemorySource { + fn apply_expressions( + &self, + f: &mut dyn FnMut(&dyn PhysicalExpr) -> 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..326c767d97610 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( + &dyn datafusion_physical_plan::PhysicalExpr, + ) -> Result, + ) -> Result { + // Visit expressions in the output ordering from equivalence properties + let mut tnr = TreeNodeRecursion::Continue; + if let Some(ordering) = self.plan_properties.output_ordering() { + for sort_expr in ordering { + tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; + } + } + Ok(tnr) + } } 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..0e95ebcd621ba 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( + &dyn datafusion::physical_plan::PhysicalExpr, + ) -> Result, + ) -> Result { + // Visit expressions in the output ordering from equivalence properties + let mut tnr = TreeNodeRecursion::Continue; + if let Some(ordering) = self.cache.output_ordering() { + for sort_expr in ordering { + tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; + } + } + Ok(tnr) + } } // A very specialized TopK implementation diff --git a/datafusion/datasource-arrow/src/source.rs b/datafusion/datasource-arrow/src/source.rs index 27533052ce03f..42165c1ee01a2 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,20 @@ impl FileSource for ArrowSource { fn projection(&self) -> Option<&ProjectionExprs> { Some(&self.projection.source) } + + fn apply_expressions( + &self, + f: &mut dyn FnMut( + &dyn datafusion_physical_plan::PhysicalExpr, + ) -> Result, + ) -> Result { + // Visit projection expressions + let mut tnr = TreeNodeRecursion::Continue; + for proj_expr in &self.projection.source { + tnr = tnr.visit_sibling(|| f(proj_expr.expr.as_ref()))?; + } + Ok(tnr) + } } /// `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..b80d4f462e425 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,20 @@ 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( + &dyn datafusion_physical_plan::PhysicalExpr, + ) -> Result, + ) -> Result { + // Visit projection expressions + let mut tnr = TreeNodeRecursion::Continue; + for proj_expr in &self.projection.source { + tnr = tnr.visit_sibling(|| f(proj_expr.expr.as_ref()))?; + } + Ok(tnr) + } } mod private { diff --git a/datafusion/datasource-csv/src/source.rs b/datafusion/datasource-csv/src/source.rs index 25ec311880405..76d2c66a6118c 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,20 @@ impl FileSource for CsvSource { DisplayFormatType::TreeRender => Ok(()), } } + + fn apply_expressions( + &self, + f: &mut dyn FnMut( + &dyn datafusion_physical_plan::PhysicalExpr, + ) -> Result, + ) -> Result { + // Visit projection expressions + let mut tnr = TreeNodeRecursion::Continue; + for proj_expr in &self.projection.source { + tnr = tnr.visit_sibling(|| f(proj_expr.expr.as_ref()))?; + } + Ok(tnr) + } } impl FileOpener for CsvOpener { diff --git a/datafusion/datasource-json/src/source.rs b/datafusion/datasource-json/src/source.rs index 8632d6b942bc1..a953e16b934de 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,20 @@ impl FileSource for JsonSource { fn file_type(&self) -> &str { "json" } + + fn apply_expressions( + &self, + f: &mut dyn FnMut( + &dyn datafusion_physical_plan::PhysicalExpr, + ) -> Result, + ) -> Result { + // Visit projection expressions + let mut tnr = TreeNodeRecursion::Continue; + for proj_expr in &self.projection.source { + tnr = tnr.visit_sibling(|| f(proj_expr.expr.as_ref()))?; + } + Ok(tnr) + } } impl FileOpener for JsonOpener { diff --git a/datafusion/datasource-parquet/src/source.rs b/datafusion/datasource-parquet/src/source.rs index 3443b08475e0d..a94c6deb5e7a9 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,26 @@ impl FileSource for ParquetSource { inner: Arc::new(new_source) as Arc, }) } + + fn apply_expressions( + &self, + f: &mut dyn FnMut( + &dyn PhysicalExpr, + ) -> datafusion_common::Result, + ) -> datafusion_common::Result { + // Visit predicate (filter) expression if present + let mut tnr = TreeNodeRecursion::Continue; + if let Some(predicate) = &self.predicate { + tnr = tnr.visit_sibling(|| f(predicate.as_ref()))?; + } + + // Visit projection expressions + for proj_expr in &self.projection { + tnr = tnr.visit_sibling(|| f(proj_expr.expr.as_ref()))?; + } + + Ok(tnr) + } } /// 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..32bee63b54f23 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(&dyn PhysicalExpr) -> 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..dfbf339b6693d 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(&dyn PhysicalExpr) -> 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(&dyn PhysicalExpr) -> 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(&dyn PhysicalExpr) -> 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(&dyn PhysicalExpr) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } } #[test] diff --git a/datafusion/datasource/src/memory.rs b/datafusion/datasource/src/memory.rs index 255dd76cbd6b4..18f3d05657015 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,20 @@ impl DataSource for MemorySourceConfig { }) .transpose() } + + fn apply_expressions( + &self, + f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + ) -> Result { + // Visit expressions in sort_information + let mut tnr = TreeNodeRecursion::Continue; + for ordering in &self.sort_information { + for sort_expr in ordering { + tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; + } + } + Ok(tnr) + } } impl MemorySourceConfig { diff --git a/datafusion/datasource/src/sink.rs b/datafusion/datasource/src/sink.rs index 18ebe80773e8a..023253de38d08 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(&dyn PhysicalExpr) -> Result, + ) -> Result { + // Apply to sort order requirements if present + if let Some(sort_order) = &self.sort_order { + for req in sort_order.iter() { + f(req.expr.as_ref())?; + } + } + Ok(TreeNodeRecursion::Continue) + } + /// 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..f2e1c84df79bd 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(&dyn PhysicalExpr) -> 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(&dyn PhysicalExpr) -> 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..19396411df1ba 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(&dyn PhysicalExpr) -> 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..0956c6c6dab17 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; @@ -442,6 +443,22 @@ impl ExecutionPlan for ForeignExecutionPlan { } } + fn apply_expressions( + &self, + f: &mut dyn FnMut( + &dyn datafusion_physical_plan::PhysicalExpr, + ) -> Result, + ) -> Result { + // Visit expressions in the output ordering from equivalence properties + let mut tnr = TreeNodeRecursion::Continue; + if let Some(ordering) = self.properties.output_ordering() { + for sort_expr in ordering { + tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; + } + } + Ok(tnr) + } + fn repartitioned( &self, target_partitions: usize, @@ -569,6 +586,22 @@ pub mod tests { Statistics::new_unknown(self.props.eq_properties.schema()) }))) } + + fn apply_expressions( + &self, + f: &mut dyn FnMut( + &dyn datafusion_physical_plan::PhysicalExpr, + ) -> Result, + ) -> Result { + // Visit expressions in the output ordering from equivalence properties + let mut tnr = TreeNodeRecursion::Continue; + if let Some(ordering) = self.props.output_ordering() { + for sort_expr in ordering { + tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; + } + } + Ok(tnr) + } } #[test] diff --git a/datafusion/ffi/src/tests/async_provider.rs b/datafusion/ffi/src/tests/async_provider.rs index 9821c3e501f67..2f7d7d44dfc82 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( + &dyn datafusion_physical_plan::PhysicalExpr, + ) -> Result, + ) -> Result { + // Visit expressions in the output ordering from equivalence properties + let mut tnr = TreeNodeRecursion::Continue; + if let Some(ordering) = self.properties.output_ordering() { + for sort_expr in ordering { + tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; + } + } + Ok(tnr) + } } impl datafusion_physical_plan::DisplayAs for AsyncTestExecutionPlan { diff --git a/datafusion/physical-optimizer/src/ensure_coop.rs b/datafusion/physical-optimizer/src/ensure_coop.rs index e7aacb2321b67..102e21a4853a4 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(&dyn PhysicalExpr) -> 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..810ba27e2b82b 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,39 @@ impl ExecutionPlan for OutputRequirementExec { fn fetch(&self) -> Option { self.fetch } + + fn apply_expressions( + &self, + f: &mut dyn FnMut( + &dyn datafusion_physical_expr_common::physical_expr::PhysicalExpr, + ) -> Result, + ) -> Result { + // Visit expressions in order_requirement + let mut tnr = TreeNodeRecursion::Continue; + if let Some(order_reqs) = &self.order_requirement { + let lexes = match order_reqs { + OrderingRequirements::Hard(alternatives) => alternatives, + OrderingRequirements::Soft(alternatives) => alternatives, + }; + for lex in lexes { + for sort_expr in lex { + tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; + } + } + } + + // Visit expressions in dist_requirement if it's key partitioned + #[expect(deprecated)] + if let Distribution::HashPartitioned(exprs) + | Distribution::KeyPartitioned(exprs) = &self.dist_requirement + { + for expr in exprs { + tnr = tnr.visit_sibling(|| f(expr.as_ref()))?; + } + } + + Ok(tnr) + } } impl PhysicalOptimizerRule for OutputRequirements { diff --git a/datafusion/physical-plan/benches/compute_statistics.rs b/datafusion/physical-plan/benches/compute_statistics.rs index 56a518c95292e..afb4ef38c82d6 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(&dyn PhysicalExpr) -> 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..b36c2096f709b 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::TreeNodeRecursion; use datafusion_common::{ Constraint, Constraints, Result, ScalarValue, assert_eq_or_internal_err, internal_err, not_impl_err, @@ -1960,6 +1961,36 @@ impl ExecutionPlan for AggregateExec { vec![&self.input] } + fn apply_expressions( + &self, + f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + ) -> Result { + // Apply to group by expressions + let mut tnr = TreeNodeRecursion::Continue; + for expr in self.group_by.input_exprs() { + tnr = tnr.visit_sibling(|| f(expr.as_ref()))?; + } + + // Apply to aggregate expressions + for aggr in self.aggr_expr.iter() { + for expr in aggr.expressions() { + tnr = tnr.visit_sibling(|| f(expr.as_ref()))?; + } + } + + // Apply to filter expressions (FILTER WHERE clauses) + for filter in self.filter_expr.iter().flatten() { + tnr = tnr.visit_sibling(|| f(filter.as_ref()))?; + } + + // Apply to dynamic filter expression if present + if let Some(dyn_filter) = &self.dynamic_filter { + tnr = tnr.visit_sibling(|| f(dyn_filter.filter.as_ref()))?; + } + + Ok(tnr) + } + fn with_new_children( self: Arc, children: Vec>, @@ -3564,6 +3595,13 @@ mod tests { vec![] } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + fn with_new_children( self: Arc, _: Vec>, diff --git a/datafusion/physical-plan/src/analyze.rs b/datafusion/physical-plan/src/analyze.rs index 31e0a27410ff9..67e221c3989d7 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(&dyn PhysicalExpr) -> 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..c421b4f4a28ab 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(&dyn PhysicalExpr) -> 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..6478b68901eb0 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(&dyn PhysicalExpr) -> 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..99a0dd6891536 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(&dyn PhysicalExpr) -> 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..b7d1206dbeabe 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(&dyn PhysicalExpr) -> 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..a74c129ac9a30 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(&dyn PhysicalExpr) -> 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..ff82c501045c2 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(&dyn PhysicalExpr) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + fn execute( &self, _: usize, @@ -1677,7 +1687,9 @@ mod tests { use crate::metrics::{Count, Metric, MetricValue, MetricsSet, Time}; use crate::{DisplayFormatType, ExecutionPlan, PlanProperties}; use datafusion_common::Result; + use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_execution::{SendableRecordBatchStream, TaskContext}; + use datafusion_physical_expr::PhysicalExpr; // Wrap `sample_plan()` with an adapter node that exposes a // hand-crafted `MetricsSet` so we can assert the PG key mapping @@ -1706,6 +1718,12 @@ mod tests { fn children(&self) -> Vec<&Arc> { vec![&self.inner] } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } 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..64abe8e8e00c8 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(&dyn PhysicalExpr) -> 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..1f24b0230adc5 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}; @@ -247,6 +249,80 @@ 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 skip any remaining expressions (though typically all expressions are visited). + /// + /// 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(&dyn PhysicalExpr) -> Result, + /// ) -> Result { + /// Ok(TreeNodeRecursion::Continue) + /// } + /// ``` + /// + /// ## Node with a single expression (e.g., FilterExec) + /// ```ignore + /// fn apply_expressions( + /// &self, + /// f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + /// ) -> Result { + /// f(self.predicate.as_ref()) + /// } + /// ``` + /// + /// ## Node with multiple expressions (e.g., ProjectionExec, JoinExec) + /// + /// Use [`TreeNodeRecursion::visit_sibling`] when iterating over multiple + /// expressions. This correctly propagates [`TreeNodeRecursion::Stop`]: if + /// `f` returns `Stop` for an earlier expression, `visit_sibling` short-circuits + /// and skips the remaining ones. + /// ```ignore + /// fn apply_expressions( + /// &self, + /// f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + /// ) -> Result { + /// let mut tnr = TreeNodeRecursion::Continue; + /// for expr in &self.expressions { + /// tnr = tnr.visit_sibling(|| f(expr.as_ref()))?; + /// } + /// Ok(tnr) + /// } + /// ``` + fn apply_expressions( + &self, + f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + ) -> Result; + /// Returns a new `ExecutionPlan` where all existing children were replaced /// by the `children`, in order fn with_new_children( @@ -1772,6 +1848,13 @@ mod tests { unimplemented!() } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + fn execute( &self, _partition: usize, @@ -1828,6 +1911,13 @@ mod tests { vec![] } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + fn with_new_children( self: Arc, _: Vec>, @@ -1878,6 +1968,13 @@ mod tests { vec![] } + fn apply_expressions( + &self, + f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + ) -> Result { + self.0.apply_expressions(f) + } + fn with_new_children( self: Arc, _: Vec>, @@ -1937,6 +2034,12 @@ mod tests { fn children(&self) -> Vec<&Arc> { vec![] } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } fn with_new_children( self: Arc, _: Vec>, @@ -2000,6 +2103,12 @@ mod tests { fn children(&self) -> Vec<&Arc> { vec![&self.input] } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } fn with_new_children( self: Arc, mut children: Vec>, @@ -2091,6 +2200,12 @@ mod tests { fn children(&self) -> Vec<&Arc> { vec![&self.input] } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } fn with_new_children( self: Arc, mut children: Vec>, @@ -2232,6 +2347,109 @@ 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>, + } + + 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> { + vec![] + } + + fn with_new_children( + self: Arc, + _: Vec>, + ) -> Result> { + unimplemented!() + } + + fn apply_expressions( + &self, + f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + ) -> Result { + let mut tnr = TreeNodeRecursion::Continue; + for expr in &self.exprs { + tnr = tnr.visit_sibling(|| f(expr.as_ref()))?; + } + Ok(tnr) + } + + 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)], + }; + 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)], + }; + 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_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..60e3fcd350aa2 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(&dyn PhysicalExpr) -> 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..ba37f5f4d014b 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(&dyn PhysicalExpr) -> Result, + ) -> Result { + f(self.predicate.as_ref()) + } + 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..99f14a49e3445 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(&dyn PhysicalExpr) -> 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..42d7acfbb4157 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,30 @@ impl ExecutionPlan for HashJoinExec { vec![&self.left, &self.right] } + fn apply_expressions( + &self, + f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + ) -> Result { + // Apply to join key expressions from both sides + let mut tnr = TreeNodeRecursion::Continue; + for (left, right) in &self.on { + tnr = tnr.visit_sibling(|| f(left.as_ref()))?; + tnr = tnr.visit_sibling(|| f(right.as_ref()))?; + } + + // Apply to join filter expression if present + if let Some(filter) = &self.filter { + tnr = tnr.visit_sibling(|| f(filter.expression().as_ref()))?; + } + + // Apply to dynamic filter expression if present + if let Some(df) = &self.dynamic_filter { + tnr = tnr.visit_sibling(|| f(df.filter.as_ref()))?; + } + + Ok(tnr) + } + /// Creates a new HashJoinExec with different children while preserving configuration. /// /// This method is called during query optimization when the optimizer creates new @@ -2461,6 +2486,13 @@ mod tests { vec![] } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&dyn PhysicalExpr) -> 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..82e8eff7b7abc 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(&dyn crate::PhysicalExpr) -> Result, + ) -> Result { + // Apply to join filter expressions if present + if let Some(filter) = &self.filter { + f(filter.expression().as_ref())?; + } + Ok(TreeNodeRecursion::Continue) + } + 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..a8e7959376ff4 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(&dyn PhysicalExpr) -> Result, + ) -> Result { + // Apply to the two expressions being compared in the range predicate + f(self.on.0.as_ref())?.visit_sibling(|| f(self.on.1.as_ref())) + } + 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..4d4d962f73467 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,23 @@ impl ExecutionPlan for SortMergeJoinExec { vec![&self.left, &self.right] } + fn apply_expressions( + &self, + f: &mut dyn FnMut(&dyn crate::PhysicalExpr) -> Result, + ) -> Result { + // Apply to join keys from both sides + let mut tnr = TreeNodeRecursion::Continue; + for (left, right) in &self.on { + tnr = tnr.visit_sibling(|| f(left.as_ref()))?; + tnr = tnr.visit_sibling(|| f(right.as_ref()))?; + } + // Apply to join filter expressions if present + if let Some(filter) = &self.filter { + tnr = tnr.visit_sibling(|| f(filter.expression().as_ref()))?; + } + Ok(tnr) + } + 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..e2f061587378c 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,23 @@ impl ExecutionPlan for SymmetricHashJoinExec { vec![&self.left, &self.right] } + fn apply_expressions( + &self, + f: &mut dyn FnMut(&dyn crate::PhysicalExpr) -> Result, + ) -> Result { + // Apply to join keys from both sides + let mut tnr = TreeNodeRecursion::Continue; + for (left, right) in &self.on { + tnr = tnr.visit_sibling(|| f(left.as_ref()))?; + tnr = tnr.visit_sibling(|| f(right.as_ref()))?; + } + // Apply to join filter expressions if present + if let Some(filter) = &self.filter { + tnr = tnr.visit_sibling(|| f(filter.expression().as_ref()))?; + } + Ok(tnr) + } + fn with_new_children( self: Arc, children: Vec>, diff --git a/datafusion/physical-plan/src/limit.rs b/datafusion/physical-plan/src/limit.rs index ddce680fc18ad..e632bcdd349a1 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,20 @@ impl ExecutionPlan for GlobalLimitExec { vec![false] } + fn apply_expressions( + &self, + f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + ) -> Result { + // Apply to required ordering expressions if present + let mut tnr = TreeNodeRecursion::Continue; + if let Some(ordering) = &self.required_ordering { + for sort_expr in ordering { + tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; + } + } + Ok(tnr) + } + fn with_new_children( self: Arc, mut children: Vec>, @@ -398,6 +413,20 @@ impl ExecutionPlan for LocalLimitExec { vec![true] } + fn apply_expressions( + &self, + f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + ) -> Result { + // Apply to required ordering expressions if present + let mut tnr = TreeNodeRecursion::Continue; + if let Some(ordering) = &self.required_ordering { + for sort_expr in ordering { + tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; + } + } + Ok(tnr) + } + fn with_new_children( self: Arc, children: Vec>, @@ -648,7 +677,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..e172ef4463ec4 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(&dyn PhysicalExpr) -> 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..33226392c7e29 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(&dyn PhysicalExpr) -> 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(&dyn PhysicalExpr) -> 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..7fcfcaf86b6cb 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(&dyn PhysicalExpr) -> 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..2e915b283f897 100644 --- a/datafusion/physical-plan/src/projection.rs +++ b/datafusion/physical-plan/src/projection.rs @@ -302,6 +302,17 @@ impl ExecutionPlan for ProjectionExec { vec![&self.input] } + fn apply_expressions( + &self, + f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + ) -> Result { + let mut tnr = TreeNodeRecursion::Continue; + for proj_expr in self.projector.projection().as_ref().iter() { + tnr = tnr.visit_sibling(|| f(proj_expr.expr.as_ref()))?; + } + Ok(tnr) + } + 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..a02ec38a7431f 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(&dyn PhysicalExpr) -> 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..63469cfd6a9cc 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,26 @@ impl ExecutionPlan for RepartitionExec { vec![&self.input] } + fn apply_expressions( + &self, + f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + ) -> Result { + let exprs: Vec<_> = match self.partitioning() { + Partitioning::Hash(exprs, _) => exprs.iter().collect(), + Partitioning::Range(range) => range + .ordering() + .iter() + .map(|sort_expr| &sort_expr.expr) + .collect(), + _ => return Ok(TreeNodeRecursion::Continue), + }; + let mut tnr = TreeNodeRecursion::Continue; + for expr in exprs { + tnr = tnr.visit_sibling(|| f(expr.as_ref()))?; + } + Ok(tnr) + } + fn with_new_children( self: Arc, mut children: Vec>, @@ -3045,6 +3066,13 @@ mod tests { vec![] } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + fn with_new_children( self: Arc, _: Vec>, diff --git a/datafusion/physical-plan/src/scalar_subquery.rs b/datafusion/physical-plan/src/scalar_subquery.rs index 73acb2ab13480..08a1a1d430770 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(&dyn PhysicalExpr) -> 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(&dyn PhysicalExpr) -> 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..724dfd3c16c74 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,17 @@ impl ExecutionPlan for PartialSortExec { vec![&self.input] } + fn apply_expressions( + &self, + f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + ) -> Result { + let mut tnr = TreeNodeRecursion::Continue; + for sort_expr in &self.expr { + tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; + } + Ok(tnr) + } + 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..897eadbb60e51 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,17 @@ impl ExecutionPlan for PartitionedTopKExec { )?)) } + fn apply_expressions( + &self, + f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + ) -> Result { + let mut tnr = TreeNodeRecursion::Continue; + for sort_expr in &self.expr { + tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; + } + Ok(tnr) + } + 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..e489369be980d 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,25 @@ impl ExecutionPlan for SortExec { vec![&self.input] } + fn apply_expressions( + &self, + f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + ) -> Result { + // Apply to sort expressions + let mut tnr = TreeNodeRecursion::Continue; + for sort_expr in &self.expr { + tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; + } + + // Apply to dynamic filter expression if present (when fetch is Some, TopK mode) + if let Some(filter) = &self.filter { + let filter_guard = filter.read(); + tnr = tnr.visit_sibling(|| f(filter_guard.expr().as_ref()))?; + } + + Ok(tnr) + } + fn benefits_from_input_partitioning(&self) -> Vec { vec![false] } @@ -1760,6 +1780,13 @@ mod tests { Ok(self) } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&dyn PhysicalExpr) -> 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..2de468455408b 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,17 @@ impl ExecutionPlan for SortPreservingMergeExec { vec![&self.input] } + fn apply_expressions( + &self, + f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + ) -> Result { + let mut tnr = TreeNodeRecursion::Continue; + for sort_expr in &self.expr { + tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; + } + Ok(tnr) + } + fn with_new_children( self: Arc, mut children: Vec>, @@ -1573,6 +1586,12 @@ mod tests { fn children(&self) -> Vec<&Arc> { vec![] } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&dyn PhysicalExpr) -> 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..4e792c8f5c504 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(&dyn PhysicalExpr) -> 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..491f38b9eccec 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,20 @@ impl ExecutionPlan for TestMemoryExec { Vec::new() } + fn apply_expressions( + &self, + f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + ) -> Result { + // Apply to all sort information orderings + let mut tnr = TreeNodeRecursion::Continue; + for ordering in &self.sort_information { + for sort_expr in ordering { + tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; + } + } + Ok(tnr) + } + 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..819259bcf8b8f 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(&dyn PhysicalExpr) -> 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(&dyn PhysicalExpr) -> 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(&dyn PhysicalExpr) -> 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(&dyn PhysicalExpr) -> 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(&dyn PhysicalExpr) -> 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(&dyn PhysicalExpr) -> 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..e1eb46e2fbe73 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(&dyn PhysicalExpr) -> 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(&dyn PhysicalExpr) -> 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..0372a7adf827c 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(&dyn PhysicalExpr) -> 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..a7beb47395652 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,19 @@ impl ExecutionPlan for BoundedWindowAggExec { vec![&self.input] } + fn apply_expressions( + &self, + f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + ) -> Result { + let mut tnr = TreeNodeRecursion::Continue; + for window_expr in &self.window_expr { + for expr in window_expr.expressions() { + tnr = tnr.visit_sibling(|| f(expr.as_ref()))?; + } + } + Ok(tnr) + } + 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..2240eaf6aff90 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,19 @@ impl ExecutionPlan for WindowAggExec { vec![&self.input] } + fn apply_expressions( + &self, + f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + ) -> Result { + let mut tnr = TreeNodeRecursion::Continue; + for window_expr in &self.window_expr { + for expr in window_expr.expressions() { + tnr = tnr.visit_sibling(|| f(expr.as_ref()))?; + } + } + Ok(tnr) + } + 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..0ae39f2dc0b6a 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(&dyn PhysicalExpr) -> 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..b003e30cd9f5c 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(&dyn PhysicalExpr) -> Result, + ) -> Result { + self.inner.apply_expressions(f) + } + fn with_new_children( self: Arc, children: Vec>, @@ -4595,6 +4603,17 @@ impl ExecutionPlan for CustomExecWithExprs { vec![&self.child] } + fn apply_expressions( + &self, + f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + ) -> Result { + let mut tnr = TreeNodeRecursion::Continue; + for expr in &self.exprs { + tnr = tnr.visit_sibling(|| f(expr.as_ref()))?; + } + Ok(tnr) + } + 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..81b2d131e65c3 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(&dyn PhysicalExpr) -> Result) -> Result { Ok(TreeNodeRecursion::Continue) } # } ``` @@ -909,6 +910,13 @@ impl ExecutionPlan for CountingExec { batch_stream, ))) } + +# fn apply_expressions( +# &self, +# _f: &mut dyn FnMut(&dyn PhysicalExpr) -> 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..349e15a2188a3 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,73 @@ 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, using `visit_sibling` to correctly propagate `TreeNodeRecursion`: + +**Node with no expressions:** + +```rust,ignore +fn apply_expressions( + &self, + _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, +) -> Result { + Ok(TreeNodeRecursion::Continue) +} +``` + +**Node with a single expression:** + +```rust,ignore +fn apply_expressions( + &self, + f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, +) -> Result { + f(self.predicate.as_ref()) +} +``` + +**Node with multiple expressions:** + +```rust,ignore +fn apply_expressions( + &self, + f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, +) -> Result { + let mut tnr = TreeNodeRecursion::Continue; + for expr in &self.expressions { + tnr = tnr.visit_sibling(|| f(expr.as_ref()))?; + } + Ok(tnr) +} +``` + +**Node whose only expressions are in `output_ordering()` (e.g. a synthetic test node with no owned expression fields):** + +```rust,ignore +fn apply_expressions( + &self, + f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, +) -> Result { + let mut tnr = TreeNodeRecursion::Continue; + if let Some(ordering) = self.cache.output_ordering() { + for sort_expr in ordering { + tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; + } + } + Ok(tnr) +} +``` + ### `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. From 380f5c7eff084592938a1b685db1efb57b366da1 Mon Sep 17 00:00:00 2001 From: Jayant Shrivastava Date: Thu, 30 Jul 2026 21:58:08 +0000 Subject: [PATCH 2/3] Refine physical expression traversal API --- .../custom_data_source/custom_datasource.rs | 13 +- .../memory_pool_execution_plan.rs | 13 +- .../proto/composed_extension_codec.rs | 4 +- .../examples/relation_planner/table_sample.rs | 18 +-- datafusion/catalog/src/memory/table.rs | 2 +- datafusion/core/src/physical_planner.rs | 17 +- .../core/tests/custom_sources_cases/mod.rs | 18 +-- .../provider_filter_pushdown.rs | 18 +-- .../tests/custom_sources_cases/statistics.rs | 18 +-- datafusion/core/tests/fuzz_cases/once_exec.rs | 18 +-- .../enforce_distribution.rs | 20 +-- .../physical_optimizer/ensure_requirements.rs | 31 ++-- .../physical_optimizer/filter_pushdown.rs | 12 +- .../physical_optimizer/join_selection.rs | 36 ++--- .../physical_optimizer/pushdown_utils.rs | 30 ++-- .../tests/physical_optimizer/test_utils.rs | 45 +++--- .../tests/user_defined/insert_operation.rs | 18 +-- .../tests/user_defined/user_defined_plan.rs | 18 +-- datafusion/datasource-arrow/src/source.rs | 15 +- datafusion/datasource-avro/src/source.rs | 15 +- datafusion/datasource-csv/src/source.rs | 15 +- datafusion/datasource-json/src/source.rs | 15 +- datafusion/datasource-parquet/src/source.rs | 20 +-- datafusion/datasource/src/file.rs | 2 +- .../datasource/src/file_scan_config/mod.rs | 8 +- datafusion/datasource/src/memory.rs | 17 +- datafusion/datasource/src/sink.rs | 16 +- datafusion/datasource/src/source.rs | 4 +- datafusion/datasource/src/test_util.rs | 2 +- datafusion/ffi/src/execution_plan.rs | 99 +++++++++--- datafusion/ffi/src/physical_expr/mod.rs | 33 +++- datafusion/ffi/src/tests/async_provider.rs | 18 +-- datafusion/ffi/src/tests/mod.rs | 13 ++ datafusion/ffi/tests/ffi_execution_plan.rs | 25 +++ .../physical-optimizer/src/ensure_coop.rs | 2 +- .../src/output_requirements.rs | 42 +++-- .../benches/compute_statistics.rs | 2 +- .../physical-plan/src/aggregates/mod.rs | 71 ++++----- datafusion/physical-plan/src/analyze.rs | 2 +- datafusion/physical-plan/src/async_func.rs | 2 +- datafusion/physical-plan/src/buffer.rs | 2 +- .../physical-plan/src/coalesce_batches.rs | 2 +- .../physical-plan/src/coalesce_partitions.rs | 2 +- datafusion/physical-plan/src/coop.rs | 2 +- datafusion/physical-plan/src/display.rs | 15 +- datafusion/physical-plan/src/empty.rs | 2 +- .../physical-plan/src/execution_plan.rs | 146 +++++++++++++----- datafusion/physical-plan/src/explain.rs | 2 +- datafusion/physical-plan/src/filter.rs | 4 +- .../physical-plan/src/joins/cross_join.rs | 2 +- .../physical-plan/src/joins/hash_join/exec.rs | 35 ++--- .../src/joins/nested_loop_join.rs | 10 +- .../src/joins/piecewise_merge_join/exec.rs | 4 +- .../src/joins/sort_merge_join/exec.rs | 16 +- .../src/joins/symmetric_hash_join.rs | 16 +- datafusion/physical-plan/src/lib.rs | 6 +- datafusion/physical-plan/src/limit.rs | 34 ++-- datafusion/physical-plan/src/memory.rs | 2 +- .../src/operator_statistics/mod.rs | 4 +- .../physical-plan/src/placeholder_row.rs | 2 +- datafusion/physical-plan/src/projection.rs | 15 +- .../physical-plan/src/recursive_query.rs | 2 +- .../physical-plan/src/repartition/mod.rs | 31 ++-- .../physical-plan/src/scalar_subquery.rs | 4 +- .../physical-plan/src/sorts/partial_sort.rs | 11 +- .../src/sorts/partitioned_topk.rs | 11 +- datafusion/physical-plan/src/sorts/sort.rs | 28 ++-- .../src/sorts/sort_preserving_merge.rs | 13 +- datafusion/physical-plan/src/streaming.rs | 2 +- datafusion/physical-plan/src/test.rs | 17 +- datafusion/physical-plan/src/test/exec.rs | 12 +- datafusion/physical-plan/src/union.rs | 4 +- datafusion/physical-plan/src/unnest.rs | 2 +- .../src/windows/bounded_window_agg_exec.rs | 18 ++- .../src/windows/window_agg_exec.rs | 18 ++- datafusion/physical-plan/src/work_table.rs | 2 +- .../tests/cases/roundtrip_physical_plan.rs | 10 +- .../custom-table-providers.md | 4 +- .../library-user-guide/upgrading/54.0.0.md | 36 ++--- 79 files changed, 714 insertions(+), 621 deletions(-) diff --git a/datafusion-examples/examples/custom_data_source/custom_datasource.rs b/datafusion-examples/examples/custom_data_source/custom_datasource.rs index e5a8838d797ac..a5a38edf0b6f5 100644 --- a/datafusion-examples/examples/custom_data_source/custom_datasource.rs +++ b/datafusion-examples/examples/custom_data_source/custom_datasource.rs @@ -318,17 +318,10 @@ impl ExecutionPlan for CustomExec { fn apply_expressions( &self, - f: &mut dyn FnMut( - &dyn datafusion::physical_plan::PhysicalExpr, + _f: &mut dyn FnMut( + &Arc, ) -> Result, ) -> Result { - // Visit expressions in the output ordering from equivalence properties - let mut tnr = TreeNodeRecursion::Continue; - if let Some(ordering) = self.cache.output_ordering() { - for sort_expr in ordering { - tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; - } - } - Ok(tnr) + 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 f761e9b6e6d50..6decb84b55be1 100644 --- a/datafusion-examples/examples/execution_monitoring/memory_pool_execution_plan.rs +++ b/datafusion-examples/examples/execution_monitoring/memory_pool_execution_plan.rs @@ -295,17 +295,10 @@ impl ExecutionPlan for BufferingExecutionPlan { fn apply_expressions( &self, - f: &mut dyn FnMut( - &dyn datafusion::physical_plan::PhysicalExpr, + _f: &mut dyn FnMut( + &Arc, ) -> Result, ) -> Result { - // Visit expressions in the output ordering from equivalence properties - let mut tnr = TreeNodeRecursion::Continue; - if let Some(ordering) = self.properties.output_ordering() { - for sort_expr in ordering { - tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; - } - } - Ok(tnr) + Ok(TreeNodeRecursion::Continue) } } diff --git a/datafusion-examples/examples/proto/composed_extension_codec.rs b/datafusion-examples/examples/proto/composed_extension_codec.rs index 3405cc6a63aca..d5197fe61bea7 100644 --- a/datafusion-examples/examples/proto/composed_extension_codec.rs +++ b/datafusion-examples/examples/proto/composed_extension_codec.rs @@ -129,7 +129,7 @@ impl ExecutionPlan for ParentExec { fn apply_expressions( &self, _f: &mut dyn FnMut( - &dyn datafusion::physical_plan::PhysicalExpr, + &Arc, ) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) @@ -216,7 +216,7 @@ impl ExecutionPlan for ChildExec { fn apply_expressions( &self, _f: &mut dyn FnMut( - &dyn datafusion::physical_plan::PhysicalExpr, + &Arc, ) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) diff --git a/datafusion-examples/examples/relation_planner/table_sample.rs b/datafusion-examples/examples/relation_planner/table_sample.rs index ba5b474b2e240..4a4dfff1da3e8 100644 --- a/datafusion-examples/examples/relation_planner/table_sample.rs +++ b/datafusion-examples/examples/relation_planner/table_sample.rs @@ -753,17 +753,17 @@ impl ExecutionPlan for SampleExec { fn apply_expressions( &self, f: &mut dyn FnMut( - &dyn datafusion::physical_plan::PhysicalExpr, + &Arc, ) -> Result, ) -> Result { - // Visit expressions in the output ordering from equivalence properties - let mut tnr = TreeNodeRecursion::Continue; - if let Some(ordering) = self.cache.output_ordering() { - for sort_expr in ordering { - tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; - } - } - Ok(tnr) + datafusion::physical_plan::apply_expression_roots( + self.cache + .output_ordering() + .into_iter() + .flatten() + .map(|sort_expr| &sort_expr.expr), + f, + ) } } diff --git a/datafusion/catalog/src/memory/table.rs b/datafusion/catalog/src/memory/table.rs index bc71e8c010353..ef5669a3a13f0 100644 --- a/datafusion/catalog/src/memory/table.rs +++ b/datafusion/catalog/src/memory/table.rs @@ -601,7 +601,7 @@ impl ExecutionPlan for DmlResultExec { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _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 14c8d2045cf5e..723617006b422 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -4714,16 +4714,9 @@ mod tests { fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - // Visit expressions in the output ordering from equivalence properties - let mut tnr = TreeNodeRecursion::Continue; - if let Some(ordering) = self.cache.output_ordering() { - for sort_expr in ordering { - tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; - } - } - Ok(tnr) + Ok(TreeNodeRecursion::Continue) } } @@ -4894,7 +4887,7 @@ digraph { } fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } @@ -4946,7 +4939,7 @@ digraph { } fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } @@ -5076,7 +5069,7 @@ digraph { } fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } diff --git a/datafusion/core/tests/custom_sources_cases/mod.rs b/datafusion/core/tests/custom_sources_cases/mod.rs index 3b8548c49e4f0..9582c13669358 100644 --- a/datafusion/core/tests/custom_sources_cases/mod.rs +++ b/datafusion/core/tests/custom_sources_cases/mod.rs @@ -214,17 +214,17 @@ impl ExecutionPlan for CustomExecutionPlan { fn apply_expressions( &self, f: &mut dyn FnMut( - &dyn datafusion::physical_plan::PhysicalExpr, + &Arc, ) -> Result, ) -> Result { - // Visit expressions in the output ordering from equivalence properties - let mut tnr = TreeNodeRecursion::Continue; - if let Some(ordering) = self.cache.output_ordering() { - for sort_expr in ordering { - tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; - } - } - Ok(tnr) + 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/custom_sources_cases/provider_filter_pushdown.rs b/datafusion/core/tests/custom_sources_cases/provider_filter_pushdown.rs index e52c559ec79ef..45fcc1b423de9 100644 --- a/datafusion/core/tests/custom_sources_cases/provider_filter_pushdown.rs +++ b/datafusion/core/tests/custom_sources_cases/provider_filter_pushdown.rs @@ -153,17 +153,17 @@ impl ExecutionPlan for CustomPlan { fn apply_expressions( &self, f: &mut dyn FnMut( - &dyn datafusion::physical_plan::PhysicalExpr, + &Arc, ) -> Result, ) -> Result { - // Visit expressions in the output ordering from equivalence properties - let mut tnr = TreeNodeRecursion::Continue; - if let Some(ordering) = self.cache.output_ordering() { - for sort_expr in ordering { - tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; - } - } - Ok(tnr) + 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/custom_sources_cases/statistics.rs b/datafusion/core/tests/custom_sources_cases/statistics.rs index da5ff444f15a8..5665b4b193aed 100644 --- a/datafusion/core/tests/custom_sources_cases/statistics.rs +++ b/datafusion/core/tests/custom_sources_cases/statistics.rs @@ -190,17 +190,17 @@ impl ExecutionPlan for StatisticsValidation { fn apply_expressions( &self, f: &mut dyn FnMut( - &dyn datafusion::physical_plan::PhysicalExpr, + &Arc, ) -> Result, ) -> Result { - // Visit expressions in the output ordering from equivalence properties - let mut tnr = TreeNodeRecursion::Continue; - if let Some(ordering) = self.cache.output_ordering() { - for sort_expr in ordering { - tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; - } - } - Ok(tnr) + 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/fuzz_cases/once_exec.rs b/datafusion/core/tests/fuzz_cases/once_exec.rs index 403e377a690e2..e3ee52a4872e5 100644 --- a/datafusion/core/tests/fuzz_cases/once_exec.rs +++ b/datafusion/core/tests/fuzz_cases/once_exec.rs @@ -110,16 +110,16 @@ impl ExecutionPlan for OnceExec { fn apply_expressions( &self, f: &mut dyn FnMut( - &dyn datafusion_physical_plan::PhysicalExpr, + &Arc, ) -> datafusion_common::Result, ) -> datafusion_common::Result { - // Visit expressions in the output ordering from equivalence properties - let mut tnr = TreeNodeRecursion::Continue; - if let Some(ordering) = self.cache.output_ordering() { - for sort_expr in ordering { - tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; - } - } - Ok(tnr) + 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 5e1182c09c1f7..ccfc8fb055eb1 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs +++ b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs @@ -207,16 +207,16 @@ impl ExecutionPlan for SortRequiredExec { fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - // Visit expressions in the output ordering from equivalence properties - let mut tnr = TreeNodeRecursion::Continue; - if let Some(ordering) = self.cache.output_ordering() { - for sort_expr in ordering { - tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; - } - } - Ok(tnr) + datafusion_physical_plan::apply_expression_roots( + self.cache + .output_ordering() + .into_iter() + .flatten() + .map(|sort_expr| &sort_expr.expr), + f, + ) } fn execute( @@ -308,7 +308,7 @@ impl ExecutionPlan for SinglePartitionMaintainsOrderExec { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } diff --git a/datafusion/core/tests/physical_optimizer/ensure_requirements.rs b/datafusion/core/tests/physical_optimizer/ensure_requirements.rs index 5bf361c8ada0c..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; @@ -30,7 +31,6 @@ use std::sync::Arc; use arrow::compute::SortOptions; use arrow::datatypes::{DataType, Field, Schema}; use datafusion_common::Result; -use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_physical_expr::expressions::Column; use datafusion_physical_expr::{ EquivalenceProperties, LexOrdering, PhysicalExpr, PhysicalSortExpr, @@ -114,7 +114,7 @@ impl ExecutionPlan for MockMultiPartitionExec { } fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } @@ -983,23 +983,24 @@ impl ExecutionPlan for MockReqExec { } fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - let mut tnr = TreeNodeRecursion::Continue; - if let Some(ordering) = &self.ord { - for sort_expr in ordering { - tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; - } - } #[expect(deprecated)] - if let Distribution::HashPartitioned(exprs) + let distribution = if let Distribution::HashPartitioned(exprs) | Distribution::KeyPartitioned(exprs) = &self.dist { - for expr in exprs { - tnr = tnr.visit_sibling(|| f(expr.as_ref()))?; - } - } - Ok(tnr) + 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, diff --git a/datafusion/core/tests/physical_optimizer/filter_pushdown.rs b/datafusion/core/tests/physical_optimizer/filter_pushdown.rs index ad6ca6fba0998..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, @@ -3092,11 +3095,6 @@ async fn test_filter_with_projection_pushdown() { /// counting nodes. Neither API is observable from SQL. #[tokio::test] async fn test_discover_dynamic_filters_via_expressions_api() { - use datafusion_common::JoinType; - use datafusion_common::tree_node::TreeNodeRecursion; - use datafusion_physical_expr::expressions::DynamicFilterPhysicalExpr; - use datafusion_physical_plan::joins::{HashJoinExec, PartitionMode}; - fn count_dynamic_filters(plan: &Arc) -> usize { let mut count = 0; diff --git a/datafusion/core/tests/physical_optimizer/join_selection.rs b/datafusion/core/tests/physical_optimizer/join_selection.rs index f8d2004bf1989..41f8960e8c21b 100644 --- a/datafusion/core/tests/physical_optimizer/join_selection.rs +++ b/datafusion/core/tests/physical_optimizer/join_selection.rs @@ -1129,16 +1129,16 @@ impl ExecutionPlan for UnboundedExec { fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - // Visit expressions in the output ordering from equivalence properties - let mut tnr = TreeNodeRecursion::Continue; - if let Some(ordering) = self.cache.output_ordering() { - for sort_expr in ordering { - tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; - } - } - Ok(tnr) + datafusion_physical_plan::apply_expression_roots( + self.cache + .output_ordering() + .into_iter() + .flatten() + .map(|sort_expr| &sort_expr.expr), + f, + ) } } @@ -1248,16 +1248,16 @@ impl ExecutionPlan for StatisticsExec { fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - // Visit expressions in the output ordering from equivalence properties - let mut tnr = TreeNodeRecursion::Continue; - if let Some(ordering) = self.cache.output_ordering() { - for sort_expr in ordering { - tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; - } - } - Ok(tnr) + 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/pushdown_utils.rs b/datafusion/core/tests/physical_optimizer/pushdown_utils.rs index 9333dac290948..4f8b9ad42b6c8 100644 --- a/datafusion/core/tests/physical_optimizer/pushdown_utils.rs +++ b/datafusion/core/tests/physical_optimizer/pushdown_utils.rs @@ -238,21 +238,17 @@ impl FileSource for TestSource { fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - // Visit predicate (filter) expression if present - if let Some(predicate) = &self.predicate { - f(predicate.as_ref())?; - } - - // Visit projection expressions if present - if let Some(projection) = &self.projection { - for proj_expr in projection { - f(proj_expr.expr.as_ref())?; - } - } - - Ok(TreeNodeRecursion::Continue) + datafusion_physical_plan::apply_expression_roots( + self.predicate.iter().chain( + self.projection + .iter() + .flatten() + .map(|proj_expr| &proj_expr.expr), + ), + f, + ) } } @@ -572,10 +568,8 @@ impl ExecutionPlan for TestNode { fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - // Visit the predicate expression - f(self.predicate.as_ref())?; - Ok(TreeNodeRecursion::Continue) + 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 09a322b47e43c..a570089b94f01 100644 --- a/datafusion/core/tests/physical_optimizer/test_utils.rs +++ b/datafusion/core/tests/physical_optimizer/test_utils.rs @@ -531,16 +531,15 @@ impl ExecutionPlan for RequirementsTestExec { fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - // Visit expressions in required_input_ordering if present - let mut tnr = TreeNodeRecursion::Continue; - if let Some(ordering) = &self.required_input_ordering { - for sort_expr in ordering { - tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; - } - } - Ok(tnr) + datafusion_physical_plan::apply_expression_roots( + self.required_input_ordering + .iter() + .flatten() + .map(|sort_expr| &sort_expr.expr), + f, + ) } } @@ -1093,24 +1092,16 @@ impl ExecutionPlan for TestScan { fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - // Visit expressions in output_ordering - let mut tnr = TreeNodeRecursion::Continue; - for ordering in &self.output_ordering { - for sort_expr in ordering { - tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; - } - } - - // Visit expressions in requested_ordering if present - if let Some(ordering) = &self.requested_ordering { - for sort_expr in ordering { - tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; - } - } - - Ok(tnr) + datafusion_physical_plan::apply_expression_roots( + self.output_ordering + .iter() + .flatten() + .chain(self.requested_ordering.iter().flatten()) + .map(|sort_expr| &sort_expr.expr), + f, + ) } } @@ -1130,7 +1121,7 @@ struct InexactMemorySource { impl DataSource for InexactMemorySource { fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { self.inner.apply_expressions(f) } diff --git a/datafusion/core/tests/user_defined/insert_operation.rs b/datafusion/core/tests/user_defined/insert_operation.rs index 326c767d97610..9efe1ec9eba15 100644 --- a/datafusion/core/tests/user_defined/insert_operation.rs +++ b/datafusion/core/tests/user_defined/insert_operation.rs @@ -181,17 +181,17 @@ impl ExecutionPlan for TestInsertExec { fn apply_expressions( &self, f: &mut dyn FnMut( - &dyn datafusion_physical_plan::PhysicalExpr, + &Arc, ) -> Result, ) -> Result { - // Visit expressions in the output ordering from equivalence properties - let mut tnr = TreeNodeRecursion::Continue; - if let Some(ordering) = self.plan_properties.output_ordering() { - for sort_expr in ordering { - tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; - } - } - Ok(tnr) + datafusion_physical_plan::apply_expression_roots( + self.plan_properties + .output_ordering() + .into_iter() + .flatten() + .map(|sort_expr| &sort_expr.expr), + f, + ) } } diff --git a/datafusion/core/tests/user_defined/user_defined_plan.rs b/datafusion/core/tests/user_defined/user_defined_plan.rs index 0e95ebcd621ba..07d44c8de39fc 100644 --- a/datafusion/core/tests/user_defined/user_defined_plan.rs +++ b/datafusion/core/tests/user_defined/user_defined_plan.rs @@ -754,17 +754,17 @@ impl ExecutionPlan for TopKExec { fn apply_expressions( &self, f: &mut dyn FnMut( - &dyn datafusion::physical_plan::PhysicalExpr, + &Arc, ) -> Result, ) -> Result { - // Visit expressions in the output ordering from equivalence properties - let mut tnr = TreeNodeRecursion::Continue; - if let Some(ordering) = self.cache.output_ordering() { - for sort_expr in ordering { - tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; - } - } - Ok(tnr) + datafusion::physical_plan::apply_expression_roots( + self.cache + .output_ordering() + .into_iter() + .flatten() + .map(|sort_expr| &sort_expr.expr), + f, + ) } } diff --git a/datafusion/datasource-arrow/src/source.rs b/datafusion/datasource-arrow/src/source.rs index 42165c1ee01a2..3da60697fb4de 100644 --- a/datafusion/datasource-arrow/src/source.rs +++ b/datafusion/datasource-arrow/src/source.rs @@ -397,15 +397,16 @@ impl FileSource for ArrowSource { fn apply_expressions( &self, f: &mut dyn FnMut( - &dyn datafusion_physical_plan::PhysicalExpr, + &Arc, ) -> Result, ) -> Result { - // Visit projection expressions - let mut tnr = TreeNodeRecursion::Continue; - for proj_expr in &self.projection.source { - tnr = tnr.visit_sibling(|| f(proj_expr.expr.as_ref()))?; - } - Ok(tnr) + datafusion_physical_plan::apply_expression_roots( + self.projection + .source + .iter() + .map(|proj_expr| &proj_expr.expr), + f, + ) } } diff --git a/datafusion/datasource-avro/src/source.rs b/datafusion/datasource-avro/src/source.rs index b80d4f462e425..a069c1b4287e6 100644 --- a/datafusion/datasource-avro/src/source.rs +++ b/datafusion/datasource-avro/src/source.rs @@ -173,15 +173,16 @@ impl FileSource for AvroSource { fn apply_expressions( &self, f: &mut dyn FnMut( - &dyn datafusion_physical_plan::PhysicalExpr, + &Arc, ) -> Result, ) -> Result { - // Visit projection expressions - let mut tnr = TreeNodeRecursion::Continue; - for proj_expr in &self.projection.source { - tnr = tnr.visit_sibling(|| f(proj_expr.expr.as_ref()))?; - } - Ok(tnr) + datafusion_physical_plan::apply_expression_roots( + self.projection + .source + .iter() + .map(|proj_expr| &proj_expr.expr), + f, + ) } } diff --git a/datafusion/datasource-csv/src/source.rs b/datafusion/datasource-csv/src/source.rs index 76d2c66a6118c..725e36e33cf32 100644 --- a/datafusion/datasource-csv/src/source.rs +++ b/datafusion/datasource-csv/src/source.rs @@ -313,15 +313,16 @@ impl FileSource for CsvSource { fn apply_expressions( &self, f: &mut dyn FnMut( - &dyn datafusion_physical_plan::PhysicalExpr, + &Arc, ) -> Result, ) -> Result { - // Visit projection expressions - let mut tnr = TreeNodeRecursion::Continue; - for proj_expr in &self.projection.source { - tnr = tnr.visit_sibling(|| f(proj_expr.expr.as_ref()))?; - } - Ok(tnr) + datafusion_physical_plan::apply_expression_roots( + self.projection + .source + .iter() + .map(|proj_expr| &proj_expr.expr), + f, + ) } } diff --git a/datafusion/datasource-json/src/source.rs b/datafusion/datasource-json/src/source.rs index a953e16b934de..f19f6a71d8f67 100644 --- a/datafusion/datasource-json/src/source.rs +++ b/datafusion/datasource-json/src/source.rs @@ -236,15 +236,16 @@ impl FileSource for JsonSource { fn apply_expressions( &self, f: &mut dyn FnMut( - &dyn datafusion_physical_plan::PhysicalExpr, + &Arc, ) -> Result, ) -> Result { - // Visit projection expressions - let mut tnr = TreeNodeRecursion::Continue; - for proj_expr in &self.projection.source { - tnr = tnr.visit_sibling(|| f(proj_expr.expr.as_ref()))?; - } - Ok(tnr) + datafusion_physical_plan::apply_expression_roots( + self.projection + .source + .iter() + .map(|proj_expr| &proj_expr.expr), + f, + ) } } diff --git a/datafusion/datasource-parquet/src/source.rs b/datafusion/datasource-parquet/src/source.rs index a94c6deb5e7a9..b2674efd404c3 100644 --- a/datafusion/datasource-parquet/src/source.rs +++ b/datafusion/datasource-parquet/src/source.rs @@ -1052,21 +1052,15 @@ impl FileSource for ParquetSource { fn apply_expressions( &self, f: &mut dyn FnMut( - &dyn PhysicalExpr, + &Arc, ) -> datafusion_common::Result, ) -> datafusion_common::Result { - // Visit predicate (filter) expression if present - let mut tnr = TreeNodeRecursion::Continue; - if let Some(predicate) = &self.predicate { - tnr = tnr.visit_sibling(|| f(predicate.as_ref()))?; - } - - // Visit projection expressions - for proj_expr in &self.projection { - tnr = tnr.visit_sibling(|| f(proj_expr.expr.as_ref()))?; - } - - Ok(tnr) + datafusion_physical_plan::apply_expression_roots( + self.predicate + .iter() + .chain(self.projection.iter().map(|proj_expr| &proj_expr.expr)), + f, + ) } } diff --git a/datafusion/datasource/src/file.rs b/datafusion/datasource/src/file.rs index 32bee63b54f23..155daaa501c4b 100644 --- a/datafusion/datasource/src/file.rs +++ b/datafusion/datasource/src/file.rs @@ -371,7 +371,7 @@ pub trait FileSource: Any + Send + Sync { /// [`ExecutionPlan::apply_expressions`]: datafusion_physical_plan::ExecutionPlan::apply_expressions fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result; } diff --git a/datafusion/datasource/src/file_scan_config/mod.rs b/datafusion/datasource/src/file_scan_config/mod.rs index dfbf339b6693d..f38a95b1ef879 100644 --- a/datafusion/datasource/src/file_scan_config/mod.rs +++ b/datafusion/datasource/src/file_scan_config/mod.rs @@ -115,7 +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(&dyn PhysicalExpr) -> Result) -> Result { Ok(TreeNodeRecursion::Continue) } +/// # 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()} } @@ -1161,7 +1161,7 @@ impl DataSource for FileScanConfig { fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { // Delegate to the file source self.file_source.apply_expressions(f) @@ -1648,7 +1648,7 @@ mod tests { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } @@ -2920,7 +2920,7 @@ mod tests { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } diff --git a/datafusion/datasource/src/memory.rs b/datafusion/datasource/src/memory.rs index 18f3d05657015..874457e249db6 100644 --- a/datafusion/datasource/src/memory.rs +++ b/datafusion/datasource/src/memory.rs @@ -260,16 +260,15 @@ impl DataSource for MemorySourceConfig { fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - // Visit expressions in sort_information - let mut tnr = TreeNodeRecursion::Continue; - for ordering in &self.sort_information { - for sort_expr in ordering { - tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; - } - } - Ok(tnr) + datafusion_physical_plan::apply_expression_roots( + self.sort_information + .iter() + .flatten() + .map(|sort_expr| &sort_expr.expr), + f, + ) } } diff --git a/datafusion/datasource/src/sink.rs b/datafusion/datasource/src/sink.rs index 023253de38d08..881c5258ca388 100644 --- a/datafusion/datasource/src/sink.rs +++ b/datafusion/datasource/src/sink.rs @@ -234,15 +234,15 @@ impl ExecutionPlan for DataSinkExec { fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - // Apply to sort order requirements if present - if let Some(sort_order) = &self.sort_order { - for req in sort_order.iter() { - f(req.expr.as_ref())?; - } - } - Ok(TreeNodeRecursion::Continue) + 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 diff --git a/datafusion/datasource/src/source.rs b/datafusion/datasource/src/source.rs index f2e1c84df79bd..a4a167137c46d 100644 --- a/datafusion/datasource/src/source.rs +++ b/datafusion/datasource/src/source.rs @@ -239,7 +239,7 @@ pub trait DataSource: Any + Send + Sync + Debug { /// [`ExecutionPlan::apply_expressions`]: datafusion_physical_plan::ExecutionPlan::apply_expressions fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result; /// Injects arbitrary run-time state into this DataSource, returning a new instance @@ -378,7 +378,7 @@ impl ExecutionPlan for DataSourceExec { fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { // Delegate to the underlying data source self.data_source.apply_expressions(f) diff --git a/datafusion/datasource/src/test_util.rs b/datafusion/datasource/src/test_util.rs index 19396411df1ba..20dfae5b3ac79 100644 --- a/datafusion/datasource/src/test_util.rs +++ b/datafusion/datasource/src/test_util.rs @@ -128,7 +128,7 @@ impl FileSource for MockSource { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } diff --git a/datafusion/ffi/src/execution_plan.rs b/datafusion/ffi/src/execution_plan.rs index 0956c6c6dab17..2ba1d0c406c6e 100644 --- a/datafusion/ffi/src/execution_plan.rs +++ b/datafusion/ffi/src/execution_plan.rs @@ -34,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; @@ -51,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, @@ -92,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, @@ -139,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, @@ -307,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, @@ -315,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, } @@ -446,17 +467,19 @@ impl ExecutionPlan for ForeignExecutionPlan { fn apply_expressions( &self, f: &mut dyn FnMut( - &dyn datafusion_physical_plan::PhysicalExpr, + &Arc, ) -> Result, ) -> Result { - // Visit expressions in the output ordering from equivalence properties - let mut tnr = TreeNodeRecursion::Continue; - if let Some(ordering) = self.properties.output_ordering() { - for sort_expr in ordering { - tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; - } - } - Ok(tnr) + 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( @@ -491,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::*; @@ -500,6 +523,7 @@ pub mod tests { pub struct EmptyExec { props: Arc, children: Vec>, + expressions: Vec>, metrics: Option, statistics: Option, } @@ -514,6 +538,7 @@ pub mod tests { Boundedness::Bounded, )), children: Vec::default(), + expressions: Vec::default(), metrics: None, statistics: None, } @@ -528,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 { @@ -560,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(), })) @@ -589,18 +623,9 @@ pub mod tests { fn apply_expressions( &self, - f: &mut dyn FnMut( - &dyn datafusion_physical_plan::PhysicalExpr, - ) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - // Visit expressions in the output ordering from equivalence properties - let mut tnr = TreeNodeRecursion::Continue; - if let Some(ordering) = self.props.output_ordering() { - for sort_expr in ordering { - tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; - } - } - Ok(tnr) + datafusion_physical_plan::apply_expression_roots(&self.expressions, f) } } @@ -633,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 2f7d7d44dfc82..f14ee2f3efd60 100644 --- a/datafusion/ffi/src/tests/async_provider.rs +++ b/datafusion/ffi/src/tests/async_provider.rs @@ -232,17 +232,17 @@ impl ExecutionPlan for AsyncTestExecutionPlan { fn apply_expressions( &self, f: &mut dyn FnMut( - &dyn datafusion_physical_plan::PhysicalExpr, + &Arc, ) -> Result, ) -> Result { - // Visit expressions in the output ordering from equivalence properties - let mut tnr = TreeNodeRecursion::Continue; - if let Some(ordering) = self.properties.output_ordering() { - for sort_expr in ordering { - tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; - } - } - Ok(tnr) + datafusion_physical_plan::apply_expression_roots( + self.properties + .output_ordering() + .into_iter() + .flatten() + .map(|sort_expr| &sort_expr.expr), + f, + ) } } 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 102e21a4853a4..10da9e4e75174 100644 --- a/datafusion/physical-optimizer/src/ensure_coop.rs +++ b/datafusion/physical-optimizer/src/ensure_coop.rs @@ -349,7 +349,7 @@ mod tests { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } diff --git a/datafusion/physical-optimizer/src/output_requirements.rs b/datafusion/physical-optimizer/src/output_requirements.rs index 810ba27e2b82b..59a92f971974e 100644 --- a/datafusion/physical-optimizer/src/output_requirements.rs +++ b/datafusion/physical-optimizer/src/output_requirements.rs @@ -330,34 +330,28 @@ impl ExecutionPlan for OutputRequirementExec { fn apply_expressions( &self, f: &mut dyn FnMut( - &dyn datafusion_physical_expr_common::physical_expr::PhysicalExpr, + &Arc, ) -> Result, ) -> Result { - // Visit expressions in order_requirement - let mut tnr = TreeNodeRecursion::Continue; - if let Some(order_reqs) = &self.order_requirement { - let lexes = match order_reqs { - OrderingRequirements::Hard(alternatives) => alternatives, - OrderingRequirements::Soft(alternatives) => alternatives, - }; - for lex in lexes { - for sort_expr in lex { - tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; - } - } - } - - // Visit expressions in dist_requirement if it's key partitioned #[expect(deprecated)] - if let Distribution::HashPartitioned(exprs) - | Distribution::KeyPartitioned(exprs) = &self.dist_requirement + let distribution = if let Distribution::HashPartitioned(exprs) + | Distribution::KeyPartitioned(exprs) = + &self.dist_requirement { - for expr in exprs { - tnr = tnr.visit_sibling(|| f(expr.as_ref()))?; - } - } - - Ok(tnr) + 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) } } diff --git a/datafusion/physical-plan/benches/compute_statistics.rs b/datafusion/physical-plan/benches/compute_statistics.rs index afb4ef38c82d6..93c95ea4ba099 100644 --- a/datafusion/physical-plan/benches/compute_statistics.rs +++ b/datafusion/physical-plan/benches/compute_statistics.rs @@ -100,7 +100,7 @@ impl ExecutionPlan for BenchLeaf { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index b36c2096f709b..a6ed7c444ff90 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -1963,32 +1963,29 @@ impl ExecutionPlan for AggregateExec { fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - // Apply to group by expressions - let mut tnr = TreeNodeRecursion::Continue; - for expr in self.group_by.input_exprs() { - tnr = tnr.visit_sibling(|| f(expr.as_ref()))?; - } - - // Apply to aggregate expressions - for aggr in self.aggr_expr.iter() { - for expr in aggr.expressions() { - tnr = tnr.visit_sibling(|| f(expr.as_ref()))?; - } - } - - // Apply to filter expressions (FILTER WHERE clauses) - for filter in self.filter_expr.iter().flatten() { - tnr = tnr.visit_sibling(|| f(filter.as_ref()))?; - } - - // Apply to dynamic filter expression if present - if let Some(dyn_filter) = &self.dynamic_filter { - tnr = tnr.visit_sibling(|| f(dyn_filter.filter.as_ref()))?; - } - - Ok(tnr) + 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( @@ -2128,25 +2125,9 @@ 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. + // HACK: A child reply of `PushedDown::No` can mean it cannot apply the + // row-level filter but still retains it for statistics pruning. + // Use the reference count to detect whether a child retained it. // Issue: let child_accepts_dyn_filter = Arc::strong_count(dyn_filter) > 1; @@ -3597,7 +3578,7 @@ mod tests { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } diff --git a/datafusion/physical-plan/src/analyze.rs b/datafusion/physical-plan/src/analyze.rs index 67e221c3989d7..0ab28f56e169c 100644 --- a/datafusion/physical-plan/src/analyze.rs +++ b/datafusion/physical-plan/src/analyze.rs @@ -223,7 +223,7 @@ impl ExecutionPlan for AnalyzeExec { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } diff --git a/datafusion/physical-plan/src/async_func.rs b/datafusion/physical-plan/src/async_func.rs index c421b4f4a28ab..1696e2a677ccf 100644 --- a/datafusion/physical-plan/src/async_func.rs +++ b/datafusion/physical-plan/src/async_func.rs @@ -155,7 +155,7 @@ impl ExecutionPlan for AsyncFuncExec { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } diff --git a/datafusion/physical-plan/src/buffer.rs b/datafusion/physical-plan/src/buffer.rs index 6478b68901eb0..a1c3c7ea01658 100644 --- a/datafusion/physical-plan/src/buffer.rs +++ b/datafusion/physical-plan/src/buffer.rs @@ -162,7 +162,7 @@ impl ExecutionPlan for BufferExec { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } diff --git a/datafusion/physical-plan/src/coalesce_batches.rs b/datafusion/physical-plan/src/coalesce_batches.rs index 99a0dd6891536..511e5d793b873 100644 --- a/datafusion/physical-plan/src/coalesce_batches.rs +++ b/datafusion/physical-plan/src/coalesce_batches.rs @@ -176,7 +176,7 @@ impl ExecutionPlan for CoalesceBatchesExec { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } diff --git a/datafusion/physical-plan/src/coalesce_partitions.rs b/datafusion/physical-plan/src/coalesce_partitions.rs index b7d1206dbeabe..5cd7a707c23b3 100644 --- a/datafusion/physical-plan/src/coalesce_partitions.rs +++ b/datafusion/physical-plan/src/coalesce_partitions.rs @@ -146,7 +146,7 @@ impl ExecutionPlan for CoalescePartitionsExec { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } diff --git a/datafusion/physical-plan/src/coop.rs b/datafusion/physical-plan/src/coop.rs index a74c129ac9a30..dc7d98891e114 100644 --- a/datafusion/physical-plan/src/coop.rs +++ b/datafusion/physical-plan/src/coop.rs @@ -270,7 +270,7 @@ impl ExecutionPlan for CooperativeExec { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } diff --git a/datafusion/physical-plan/src/display.rs b/datafusion/physical-plan/src/display.rs index ff82c501045c2..89029aea87546 100644 --- a/datafusion/physical-plan/src/display.rs +++ b/datafusion/physical-plan/src/display.rs @@ -1554,7 +1554,7 @@ mod tests { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } @@ -1687,9 +1687,7 @@ mod tests { use crate::metrics::{Count, Metric, MetricValue, MetricsSet, Time}; use crate::{DisplayFormatType, ExecutionPlan, PlanProperties}; use datafusion_common::Result; - use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_execution::{SendableRecordBatchStream, TaskContext}; - use datafusion_physical_expr::PhysicalExpr; // Wrap `sample_plan()` with an adapter node that exposes a // hand-crafted `MetricsSet` so we can assert the PG key mapping @@ -1720,9 +1718,14 @@ mod tests { } fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, - ) -> Result { - Ok(TreeNodeRecursion::Continue) + f: &mut dyn FnMut( + &Arc, + ) -> Result< + datafusion_common::tree_node::TreeNodeRecursion, + >, + ) -> Result + { + self.inner.apply_expressions(f) } fn with_new_children( self: Arc, diff --git a/datafusion/physical-plan/src/empty.rs b/datafusion/physical-plan/src/empty.rs index 64abe8e8e00c8..bd91ec742d48c 100644 --- a/datafusion/physical-plan/src/empty.rs +++ b/datafusion/physical-plan/src/empty.rs @@ -122,7 +122,7 @@ impl ExecutionPlan for EmptyExec { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } diff --git a/datafusion/physical-plan/src/execution_plan.rs b/datafusion/physical-plan/src/execution_plan.rs index 1f24b0230adc5..0b85370f9d88c 100644 --- a/datafusion/physical-plan/src/execution_plan.rs +++ b/datafusion/physical-plan/src/execution_plan.rs @@ -41,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}; @@ -69,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 @@ -255,7 +278,8 @@ pub trait ExecutionPlan: Any + Debug + DisplayAs + Send + Sync { /// 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 skip any remaining expressions (though typically all expressions are visited). + /// 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 @@ -284,43 +308,29 @@ pub trait ExecutionPlan: Any + Debug + DisplayAs + Send + Sync { /// ```ignore /// fn apply_expressions( /// &self, - /// _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + /// _f: &mut dyn FnMut(&Arc) -> Result, /// ) -> Result { /// Ok(TreeNodeRecursion::Continue) /// } /// ``` /// - /// ## Node with a single expression (e.g., FilterExec) - /// ```ignore - /// fn apply_expressions( - /// &self, - /// f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, - /// ) -> Result { - /// f(self.predicate.as_ref()) - /// } - /// ``` - /// - /// ## Node with multiple expressions (e.g., ProjectionExec, JoinExec) + /// ## Node with expressions (e.g., FilterExec, ProjectionExec) /// - /// Use [`TreeNodeRecursion::visit_sibling`] when iterating over multiple - /// expressions. This correctly propagates [`TreeNodeRecursion::Stop`]: if - /// `f` returns `Stop` for an earlier expression, `visit_sibling` short-circuits - /// and skips the remaining ones. + /// 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(&dyn PhysicalExpr) -> Result, + /// f: &mut dyn FnMut(&Arc) -> Result, /// ) -> Result { - /// let mut tnr = TreeNodeRecursion::Continue; - /// for expr in &self.expressions { - /// tnr = tnr.visit_sibling(|| f(expr.as_ref()))?; - /// } - /// Ok(tnr) + /// apply_expression_roots([&self.predicate], f) /// } /// ``` fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result; /// Returns a new `ExecutionPlan` where all existing children were replaced @@ -1850,7 +1860,7 @@ mod tests { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } @@ -1913,7 +1923,7 @@ mod tests { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } @@ -1970,7 +1980,7 @@ mod tests { fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { self.0.apply_expressions(f) } @@ -2036,7 +2046,7 @@ mod tests { } fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } @@ -2105,7 +2115,7 @@ mod tests { } fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } @@ -2202,7 +2212,7 @@ mod tests { } fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } @@ -2352,6 +2362,7 @@ mod tests { #[derive(Debug)] struct MultiExprExec { exprs: Vec>, + children: Vec>, } impl DisplayAs for MultiExprExec { @@ -2374,7 +2385,7 @@ mod tests { } fn children(&self) -> Vec<&Arc> { - vec![] + self.children.iter().collect() } fn with_new_children( @@ -2386,13 +2397,9 @@ mod tests { fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - let mut tnr = TreeNodeRecursion::Continue; - for expr in &self.exprs { - tnr = tnr.visit_sibling(|| f(expr.as_ref()))?; - } - Ok(tnr) + apply_expression_roots(&self.exprs, f) } fn execute( @@ -2424,6 +2431,7 @@ mod tests { 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| { @@ -2438,6 +2446,7 @@ mod tests { 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| { @@ -2450,6 +2459,69 @@ mod tests { 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 60e3fcd350aa2..72ceae81f3724 100644 --- a/datafusion/physical-plan/src/explain.rs +++ b/datafusion/physical-plan/src/explain.rs @@ -119,7 +119,7 @@ impl ExecutionPlan for ExplainExec { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } diff --git a/datafusion/physical-plan/src/filter.rs b/datafusion/physical-plan/src/filter.rs index ba37f5f4d014b..c6800283c5c23 100644 --- a/datafusion/physical-plan/src/filter.rs +++ b/datafusion/physical-plan/src/filter.rs @@ -538,9 +538,9 @@ impl ExecutionPlan for FilterExec { fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - f(self.predicate.as_ref()) + crate::apply_expression_roots([&self.predicate], f) } fn maintains_input_order(&self) -> Vec { diff --git a/datafusion/physical-plan/src/joins/cross_join.rs b/datafusion/physical-plan/src/joins/cross_join.rs index 99f14a49e3445..d56ef7d0881dc 100644 --- a/datafusion/physical-plan/src/joins/cross_join.rs +++ b/datafusion/physical-plan/src/joins/cross_join.rs @@ -294,7 +294,7 @@ impl ExecutionPlan for CrossJoinExec { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { // CrossJoin has no join conditions or expressions Ok(TreeNodeRecursion::Continue) diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index 42d7acfbb4157..190311979f133 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -1333,26 +1333,21 @@ impl ExecutionPlan for HashJoinExec { fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - // Apply to join key expressions from both sides - let mut tnr = TreeNodeRecursion::Continue; - for (left, right) in &self.on { - tnr = tnr.visit_sibling(|| f(left.as_ref()))?; - tnr = tnr.visit_sibling(|| f(right.as_ref()))?; - } - - // Apply to join filter expression if present - if let Some(filter) = &self.filter { - tnr = tnr.visit_sibling(|| f(filter.expression().as_ref()))?; - } - - // Apply to dynamic filter expression if present - if let Some(df) = &self.dynamic_filter { - tnr = tnr.visit_sibling(|| f(df.filter.as_ref()))?; - } - - Ok(tnr) + 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. @@ -2488,7 +2483,7 @@ mod tests { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } diff --git a/datafusion/physical-plan/src/joins/nested_loop_join.rs b/datafusion/physical-plan/src/joins/nested_loop_join.rs index 82e8eff7b7abc..7069a8b44805c 100644 --- a/datafusion/physical-plan/src/joins/nested_loop_join.rs +++ b/datafusion/physical-plan/src/joins/nested_loop_join.rs @@ -565,13 +565,13 @@ impl ExecutionPlan for NestedLoopJoinExec { fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn crate::PhysicalExpr) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { // Apply to join filter expressions if present - if let Some(filter) = &self.filter { - f(filter.expression().as_ref())?; - } - Ok(TreeNodeRecursion::Continue) + crate::apply_expression_roots( + self.filter.iter().map(|filter| filter.expression()), + f, + ) } fn with_new_children( 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 a8e7959376ff4..b60ec1c784de5 100644 --- a/datafusion/physical-plan/src/joins/piecewise_merge_join/exec.rs +++ b/datafusion/physical-plan/src/joins/piecewise_merge_join/exec.rs @@ -485,10 +485,10 @@ impl ExecutionPlan for PiecewiseMergeJoinExec { fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { // Apply to the two expressions being compared in the range predicate - f(self.on.0.as_ref())?.visit_sibling(|| f(self.on.1.as_ref())) + crate::apply_expression_roots([&self.on.0, &self.on.1], f) } fn required_input_distribution(&self) -> Vec { 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 4d4d962f73467..4d82eed547a46 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs @@ -444,19 +444,11 @@ impl ExecutionPlan for SortMergeJoinExec { fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn crate::PhysicalExpr) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - // Apply to join keys from both sides - let mut tnr = TreeNodeRecursion::Continue; - for (left, right) in &self.on { - tnr = tnr.visit_sibling(|| f(left.as_ref()))?; - tnr = tnr.visit_sibling(|| f(right.as_ref()))?; - } - // Apply to join filter expressions if present - if let Some(filter) = &self.filter { - tnr = tnr.visit_sibling(|| f(filter.expression().as_ref()))?; - } - Ok(tnr) + 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( diff --git a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs index e2f061587378c..4b161a7b47473 100644 --- a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs +++ b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs @@ -455,19 +455,11 @@ impl ExecutionPlan for SymmetricHashJoinExec { fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn crate::PhysicalExpr) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - // Apply to join keys from both sides - let mut tnr = TreeNodeRecursion::Continue; - for (left, right) in &self.on { - tnr = tnr.visit_sibling(|| f(left.as_ref()))?; - tnr = tnr.visit_sibling(|| f(right.as_ref()))?; - } - // Apply to join filter expressions if present - if let Some(filter) = &self.filter { - tnr = tnr.visit_sibling(|| f(filter.expression().as_ref()))?; - } - Ok(tnr) + 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( 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 e632bcdd349a1..0e588ba1aacc3 100644 --- a/datafusion/physical-plan/src/limit.rs +++ b/datafusion/physical-plan/src/limit.rs @@ -170,16 +170,15 @@ impl ExecutionPlan for GlobalLimitExec { fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - // Apply to required ordering expressions if present - let mut tnr = TreeNodeRecursion::Continue; - if let Some(ordering) = &self.required_ordering { - for sort_expr in ordering { - tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; - } - } - Ok(tnr) + crate::apply_expression_roots( + self.required_ordering + .iter() + .flatten() + .map(|sort_expr| &sort_expr.expr), + f, + ) } fn with_new_children( @@ -415,16 +414,15 @@ impl ExecutionPlan for LocalLimitExec { fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - // Apply to required ordering expressions if present - let mut tnr = TreeNodeRecursion::Continue; - if let Some(ordering) = &self.required_ordering { - for sort_expr in ordering { - tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; - } - } - Ok(tnr) + crate::apply_expression_roots( + self.required_ordering + .iter() + .flatten() + .map(|sort_expr| &sort_expr.expr), + f, + ) } fn with_new_children( diff --git a/datafusion/physical-plan/src/memory.rs b/datafusion/physical-plan/src/memory.rs index e172ef4463ec4..eb141b8c70d5e 100644 --- a/datafusion/physical-plan/src/memory.rs +++ b/datafusion/physical-plan/src/memory.rs @@ -314,7 +314,7 @@ impl ExecutionPlan for LazyMemoryExec { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } diff --git a/datafusion/physical-plan/src/operator_statistics/mod.rs b/datafusion/physical-plan/src/operator_statistics/mod.rs index 33226392c7e29..ec54201e7b3d5 100644 --- a/datafusion/physical-plan/src/operator_statistics/mod.rs +++ b/datafusion/physical-plan/src/operator_statistics/mod.rs @@ -1120,7 +1120,7 @@ mod tests { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } @@ -1229,7 +1229,7 @@ mod tests { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } diff --git a/datafusion/physical-plan/src/placeholder_row.rs b/datafusion/physical-plan/src/placeholder_row.rs index 7fcfcaf86b6cb..de07529bab70c 100644 --- a/datafusion/physical-plan/src/placeholder_row.rs +++ b/datafusion/physical-plan/src/placeholder_row.rs @@ -140,7 +140,7 @@ impl ExecutionPlan for PlaceholderRowExec { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } diff --git a/datafusion/physical-plan/src/projection.rs b/datafusion/physical-plan/src/projection.rs index 2e915b283f897..cf1db9328ada2 100644 --- a/datafusion/physical-plan/src/projection.rs +++ b/datafusion/physical-plan/src/projection.rs @@ -304,13 +304,16 @@ impl ExecutionPlan for ProjectionExec { fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - let mut tnr = TreeNodeRecursion::Continue; - for proj_expr in self.projector.projection().as_ref().iter() { - tnr = tnr.visit_sibling(|| f(proj_expr.expr.as_ref()))?; - } - Ok(tnr) + crate::apply_expression_roots( + self.projector + .projection() + .as_ref() + .iter() + .map(|proj_expr| &proj_expr.expr), + f, + ) } fn with_new_children( diff --git a/datafusion/physical-plan/src/recursive_query.rs b/datafusion/physical-plan/src/recursive_query.rs index a02ec38a7431f..4c6f0493adf40 100644 --- a/datafusion/physical-plan/src/recursive_query.rs +++ b/datafusion/physical-plan/src/recursive_query.rs @@ -157,7 +157,7 @@ impl ExecutionPlan for RecursiveQueryExec { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs index 63469cfd6a9cc..dce5805e16bd4 100644 --- a/datafusion/physical-plan/src/repartition/mod.rs +++ b/datafusion/physical-plan/src/repartition/mod.rs @@ -1340,22 +1340,16 @@ impl ExecutionPlan for RepartitionExec { fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - let exprs: Vec<_> = match self.partitioning() { - Partitioning::Hash(exprs, _) => exprs.iter().collect(), - Partitioning::Range(range) => range - .ordering() - .iter() - .map(|sort_expr| &sort_expr.expr) - .collect(), - _ => return Ok(TreeNodeRecursion::Continue), - }; - let mut tnr = TreeNodeRecursion::Continue; - for expr in exprs { - tnr = tnr.visit_sibling(|| f(expr.as_ref()))?; + 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), } - Ok(tnr) } fn with_new_children( @@ -3068,7 +3062,7 @@ mod tests { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } @@ -4307,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 08a1a1d430770..ee3c2e5d077f5 100644 --- a/datafusion/physical-plan/src/scalar_subquery.rs +++ b/datafusion/physical-plan/src/scalar_subquery.rs @@ -228,7 +228,7 @@ impl ExecutionPlan for ScalarSubqueryExec { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } @@ -464,7 +464,7 @@ mod tests { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } diff --git a/datafusion/physical-plan/src/sorts/partial_sort.rs b/datafusion/physical-plan/src/sorts/partial_sort.rs index 724dfd3c16c74..5a7c2b2e4557e 100644 --- a/datafusion/physical-plan/src/sorts/partial_sort.rs +++ b/datafusion/physical-plan/src/sorts/partial_sort.rs @@ -310,13 +310,12 @@ impl ExecutionPlan for PartialSortExec { fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - let mut tnr = TreeNodeRecursion::Continue; - for sort_expr in &self.expr { - tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; - } - Ok(tnr) + crate::apply_expression_roots( + self.expr.iter().map(|sort_expr| &sort_expr.expr), + f, + ) } fn with_new_children( diff --git a/datafusion/physical-plan/src/sorts/partitioned_topk.rs b/datafusion/physical-plan/src/sorts/partitioned_topk.rs index 897eadbb60e51..78dd9b9696d45 100644 --- a/datafusion/physical-plan/src/sorts/partitioned_topk.rs +++ b/datafusion/physical-plan/src/sorts/partitioned_topk.rs @@ -382,13 +382,12 @@ impl ExecutionPlan for PartitionedTopKExec { fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - let mut tnr = TreeNodeRecursion::Continue; - for sort_expr in &self.expr { - tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; - } - Ok(tnr) + crate::apply_expression_roots( + self.expr.iter().map(|sort_expr| &sort_expr.expr), + f, + ) } fn execute( diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs index e489369be980d..b11db4633e979 100644 --- a/datafusion/physical-plan/src/sorts/sort.rs +++ b/datafusion/physical-plan/src/sorts/sort.rs @@ -1277,21 +1277,19 @@ impl ExecutionPlan for SortExec { fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - // Apply to sort expressions - let mut tnr = TreeNodeRecursion::Continue; - for sort_expr in &self.expr { - tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; - } - - // Apply to dynamic filter expression if present (when fetch is Some, TopK mode) - if let Some(filter) = &self.filter { - let filter_guard = filter.read(); - tnr = tnr.visit_sibling(|| f(filter_guard.expr().as_ref()))?; - } - - Ok(tnr) + 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 { @@ -1782,7 +1780,7 @@ mod tests { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } diff --git a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs index 2de468455408b..ac6f5d18cd2ff 100644 --- a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs +++ b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs @@ -285,13 +285,12 @@ impl ExecutionPlan for SortPreservingMergeExec { fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - let mut tnr = TreeNodeRecursion::Continue; - for sort_expr in &self.expr { - tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; - } - Ok(tnr) + crate::apply_expression_roots( + self.expr.iter().map(|sort_expr| &sort_expr.expr), + f, + ) } fn with_new_children( @@ -1588,7 +1587,7 @@ mod tests { } fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } diff --git a/datafusion/physical-plan/src/streaming.rs b/datafusion/physical-plan/src/streaming.rs index 4e792c8f5c504..a82f8d9441e95 100644 --- a/datafusion/physical-plan/src/streaming.rs +++ b/datafusion/physical-plan/src/streaming.rs @@ -275,7 +275,7 @@ impl ExecutionPlan for StreamingTableExec { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } diff --git a/datafusion/physical-plan/src/test.rs b/datafusion/physical-plan/src/test.rs index 491f38b9eccec..fd66cc9e05b2f 100644 --- a/datafusion/physical-plan/src/test.rs +++ b/datafusion/physical-plan/src/test.rs @@ -143,16 +143,15 @@ impl ExecutionPlan for TestMemoryExec { fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - // Apply to all sort information orderings - let mut tnr = TreeNodeRecursion::Continue; - for ordering in &self.sort_information { - for sort_expr in ordering { - tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; - } - } - Ok(tnr) + crate::apply_expression_roots( + self.sort_information + .iter() + .flatten() + .map(|sort_expr| &sort_expr.expr), + f, + ) } fn with_new_children( diff --git a/datafusion/physical-plan/src/test/exec.rs b/datafusion/physical-plan/src/test/exec.rs index 819259bcf8b8f..22e6859e5ee06 100644 --- a/datafusion/physical-plan/src/test/exec.rs +++ b/datafusion/physical-plan/src/test/exec.rs @@ -198,7 +198,7 @@ impl ExecutionPlan for MockExec { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } @@ -442,7 +442,7 @@ impl ExecutionPlan for BarrierExec { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } @@ -585,7 +585,7 @@ impl ExecutionPlan for ErrorExec { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } @@ -671,7 +671,7 @@ impl ExecutionPlan for StatisticsExec { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } @@ -789,7 +789,7 @@ impl ExecutionPlan for BlockingExec { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } @@ -931,7 +931,7 @@ impl ExecutionPlan for PanicExec { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } diff --git a/datafusion/physical-plan/src/union.rs b/datafusion/physical-plan/src/union.rs index e1eb46e2fbe73..9be852e6cc649 100644 --- a/datafusion/physical-plan/src/union.rs +++ b/datafusion/physical-plan/src/union.rs @@ -255,7 +255,7 @@ impl ExecutionPlan for UnionExec { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } @@ -634,7 +634,7 @@ impl ExecutionPlan for InterleaveExec { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } diff --git a/datafusion/physical-plan/src/unnest.rs b/datafusion/physical-plan/src/unnest.rs index 0372a7adf827c..f3d75b7e7d263 100644 --- a/datafusion/physical-plan/src/unnest.rs +++ b/datafusion/physical-plan/src/unnest.rs @@ -230,7 +230,7 @@ impl ExecutionPlan for UnnestExec { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } 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 a7beb47395652..938f6138cd760 100644 --- a/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs +++ b/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs @@ -314,15 +314,17 @@ impl ExecutionPlan for BoundedWindowAggExec { fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - let mut tnr = TreeNodeRecursion::Continue; - for window_expr in &self.window_expr { - for expr in window_expr.expressions() { - tnr = tnr.visit_sibling(|| f(expr.as_ref()))?; - } - } - Ok(tnr) + 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> { diff --git a/datafusion/physical-plan/src/windows/window_agg_exec.rs b/datafusion/physical-plan/src/windows/window_agg_exec.rs index 2240eaf6aff90..31ff5189253cf 100644 --- a/datafusion/physical-plan/src/windows/window_agg_exec.rs +++ b/datafusion/physical-plan/src/windows/window_agg_exec.rs @@ -217,15 +217,17 @@ impl ExecutionPlan for WindowAggExec { fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - let mut tnr = TreeNodeRecursion::Continue; - for window_expr in &self.window_expr { - for expr in window_expr.expressions() { - tnr = tnr.visit_sibling(|| f(expr.as_ref()))?; - } - } - Ok(tnr) + 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 { diff --git a/datafusion/physical-plan/src/work_table.rs b/datafusion/physical-plan/src/work_table.rs index 0ae39f2dc0b6a..83cd0a15a6d26 100644 --- a/datafusion/physical-plan/src/work_table.rs +++ b/datafusion/physical-plan/src/work_table.rs @@ -189,7 +189,7 @@ impl ExecutionPlan for WorkTableExec { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } diff --git a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs index b003e30cd9f5c..5f2ed3573fbb0 100644 --- a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs @@ -324,7 +324,7 @@ impl ExecutionPlan for DowncastDelegatingExec { fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { self.inner.apply_expressions(f) } @@ -4605,13 +4605,9 @@ impl ExecutionPlan for CustomExecWithExprs { fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - let mut tnr = TreeNodeRecursion::Continue; - for expr in &self.exprs { - tnr = tnr.visit_sibling(|| f(expr.as_ref()))?; - } - Ok(tnr) + datafusion_physical_plan::apply_expression_roots(&self.exprs, f) } fn with_new_children( diff --git a/docs/source/library-user-guide/custom-table-providers.md b/docs/source/library-user-guide/custom-table-providers.md index 81b2d131e65c3..c6a316aa74b94 100644 --- a/docs/source/library-user-guide/custom-table-providers.md +++ b/docs/source/library-user-guide/custom-table-providers.md @@ -766,7 +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(&dyn PhysicalExpr) -> Result) -> Result { Ok(TreeNodeRecursion::Continue) } +# fn apply_expressions(&self, _f: &mut dyn FnMut(&Arc) -> Result) -> Result { Ok(TreeNodeRecursion::Continue) } # } ``` @@ -913,7 +913,7 @@ impl ExecutionPlan for CountingExec { # fn apply_expressions( # &self, -# _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, +# _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 349e15a2188a3..b067f6e782fd3 100644 --- a/docs/source/library-user-guide/upgrading/54.0.0.md +++ b/docs/source/library-user-guide/upgrading/54.0.0.md @@ -176,14 +176,17 @@ where string types are preferred (`UNION`, `CASE THEN/ELSE`, `NVL2`). **Migration guide:** -Add `apply_expressions` to your implementation. Call `f` on each top-level `PhysicalExpr` your node owns, using `visit_sibling` to correctly propagate `TreeNodeRecursion`: +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(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } @@ -194,9 +197,9 @@ fn apply_expressions( ```rust,ignore fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - f(self.predicate.as_ref()) + apply_expression_roots([&self.predicate], f) } ``` @@ -205,30 +208,9 @@ fn apply_expressions( ```rust,ignore fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - let mut tnr = TreeNodeRecursion::Continue; - for expr in &self.expressions { - tnr = tnr.visit_sibling(|| f(expr.as_ref()))?; - } - Ok(tnr) -} -``` - -**Node whose only expressions are in `output_ordering()` (e.g. a synthetic test node with no owned expression fields):** - -```rust,ignore -fn apply_expressions( - &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, -) -> Result { - let mut tnr = TreeNodeRecursion::Continue; - if let Some(ordering) = self.cache.output_ordering() { - for sort_expr in ordering { - tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; - } - } - Ok(tnr) + apply_expression_roots(&self.expressions, f) } ``` From 5514194fa7362e7fbc2e3f891779e5df90f57322 Mon Sep 17 00:00:00 2001 From: Jayant Shrivastava Date: Thu, 30 Jul 2026 22:06:56 +0000 Subject: [PATCH 3/3] Use apply_expressions for aggregate dynamic filters --- .../physical-plan/src/aggregates/mod.rs | 76 +++++++++++++++++-- .../tests/cases/roundtrip_physical_plan.rs | 55 ++++++++++++++ 2 files changed, 124 insertions(+), 7 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index a6ed7c444ff90..27dbd62884e02 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -176,7 +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::TreeNodeRecursion; +use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion}; use datafusion_common::{ Constraint, Constraints, Result, ScalarValue, assert_eq_or_internal_err, internal_err, not_impl_err, @@ -764,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. @@ -2125,11 +2151,12 @@ impl ExecutionPlan for AggregateExec { if phase == FilterPushdownPhase::Post && let Some(dyn_filter) = &self.dynamic_filter { - // HACK: A child reply of `PushedDown::No` can mean it cannot apply the - // row-level filter but still retains it for statistics pruning. - // Use the reference count to detect whether a child retained it. - // 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 @@ -2518,6 +2545,8 @@ impl AggregateExec { })?; aggregate.with_dynamic_filter_expr(dynamic_filter)? } else { + let mut aggregate = aggregate; + aggregate.dynamic_filter = None; aggregate }; @@ -3021,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; @@ -3056,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; @@ -7598,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/proto/tests/cases/roundtrip_physical_plan.rs b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs index 5f2ed3573fbb0..c218a8c1810bb 100644 --- a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs @@ -4485,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]