From 052f2b97e495514166bd2167ca3ddd9f41768ee8 Mon Sep 17 00:00:00 2001 From: Jayant Shrivastava Date: Mon, 3 Aug 2026 20:37:11 +0000 Subject: [PATCH] add ExecutionPlan::dynamic_filters() method --- .../physical_optimizer/filter_pushdown.rs | 73 +++++++++++++++- datafusion/ffi/src/execution_plan.rs | 69 ++++++++++++++- datafusion/ffi/src/tests/mod.rs | 14 +++ datafusion/ffi/tests/ffi_execution_plan.rs | 18 ++++ .../physical-plan/src/aggregates/mod.rs | 13 +++ .../physical-plan/src/execution_plan.rs | 85 ++++++++++++++++++- .../physical-plan/src/joins/hash_join/exec.rs | 14 +++ datafusion/physical-plan/src/sorts/sort.rs | 13 +++ 8 files changed, 295 insertions(+), 4 deletions(-) diff --git a/datafusion/core/tests/physical_optimizer/filter_pushdown.rs b/datafusion/core/tests/physical_optimizer/filter_pushdown.rs index 7593fe351548e..3020305e38fab 100644 --- a/datafusion/core/tests/physical_optimizer/filter_pushdown.rs +++ b/datafusion/core/tests/physical_optimizer/filter_pushdown.rs @@ -34,7 +34,11 @@ use datafusion::{ scalar::ScalarValue, }; use datafusion_catalog::memory::DataSourceExec; -use datafusion_common::config::ConfigOptions; +use datafusion_common::{ + JoinType, + config::ConfigOptions, + tree_node::{TreeNode, TreeNodeRecursion}, +}; use datafusion_datasource::{ PartitionedFile, file_groups::FileGroup, file_scan_config::FileScanConfigBuilder, }; @@ -60,6 +64,7 @@ use datafusion_physical_plan::{ coalesce_partitions::CoalescePartitionsExec, collect, filter::{FilterExec, FilterExecBuilder}, + joins::{HashJoinExec, PartitionMode}, projection::ProjectionExec, repartition::RepartitionExec, sorts::sort::SortExec, @@ -3085,6 +3090,72 @@ async fn test_filter_with_projection_pushdown() { assert_batches_eq!(expected, &result); } +#[test] +fn test_discover_dynamic_expression_producers() { + fn producer_count(plan: &Arc) -> usize { + let mut count = 0; + plan.apply(|node| { + count += node.dynamic_expressions().len(); + Ok(TreeNodeRecursion::Continue) + }) + .expect("plan traversal should succeed"); + count + } + + 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(Arc::clone(&build_schema)) + .with_support(true) + .with_batches(vec![ + record_batch!(("a", Utf8, ["foo", "bar"]), ("b", Int32, [1, 2])).unwrap(), + ]) + .build(); + + 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(Arc::clone(&probe_schema)) + .with_support(true) + .with_batches(vec![ + record_batch!( + ("a", Utf8, ["foo", "bar", "baz", "qux"]), + ("c", Float64, [1.0, 2.0, 3.0, 4.0]) + ) + .unwrap(), + ]) + .build(); + + 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; + assert_eq!(producer_count(&plan), 0); + + 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(); + assert_eq!(producer_count(&optimized_plan), 1); +} + // ==== Filter pushdown through SortExec tests ==== /// FilterExec above a plain SortExec (no fetch) should be pushed below it. diff --git a/datafusion/ffi/src/execution_plan.rs b/datafusion/ffi/src/execution_plan.rs index 087a351b697cc..320c87c96f01b 100644 --- a/datafusion/ffi/src/execution_plan.rs +++ b/datafusion/ffi/src/execution_plan.rs @@ -33,6 +33,7 @@ use tokio::runtime::Handle; use crate::config::FFI_ConfigOptions; use crate::execution::FFI_TaskContext; +use crate::physical_expr::FFI_PhysicalExpr; use crate::physical_expr::metrics::FFI_MetricsSet; use crate::plan_properties::FFI_PlanProperties; use crate::record_batch_stream::FFI_RecordBatchStream; @@ -50,6 +51,9 @@ pub struct FFI_ExecutionPlan { /// Return a vector of children plans pub children: unsafe extern "C" fn(plan: &Self) -> SVec, + /// Return the dynamic expressions produced by this plan node. + pub dynamic_expressions: unsafe extern "C" fn(plan: &Self) -> SVec, + pub with_new_children: unsafe extern "C" fn(plan: &Self, children: SVec) -> FFI_Result, @@ -138,6 +142,16 @@ unsafe extern "C" fn children_fn_wrapper( .collect() } +unsafe extern "C" fn dynamic_expressions_fn_wrapper( + plan: &FFI_ExecutionPlan, +) -> SVec { + plan.inner() + .dynamic_expressions() + .into_iter() + .map(FFI_PhysicalExpr::from) + .collect() +} + unsafe extern "C" fn with_new_children_fn_wrapper( plan: &FFI_ExecutionPlan, children: SVec, @@ -306,6 +320,7 @@ impl FFI_ExecutionPlan { Self { properties: properties_fn_wrapper, children: children_fn_wrapper, + dynamic_expressions: dynamic_expressions_fn_wrapper, with_new_children: with_new_children_fn_wrapper, name: name_fn_wrapper, execute: execute_fn_wrapper, @@ -442,6 +457,15 @@ impl ExecutionPlan for ForeignExecutionPlan { } } + fn dynamic_expressions( + &self, + ) -> Vec> { + unsafe { (self.plan.dynamic_expressions)(&self.plan) } + .iter() + .map(>::from) + .collect() + } + fn repartitioned( &self, target_partitions: usize, @@ -474,8 +498,10 @@ impl ExecutionPlan for ForeignExecutionPlan { #[cfg(any(test, feature = "integration-tests"))] pub mod tests { - use datafusion_physical_plan::Partitioning; + #[cfg(test)] + use datafusion_physical_expr::expressions::{DynamicFilterPhysicalExpr, lit}; use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType}; + use datafusion_physical_plan::{Partitioning, PhysicalExpr}; use super::*; @@ -483,6 +509,7 @@ pub mod tests { pub struct EmptyExec { props: Arc, children: Vec>, + dynamic_expressions: Vec>, metrics: Option, statistics: Option, } @@ -497,6 +524,7 @@ pub mod tests { Boundedness::Bounded, )), children: Vec::default(), + dynamic_expressions: Vec::default(), metrics: None, statistics: None, } @@ -511,6 +539,14 @@ pub mod tests { self.statistics = Some(statistics); self } + + pub fn with_dynamic_expressions( + mut self, + dynamic_expressions: Vec>, + ) -> Self { + self.dynamic_expressions = dynamic_expressions; + self + } } impl DisplayAs for EmptyExec { @@ -543,6 +579,7 @@ pub mod tests { Ok(Arc::new(EmptyExec { props: Arc::clone(&self.props), children, + dynamic_expressions: self.dynamic_expressions.clone(), metrics: self.metrics.clone(), statistics: self.statistics.clone(), })) @@ -556,6 +593,10 @@ pub mod tests { unimplemented!() } + fn dynamic_expressions(&self) -> Vec> { + self.dynamic_expressions.iter().map(Arc::clone).collect() + } + fn metrics(&self) -> Option { self.metrics.clone() } @@ -600,6 +641,32 @@ pub mod tests { Ok(()) } + #[test] + fn test_ffi_execution_plan_dynamic_expressions() -> Result<()> { + 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) as _; + let original_plan = + Arc::new(EmptyExec::new(schema).with_dynamic_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()?; + foreign_plan.check_invariants( + datafusion_physical_plan::execution_plan::InvariantLevel::Always, + )?; + + let produced = foreign_plan.dynamic_expressions(); + assert_eq!(produced.len(), 1); + assert_eq!(produced[0].expression_id(), Some(expected_id)); + drop(foreign_plan); + assert_eq!(produced[0].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/tests/mod.rs b/datafusion/ffi/src/tests/mod.rs index d372dcf9177e6..c116fd5ddb3db 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_dynamic_expressions: extern "C" fn() -> FFI_ExecutionPlan, + pub create_exec_with_statistics: extern "C" fn() -> FFI_ExecutionPlan, pub create_table_with_statistics: @@ -161,6 +165,15 @@ pub(crate) extern "C" fn create_empty_exec() -> FFI_ExecutionPlan { FFI_ExecutionPlan::new(plan, None) } +pub(crate) extern "C" fn create_exec_with_dynamic_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_dynamic_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 +272,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_dynamic_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..aa797af14b038 100644 --- a/datafusion/ffi/tests/ffi_execution_plan.rs +++ b/datafusion/ffi/tests/ffi_execution_plan.rs @@ -26,6 +26,7 @@ mod tests { use datafusion_ffi::execution_plan::{ExecutionPlanPrivateData, tests::EmptyExec}; use datafusion_ffi::tests::utils::get_module; use datafusion_physical_plan::ExecutionPlan; + use datafusion_physical_plan::execution_plan::InvariantLevel; use std::sync::Arc; #[test] @@ -63,6 +64,23 @@ mod tests { Ok(()) } + #[test] + fn test_ffi_execution_plan_dynamic_expressions_cross_library() + -> Result<(), DataFusionError> { + let module = get_module()?; + let plan = (module.create_exec_with_dynamic_expressions)(); + let plan: Arc = (&plan).try_into()?; + assert!(plan.is::()); + plan.check_invariants(InvariantLevel::Always)?; + + let produced = plan.dynamic_expressions(); + assert_eq!(produced.len(), 1); + assert!(produced[0].expression_id().is_some()); + drop(plan); + assert!(produced[0].expression_id().is_some()); + Ok(()) + } + #[test] fn test_ffi_execution_plan_new_sets_runtimes_on_children() -> Result<(), DataFusionError> { diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index db1eb951d6fbc..08030cf37b8aa 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -1960,6 +1960,16 @@ impl ExecutionPlan for AggregateExec { vec![&self.input] } + fn dynamic_expressions(&self) -> Vec> { + self.dynamic_filter + .iter() + .map(|dynamic_filter| { + Arc::::clone(&dynamic_filter.filter) + as Arc + }) + .collect() + } + fn with_new_children( self: Arc, children: Vec>, @@ -7554,6 +7564,9 @@ mod tests { lit(false), )); let agg = agg.with_dynamic_filter_expr(Arc::clone(&new_df))?; + let produced = agg.dynamic_expressions(); + assert_eq!(produced.len(), 1); + assert_eq!(produced[0].expression_id(), new_df.expression_id()); // The aggregate's filter should now resolve to the new inner expression. let swapped = agg diff --git a/datafusion/physical-plan/src/execution_plan.rs b/datafusion/physical-plan/src/execution_plan.rs index 11a8d69a37669..f459ff26d7515 100644 --- a/datafusion/physical-plan/src/execution_plan.rs +++ b/datafusion/physical-plan/src/execution_plan.rs @@ -39,6 +39,7 @@ pub use datafusion_physical_expr::{ }; use std::any::Any; +use std::collections::HashSet; use std::fmt::Debug; use std::sync::{Arc, LazyLock}; @@ -165,6 +166,22 @@ pub trait ExecutionPlan: Any + Debug + DisplayAs + Send + Sync { check_default_invariants(self, check) } + /// Returns the dynamic expressions produced by this plan node. + /// + /// A dynamic expression is produced when this node updates or completes its + /// runtime state during execution. Expressions that this node only consumes + /// must not be returned. This method is shallow and does not include dynamic + /// expressions produced by child plans. + /// + /// Each returned expression must have a stable, node-unique + /// [`PhysicalExpr::expression_id`]. The returned [`Arc`]s allow callers to + /// retain access to the live runtime state after this plan is no longer + /// borrowed. The default implementation reports that this node produces no + /// dynamic expressions. + fn dynamic_expressions(&self) -> Vec> { + Vec::new() + } + /// Specifies simple per-child input distribution requirements. /// /// Deprecated: override [`Self::input_distribution_requirements`] instead. @@ -1317,6 +1334,26 @@ macro_rules! check_len { }; } +fn check_dynamic_expression_invariants( + plan: &P, +) -> Result<()> { + let mut produced_ids = HashSet::new(); + for expr in plan.dynamic_expressions() { + let Some(expression_id) = expr.expression_id() else { + return internal_err!( + "{}::dynamic_expressions returned an expression without an expression ID", + plan.name() + ); + }; + assert_or_internal_err!( + produced_ids.insert(expression_id), + "{}::dynamic_expressions returned duplicate expression ID {expression_id}", + plan.name() + ); + } + Ok(()) +} + /// Checks a set of invariants that apply to all ExecutionPlan implementations. /// Returns an error if the given node does not conform. pub fn check_default_invariants( @@ -1330,6 +1367,7 @@ pub fn check_default_invariants( check_len!(plan, benefits_from_input_partitioning, children_len); plan.input_distribution_requirements() .check_invariants(plan, check)?; + check_dynamic_expression_invariants(plan)?; Ok(()) } @@ -1732,13 +1770,26 @@ mod tests { use arrow::array::{DictionaryArray, Int32Array, NullArray, RunArray}; use arrow::datatypes::{DataType, Field, Schema}; + use datafusion_physical_expr::expressions::{DynamicFilterPhysicalExpr, lit}; #[derive(Debug)] - pub struct EmptyExec; + pub struct EmptyExec { + dynamic_expressions: Vec>, + } impl EmptyExec { pub fn new(_schema: SchemaRef) -> Self { - Self + Self { + dynamic_expressions: vec![], + } + } + + fn with_dynamic_expressions( + mut self, + dynamic_expressions: Vec>, + ) -> Self { + self.dynamic_expressions = dynamic_expressions; + self } } @@ -1772,6 +1823,10 @@ mod tests { unimplemented!() } + fn dynamic_expressions(&self) -> Vec> { + self.dynamic_expressions.iter().map(Arc::clone).collect() + } + fn execute( &self, _partition: usize, @@ -1789,6 +1844,32 @@ mod tests { } } + #[test] + fn test_dynamic_expression_invariants() -> Result<()> { + let schema = Arc::new(Schema::empty()); + let dynamic: Arc = + Arc::new(DynamicFilterPhysicalExpr::new(vec![], lit(true))); + let valid = EmptyExec::new(Arc::clone(&schema)) + .with_dynamic_expressions(vec![Arc::clone(&dynamic)]); + check_default_invariants(&valid, InvariantLevel::Always)?; + + let missing_id = + EmptyExec::new(Arc::clone(&schema)).with_dynamic_expressions(vec![lit(true)]); + let error = check_default_invariants(&missing_id, InvariantLevel::Always) + .unwrap_err() + .strip_backtrace(); + assert!(error.contains("without an expression ID"), "{error}"); + + let duplicate = EmptyExec::new(schema) + .with_dynamic_expressions(vec![Arc::clone(&dynamic), dynamic]); + let error = check_default_invariants(&duplicate, InvariantLevel::Always) + .unwrap_err() + .strip_backtrace(); + assert!(error.contains("duplicate expression ID"), "{error}"); + + Ok(()) + } + #[derive(Debug)] pub struct RenamedEmptyExec; diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index 9d9c867c2724b..85bd060225d66 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -1330,6 +1330,16 @@ impl ExecutionPlan for HashJoinExec { vec![&self.left, &self.right] } + fn dynamic_expressions(&self) -> Vec> { + self.dynamic_filter + .iter() + .map(|dynamic_filter| { + Arc::::clone(&dynamic_filter.filter) + as Arc + }) + .collect() + } + /// Creates a new HashJoinExec with different children while preserving configuration. /// /// This method is called during query optimization when the optimizer creates new @@ -6873,6 +6883,7 @@ mod tests { false, )?; assert!(join.dynamic_filter_expr().is_none()); + assert!(join.dynamic_expressions().is_empty()); let df = Arc::new(DynamicFilterPhysicalExpr::new( vec![Arc::new(Column::new("b1", 1)) as _], @@ -6890,6 +6901,9 @@ mod tests { df.expression_id() .expect("DynamicFilterPhysicalExpr always has an expression_id"), ); + let produced = join.dynamic_expressions(); + assert_eq!(produced.len(), 1); + assert_eq!(produced[0].expression_id(), df.expression_id()); Ok(()) } diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs index 4b30aede7d02a..63c7205d2e036 100644 --- a/datafusion/physical-plan/src/sorts/sort.rs +++ b/datafusion/physical-plan/src/sorts/sort.rs @@ -1274,6 +1274,13 @@ impl ExecutionPlan for SortExec { vec![&self.input] } + fn dynamic_expressions(&self) -> Vec> { + self.dynamic_filter_expr() + .into_iter() + .map(|expr| expr as Arc) + .collect() + } + fn benefits_from_input_partitioning(&self) -> Vec { vec![false] } @@ -3194,6 +3201,9 @@ mod tests { .expect("should have dynamic filter with fetch") .expression_id() .expect("DynamicFilterPhysicalExpr always has an expression_id"); + let produced = sort.dynamic_expressions(); + assert_eq!(produced.len(), 1); + assert_eq!(produced[0].expression_id(), Some(original_id)); // with_dynamic_filter replaces it with a new TopKDynamicFilters. let new_df = Arc::new(DynamicFilterPhysicalExpr::new( @@ -3211,6 +3221,9 @@ mod tests { .expect("DynamicFilterPhysicalExpr always has an expression_id"); assert_eq!(restored_id, new_id); assert_ne!(restored_id, original_id); + let produced = sort.dynamic_expressions(); + assert_eq!(produced.len(), 1); + assert_eq!(produced[0].expression_id(), Some(new_id)); Ok(()) }