From e2c33e04fac13104e8fefa6ac2c693250630f33f Mon Sep 17 00:00:00 2001 From: Ping Zhang Date: Tue, 15 Sep 2026 00:05:40 +0000 Subject: [PATCH 1/4] refactor: separate join runtime-filter wiring --- .../src/execution/operators/dynamic_filter.rs | 403 +---- .../operators/dynamic_filter/tests.rs | 1541 +--------------- .../operators/dynamic_filter_join.rs | 415 +++++ .../operators/dynamic_filter_join/tests.rs | 1610 +++++++++++++++++ native/core/src/execution/operators/mod.rs | 3 +- 5 files changed, 2055 insertions(+), 1917 deletions(-) create mode 100644 native/core/src/execution/operators/dynamic_filter_join.rs create mode 100644 native/core/src/execution/operators/dynamic_filter_join/tests.rs diff --git a/native/core/src/execution/operators/dynamic_filter.rs b/native/core/src/execution/operators/dynamic_filter.rs index 27e650a2b78..8a294201965 100644 --- a/native/core/src/execution/operators/dynamic_filter.rs +++ b/native/core/src/execution/operators/dynamic_filter.rs @@ -15,47 +15,29 @@ // specific language governing permissions and limitations // under the License. -//! Connect a hash join's completed build domain to its probe input. -//! -//! Comet does not run DataFusion's physical optimizer, which normally connects -//! dynamic-filter producers and consumers. This targeted wiring filters probe -//! batches and lets a direct Parquet reader use the same live predicate for -//! pruning. The original join verifies matches, including hash collisions. -//! This leaves Spark's operator tree and partitioning intact and does -//! not cross Spark exchanges or JVM/Arrow boundaries. +//! Apply a live runtime predicate to decoded batches. use std::fmt::Formatter; use std::sync::Arc; use arrow::compute::filter_record_batch; -use arrow::datatypes::DataType; use datafusion::common::cast::as_boolean_array; -use datafusion::common::config::ConfigOptions; use datafusion::common::tree_node::TreeNodeRecursion; -use datafusion::common::{internal_err, JoinType, NullEquality, Result, ScalarValue, Statistics}; -use datafusion::datasource::physical_plan::ParquetSource; -use datafusion::datasource::source::DataSourceExec; +use datafusion::common::{internal_err, Result, ScalarValue}; use datafusion::execution::TaskContext; -use datafusion::logical_expr::{ColumnarValue, Operator}; -use datafusion::physical_expr::expressions::{ - lit, BinaryExpr, Column, DynamicFilterPhysicalExpr, IsNotNullExpr, -}; +use datafusion::logical_expr::ColumnarValue; +use datafusion::physical_expr::expressions::{lit, Column, DynamicFilterPhysicalExpr}; use datafusion::physical_expr::PhysicalExpr; -use datafusion::physical_plan::distribution_requirements::InputDistributionRequirements; use datafusion::physical_plan::execution_plan::CardinalityEffect; -use datafusion::physical_plan::joins::{HashJoinExec, PartitionMode}; use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, MetricBuilder, MetricsSet}; -use datafusion::physical_plan::statistics::{ChildStats, StatisticsArgs}; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use datafusion::physical_plan::{ apply_expression_roots, ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, - ExecutionPlanProperties, PlanProperties, ReplaceChildrenOptions, SendableRecordBatchStream, + PlanProperties, ReplaceChildrenOptions, SendableRecordBatchStream, }; use futures::StreamExt; -use super::CometFilterExec; - -/// A task-local consumer of DataFusion's build-side runtime filter. +/// A task-local consumer of a live runtime predicate. #[derive(Debug)] pub(crate) struct DynamicFilterExec { input: Arc, @@ -64,11 +46,15 @@ pub(crate) struct DynamicFilterExec { } impl DynamicFilterExec { - fn new(input: Arc, predicate: Arc) -> Self { + pub(super) fn new( + input: Arc, + predicate: Arc, + metrics: ExecutionPlanMetricsSet, + ) -> Self { Self { input, predicate, - metrics: ExecutionPlanMetricsSet::new(), + metrics, } } } @@ -129,6 +115,7 @@ impl ExecutionPlan for DynamicFilterExec { Ok(Arc::new(Self::new( children.remove(0), Arc::clone(&self.predicate), + ExecutionPlanMetricsSet::new(), ))) } @@ -140,7 +127,11 @@ impl ExecutionPlan for DynamicFilterExec { self.predicate.children().into_iter().cloned().collect(), lit(true), )); - Ok(Arc::new(Self::new(Arc::clone(&self.input), predicate))) + Ok(Arc::new(Self::new( + Arc::clone(&self.input), + predicate, + ExecutionPlanMetricsSet::new(), + ))) } fn execute( @@ -210,363 +201,5 @@ impl ExecutionPlan for DynamicFilterExec { } } -/// A permanent plan must not own a completed join's filter or build accumulator: -/// those can retain the hash map after its stream-owned reservation is released. -/// Keep an unexecuted template here and create the producer and consumer together -/// for each stream. Only their metric handles are retained by the Spark plan. -#[derive(Debug)] -pub(crate) struct DynamicFilterJoinExec { - template: HashJoinExec, - config: ConfigOptions, - metrics: ExecutionPlanMetricsSet, -} - -/// Per-execution join state. The permanent plan keeps no live filter; this value -/// records whether this execution also connected its filter to the Parquet reader. -struct RuntimeDynamicFilterJoin { - join: HashJoinExec, - reader_filter_attached: bool, -} - -/// Recognize only direct-column null checks joined by AND, without evaluating -/// or changing the predicate. Every accepted leaf is deterministic, infallible, -/// and only discards rows, so reader pruning cannot suppress expression errors -/// or alter stateful evaluation. All other expressions remain a boundary. -fn is_direct_column_null_checks(predicate: &Arc) -> bool { - if let Some(binary) = predicate.downcast_ref::() { - return binary.op() == &Operator::And - && is_direct_column_null_checks(binary.left()) - && is_direct_column_null_checks(binary.right()); - } - predicate - .downcast_ref::() - .is_some_and(|is_not_null| is_not_null.arg().is::()) -} - -fn try_attach_parquet_reader_filter( - input: &Arc, - predicate: Arc, - config: &ConfigOptions, -) -> Result>> { - // Filtering before a fetch can change which rows are selected by its limit. - if input.fetch().is_some() { - log::debug!("Join dynamic filter reader pushdown skipped: probe has a fetch limit"); - return Ok(None); - } - // Spark inserts IS NOT NULL residuals above equijoin inputs, including AND - // chains of inferred null checks. A reader predicate can cross those direct - // checks because both operations only discard rows. Keep every other filter - // as a boundary: reader pruning would change which rows reach stateful - // expressions and can suppress expression errors. - if let Some(filter) = input.downcast_ref::() { - if filter.has_projection() { - log::debug!( - "Join dynamic filter reader pushdown skipped: probe FilterExec has a projection" - ); - return Ok(None); - } - if !is_direct_column_null_checks(filter.predicate()) { - log::debug!( - "Join dynamic filter reader pushdown skipped: probe filter is not direct column IS NOT NULL checks" - ); - return Ok(None); - } - let Some(reader) = - try_attach_parquet_reader_filter(filter.input(), Arc::clone(&predicate), config)? - else { - return Ok(None); - }; - return match filter.with_execution_input(reader) { - Ok(updated) => Ok(Some(updated)), - Err(error) => { - log::debug!( - "Join dynamic filter reader pushdown skipped: probe filter rebuild failed: {error}" - ); - Ok(None) - } - }; - } - let Some(scan) = input.downcast_ref::() else { - log::debug!( - "Join dynamic filter reader pushdown skipped: probe root is {}", - input.name() - ); - return Ok(None); - }; - if scan.downcast_to_file_source::().is_none() { - log::debug!("Join dynamic filter reader pushdown skipped: probe is not Parquet"); - return Ok(None); - } - - let predicate: Arc = predicate; - let propagation = match scan - .data_source() - .try_pushdown_filters(vec![predicate], config) - { - Ok(propagation) => propagation, - Err(error) => { - log::debug!( - "Join dynamic filter reader pushdown skipped: predicate remapping failed: {error}" - ); - return Ok(None); - } - }; - let Some(data_source) = propagation.updated_node else { - log::debug!("Join dynamic filter reader pushdown skipped: Parquet declined the predicate"); - return Ok(None); - }; - Ok(Some(Arc::new(scan.clone().with_data_source(data_source)))) -} - -impl DynamicFilterJoinExec { - /// Return no wrapper when the join cannot safely use a runtime filter. - pub(crate) fn try_new(join: &HashJoinExec, config: &ConfigOptions) -> Result> { - if let Some(reason) = ineligible_reason(join, config)? { - log::debug!("Join dynamic filter skipped: {reason}"); - return Ok(None); - } - Ok(Some(Self::new(join, config.clone())?)) - } - - fn new(join: &HashJoinExec, config: ConfigOptions) -> Result { - Ok(Self { - template: join.builder().reset_state().build()?, - config, - metrics: ExecutionPlanMetricsSet::new(), - }) - } - - fn build_runtime_join(&self) -> Result { - let predicate = Arc::new(DynamicFilterPhysicalExpr::new( - vec![Arc::clone(&self.template.on()[0].1)], - lit(true), - )); - let reader = try_attach_parquet_reader_filter( - self.template.right(), - Arc::clone(&predicate), - &self.config, - )?; - let reader_filter_attached = reader.is_some(); - let consumer = Arc::new(DynamicFilterExec { - input: reader.unwrap_or_else(|| Arc::clone(self.template.right())), - predicate: Arc::clone(&predicate), - metrics: self.metrics.clone(), - }); - // In particular, do not share CollectLeft's cached build future with the - // template, another execution, or a reset plan. - let join = self - .template - .builder() - .reset_state() - .with_new_children(vec![Arc::clone(self.template.left()), consumer])? - .build()? - .with_dynamic_filter_expr(predicate)?; - Ok(RuntimeDynamicFilterJoin { - join, - reader_filter_attached, - }) - } - - fn execute_runtime_join( - &self, - join: HashJoinExec, - partition: usize, - context: Arc, - ) -> Result { - // DataFusion can materialize one IN-list literal per build row, - // despite admitting the list by packed-array bytes and distinct-key count. - // Avoid that unreserved allocation. Map membership - // reuses the join's already-reserved hash table and preserves duplicates. - let mut config = context.session_config().clone(); - config - .options_mut() - .optimizer - .hash_join_inlist_pushdown_max_size = 0; - config - .options_mut() - .optimizer - .hash_join_inlist_pushdown_max_distinct_values = 0; - let context = Arc::new(TaskContext::new( - context.task_id(), - context.session_id(), - config, - context.scalar_functions().clone(), - context.higher_order_functions().clone(), - context.aggregate_functions().clone(), - context.window_functions().clone(), - context.runtime_env(), - )); - let result = join.execute(partition, context); - // HashJoinExec registers its metrics synchronously in execute(). Keep the - // live counters, including on error, without retaining the producer plan. - for metric in join.metrics().unwrap_or_default().iter() { - self.metrics.register(Arc::clone(metric)); - } - drop(join); - let input = result?; - // Drop execution state at EOF or error even if the caller retains the - // exhausted stream. Dropping a pending stream also drops all of its state. - let stream = futures::stream::unfold(Some(input), |input| async move { - let mut input = input?; - let batch = input.next().await?; - let remaining = if batch.is_ok() { Some(input) } else { None }; - Some((batch, remaining)) - }); - Ok(Box::pin(RecordBatchStreamAdapter::new( - self.schema(), - stream, - ))) - } -} - -impl DisplayAs for DynamicFilterJoinExec { - fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { - write!(f, "CometDynamicFilterJoinExec: ")?; - self.template.fmt_as(t, f) - } -} - -impl ExecutionPlan for DynamicFilterJoinExec { - fn name(&self) -> &str { - "CometDynamicFilterJoinExec" - } - - fn properties(&self) -> &Arc { - self.template.properties() - } - - fn input_distribution_requirements(&self) -> InputDistributionRequirements { - self.template.input_distribution_requirements() - } - - fn maintains_input_order(&self) -> Vec { - self.template.maintains_input_order() - } - - fn children(&self) -> Vec<&Arc> { - self.template.children() - } - - fn apply_expressions( - &self, - f: &mut dyn FnMut(&Arc) -> Result, - ) -> Result { - self.template.apply_expressions(f) - } - - fn with_new_children( - self: Arc, - children: Vec>, - ) -> Result> { - self.replace_children( - children, - ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), - ) - } - - fn replace_children( - self: Arc, - children: Vec>, - _options: ReplaceChildrenOptions, - ) -> Result> { - let join = self - .template - .builder() - .reset_state() - .with_new_children(children)? - .build()?; - match Self::try_new(&join, &self.config)? { - Some(wrapper) => Ok(Arc::new(wrapper)), - None => Ok(Arc::new(join)), - } - } - - fn reset_state(self: Arc) -> Result> { - Ok(Arc::new(Self::new(&self.template, self.config.clone())?)) - } - - fn execute( - &self, - partition: usize, - context: Arc, - ) -> Result { - let runtime = self.build_runtime_join()?; - let attachment_metric = if runtime.reader_filter_attached { - "dynamic_filter_reader_filters_attached" - } else { - "dynamic_filter_reader_filters_skipped" - }; - MetricBuilder::new(&self.metrics) - .counter(attachment_metric, partition) - .add(1); - self.execute_runtime_join(runtime.join, partition, context) - } - - fn metrics(&self) -> Option { - Some(self.metrics.clone_inner()) - } - - fn child_stats_requests(&self, partition: Option) -> Vec { - self.template.child_stats_requests(partition) - } - - fn statistics_from_inputs( - &self, - input_stats: &[Arc], - args: &StatisticsArgs, - ) -> Result> { - self.template.statistics_from_inputs(input_stats, args) - } -} - -fn ineligible_reason(join: &HashJoinExec, config: &ConfigOptions) -> Result> { - if !config.optimizer.enable_dynamic_filter_pushdown - || !config.optimizer.enable_join_dynamic_filter_pushdown - { - return Ok(Some("disabled by DataFusion session options")); - } - if join.join_type() != &JoinType::Inner - || join.null_equality() != NullEquality::NullEqualsNothing - { - return Ok(Some("only ordinary inner equijoins are supported")); - } - if !matches!( - join.partition_mode(), - PartitionMode::Partitioned | PartitionMode::CollectLeft - ) { - return Ok(Some("unresolved hash join partition mode")); - } - if config.optimizer.preserve_file_partitions > 0 - && matches!(join.partition_mode(), PartitionMode::Partitioned) - { - return Ok(Some("DataFusion preserve_file_partitions is enabled")); - } - // Spark, not DataFusion, routes rows across tasks. Restrict the filter to the - // single native partition executed by this Spark task: no shared domains or - // assumptions about DataFusion's repartition hash across Spark partitions. - if join.left().output_partitioning().partition_count() != 1 - || join.right().output_partitioning().partition_count() != 1 - { - return Ok(Some("requires one native partition per input")); - } - let [(build_key, probe_key)] = join.on() else { - return Ok(Some("requires one join key")); - }; - if !build_key.is::() || !probe_key.is::() { - return Ok(Some("computed join keys are not supported")); - } - let build_type = build_key.data_type(join.left().schema().as_ref())?; - let probe_type = probe_key.data_type(join.right().schema().as_ref())?; - if build_type != probe_type - || !matches!( - build_type, - DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 - ) - { - return Ok(Some("requires matching signed integer keys")); - } - Ok(None) -} - #[cfg(test)] mod tests; diff --git a/native/core/src/execution/operators/dynamic_filter/tests.rs b/native/core/src/execution/operators/dynamic_filter/tests.rs index fe6b0d06539..9460272c27a 100644 --- a/native/core/src/execution/operators/dynamic_filter/tests.rs +++ b/native/core/src/execution/operators/dynamic_filter/tests.rs @@ -16,33 +16,16 @@ // under the License. use super::*; -use std::fmt::Display; -use std::hash::{Hash, Hasher}; -use crate::execution::planner::PhysicalPlanner; -use crate::parquet::parquet_exec::init_datasource_exec; -use crate::parquet::parquet_support::ObjectStoreBackend; -use arrow::array::{ArrayRef, BooleanArray, Int32Array, Int64Array, Int8Array, RecordBatch}; +use arrow::array::{ArrayRef, Int32Array, RecordBatch}; use arrow::compute::cast; -use arrow::datatypes::{Field, Schema}; -use datafusion::common::test_util::batches_to_sort_string; -use datafusion::common::tree_node::TreeNodeRecursion; -use datafusion::datasource::listing::PartitionedFile; +use arrow::datatypes::{DataType, Field, Schema}; use datafusion::datasource::memory::MemorySourceConfig; -use datafusion::execution::memory_pool::{GreedyMemoryPool, MemoryPool}; -use datafusion::execution::object_store::ObjectStoreUrl; -use datafusion::execution::runtime_env::RuntimeEnvBuilder; +use datafusion::logical_expr::Operator; +use datafusion::physical_expr::expressions::BinaryExpr; use datafusion::physical_plan::collect; -use datafusion::physical_plan::filter::FilterExec; -use datafusion::physical_plan::metrics::MetricValue; use datafusion::physical_plan::projection::ProjectionExec; -use datafusion::physical_plan::statistics::{StatisticsArgs, StatisticsContext}; -use datafusion::physical_plan::{ChildrenPropertiesMode, Distribution, ReplaceChildrenOptions}; -use datafusion::prelude::{SessionConfig, SessionContext}; -use datafusion_comet_spark_expr::RandExpr; -use datafusion_datasource::file::FileSource; -use parquet::arrow::ArrowWriter; -use parquet::file::properties::{EnabledStatistics, WriterProperties}; +use datafusion::prelude::SessionContext; fn input( values: Vec>, @@ -78,44 +61,6 @@ fn memory_exec(batches: Vec) -> Arc { .unwrap() } -fn produced_join_filter(join: &HashJoinExec) -> Arc { - let mut expressions = join.dynamic_expressions_produced(); - assert_eq!(expressions.len(), 1); - let expression: Arc = expressions.pop().unwrap(); - expression.downcast::().unwrap() -} - -fn join( - build: Arc, - probe: Arc, - swap: bool, -) -> Arc { - let build_key = Arc::new(Column::new("key", 1)) as Arc; - let probe_key = Arc::new(Column::new("key", 0)) as Arc; - let (left, right, on) = if swap { - (probe, build, vec![(probe_key, build_key)]) - } else { - (build, probe, vec![(build_key, probe_key)]) - }; - let join = HashJoinExec::try_new( - left, - right, - on, - None, - &JoinType::Inner, - None, - PartitionMode::Partitioned, - NullEquality::NullEqualsNothing, - false, - ) - .unwrap(); - if swap { - join.swap_inputs(PartitionMode::Partitioned).unwrap() - } else { - Arc::new(join) - } -} - fn metric(plan: &Arc, name: &str) -> usize { if let Some(projection) = plan.downcast_ref::() { return metric(projection.input(), name); @@ -131,231 +76,6 @@ fn row_count(batches: &[RecordBatch]) -> usize { batches.iter().map(RecordBatch::num_rows).sum() } -/// Inspect the batch given to the real completed predicate, without changing its result. -#[derive(Debug, Eq)] -struct AssertKeyOnlyBatch { - child: Arc, - key_values_ptr: usize, -} - -impl PartialEq for AssertKeyOnlyBatch { - fn eq(&self, other: &Self) -> bool { - self.child.eq(&other.child) && self.key_values_ptr == other.key_values_ptr - } -} - -impl Hash for AssertKeyOnlyBatch { - fn hash(&self, state: &mut H) { - self.child.hash(state); - self.key_values_ptr.hash(state); - } -} - -impl Display for AssertKeyOnlyBatch { - fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - write!(f, "AssertKeyOnlyBatch({})", self.child) - } -} - -impl PhysicalExpr for AssertKeyOnlyBatch { - fn data_type(&self, input_schema: &Schema) -> Result { - self.child.data_type(input_schema) - } - - fn nullable(&self, input_schema: &Schema) -> Result { - self.child.nullable(input_schema) - } - - fn evaluate(&self, batch: &RecordBatch) -> Result { - assert_eq!( - batch.num_columns(), - 1, - "predicate must not receive payload columns" - ); - let keys = batch - .column(0) - .as_any() - .downcast_ref::() - .unwrap(); - assert_eq!( - keys.values().as_ptr() as usize, - self.key_values_ptr, - "projecting the join key must not copy its values" - ); - self.child.evaluate(batch) - } - - fn children(&self) -> Vec<&Arc> { - vec![&self.child] - } - - fn with_new_children( - self: Arc, - mut children: Vec>, - ) -> Result> { - assert_eq!(children.len(), 1); - Ok(Arc::new(Self { - child: children.remove(0), - key_values_ptr: self.key_values_ptr, - })) - } - - fn fmt_sql(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - Display::fmt(self, f) - } -} - -#[tokio::test] -async fn completed_filter_evaluates_only_the_shared_probe_key() { - // With no nulls and only 1/8 of rows inside the build bounds, DataFusion's AND - // evaluation preselects those rows before evaluating hash_lookup. A permutation - // spreads the selected rows throughout the batch, forcing payload copies if the - // full probe batch reaches the predicate. - let keys = Arc::new(Int32Array::from_iter_values( - (0..8192).map(|row| (row * 641) % 8192), - )); - let mut fields = (0..32) - .map(|column| Field::new(format!("payload_{column}"), DataType::Int64, false)) - .collect::>(); - let mut columns = (0..32) - .map(|column| { - Arc::new(Int64Array::from_iter_values( - (0..8192).map(move |row| i64::from(row) * 32 + i64::from(column)), - )) as ArrayRef - }) - .collect::>(); - let key_index = 17; - fields.insert(key_index, Field::new("key", DataType::Int32, false)); - columns.insert(key_index, Arc::clone(&keys) as ArrayRef); - let batch = RecordBatch::try_new(Arc::new(Schema::new(fields)), columns).unwrap(); - let probe = memory_exec(vec![batch.clone()]); - let join = HashJoinExec::try_new( - input((0..1024).map(Some).collect(), &DataType::Int32, 0), - Arc::clone(&probe), - vec![( - Arc::new(Column::new("key", 0)), - Arc::new(Column::new("key", key_index)), - )], - None, - &JoinType::Inner, - None, - PartitionMode::Partitioned, - NullEquality::NullEqualsNothing, - false, - ) - .unwrap(); - let session = SessionContext::new(); - let expected = collect( - Arc::new(join.builder().build().unwrap()), - session.task_ctx(), - ) - .await - .unwrap(); - let wrapper = - DynamicFilterJoinExec::new(&join, session.copied_config().options().as_ref().clone()) - .unwrap(); - let runtime = wrapper.build_runtime_join().unwrap(); - let predicate = produced_join_filter(&runtime.join); - let actual = datafusion::physical_plan::common::collect( - wrapper - .execute_runtime_join(runtime.join, 0, session.task_ctx()) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!( - batches_to_sort_string(&actual), - batches_to_sort_string(&expected) - ); - - // Inspect the actual build-generated bounds AND hash-membership expression. Its key - // remains at index 17 here, so the consumer must also remap every nested reference. - let completed = predicate.current().unwrap(); - assert!(completed.to_string().contains("hash_lookup")); - assert!(completed.to_string().contains("AND")); - predicate - .update(Arc::new(AssertKeyOnlyBatch { - child: completed, - key_values_ptr: keys.values().as_ptr() as usize, - })) - .unwrap(); - let consumer: Arc = Arc::new(DynamicFilterExec::new(probe, predicate)); - let filtered = collect(consumer, session.task_ctx()).await.unwrap(); - let selected = BooleanArray::from( - keys.values() - .iter() - .map(|key| *key < 1024) - .collect::>(), - ); - let expected = filter_record_batch(&batch, &selected).unwrap(); - assert_eq!(expected.num_rows(), 1024); - assert_eq!(filtered, vec![expected]); -} - -#[tokio::test] -async fn completed_build_filters_both_sides_and_session_inlist_settings() { - for key_type in [ - DataType::Int8, - DataType::Int16, - DataType::Int32, - DataType::Int64, - ] { - for swap in [false, true] { - for max_inlist_size in [0, 1024 * 1024] { - let mut config = SessionConfig::new(); - config - .options_mut() - .optimizer - .hash_join_inlist_pushdown_max_size = max_inlist_size; - let session = SessionContext::new_with_config(config); - for build_values in [ - vec![Some(-5), Some(20), None, Some(20), Some(90)], - vec![], - vec![None, None], - ] { - let build = input(build_values.clone(), &key_type, 1); - let probe = input((-100..=100).map(Some).chain([None]).collect(), &key_type, 0); - let plain = join(Arc::clone(&build), Arc::clone(&probe), swap); - let attached = PhysicalPlanner::apply_join_dynamic_filter( - join(build, probe, swap), - true, - session.copied_config().options(), - ) - .unwrap(); - let native_join = - if let Some(projection) = attached.downcast_ref::() { - projection.input() - } else { - &attached - }; - assert!(native_join.is::()); - assert_eq!(plain.schema(), attached.schema()); - let expected = collect(plain, session.task_ctx()).await.unwrap(); - let actual = collect(Arc::clone(&attached), session.task_ctx()) - .await - .unwrap(); - assert_eq!( - batches_to_sort_string(&actual), - batches_to_sort_string(&expected) - ); - if build_values.iter().any(Option::is_some) { - assert_eq!( - row_count(&actual), - 4, - "duplicates and late build keys must match" - ); - assert_eq!(metric(&attached, "dynamic_filter_rows_evaluated"), 202); - assert!(metric(&attached, "dynamic_filter_rows_pruned") >= 199); - assert_eq!(metric(&attached, "dynamic_filter_rows_bypassed"), 0); - } else { - assert_eq!(row_count(&actual), 0); - } - } - } - } - } -} - #[tokio::test] async fn placeholder_updates_and_errors_are_not_hidden() { let source = input((0..10).map(Some).collect(), &DataType::Int32, 1); @@ -363,8 +83,11 @@ async fn placeholder_updates_and_errors_are_not_hidden() { vec![Arc::new(Column::new("key", 1))], lit(true), )); - let wrapper: Arc = - Arc::new(DynamicFilterExec::new(source, Arc::clone(&predicate))); + let wrapper: Arc = Arc::new(DynamicFilterExec::new( + source, + Arc::clone(&predicate), + ExecutionPlanMetricsSet::new(), + )); let task = SessionContext::new().task_ctx(); let mut stream = wrapper.execute(0, Arc::clone(&task)).unwrap(); let first = stream.next().await.unwrap().unwrap(); @@ -399,1247 +122,3 @@ async fn placeholder_updates_and_errors_are_not_hidden() { let error = collect(wrapper, task).await.unwrap_err(); assert!(error.to_string().contains("must evaluate to a Boolean")); } - -fn plain_join() -> HashJoinExec { - let plan = join( - input(vec![Some(10)], &DataType::Int32, 1), - input(vec![Some(10), Some(20)], &DataType::Int32, 0), - false, - ); - plan.downcast_ref::() - .unwrap() - .builder() - .build() - .unwrap() -} - -fn assert_skipped(join: HashJoinExec, config: &ConfigOptions) { - let plain: Arc = Arc::new(join); - let attached = - PhysicalPlanner::apply_join_dynamic_filter(Arc::clone(&plain), true, config).unwrap(); - assert!( - Arc::ptr_eq(&plain, &attached), - "fallback must preserve the original plan" - ); - assert!(attached - .downcast_ref::() - .unwrap() - .dynamic_expressions_produced() - .is_empty()); -} - -#[test] -fn skips_unsupported_joins_and_session_disables() { - let default = ConfigOptions::default(); - for join_type in [ - JoinType::Left, - JoinType::Right, - JoinType::Full, - JoinType::LeftSemi, - JoinType::RightSemi, - JoinType::LeftAnti, - JoinType::RightAnti, - JoinType::LeftMark, - JoinType::RightMark, - ] { - assert_skipped( - plain_join().builder().with_type(join_type).build().unwrap(), - &default, - ); - } - assert_skipped( - plain_join() - .builder() - .with_null_equality(NullEquality::NullEqualsNull) - .build() - .unwrap(), - &default, - ); - assert_skipped( - plain_join() - .builder() - .with_partition_mode(PartitionMode::Auto) - .build() - .unwrap(), - &default, - ); - assert_skipped( - plain_join() - .builder() - .with_type(JoinType::LeftAnti) - .with_null_aware(true) - .with_partition_mode(PartitionMode::CollectLeft) - .build() - .unwrap(), - &default, - ); - for option in [ - "enable_dynamic_filter_pushdown", - "enable_join_dynamic_filter_pushdown", - "preserve_file_partitions", - ] { - let mut config = ConfigOptions::default(); - match option { - "enable_dynamic_filter_pushdown" => { - config.optimizer.enable_dynamic_filter_pushdown = false - } - "enable_join_dynamic_filter_pushdown" => { - config.optimizer.enable_join_dynamic_filter_pushdown = false - } - _ => config.optimizer.preserve_file_partitions = 1, - } - assert_skipped(plain_join(), &config); - } -} - -#[test] -fn skips_unsupported_keys_and_multiple_native_partitions() { - for key_type in [ - DataType::Float32, - DataType::Float64, - DataType::Utf8, - DataType::Decimal128(10, 0), - ] { - let plan = join( - input(vec![Some(10)], &key_type, 1), - input(vec![Some(10)], &key_type, 0), - false, - ); - let same = PhysicalPlanner::apply_join_dynamic_filter( - Arc::clone(&plan), - true, - &ConfigOptions::default(), - ) - .unwrap(); - assert!(Arc::ptr_eq(&plan, &same)); - } - let plain = plain_join(); - let (build, probe) = plain.on()[0].clone(); - let computed = Arc::new(BinaryExpr::new( - Arc::clone(&probe), - Operator::Plus, - lit(1i32), - )); - for keys in [ - vec![(Arc::clone(&build), computed as Arc)], - vec![(Arc::clone(&build), Arc::clone(&probe)), (build, probe)], - ] { - assert_skipped( - plain.builder().with_on(keys).build().unwrap(), - &ConfigOptions::default(), - ); - } - for replace_build in [false, true] { - let schema = if replace_build { - plain.left().schema() - } else { - plain.right().schema() - }; - let two_partitions: Arc = MemorySourceConfig::try_new_exec( - &[ - vec![RecordBatch::new_empty(Arc::clone(&schema))], - vec![RecordBatch::new_empty(Arc::clone(&schema))], - ], - schema, - None, - ) - .unwrap(); - let children = if replace_build { - vec![two_partitions, Arc::clone(plain.right())] - } else { - vec![Arc::clone(plain.left()), two_partitions] - }; - assert_skipped( - plain - .builder() - .with_new_children(children) - .unwrap() - .build() - .unwrap(), - &ConfigOptions::default(), - ); - } -} - -#[tokio::test] -async fn independent_attempts_do_not_share_build_domains() { - let session = SessionContext::new(); - for build_key in [5, 90] { - let plan = join( - input(vec![Some(build_key)], &DataType::Int32, 1), - input(vec![Some(5), Some(90)], &DataType::Int32, 0), - false, - ); - let attached = PhysicalPlanner::apply_join_dynamic_filter( - plan, - true, - session.copied_config().options(), - ) - .unwrap(); - let output = collect(Arc::clone(&attached), session.task_ctx()) - .await - .unwrap(); - assert_eq!(row_count(&output), 1); - let keys = output - .iter() - .flat_map(|batch| { - batch - .column(1) - .as_any() - .downcast_ref::() - .unwrap() - .values() - .iter() - .copied() - }) - .collect::>(); - assert_eq!(keys, vec![build_key]); - assert_eq!(metric(&attached, "dynamic_filter_rows_pruned"), 1); - } -} - -fn single_key_join_plans( - build: Arc, - probe: Arc, - mode: PartitionMode, -) -> HashJoinExec { - HashJoinExec::try_new( - build, - probe, - vec![( - Arc::new(Column::new("key", 0)), - Arc::new(Column::new("key", 0)), - )], - None, - &JoinType::Inner, - None, - mode, - NullEquality::NullEqualsNothing, - false, - ) - .unwrap() -} - -fn single_key_join(build: ArrayRef, probe: ArrayRef, mode: PartitionMode) -> HashJoinExec { - let schema = Arc::new(Schema::new(vec![Field::new( - "key", - build.data_type().clone(), - true, - )])); - let build = memory_exec(vec![ - RecordBatch::try_new(Arc::clone(&schema), vec![build]).unwrap() - ]); - let probe = memory_exec(vec![RecordBatch::try_new(schema, vec![probe]).unwrap()]); - single_key_join_plans(build, probe, mode) -} - -fn parquet_probe( - values: Vec, - session: &Arc, - max_row_group_rows: usize, -) -> (tempfile::NamedTempFile, Arc) { - // Put the key after an unused physical column and project only the key. - // Reader attachment must remap key@0 in the join to key@1 in the file. - let file_schema = Arc::new(Schema::new(vec![ - Field::new("payload", DataType::Int32, false), - Field::new("key", DataType::Int32, false), - ])); - let required_schema = Arc::new(Schema::new(vec![file_schema.field(1).clone()])); - let row_count = values.len(); - let batch = RecordBatch::try_new( - Arc::clone(&file_schema), - vec![ - Arc::new(Int32Array::from_iter_values(0..row_count as i32)), - Arc::new(Int32Array::from(values)), - ], - ) - .unwrap(); - let file = tempfile::NamedTempFile::new().unwrap(); - let props = WriterProperties::builder() - .set_max_row_group_row_count(Some(max_row_group_rows)) - .set_statistics_enabled(EnabledStatistics::Chunk) - .set_dictionary_enabled(false) - .build(); - let mut writer = ArrowWriter::try_new( - file.reopen().unwrap(), - Arc::clone(&file_schema), - Some(props), - ) - .unwrap(); - writer.write(&batch).unwrap(); - let metadata = writer.close().unwrap(); - assert_eq!( - metadata.num_row_groups(), - row_count.div_ceil(max_row_group_rows) - ); - assert!(metadata - .row_groups() - .iter() - .all(|group| { group.num_rows() > 0 && group.num_rows() as usize <= max_row_group_rows })); - - let non_null: Arc = - Arc::new(IsNotNullExpr::new(Arc::new(Column::new("key", 0)))); - let scan = init_datasource_exec( - required_schema, - Some(file_schema), - None, - ObjectStoreUrl::local_filesystem(), - ObjectStoreBackend::Local, - vec![vec![PartitionedFile::from_path( - file.path().to_string_lossy().into_owned(), - ) - .unwrap()]], - Some(vec![1]), - Some(vec![non_null]), - None, - "UTC", - true, - false, - false, - false, - session, - false, - false, - false, - ) - .unwrap(); - (file, scan) -} - -fn filtered_probe(scan: &Arc) -> Arc { - let predicate: Arc = - Arc::new(IsNotNullExpr::new(Arc::new(Column::new("key", 0)))); - Arc::new(CometFilterExec::from_datafusion( - FilterExec::try_new(predicate, Arc::clone(scan) as Arc).unwrap(), - )) -} - -fn find_dynamic_filter(expr: &Arc) -> Option<&DynamicFilterPhysicalExpr> { - if let Some(filter) = expr.downcast_ref::() { - return Some(filter); - } - expr.children() - .into_iter() - .find_map(|child| find_dynamic_filter(child)) -} - -/// Accept nested null-check conjunctions while retaining all other filter -/// boundaries, including OR and computed or potentially failing expressions. -#[test] -fn reader_filter_crosses_only_direct_column_null_checks() { - let key: Arc = Arc::new(Column::new("key", 0)); - let direct_null_check: Arc = Arc::new(IsNotNullExpr::new(Arc::clone(&key))); - assert!(is_direct_column_null_checks(&direct_null_check)); - let other_null_check: Arc = - Arc::new(IsNotNullExpr::new(Arc::new(Column::new("other", 1)))); - let conjunction: Arc = Arc::new(BinaryExpr::new( - Arc::clone(&direct_null_check), - Operator::And, - Arc::clone(&other_null_check), - )); - let nested: Arc = Arc::new(BinaryExpr::new( - conjunction, - Operator::And, - Arc::clone(&direct_null_check), - )); - assert!(is_direct_column_null_checks(&nested)); - let right_nested: Arc = Arc::new(BinaryExpr::new( - Arc::clone(&direct_null_check), - Operator::And, - Arc::new(BinaryExpr::new( - Arc::clone(&other_null_check), - Operator::And, - Arc::clone(&direct_null_check), - )), - )); - assert!(is_direct_column_null_checks(&right_nested)); - let disjunction: Arc = Arc::new(BinaryExpr::new( - Arc::clone(&direct_null_check), - Operator::Or, - other_null_check, - )); - assert!(!is_direct_column_null_checks(&disjunction)); - - let comparison: Arc = - Arc::new(BinaryExpr::new(Arc::clone(&key), Operator::Gt, lit(0_i32))); - let conjunction: Arc = Arc::new(BinaryExpr::new( - Arc::clone(&direct_null_check), - Operator::And, - comparison, - )); - assert!(!is_direct_column_null_checks(&conjunction)); - - let computed: Arc = - Arc::new(BinaryExpr::new(key, Operator::Plus, lit(1_i32))); - let computed_null_check: Arc = Arc::new(IsNotNullExpr::new(computed)); - assert!(!is_direct_column_null_checks(&computed_null_check)); -} - -/// Exercise a real Parquet reader with three distinct nullable columns and -/// remapped projection. Both statistics-only and row-filter reads must retain -/// the AND residual: build keys 2 and 3 fail its payload checks, leaving key 1. -#[tokio::test] -async fn reader_filter_crosses_null_check_conjunction_and_retains_residual() { - let file_schema = Arc::new(Schema::new(vec![ - Field::new("payload", DataType::Int32, true), - Field::new("key", DataType::Int32, true), - Field::new("other", DataType::Int32, true), - ])); - let batch = RecordBatch::try_new( - Arc::clone(&file_schema), - vec![ - Arc::new(Int32Array::from(vec![ - Some(10), - None, - Some(30), - Some(40), - Some(50), - ])), - Arc::new(Int32Array::from(vec![ - Some(1), - Some(2), - Some(3), - Some(4), - None, - ])), - Arc::new(Int32Array::from(vec![ - Some(1), - Some(1), - None, - Some(1), - Some(1), - ])), - ], - ) - .unwrap(); - let file = tempfile::NamedTempFile::new().unwrap(); - let props = WriterProperties::builder() - .set_dictionary_enabled(false) - .build(); - let mut writer = ArrowWriter::try_new( - file.reopen().unwrap(), - Arc::clone(&file_schema), - Some(props), - ) - .unwrap(); - writer.write(&batch).unwrap(); - writer.close().unwrap(); - let required_schema = Arc::new(Schema::new(vec![ - file_schema.field(1).clone(), - file_schema.field(0).clone(), - file_schema.field(2).clone(), - ])); - let mut outputs = Vec::new(); - for (row_filter, enabled) in [(false, false), (false, true), (true, false), (true, true)] { - let mut config = SessionConfig::new() - .with_target_partitions(1) - .with_parquet_page_index_pruning(false); - config.options_mut().execution.parquet.pushdown_filters = row_filter; - let session = Arc::new(SessionContext::new_with_config(config)); - let scan = init_datasource_exec( - Arc::clone(&required_schema), - Some(Arc::clone(&file_schema)), - None, - ObjectStoreUrl::local_filesystem(), - ObjectStoreBackend::Local, - vec![vec![PartitionedFile::from_path( - file.path().to_string_lossy().into_owned(), - ) - .unwrap()]], - Some(vec![1, 0, 2]), - None, - None, - "UTC", - true, - false, - false, - false, - &session, - false, - false, - false, - ) - .unwrap(); - let checks = [("key", 0), ("payload", 1), ("other", 2)].map(|(name, index)| { - Arc::new(IsNotNullExpr::new(Arc::new(Column::new(name, index)))) - as Arc - }); - let predicate: Arc = Arc::new(BinaryExpr::new( - Arc::new(BinaryExpr::new( - Arc::clone(&checks[0]), - Operator::And, - Arc::clone(&checks[1]), - )), - Operator::And, - Arc::clone(&checks[2]), - )); - let filter = Arc::new(CometFilterExec::from_datafusion( - FilterExec::try_new( - Arc::clone(&predicate), - Arc::clone(&scan) as Arc, - ) - .unwrap(), - )); - let build_schema = Arc::new(Schema::new(vec![Field::new("key", DataType::Int32, true)])); - let build = memory_exec(vec![RecordBatch::try_new( - build_schema, - vec![Arc::new(Int32Array::from(vec![1, 2, 3]))], - ) - .unwrap()]); - let join = single_key_join_plans(build, filter, PartitionMode::Partitioned); - let plan: Arc = if enabled { - let wrapper = DynamicFilterJoinExec::new( - &join, - session.copied_config().options().as_ref().clone(), - ) - .unwrap(); - let runtime = wrapper.build_runtime_join().unwrap(); - assert!(runtime.reader_filter_attached); - let consumer = runtime - .join - .right() - .downcast_ref::() - .unwrap(); - let retained = consumer.input.downcast_ref::().unwrap(); - assert_eq!(retained.predicate().to_string(), predicate.to_string()); - Arc::new(wrapper) - } else { - Arc::new(join) - }; - let output = collect(Arc::clone(&plan), session.task_ctx()) - .await - .unwrap(); - assert_eq!(row_count(&output), 1); - if enabled { - if row_filter { - let metrics = scan.metrics().unwrap(); - assert!( - metrics - .sum_by_name("pushdown_rows_pruned") - .unwrap() - .as_usize() - > 0, - "{metrics}" - ); - } - assert_eq!( - plan.metrics() - .unwrap() - .sum_by_name("dynamic_filter_reader_filters_attached") - .unwrap() - .as_usize(), - 1 - ); - } - outputs.push(batches_to_sort_string(&output)); - } - assert!(outputs.windows(2).all(|pair| pair[0] == pair[1])); -} - -#[tokio::test] -async fn reader_filter_does_not_cross_fetch_limits() { - for limit_filter in [false, true] { - let mut outputs = Vec::new(); - for enabled in [false, true] { - let session = Arc::new(SessionContext::new_with_config( - SessionConfig::new().with_target_partitions(1), - )); - let (_file, scan) = parquet_probe((0..4).collect(), &session, 1); - let probe = if limit_filter { - filtered_probe(&scan).with_fetch(Some(1)).unwrap() - } else { - scan.with_fetch(Some(1)).unwrap() - }; - let schema = Arc::new(Schema::new(vec![Field::new("key", DataType::Int32, false)])); - let build = memory_exec(vec![RecordBatch::try_new( - schema, - vec![Arc::new(Int32Array::from(vec![3]))], - ) - .unwrap()]); - let join = single_key_join_plans(build, probe, PartitionMode::Partitioned); - let plan: Arc = if enabled { - PhysicalPlanner::apply_join_dynamic_filter( - Arc::new(join), - true, - session.copied_config().options(), - ) - .unwrap() - } else { - Arc::new(join) - }; - let output = collect(Arc::clone(&plan), session.task_ctx()) - .await - .unwrap(); - outputs.push(row_count(&output)); - if enabled { - assert_eq!(metric(&plan, "dynamic_filter_reader_filters_skipped"), 1); - } - } - assert_eq!(outputs, vec![0, 0]); - } -} - -#[tokio::test] -async fn reader_filter_does_not_cross_seeded_rand_probe_filter() { - let mut outputs = Vec::new(); - for enabled in [false, true] { - let mut config = SessionConfig::new() - .with_target_partitions(1) - .with_parquet_page_index_pruning(false); - config.options_mut().execution.parquet.pushdown_filters = false; - let session = Arc::new(SessionContext::new_with_config(config)); - - // Four one-row groups make reader pruning observable in Rand's state. Without the - // filter boundary, pruning keys 0-2 makes key 3 receive the first draw (0.619...) - // instead of the fourth (0.263...), changing whether it passes rand(42) < 0.5. - let (_file, scan) = parquet_probe((0..4).collect(), &session, 1); - let key: Arc = Arc::new(Column::new("key", 0)); - let is_not_null: Arc = Arc::new(IsNotNullExpr::new(Arc::clone(&key))); - let random_below_half: Arc = Arc::new(BinaryExpr::new( - Arc::new(RandExpr::new(42)), - Operator::Lt, - lit(0.5_f64), - )); - let predicate: Arc = Arc::new(BinaryExpr::new( - is_not_null, - Operator::And, - random_below_half, - )); - let filter = Arc::new(CometFilterExec::from_datafusion( - FilterExec::try_new(predicate, Arc::clone(&scan) as Arc).unwrap(), - )); - - let build_schema = Arc::new(Schema::new(vec![Field::new("key", DataType::Int32, false)])); - let build = memory_exec(vec![RecordBatch::try_new( - build_schema, - vec![Arc::new(Int32Array::from(vec![3]))], - ) - .unwrap()]); - let join = single_key_join_plans( - build, - Arc::clone(&filter) as Arc, - PartitionMode::Partitioned, - ); - let plan: Arc = if enabled { - PhysicalPlanner::apply_join_dynamic_filter( - Arc::new(join), - true, - session.copied_config().options(), - ) - .unwrap() - } else { - Arc::new(join) - }; - let batches = collect(Arc::clone(&plan), session.task_ctx()) - .await - .unwrap(); - let probe_keys = batches - .iter() - .flat_map(|batch| { - batch - .column(1) - .as_any() - .downcast_ref::() - .unwrap() - .values() - .iter() - .copied() - }) - .collect::>(); - outputs.push(probe_keys); - - if enabled { - let dynamic_metric = |name| { - plan.metrics() - .and_then(|metrics| metrics.sum_by_name(name)) - .map_or(0, |value| value.as_usize()) - }; - assert_eq!(dynamic_metric("dynamic_filter_reader_filters_attached"), 0); - assert_eq!(dynamic_metric("dynamic_filter_reader_filters_skipped"), 1); - assert_eq!(dynamic_metric("dynamic_filter_rows_evaluated"), 1); - assert_eq!(dynamic_metric("dynamic_filter_rows_pruned"), 0); - assert_eq!(pruning_metric(&scan, "row_groups_pruned_statistics"), 0); - assert_eq!(filter.metrics().unwrap().output_rows().unwrap(), 1); - } - } - assert_eq!(outputs, vec![vec![3], vec![3]]); -} - -fn two_batch_build() -> Arc { - let schema = Arc::new(Schema::new(vec![Field::new("key", DataType::Int32, false)])); - memory_exec( - [150, 250] - .into_iter() - .map(|key| { - RecordBatch::try_new( - Arc::clone(&schema), - vec![Arc::new(Int32Array::from(vec![key]))], - ) - .unwrap() - }) - .collect(), - ) -} - -fn pruning_metric(plan: &Arc, name: &str) -> usize { - let metrics = plan.metrics().unwrap(); - let Some(value) = metrics.sum_by_name(name) else { - return 0; - }; - let MetricValue::PruningMetrics { - pruning_metrics, .. - } = value - else { - panic!("expected pruning metric {name}: {metrics}"); - }; - pruning_metrics.pruned() -} - -async fn run_parquet_join(values: Vec, enabled: bool) -> (usize, usize, usize, usize, usize) { - let mut config = SessionConfig::new() - .with_target_partitions(1) - .with_parquet_page_index_pruning(false); - // Isolate row-group pruning. The residual batch filter remains responsible - // for membership filtering after the reader applies conservative statistics. - config.options_mut().execution.parquet.pushdown_filters = false; - let session = Arc::new(SessionContext::new_with_config(config)); - let (_file, scan) = parquet_probe(values, &session, 100); - let filter = filtered_probe(&scan); - let probe = Arc::clone(&filter) as Arc; - let join = single_key_join_plans(two_batch_build(), probe, PartitionMode::Partitioned); - let plan: Arc = if enabled { - PhysicalPlanner::apply_join_dynamic_filter( - Arc::new(join), - true, - session.copied_config().options(), - ) - .unwrap() - } else { - Arc::new(join) - }; - let output = collect(Arc::clone(&plan), session.task_ctx()) - .await - .unwrap(); - let attached = plan - .metrics() - .and_then(|metrics| metrics.sum_by_name("dynamic_filter_reader_filters_attached")) - .map_or(0, |metric| metric.as_usize()); - ( - row_count(&output), - pruning_metric(&scan, "row_groups_pruned_statistics"), - scan.metrics() - .unwrap() - .sum_by_name("bytes_scanned") - .unwrap() - .as_usize(), - attached, - filter.metrics().unwrap().output_rows().unwrap_or_default(), - ) -} - -#[tokio::test] -async fn broadcast_filter_reaches_parquet_reader_after_complete_build() { - let mut config = SessionConfig::new() - .with_target_partitions(1) - .with_parquet_page_index_pruning(false); - config.options_mut().execution.parquet.pushdown_filters = false; - let session = Arc::new(SessionContext::new_with_config(config)); - let (_file, scan) = parquet_probe((0..400).collect(), &session, 100); - let join = single_key_join_plans( - two_batch_build(), - filtered_probe(&scan), - PartitionMode::Partitioned, - ); - let wrapper = - DynamicFilterJoinExec::new(&join, session.copied_config().options().as_ref().clone()) - .unwrap(); - let runtime = wrapper.build_runtime_join().unwrap(); - assert!(runtime.reader_filter_attached); - let consumer = runtime - .join - .right() - .downcast_ref::() - .unwrap(); - let filter = consumer.input.downcast_ref::().unwrap(); - let reader = filter.input().downcast_ref::().unwrap(); - let (_, source) = reader.downcast_to_file_source::().unwrap(); - let reader_filter = source.filter().unwrap(); - let reader_filter = find_dynamic_filter(&reader_filter).unwrap(); - let join_filter = produced_join_filter(&runtime.join); - assert_eq!( - reader_filter.expression_id().unwrap(), - join_filter.expression_id().unwrap() - ); - let reader_keys = reader_filter.children(); - assert_eq!(reader_keys.len(), 1); - let remapped_key = reader_keys[0].downcast_ref::().unwrap(); - assert_eq!(remapped_key.name(), "key"); - assert_eq!(remapped_key.index(), 1); - - let (disabled_rows, _, disabled_bytes, _, disabled_filter_rows) = - run_parquet_join((0..400).collect(), false).await; - let (enabled_rows, enabled_pruned, enabled_bytes, attached, enabled_filter_rows) = - run_parquet_join((0..400).collect(), true).await; - assert_eq!(disabled_rows, 2); - assert_eq!(enabled_rows, disabled_rows); - assert_eq!(enabled_pruned, 2); - assert!( - enabled_bytes < disabled_bytes, - "reader filter should avoid data reads: enabled={enabled_bytes}, disabled={disabled_bytes}" - ); - assert_eq!(attached, 1); - assert_eq!(disabled_filter_rows, 400); - assert_eq!(enabled_filter_rows, 200); - - // Every row group's bounds span both build keys, so the same attached - // filter is safe but cannot avoid I/O on this deliberately poor layout. - let mut unclustered = Vec::with_capacity(400); - for group in 0..4 { - unclustered.extend(0..50); - unclustered.extend(350..400); - if group == 0 { - unclustered[group * 100 + 1] = 150; - unclustered[group * 100 + 51] = 250; - } - } - let ( - unclustered_disabled_rows, - _, - unclustered_disabled_bytes, - _, - unclustered_disabled_filter_rows, - ) = run_parquet_join(unclustered.clone(), false).await; - let ( - unclustered_enabled_rows, - unclustered_pruned, - unclustered_enabled_bytes, - attached, - unclustered_enabled_filter_rows, - ) = run_parquet_join(unclustered, true).await; - assert_eq!(unclustered_enabled_rows, 2); - assert_eq!(unclustered_enabled_rows, unclustered_disabled_rows); - assert_eq!(unclustered_pruned, 0); - assert_eq!(unclustered_enabled_bytes, unclustered_disabled_bytes); - assert_eq!(attached, 1); - assert_eq!(unclustered_disabled_filter_rows, 400); - assert_eq!(unclustered_enabled_filter_rows, 400); -} - -fn limited_session(bytes: usize) -> (SessionContext, Arc) { - let pool: Arc = Arc::new(GreedyMemoryPool::new(bytes)); - let runtime = RuntimeEnvBuilder::new() - .with_memory_pool(Arc::clone(&pool)) - .build_arc() - .unwrap(); - ( - SessionContext::new_with_config_rt(SessionConfig::new(), runtime), - pool, - ) -} - -#[tokio::test] -async fn runtime_domains_release_with_streams_while_plans_remain_alive() { - for mode in [PartitionMode::Partitioned, PartitionMode::CollectLeft] { - for finish in [ - "eof", - "cancel_before_poll", - "cancel_after_output", - "build_error", - ] { - let (session, pool) = limited_session(if finish == "build_error" { - 1 - } else { - 8 * 1024 * 1024 - }); - // Sparse keys and more than 150 distinct values require the map - // strategy even without the IN-list override. - let join = single_key_join( - Arc::new(Int64Array::from_iter_values( - (0..4096).map(|i| i * 1_000_003), - )), - Arc::new(Int64Array::from(vec![0, 1])), - mode, - ); - let plan = DynamicFilterJoinExec::new(&join, ConfigOptions::default()).unwrap(); - let producer = plan.build_runtime_join().unwrap(); - let predicate = Arc::downgrade(&produced_join_filter(&producer.join)); - let mut stream = plan - .execute_runtime_join(producer.join, 0, session.task_ctx()) - .unwrap(); - assert!(predicate.upgrade().is_some()); - - match finish { - "cancel_before_poll" => drop(stream), - "build_error" => { - assert!(stream.next().await.unwrap().is_err()); - assert!(stream.next().await.is_none()); - } - _ => { - assert_eq!(stream.next().await.unwrap().unwrap().num_rows(), 1); - assert!(pool.reserved() > 0); - assert!(predicate.upgrade().is_some()); - if finish == "eof" { - assert!(stream.next().await.is_none()); - // Retain the exhausted stream as well as the plan. - assert_eq!(pool.reserved(), 0); - assert!(predicate.upgrade().is_none()); - } else { - drop(stream); - } - let metrics = plan.metrics().unwrap(); - assert_eq!(metrics.output_rows(), Some(1)); - assert_eq!( - metrics - .sum_by_name("dynamic_filter_rows_pruned") - .unwrap() - .as_usize(), - 1 - ); - } - } - // The accumulator also owns this predicate. A zero reservation - // alone would miss the old plan-owned, unaccounted map retention. - assert_eq!(pool.reserved(), 0, "{mode:?}: {finish}"); - assert!(predicate.upgrade().is_none(), "{mode:?}: {finish}"); - assert!(plan.template.dynamic_expressions_produced().is_empty()); - } - } -} - -fn expression_nodes(expr: &Arc) -> usize { - 1 + expr - .children() - .into_iter() - .map(expression_nodes) - .sum::() -} - -#[tokio::test] -async fn duplicate_heavy_builds_do_not_materialize_unreserved_inlists() { - for mode in [PartitionMode::Partitioned, PartitionMode::CollectLeft] { - let (session, pool) = limited_session(1024 * 1024); - let join = single_key_join( - Arc::new(Int8Array::from(vec![1; 65_536])), - Arc::new(Int8Array::from(vec![1, 2])), - mode, - ); - let config = session.copied_config(); - let configured_limit = config - .options() - .optimizer - .hash_join_inlist_pushdown_max_size; - assert!(configured_limit > 0); - let plan = DynamicFilterJoinExec::new(&join, config.options().as_ref().clone()).unwrap(); - let producer = plan.build_runtime_join().unwrap(); - let predicate = Arc::downgrade(&produced_join_filter(&producer.join)); - let mut stream = plan - .execute_runtime_join(producer.join, 0, session.task_ctx()) - .unwrap(); - let mut rows = stream.next().await.unwrap().unwrap().num_rows(); - { - let predicate = predicate.upgrade().unwrap(); - let current = predicate.current().unwrap(); - // Bound the published expression, rather than relying on pool - // reservations, which do not account for IN-list construction. - assert!(expression_nodes(¤t) < 32); - } - while let Some(batch) = stream.next().await { - rows += batch.unwrap().num_rows(); - } - assert_eq!(rows, 65_536, "duplicate build rows must still join"); - assert_eq!(pool.reserved(), 0); - assert!(predicate.upgrade().is_none()); - assert_eq!( - session - .copied_config() - .options() - .optimizer - .hash_join_inlist_pushdown_max_size, - configured_limit, - "the execution override must not change the session" - ); - assert_eq!( - plan.metrics() - .unwrap() - .sum_by_name("dynamic_filter_rows_pruned") - .unwrap() - .as_usize(), - 1 - ); - } -} - -#[tokio::test] -async fn executions_and_resets_have_independent_producers() { - for mode in [PartitionMode::Partitioned, PartitionMode::CollectLeft] { - let (session, pool) = limited_session(1024 * 1024); - let join = single_key_join( - Arc::new(Int64Array::from(vec![5])), - Arc::new(Int64Array::from(vec![5, 90])), - mode, - ); - let plan = Arc::new(DynamicFilterJoinExec::new(&join, ConfigOptions::default()).unwrap()); - let first = plan.build_runtime_join().unwrap(); - let first_filter = Arc::downgrade(&produced_join_filter(&first.join)); - let second = plan.build_runtime_join().unwrap(); - let second_filter = Arc::downgrade(&produced_join_filter(&second.join)); - assert!(!first_filter.ptr_eq(&second_filter)); - let mut first = plan - .execute_runtime_join(first.join, 0, session.task_ctx()) - .unwrap(); - let mut second = plan - .execute_runtime_join(second.join, 0, session.task_ctx()) - .unwrap(); - assert_eq!(first.next().await.unwrap().unwrap().num_rows(), 1); - assert_eq!(second.next().await.unwrap().unwrap().num_rows(), 1); - drop(first); - assert!(first_filter.upgrade().is_none()); - assert!(second_filter.upgrade().is_some()); - assert!(pool.reserved() > 0); - assert!(second.next().await.is_none()); - assert!(second_filter.upgrade().is_none()); - assert_eq!(pool.reserved(), 0); - - let reset = Arc::clone(&plan).reset_state().unwrap(); - assert_eq!( - row_count(&collect(reset, session.task_ctx()).await.unwrap()), - 1 - ); - let replacement = single_key_join( - Arc::new(Int64Array::from(vec![90])), - Arc::new(Int64Array::from(vec![5, 90])), - mode, - ); - let rewritten = Arc::clone(&plan) - .replace_children( - vec![ - Arc::clone(replacement.left()), - Arc::clone(replacement.right()), - ], - ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), - ) - .unwrap(); - let output = collect(rewritten, session.task_ctx()).await.unwrap(); - assert_eq!(row_count(&output), 1); - assert_eq!( - output[0] - .column(0) - .as_any() - .downcast_ref::() - .unwrap() - .value(0), - 90 - ); - assert_eq!(pool.reserved(), 0); - } -} - -#[tokio::test] -async fn child_replacement_rechecks_join_key_types() { - let session = SessionContext::new(); - for mode in [PartitionMode::Partitioned, PartitionMode::CollectLeft] { - let join = plain_join() - .builder() - .with_partition_mode(mode) - .build() - .unwrap(); - let plan = Arc::new( - DynamicFilterJoinExec::try_new(&join, &ConfigOptions::default()) - .unwrap() - .unwrap(), - ); - for (key_type, supported) in [ - (DataType::Int64, true), - (DataType::Float64, false), - (DataType::Utf8, false), - ] { - let rewritten = Arc::clone(&plan) - .replace_children( - vec![ - input(vec![Some(90)], &key_type, 1), - input(vec![Some(5), Some(90)], &key_type, 0), - ], - ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), - ) - .unwrap(); - assert_eq!(rewritten.is::(), supported); - if !supported { - assert!(rewritten - .downcast_ref::() - .unwrap() - .dynamic_expressions_produced() - .is_empty()); - } - let output = collect(Arc::clone(&rewritten), session.task_ctx()) - .await - .unwrap(); - assert_eq!(row_count(&output), 1); - if supported { - assert_eq!(metric(&rewritten, "dynamic_filter_rows_pruned"), 1); - let reset = rewritten.reset_state().unwrap(); - assert!(reset.is::()); - let reset_output = collect(Arc::clone(&reset), session.task_ctx()) - .await - .unwrap(); - assert_eq!(row_count(&reset_output), 1); - assert_eq!(metric(&reset, "dynamic_filter_rows_pruned"), 1); - } - } - } -} - -#[test] -fn child_replacement_rechecks_native_partition_counts() { - for mode in [PartitionMode::Partitioned, PartitionMode::CollectLeft] { - let join = plain_join() - .builder() - .with_partition_mode(mode) - .build() - .unwrap(); - let plan = Arc::new( - DynamicFilterJoinExec::try_new(&join, &ConfigOptions::default()) - .unwrap() - .unwrap(), - ); - for replace_build in [false, true] { - let schema = if replace_build { - join.left().schema() - } else { - join.right().schema() - }; - let two_partitions: Arc = MemorySourceConfig::try_new_exec( - &[ - vec![RecordBatch::new_empty(Arc::clone(&schema))], - vec![RecordBatch::new_empty(Arc::clone(&schema))], - ], - schema, - None, - ) - .unwrap(); - let children = if replace_build { - vec![two_partitions, Arc::clone(join.right())] - } else { - vec![Arc::clone(join.left()), two_partitions] - }; - let rewritten = Arc::clone(&plan) - .replace_children( - children, - ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), - ) - .unwrap(); - assert!(rewritten - .downcast_ref::() - .unwrap() - .dynamic_expressions_produced() - .is_empty()); - } - } -} - -#[test] -fn runtime_filter_is_visible_to_expression_visitors() { - let wrapper = DynamicFilterJoinExec::try_new(&plain_join(), &ConfigOptions::default()) - .unwrap() - .unwrap(); - let runtime = wrapper.build_runtime_join().unwrap(); - let predicate = produced_join_filter(&runtime.join); - let consumer = runtime.join.right(); - let mut consumer_expression_ids = Vec::new(); - consumer - .apply_expressions(&mut |expression| { - consumer_expression_ids.push(expression.expression_id()); - Ok(TreeNodeRecursion::Continue) - }) - .unwrap(); - assert_eq!( - consumer_expression_ids, - vec![Some(predicate.expression_id().unwrap())] - ); - assert_eq!( - consumer - .apply_expressions(&mut |_| Ok(TreeNodeRecursion::Stop)) - .unwrap(), - TreeNodeRecursion::Stop - ); - assert!(consumer.dynamic_expressions_produced().is_empty()); - assert!(wrapper.dynamic_expressions_produced().is_empty()); - - let mut join_key_count = 0; - wrapper - .apply_expressions(&mut |expression| { - assert!(expression.expression_id().is_none()); - join_key_count += 1; - Ok(TreeNodeRecursion::Continue) - }) - .unwrap(); - assert_eq!(join_key_count, 2); -} - -#[test] -fn wrapper_preserves_join_statistics_and_distribution() { - for mode in [PartitionMode::Partitioned, PartitionMode::CollectLeft] { - let join = plain_join() - .builder() - .with_partition_mode(mode) - .build() - .unwrap(); - let wrapper = DynamicFilterJoinExec::try_new(&join, &ConfigOptions::default()) - .unwrap() - .unwrap(); - for partition in [None, Some(0)] { - let context = StatisticsContext::new(); - let args = StatisticsArgs::new().with_partition(partition); - let expected = context.compute(&join, &args).unwrap(); - let actual = context.compute(&wrapper, &args).unwrap(); - assert!(expected.num_rows.get_value().is_some()); - assert_eq!(actual, expected, "{mode:?}: {partition:?}"); - } - - let expected = join.input_distribution_requirements(); - let actual = wrapper.input_distribution_requirements(); - assert_eq!( - actual.per_child_distributions().len(), - expected.per_child_distributions().len() - ); - for (actual, expected) in actual - .per_child_distributions() - .zip(expected.per_child_distributions()) - { - match (actual, expected) { - (Distribution::SinglePartition, Distribution::SinglePartition) - | (Distribution::UnspecifiedDistribution, Distribution::UnspecifiedDistribution) => { - } - (Distribution::KeyPartitioned(actual), Distribution::KeyPartitioned(expected)) => { - assert_eq!(actual, expected); - } - _ => panic!("distribution changed: {expected:?} -> {actual:?}"), - } - } - let schema = join.right().schema(); - let two_partitions = MemorySourceConfig::try_new_exec( - &[ - vec![RecordBatch::new_empty(Arc::clone(&schema))], - vec![RecordBatch::new_empty(Arc::clone(&schema))], - ], - schema, - None, - ) - .unwrap(); - let candidate_children = [join.left().as_ref(), two_partitions.as_ref()]; - let expected_unsatisfied = expected - .unsatisfied_co_partitioned_children(join.name(), &candidate_children) - .unwrap(); - let actual_unsatisfied = actual - .unsatisfied_co_partitioned_children(wrapper.name(), &candidate_children) - .unwrap(); - assert_eq!(actual_unsatisfied, expected_unsatisfied); - assert_eq!( - actual_unsatisfied, - if matches!(mode, PartitionMode::Partitioned) { - vec![0, 1] - } else { - vec![] - } - ); - } -} diff --git a/native/core/src/execution/operators/dynamic_filter_join.rs b/native/core/src/execution/operators/dynamic_filter_join.rs new file mode 100644 index 00000000000..049739116bf --- /dev/null +++ b/native/core/src/execution/operators/dynamic_filter_join.rs @@ -0,0 +1,415 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Connect a hash join's completed build domain to its probe input. +//! +//! Comet does not run DataFusion's physical optimizer, which normally connects +//! dynamic-filter producers and consumers. This targeted wiring filters probe +//! batches and lets a direct Parquet reader use the same live predicate for +//! pruning. The original join verifies matches, including hash collisions. +//! This leaves Spark's operator tree and partitioning intact and does +//! not cross Spark exchanges or JVM/Arrow boundaries. + +use std::fmt::Formatter; +use std::sync::Arc; + +use arrow::datatypes::DataType; +use datafusion::common::config::ConfigOptions; +use datafusion::common::tree_node::TreeNodeRecursion; +use datafusion::common::{JoinType, NullEquality, Result, Statistics}; +use datafusion::datasource::physical_plan::ParquetSource; +use datafusion::datasource::source::DataSourceExec; +use datafusion::execution::TaskContext; +use datafusion::logical_expr::Operator; +use datafusion::physical_expr::expressions::{ + lit, BinaryExpr, Column, DynamicFilterPhysicalExpr, IsNotNullExpr, +}; +use datafusion::physical_expr::PhysicalExpr; +use datafusion::physical_plan::distribution_requirements::InputDistributionRequirements; +use datafusion::physical_plan::joins::{HashJoinExec, PartitionMode}; +use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, MetricBuilder, MetricsSet}; +use datafusion::physical_plan::statistics::{ChildStats, StatisticsArgs}; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; +use datafusion::physical_plan::{ + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, + PlanProperties, ReplaceChildrenOptions, SendableRecordBatchStream, +}; +use futures::StreamExt; + +use super::dynamic_filter::DynamicFilterExec; +use super::CometFilterExec; + +/// A permanent plan must not own a completed join's filter or build accumulator: +/// those can retain the hash map after its stream-owned reservation is released. +/// Keep an unexecuted template here and create the producer and consumer together +/// for each stream. Only their metric handles are retained by the Spark plan. +#[derive(Debug)] +pub(crate) struct DynamicFilterJoinExec { + template: HashJoinExec, + config: ConfigOptions, + metrics: ExecutionPlanMetricsSet, +} + +/// Per-execution join state. The permanent plan keeps no live filter; this value +/// records whether this execution also connected its filter to the Parquet reader. +struct RuntimeDynamicFilterJoin { + join: HashJoinExec, + reader_filter_attached: bool, +} + +/// Recognize only direct-column null checks joined by AND, without evaluating +/// or changing the predicate. Every accepted leaf is deterministic, infallible, +/// and only discards rows, so reader pruning cannot suppress expression errors +/// or alter stateful evaluation. All other expressions remain a boundary. +fn is_direct_column_null_checks(predicate: &Arc) -> bool { + if let Some(binary) = predicate.downcast_ref::() { + return binary.op() == &Operator::And + && is_direct_column_null_checks(binary.left()) + && is_direct_column_null_checks(binary.right()); + } + predicate + .downcast_ref::() + .is_some_and(|is_not_null| is_not_null.arg().is::()) +} + +fn try_attach_parquet_reader_filter( + input: &Arc, + predicate: Arc, + config: &ConfigOptions, +) -> Result>> { + // Filtering before a fetch can change which rows are selected by its limit. + if input.fetch().is_some() { + log::debug!("Join dynamic filter reader pushdown skipped: probe has a fetch limit"); + return Ok(None); + } + // Spark inserts IS NOT NULL residuals above equijoin inputs, including AND + // chains of inferred null checks. A reader predicate can cross those direct + // checks because both operations only discard rows. Keep every other filter + // as a boundary: reader pruning would change which rows reach stateful + // expressions and can suppress expression errors. + if let Some(filter) = input.downcast_ref::() { + if filter.has_projection() { + log::debug!( + "Join dynamic filter reader pushdown skipped: probe FilterExec has a projection" + ); + return Ok(None); + } + if !is_direct_column_null_checks(filter.predicate()) { + log::debug!( + "Join dynamic filter reader pushdown skipped: probe filter is not direct column IS NOT NULL checks" + ); + return Ok(None); + } + let Some(reader) = + try_attach_parquet_reader_filter(filter.input(), Arc::clone(&predicate), config)? + else { + return Ok(None); + }; + return match filter.with_execution_input(reader) { + Ok(updated) => Ok(Some(updated)), + Err(error) => { + log::debug!( + "Join dynamic filter reader pushdown skipped: probe filter rebuild failed: {error}" + ); + Ok(None) + } + }; + } + let Some(scan) = input.downcast_ref::() else { + log::debug!( + "Join dynamic filter reader pushdown skipped: probe root is {}", + input.name() + ); + return Ok(None); + }; + if scan.downcast_to_file_source::().is_none() { + log::debug!("Join dynamic filter reader pushdown skipped: probe is not Parquet"); + return Ok(None); + } + + let predicate: Arc = predicate; + let propagation = match scan + .data_source() + .try_pushdown_filters(vec![predicate], config) + { + Ok(propagation) => propagation, + Err(error) => { + log::debug!( + "Join dynamic filter reader pushdown skipped: predicate remapping failed: {error}" + ); + return Ok(None); + } + }; + let Some(data_source) = propagation.updated_node else { + log::debug!("Join dynamic filter reader pushdown skipped: Parquet declined the predicate"); + return Ok(None); + }; + Ok(Some(Arc::new(scan.clone().with_data_source(data_source)))) +} + +impl DynamicFilterJoinExec { + /// Return no wrapper when the join cannot safely use a runtime filter. + pub(crate) fn try_new(join: &HashJoinExec, config: &ConfigOptions) -> Result> { + if let Some(reason) = ineligible_reason(join, config)? { + log::debug!("Join dynamic filter skipped: {reason}"); + return Ok(None); + } + Ok(Some(Self::new(join, config.clone())?)) + } + + fn new(join: &HashJoinExec, config: ConfigOptions) -> Result { + Ok(Self { + template: join.builder().reset_state().build()?, + config, + metrics: ExecutionPlanMetricsSet::new(), + }) + } + + fn build_runtime_join(&self) -> Result { + let predicate = Arc::new(DynamicFilterPhysicalExpr::new( + vec![Arc::clone(&self.template.on()[0].1)], + lit(true), + )); + let reader = try_attach_parquet_reader_filter( + self.template.right(), + Arc::clone(&predicate), + &self.config, + )?; + let reader_filter_attached = reader.is_some(); + let consumer = Arc::new(DynamicFilterExec::new( + reader.unwrap_or_else(|| Arc::clone(self.template.right())), + Arc::clone(&predicate), + self.metrics.clone(), + )); + // In particular, do not share CollectLeft's cached build future with the + // template, another execution, or a reset plan. + let join = self + .template + .builder() + .reset_state() + .with_new_children(vec![Arc::clone(self.template.left()), consumer])? + .build()? + .with_dynamic_filter_expr(predicate)?; + Ok(RuntimeDynamicFilterJoin { + join, + reader_filter_attached, + }) + } + + fn execute_runtime_join( + &self, + join: HashJoinExec, + partition: usize, + context: Arc, + ) -> Result { + // DataFusion can materialize one IN-list literal per build row, + // despite admitting the list by packed-array bytes and distinct-key count. + // Avoid that unreserved allocation. Map membership + // reuses the join's already-reserved hash table and preserves duplicates. + let mut config = context.session_config().clone(); + config + .options_mut() + .optimizer + .hash_join_inlist_pushdown_max_size = 0; + config + .options_mut() + .optimizer + .hash_join_inlist_pushdown_max_distinct_values = 0; + let context = Arc::new(TaskContext::new( + context.task_id(), + context.session_id(), + config, + context.scalar_functions().clone(), + context.higher_order_functions().clone(), + context.aggregate_functions().clone(), + context.window_functions().clone(), + context.runtime_env(), + )); + let result = join.execute(partition, context); + // HashJoinExec registers its metrics synchronously in execute(). Keep the + // live counters, including on error, without retaining the producer plan. + for metric in join.metrics().unwrap_or_default().iter() { + self.metrics.register(Arc::clone(metric)); + } + drop(join); + let input = result?; + // Drop execution state at EOF or error even if the caller retains the + // exhausted stream. Dropping a pending stream also drops all of its state. + let stream = futures::stream::unfold(Some(input), |input| async move { + let mut input = input?; + let batch = input.next().await?; + let remaining = if batch.is_ok() { Some(input) } else { None }; + Some((batch, remaining)) + }); + Ok(Box::pin(RecordBatchStreamAdapter::new( + self.schema(), + stream, + ))) + } +} + +impl DisplayAs for DynamicFilterJoinExec { + fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { + write!(f, "CometDynamicFilterJoinExec: ")?; + self.template.fmt_as(t, f) + } +} + +impl ExecutionPlan for DynamicFilterJoinExec { + fn name(&self) -> &str { + "CometDynamicFilterJoinExec" + } + + fn properties(&self) -> &Arc { + self.template.properties() + } + + fn input_distribution_requirements(&self) -> InputDistributionRequirements { + self.template.input_distribution_requirements() + } + + fn maintains_input_order(&self) -> Vec { + self.template.maintains_input_order() + } + + fn children(&self) -> Vec<&Arc> { + self.template.children() + } + + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + self.template.apply_expressions(f) + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn replace_children( + self: Arc, + children: Vec>, + _options: ReplaceChildrenOptions, + ) -> Result> { + let join = self + .template + .builder() + .reset_state() + .with_new_children(children)? + .build()?; + match Self::try_new(&join, &self.config)? { + Some(wrapper) => Ok(Arc::new(wrapper)), + None => Ok(Arc::new(join)), + } + } + + fn reset_state(self: Arc) -> Result> { + Ok(Arc::new(Self::new(&self.template, self.config.clone())?)) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + let runtime = self.build_runtime_join()?; + let attachment_metric = if runtime.reader_filter_attached { + "dynamic_filter_reader_filters_attached" + } else { + "dynamic_filter_reader_filters_skipped" + }; + MetricBuilder::new(&self.metrics) + .counter(attachment_metric, partition) + .add(1); + self.execute_runtime_join(runtime.join, partition, context) + } + + fn metrics(&self) -> Option { + Some(self.metrics.clone_inner()) + } + + fn child_stats_requests(&self, partition: Option) -> Vec { + self.template.child_stats_requests(partition) + } + + fn statistics_from_inputs( + &self, + input_stats: &[Arc], + args: &StatisticsArgs, + ) -> Result> { + self.template.statistics_from_inputs(input_stats, args) + } +} + +fn ineligible_reason(join: &HashJoinExec, config: &ConfigOptions) -> Result> { + if !config.optimizer.enable_dynamic_filter_pushdown + || !config.optimizer.enable_join_dynamic_filter_pushdown + { + return Ok(Some("disabled by DataFusion session options")); + } + if join.join_type() != &JoinType::Inner + || join.null_equality() != NullEquality::NullEqualsNothing + { + return Ok(Some("only ordinary inner equijoins are supported")); + } + if !matches!( + join.partition_mode(), + PartitionMode::Partitioned | PartitionMode::CollectLeft + ) { + return Ok(Some("unresolved hash join partition mode")); + } + if config.optimizer.preserve_file_partitions > 0 + && matches!(join.partition_mode(), PartitionMode::Partitioned) + { + return Ok(Some("DataFusion preserve_file_partitions is enabled")); + } + // Spark, not DataFusion, routes rows across tasks. Restrict the filter to the + // single native partition executed by this Spark task: no shared domains or + // assumptions about DataFusion's repartition hash across Spark partitions. + if join.left().output_partitioning().partition_count() != 1 + || join.right().output_partitioning().partition_count() != 1 + { + return Ok(Some("requires one native partition per input")); + } + let [(build_key, probe_key)] = join.on() else { + return Ok(Some("requires one join key")); + }; + if !build_key.is::() || !probe_key.is::() { + return Ok(Some("computed join keys are not supported")); + } + let build_type = build_key.data_type(join.left().schema().as_ref())?; + let probe_type = probe_key.data_type(join.right().schema().as_ref())?; + if build_type != probe_type + || !matches!( + build_type, + DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 + ) + { + return Ok(Some("requires matching signed integer keys")); + } + Ok(None) +} + +#[cfg(test)] +mod tests; diff --git a/native/core/src/execution/operators/dynamic_filter_join/tests.rs b/native/core/src/execution/operators/dynamic_filter_join/tests.rs new file mode 100644 index 00000000000..5b3e33f60d0 --- /dev/null +++ b/native/core/src/execution/operators/dynamic_filter_join/tests.rs @@ -0,0 +1,1610 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use super::*; +use std::fmt::Display; +use std::hash::{Hash, Hasher}; + +use crate::execution::planner::PhysicalPlanner; +use crate::parquet::parquet_exec::init_datasource_exec; +use crate::parquet::parquet_support::ObjectStoreBackend; +use arrow::array::{ArrayRef, BooleanArray, Int32Array, Int64Array, Int8Array, RecordBatch}; +use arrow::compute::{cast, filter_record_batch}; +use arrow::datatypes::{Field, Schema}; +use datafusion::common::test_util::batches_to_sort_string; +use datafusion::common::tree_node::TreeNodeRecursion; +use datafusion::datasource::listing::PartitionedFile; +use datafusion::datasource::memory::MemorySourceConfig; +use datafusion::execution::memory_pool::{GreedyMemoryPool, MemoryPool}; +use datafusion::execution::object_store::ObjectStoreUrl; +use datafusion::execution::runtime_env::RuntimeEnvBuilder; +use datafusion::logical_expr::ColumnarValue; +use datafusion::physical_plan::collect; +use datafusion::physical_plan::filter::FilterExec; +use datafusion::physical_plan::metrics::MetricValue; +use datafusion::physical_plan::projection::ProjectionExec; +use datafusion::physical_plan::statistics::{StatisticsArgs, StatisticsContext}; +use datafusion::physical_plan::{ChildrenPropertiesMode, Distribution, ReplaceChildrenOptions}; +use datafusion::prelude::{SessionConfig, SessionContext}; +use datafusion_comet_spark_expr::RandExpr; +use datafusion_datasource::file::FileSource; +use parquet::arrow::ArrowWriter; +use parquet::file::properties::{EnabledStatistics, WriterProperties}; + +fn input( + values: Vec>, + key_type: &DataType, + key_index: usize, +) -> Arc { + let payload = Arc::new(Int32Array::from_iter_values(0..values.len() as i32)) as ArrayRef; + let key = cast(&Int32Array::from(values), key_type).unwrap(); + let mut fields = vec![ + Field::new("key", key_type.clone(), true), + Field::new("payload", DataType::Int32, false), + ]; + let mut columns = vec![key, payload]; + fields.swap(0, key_index); + columns.swap(0, key_index); + let schema = Arc::new(Schema::new(fields)); + let batch = RecordBatch::try_new(schema, columns).unwrap(); + // Multiple build batches prove that an early subset of keys cannot prune + // matches belonging to a later batch. + let batches = if batch.num_rows() == 0 { + vec![batch] + } else { + (0..batch.num_rows()) + .step_by(2) + .map(|offset| batch.slice(offset, 2.min(batch.num_rows() - offset))) + .collect() + }; + memory_exec(batches) +} + +fn memory_exec(batches: Vec) -> Arc { + MemorySourceConfig::try_new_exec(std::slice::from_ref(&batches), batches[0].schema(), None) + .unwrap() +} + +fn produced_join_filter(join: &HashJoinExec) -> Arc { + let mut expressions = join.dynamic_expressions_produced(); + assert_eq!(expressions.len(), 1); + let expression: Arc = expressions.pop().unwrap(); + expression.downcast::().unwrap() +} + +fn join( + build: Arc, + probe: Arc, + swap: bool, +) -> Arc { + let build_key = Arc::new(Column::new("key", 1)) as Arc; + let probe_key = Arc::new(Column::new("key", 0)) as Arc; + let (left, right, on) = if swap { + (probe, build, vec![(probe_key, build_key)]) + } else { + (build, probe, vec![(build_key, probe_key)]) + }; + let join = HashJoinExec::try_new( + left, + right, + on, + None, + &JoinType::Inner, + None, + PartitionMode::Partitioned, + NullEquality::NullEqualsNothing, + false, + ) + .unwrap(); + if swap { + join.swap_inputs(PartitionMode::Partitioned).unwrap() + } else { + Arc::new(join) + } +} + +fn metric(plan: &Arc, name: &str) -> usize { + if let Some(projection) = plan.downcast_ref::() { + return metric(projection.input(), name); + } + plan.metrics() + .unwrap() + .sum_by_name(name) + .unwrap() + .as_usize() +} + +fn row_count(batches: &[RecordBatch]) -> usize { + batches.iter().map(RecordBatch::num_rows).sum() +} + +/// Inspect the batch given to the real completed predicate, without changing its result. +#[derive(Debug, Eq)] +struct AssertKeyOnlyBatch { + child: Arc, + key_values_ptr: usize, +} + +impl PartialEq for AssertKeyOnlyBatch { + fn eq(&self, other: &Self) -> bool { + self.child.eq(&other.child) && self.key_values_ptr == other.key_values_ptr + } +} + +impl Hash for AssertKeyOnlyBatch { + fn hash(&self, state: &mut H) { + self.child.hash(state); + self.key_values_ptr.hash(state); + } +} + +impl Display for AssertKeyOnlyBatch { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "AssertKeyOnlyBatch({})", self.child) + } +} + +impl PhysicalExpr for AssertKeyOnlyBatch { + fn data_type(&self, input_schema: &Schema) -> Result { + self.child.data_type(input_schema) + } + + fn nullable(&self, input_schema: &Schema) -> Result { + self.child.nullable(input_schema) + } + + fn evaluate(&self, batch: &RecordBatch) -> Result { + assert_eq!( + batch.num_columns(), + 1, + "predicate must not receive payload columns" + ); + let keys = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!( + keys.values().as_ptr() as usize, + self.key_values_ptr, + "projecting the join key must not copy its values" + ); + self.child.evaluate(batch) + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.child] + } + + fn with_new_children( + self: Arc, + mut children: Vec>, + ) -> Result> { + assert_eq!(children.len(), 1); + Ok(Arc::new(Self { + child: children.remove(0), + key_values_ptr: self.key_values_ptr, + })) + } + + fn fmt_sql(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + Display::fmt(self, f) + } +} + +#[tokio::test] +async fn completed_filter_evaluates_only_the_shared_probe_key() { + // With no nulls and only 1/8 of rows inside the build bounds, DataFusion's AND + // evaluation preselects those rows before evaluating hash_lookup. A permutation + // spreads the selected rows throughout the batch, forcing payload copies if the + // full probe batch reaches the predicate. + let keys = Arc::new(Int32Array::from_iter_values( + (0..8192).map(|row| (row * 641) % 8192), + )); + let mut fields = (0..32) + .map(|column| Field::new(format!("payload_{column}"), DataType::Int64, false)) + .collect::>(); + let mut columns = (0..32) + .map(|column| { + Arc::new(Int64Array::from_iter_values( + (0..8192).map(move |row| i64::from(row) * 32 + i64::from(column)), + )) as ArrayRef + }) + .collect::>(); + let key_index = 17; + fields.insert(key_index, Field::new("key", DataType::Int32, false)); + columns.insert(key_index, Arc::clone(&keys) as ArrayRef); + let batch = RecordBatch::try_new(Arc::new(Schema::new(fields)), columns).unwrap(); + let probe = memory_exec(vec![batch.clone()]); + let join = HashJoinExec::try_new( + input((0..1024).map(Some).collect(), &DataType::Int32, 0), + Arc::clone(&probe), + vec![( + Arc::new(Column::new("key", 0)), + Arc::new(Column::new("key", key_index)), + )], + None, + &JoinType::Inner, + None, + PartitionMode::Partitioned, + NullEquality::NullEqualsNothing, + false, + ) + .unwrap(); + let session = SessionContext::new(); + let expected = collect( + Arc::new(join.builder().build().unwrap()), + session.task_ctx(), + ) + .await + .unwrap(); + let wrapper = + DynamicFilterJoinExec::new(&join, session.copied_config().options().as_ref().clone()) + .unwrap(); + let runtime = wrapper.build_runtime_join().unwrap(); + let predicate = produced_join_filter(&runtime.join); + let actual = datafusion::physical_plan::common::collect( + wrapper + .execute_runtime_join(runtime.join, 0, session.task_ctx()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!( + batches_to_sort_string(&actual), + batches_to_sort_string(&expected) + ); + + // Inspect the actual build-generated bounds AND hash-membership expression. Its key + // remains at index 17 here, so the consumer must also remap every nested reference. + let completed = predicate.current().unwrap(); + assert!(completed.to_string().contains("hash_lookup")); + assert!(completed.to_string().contains("AND")); + predicate + .update(Arc::new(AssertKeyOnlyBatch { + child: completed, + key_values_ptr: keys.values().as_ptr() as usize, + })) + .unwrap(); + let consumer: Arc = Arc::new(DynamicFilterExec::new( + probe, + predicate, + ExecutionPlanMetricsSet::new(), + )); + let filtered = collect(consumer, session.task_ctx()).await.unwrap(); + let selected = BooleanArray::from( + keys.values() + .iter() + .map(|key| *key < 1024) + .collect::>(), + ); + let expected = filter_record_batch(&batch, &selected).unwrap(); + assert_eq!(expected.num_rows(), 1024); + assert_eq!(filtered, vec![expected]); +} + +#[tokio::test] +async fn completed_build_filters_both_sides_and_session_inlist_settings() { + for key_type in [ + DataType::Int8, + DataType::Int16, + DataType::Int32, + DataType::Int64, + ] { + for swap in [false, true] { + for max_inlist_size in [0, 1024 * 1024] { + let mut config = SessionConfig::new(); + config + .options_mut() + .optimizer + .hash_join_inlist_pushdown_max_size = max_inlist_size; + let session = SessionContext::new_with_config(config); + for build_values in [ + vec![Some(-5), Some(20), None, Some(20), Some(90)], + vec![], + vec![None, None], + ] { + let build = input(build_values.clone(), &key_type, 1); + let probe = input((-100..=100).map(Some).chain([None]).collect(), &key_type, 0); + let plain = join(Arc::clone(&build), Arc::clone(&probe), swap); + let attached = PhysicalPlanner::apply_join_dynamic_filter( + join(build, probe, swap), + true, + session.copied_config().options(), + ) + .unwrap(); + let native_join = + if let Some(projection) = attached.downcast_ref::() { + projection.input() + } else { + &attached + }; + assert!(native_join.is::()); + assert_eq!(plain.schema(), attached.schema()); + let expected = collect(plain, session.task_ctx()).await.unwrap(); + let actual = collect(Arc::clone(&attached), session.task_ctx()) + .await + .unwrap(); + assert_eq!( + batches_to_sort_string(&actual), + batches_to_sort_string(&expected) + ); + if build_values.iter().any(Option::is_some) { + assert_eq!( + row_count(&actual), + 4, + "duplicates and late build keys must match" + ); + assert_eq!(metric(&attached, "dynamic_filter_rows_evaluated"), 202); + assert!(metric(&attached, "dynamic_filter_rows_pruned") >= 199); + assert_eq!(metric(&attached, "dynamic_filter_rows_bypassed"), 0); + } else { + assert_eq!(row_count(&actual), 0); + } + } + } + } + } +} + +fn plain_join() -> HashJoinExec { + let plan = join( + input(vec![Some(10)], &DataType::Int32, 1), + input(vec![Some(10), Some(20)], &DataType::Int32, 0), + false, + ); + plan.downcast_ref::() + .unwrap() + .builder() + .build() + .unwrap() +} + +fn assert_skipped(join: HashJoinExec, config: &ConfigOptions) { + let plain: Arc = Arc::new(join); + let attached = + PhysicalPlanner::apply_join_dynamic_filter(Arc::clone(&plain), true, config).unwrap(); + assert!( + Arc::ptr_eq(&plain, &attached), + "fallback must preserve the original plan" + ); + assert!(attached + .downcast_ref::() + .unwrap() + .dynamic_expressions_produced() + .is_empty()); +} + +#[test] +fn skips_unsupported_joins_and_session_disables() { + let default = ConfigOptions::default(); + for join_type in [ + JoinType::Left, + JoinType::Right, + JoinType::Full, + JoinType::LeftSemi, + JoinType::RightSemi, + JoinType::LeftAnti, + JoinType::RightAnti, + JoinType::LeftMark, + JoinType::RightMark, + ] { + assert_skipped( + plain_join().builder().with_type(join_type).build().unwrap(), + &default, + ); + } + assert_skipped( + plain_join() + .builder() + .with_null_equality(NullEquality::NullEqualsNull) + .build() + .unwrap(), + &default, + ); + assert_skipped( + plain_join() + .builder() + .with_partition_mode(PartitionMode::Auto) + .build() + .unwrap(), + &default, + ); + assert_skipped( + plain_join() + .builder() + .with_type(JoinType::LeftAnti) + .with_null_aware(true) + .with_partition_mode(PartitionMode::CollectLeft) + .build() + .unwrap(), + &default, + ); + for option in [ + "enable_dynamic_filter_pushdown", + "enable_join_dynamic_filter_pushdown", + "preserve_file_partitions", + ] { + let mut config = ConfigOptions::default(); + match option { + "enable_dynamic_filter_pushdown" => { + config.optimizer.enable_dynamic_filter_pushdown = false + } + "enable_join_dynamic_filter_pushdown" => { + config.optimizer.enable_join_dynamic_filter_pushdown = false + } + _ => config.optimizer.preserve_file_partitions = 1, + } + assert_skipped(plain_join(), &config); + } +} + +#[test] +fn skips_unsupported_keys_and_multiple_native_partitions() { + for key_type in [ + DataType::Float32, + DataType::Float64, + DataType::Utf8, + DataType::Decimal128(10, 0), + ] { + let plan = join( + input(vec![Some(10)], &key_type, 1), + input(vec![Some(10)], &key_type, 0), + false, + ); + let same = PhysicalPlanner::apply_join_dynamic_filter( + Arc::clone(&plan), + true, + &ConfigOptions::default(), + ) + .unwrap(); + assert!(Arc::ptr_eq(&plan, &same)); + } + let plain = plain_join(); + let (build, probe) = plain.on()[0].clone(); + let computed = Arc::new(BinaryExpr::new( + Arc::clone(&probe), + Operator::Plus, + lit(1i32), + )); + for keys in [ + vec![(Arc::clone(&build), computed as Arc)], + vec![(Arc::clone(&build), Arc::clone(&probe)), (build, probe)], + ] { + assert_skipped( + plain.builder().with_on(keys).build().unwrap(), + &ConfigOptions::default(), + ); + } + for replace_build in [false, true] { + let schema = if replace_build { + plain.left().schema() + } else { + plain.right().schema() + }; + let two_partitions: Arc = MemorySourceConfig::try_new_exec( + &[ + vec![RecordBatch::new_empty(Arc::clone(&schema))], + vec![RecordBatch::new_empty(Arc::clone(&schema))], + ], + schema, + None, + ) + .unwrap(); + let children = if replace_build { + vec![two_partitions, Arc::clone(plain.right())] + } else { + vec![Arc::clone(plain.left()), two_partitions] + }; + assert_skipped( + plain + .builder() + .with_new_children(children) + .unwrap() + .build() + .unwrap(), + &ConfigOptions::default(), + ); + } +} + +#[tokio::test] +async fn independent_attempts_do_not_share_build_domains() { + let session = SessionContext::new(); + for build_key in [5, 90] { + let plan = join( + input(vec![Some(build_key)], &DataType::Int32, 1), + input(vec![Some(5), Some(90)], &DataType::Int32, 0), + false, + ); + let attached = PhysicalPlanner::apply_join_dynamic_filter( + plan, + true, + session.copied_config().options(), + ) + .unwrap(); + let output = collect(Arc::clone(&attached), session.task_ctx()) + .await + .unwrap(); + assert_eq!(row_count(&output), 1); + let keys = output + .iter() + .flat_map(|batch| { + batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .iter() + .copied() + }) + .collect::>(); + assert_eq!(keys, vec![build_key]); + assert_eq!(metric(&attached, "dynamic_filter_rows_pruned"), 1); + } +} + +fn single_key_join_plans( + build: Arc, + probe: Arc, + mode: PartitionMode, +) -> HashJoinExec { + HashJoinExec::try_new( + build, + probe, + vec![( + Arc::new(Column::new("key", 0)), + Arc::new(Column::new("key", 0)), + )], + None, + &JoinType::Inner, + None, + mode, + NullEquality::NullEqualsNothing, + false, + ) + .unwrap() +} + +fn single_key_join(build: ArrayRef, probe: ArrayRef, mode: PartitionMode) -> HashJoinExec { + let schema = Arc::new(Schema::new(vec![Field::new( + "key", + build.data_type().clone(), + true, + )])); + let build = memory_exec(vec![ + RecordBatch::try_new(Arc::clone(&schema), vec![build]).unwrap() + ]); + let probe = memory_exec(vec![RecordBatch::try_new(schema, vec![probe]).unwrap()]); + single_key_join_plans(build, probe, mode) +} + +fn parquet_probe( + values: Vec, + session: &Arc, + max_row_group_rows: usize, +) -> (tempfile::NamedTempFile, Arc) { + // Put the key after an unused physical column and project only the key. + // Reader attachment must remap key@0 in the join to key@1 in the file. + let file_schema = Arc::new(Schema::new(vec![ + Field::new("payload", DataType::Int32, false), + Field::new("key", DataType::Int32, false), + ])); + let required_schema = Arc::new(Schema::new(vec![file_schema.field(1).clone()])); + let row_count = values.len(); + let batch = RecordBatch::try_new( + Arc::clone(&file_schema), + vec![ + Arc::new(Int32Array::from_iter_values(0..row_count as i32)), + Arc::new(Int32Array::from(values)), + ], + ) + .unwrap(); + let file = tempfile::NamedTempFile::new().unwrap(); + let props = WriterProperties::builder() + .set_max_row_group_row_count(Some(max_row_group_rows)) + .set_statistics_enabled(EnabledStatistics::Chunk) + .set_dictionary_enabled(false) + .build(); + let mut writer = ArrowWriter::try_new( + file.reopen().unwrap(), + Arc::clone(&file_schema), + Some(props), + ) + .unwrap(); + writer.write(&batch).unwrap(); + let metadata = writer.close().unwrap(); + assert_eq!( + metadata.num_row_groups(), + row_count.div_ceil(max_row_group_rows) + ); + assert!(metadata + .row_groups() + .iter() + .all(|group| { group.num_rows() > 0 && group.num_rows() as usize <= max_row_group_rows })); + + let non_null: Arc = + Arc::new(IsNotNullExpr::new(Arc::new(Column::new("key", 0)))); + let scan = init_datasource_exec( + required_schema, + Some(file_schema), + None, + ObjectStoreUrl::local_filesystem(), + ObjectStoreBackend::Local, + vec![vec![PartitionedFile::from_path( + file.path().to_string_lossy().into_owned(), + ) + .unwrap()]], + Some(vec![1]), + Some(vec![non_null]), + None, + "UTC", + true, + false, + false, + false, + session, + false, + false, + false, + ) + .unwrap(); + (file, scan) +} + +fn filtered_probe(scan: &Arc) -> Arc { + let predicate: Arc = + Arc::new(IsNotNullExpr::new(Arc::new(Column::new("key", 0)))); + Arc::new(CometFilterExec::from_datafusion( + FilterExec::try_new(predicate, Arc::clone(scan) as Arc).unwrap(), + )) +} + +fn find_dynamic_filter(expr: &Arc) -> Option<&DynamicFilterPhysicalExpr> { + if let Some(filter) = expr.downcast_ref::() { + return Some(filter); + } + expr.children() + .into_iter() + .find_map(|child| find_dynamic_filter(child)) +} + +/// Accept nested null-check conjunctions while retaining all other filter +/// boundaries, including OR and computed or potentially failing expressions. +#[test] +fn reader_filter_crosses_only_direct_column_null_checks() { + let key: Arc = Arc::new(Column::new("key", 0)); + let direct_null_check: Arc = Arc::new(IsNotNullExpr::new(Arc::clone(&key))); + assert!(is_direct_column_null_checks(&direct_null_check)); + let other_null_check: Arc = + Arc::new(IsNotNullExpr::new(Arc::new(Column::new("other", 1)))); + let conjunction: Arc = Arc::new(BinaryExpr::new( + Arc::clone(&direct_null_check), + Operator::And, + Arc::clone(&other_null_check), + )); + let nested: Arc = Arc::new(BinaryExpr::new( + conjunction, + Operator::And, + Arc::clone(&direct_null_check), + )); + assert!(is_direct_column_null_checks(&nested)); + let right_nested: Arc = Arc::new(BinaryExpr::new( + Arc::clone(&direct_null_check), + Operator::And, + Arc::new(BinaryExpr::new( + Arc::clone(&other_null_check), + Operator::And, + Arc::clone(&direct_null_check), + )), + )); + assert!(is_direct_column_null_checks(&right_nested)); + let disjunction: Arc = Arc::new(BinaryExpr::new( + Arc::clone(&direct_null_check), + Operator::Or, + other_null_check, + )); + assert!(!is_direct_column_null_checks(&disjunction)); + + let comparison: Arc = + Arc::new(BinaryExpr::new(Arc::clone(&key), Operator::Gt, lit(0_i32))); + let conjunction: Arc = Arc::new(BinaryExpr::new( + Arc::clone(&direct_null_check), + Operator::And, + comparison, + )); + assert!(!is_direct_column_null_checks(&conjunction)); + + let computed: Arc = + Arc::new(BinaryExpr::new(key, Operator::Plus, lit(1_i32))); + let computed_null_check: Arc = Arc::new(IsNotNullExpr::new(computed)); + assert!(!is_direct_column_null_checks(&computed_null_check)); +} + +/// Exercise a real Parquet reader with three distinct nullable columns and +/// remapped projection. Both statistics-only and row-filter reads must retain +/// the AND residual: build keys 2 and 3 fail its payload checks, leaving key 1. +#[tokio::test] +async fn reader_filter_crosses_null_check_conjunction_and_retains_residual() { + let file_schema = Arc::new(Schema::new(vec![ + Field::new("payload", DataType::Int32, true), + Field::new("key", DataType::Int32, true), + Field::new("other", DataType::Int32, true), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&file_schema), + vec![ + Arc::new(Int32Array::from(vec![ + Some(10), + None, + Some(30), + Some(40), + Some(50), + ])), + Arc::new(Int32Array::from(vec![ + Some(1), + Some(2), + Some(3), + Some(4), + None, + ])), + Arc::new(Int32Array::from(vec![ + Some(1), + Some(1), + None, + Some(1), + Some(1), + ])), + ], + ) + .unwrap(); + let file = tempfile::NamedTempFile::new().unwrap(); + let props = WriterProperties::builder() + .set_dictionary_enabled(false) + .build(); + let mut writer = ArrowWriter::try_new( + file.reopen().unwrap(), + Arc::clone(&file_schema), + Some(props), + ) + .unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + let required_schema = Arc::new(Schema::new(vec![ + file_schema.field(1).clone(), + file_schema.field(0).clone(), + file_schema.field(2).clone(), + ])); + let mut outputs = Vec::new(); + for (row_filter, enabled) in [(false, false), (false, true), (true, false), (true, true)] { + let mut config = SessionConfig::new() + .with_target_partitions(1) + .with_parquet_page_index_pruning(false); + config.options_mut().execution.parquet.pushdown_filters = row_filter; + let session = Arc::new(SessionContext::new_with_config(config)); + let scan = init_datasource_exec( + Arc::clone(&required_schema), + Some(Arc::clone(&file_schema)), + None, + ObjectStoreUrl::local_filesystem(), + ObjectStoreBackend::Local, + vec![vec![PartitionedFile::from_path( + file.path().to_string_lossy().into_owned(), + ) + .unwrap()]], + Some(vec![1, 0, 2]), + None, + None, + "UTC", + true, + false, + false, + false, + &session, + false, + false, + false, + ) + .unwrap(); + let checks = [("key", 0), ("payload", 1), ("other", 2)].map(|(name, index)| { + Arc::new(IsNotNullExpr::new(Arc::new(Column::new(name, index)))) + as Arc + }); + let predicate: Arc = Arc::new(BinaryExpr::new( + Arc::new(BinaryExpr::new( + Arc::clone(&checks[0]), + Operator::And, + Arc::clone(&checks[1]), + )), + Operator::And, + Arc::clone(&checks[2]), + )); + let filter = Arc::new(CometFilterExec::from_datafusion( + FilterExec::try_new( + Arc::clone(&predicate), + Arc::clone(&scan) as Arc, + ) + .unwrap(), + )); + let build_schema = Arc::new(Schema::new(vec![Field::new("key", DataType::Int32, true)])); + let build = memory_exec(vec![RecordBatch::try_new( + build_schema, + vec![Arc::new(Int32Array::from(vec![1, 2, 3]))], + ) + .unwrap()]); + let join = single_key_join_plans(build, filter, PartitionMode::Partitioned); + let plan: Arc = if enabled { + let wrapper = DynamicFilterJoinExec::new( + &join, + session.copied_config().options().as_ref().clone(), + ) + .unwrap(); + let runtime = wrapper.build_runtime_join().unwrap(); + assert!(runtime.reader_filter_attached); + let consumer = runtime + .join + .right() + .downcast_ref::() + .unwrap(); + let retained = consumer.children()[0] + .downcast_ref::() + .unwrap(); + assert_eq!(retained.predicate().to_string(), predicate.to_string()); + Arc::new(wrapper) + } else { + Arc::new(join) + }; + let output = collect(Arc::clone(&plan), session.task_ctx()) + .await + .unwrap(); + assert_eq!(row_count(&output), 1); + if enabled { + if row_filter { + let metrics = scan.metrics().unwrap(); + assert!( + metrics + .sum_by_name("pushdown_rows_pruned") + .unwrap() + .as_usize() + > 0, + "{metrics}" + ); + } + assert_eq!( + plan.metrics() + .unwrap() + .sum_by_name("dynamic_filter_reader_filters_attached") + .unwrap() + .as_usize(), + 1 + ); + } + outputs.push(batches_to_sort_string(&output)); + } + assert!(outputs.windows(2).all(|pair| pair[0] == pair[1])); +} + +#[tokio::test] +async fn reader_filter_does_not_cross_fetch_limits() { + for limit_filter in [false, true] { + let mut outputs = Vec::new(); + for enabled in [false, true] { + let session = Arc::new(SessionContext::new_with_config( + SessionConfig::new().with_target_partitions(1), + )); + let (_file, scan) = parquet_probe((0..4).collect(), &session, 1); + let probe = if limit_filter { + filtered_probe(&scan).with_fetch(Some(1)).unwrap() + } else { + scan.with_fetch(Some(1)).unwrap() + }; + let schema = Arc::new(Schema::new(vec![Field::new("key", DataType::Int32, false)])); + let build = memory_exec(vec![RecordBatch::try_new( + schema, + vec![Arc::new(Int32Array::from(vec![3]))], + ) + .unwrap()]); + let join = single_key_join_plans(build, probe, PartitionMode::Partitioned); + let plan: Arc = if enabled { + PhysicalPlanner::apply_join_dynamic_filter( + Arc::new(join), + true, + session.copied_config().options(), + ) + .unwrap() + } else { + Arc::new(join) + }; + let output = collect(Arc::clone(&plan), session.task_ctx()) + .await + .unwrap(); + outputs.push(row_count(&output)); + if enabled { + assert_eq!(metric(&plan, "dynamic_filter_reader_filters_skipped"), 1); + } + } + assert_eq!(outputs, vec![0, 0]); + } +} + +#[tokio::test] +async fn reader_filter_does_not_cross_seeded_rand_probe_filter() { + let mut outputs = Vec::new(); + for enabled in [false, true] { + let mut config = SessionConfig::new() + .with_target_partitions(1) + .with_parquet_page_index_pruning(false); + config.options_mut().execution.parquet.pushdown_filters = false; + let session = Arc::new(SessionContext::new_with_config(config)); + + // Four one-row groups make reader pruning observable in Rand's state. Without the + // filter boundary, pruning keys 0-2 makes key 3 receive the first draw (0.619...) + // instead of the fourth (0.263...), changing whether it passes rand(42) < 0.5. + let (_file, scan) = parquet_probe((0..4).collect(), &session, 1); + let key: Arc = Arc::new(Column::new("key", 0)); + let is_not_null: Arc = Arc::new(IsNotNullExpr::new(Arc::clone(&key))); + let random_below_half: Arc = Arc::new(BinaryExpr::new( + Arc::new(RandExpr::new(42)), + Operator::Lt, + lit(0.5_f64), + )); + let predicate: Arc = Arc::new(BinaryExpr::new( + is_not_null, + Operator::And, + random_below_half, + )); + let filter = Arc::new(CometFilterExec::from_datafusion( + FilterExec::try_new(predicate, Arc::clone(&scan) as Arc).unwrap(), + )); + + let build_schema = Arc::new(Schema::new(vec![Field::new("key", DataType::Int32, false)])); + let build = memory_exec(vec![RecordBatch::try_new( + build_schema, + vec![Arc::new(Int32Array::from(vec![3]))], + ) + .unwrap()]); + let join = single_key_join_plans( + build, + Arc::clone(&filter) as Arc, + PartitionMode::Partitioned, + ); + let plan: Arc = if enabled { + PhysicalPlanner::apply_join_dynamic_filter( + Arc::new(join), + true, + session.copied_config().options(), + ) + .unwrap() + } else { + Arc::new(join) + }; + let batches = collect(Arc::clone(&plan), session.task_ctx()) + .await + .unwrap(); + let probe_keys = batches + .iter() + .flat_map(|batch| { + batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .iter() + .copied() + }) + .collect::>(); + outputs.push(probe_keys); + + if enabled { + let dynamic_metric = |name| { + plan.metrics() + .and_then(|metrics| metrics.sum_by_name(name)) + .map_or(0, |value| value.as_usize()) + }; + assert_eq!(dynamic_metric("dynamic_filter_reader_filters_attached"), 0); + assert_eq!(dynamic_metric("dynamic_filter_reader_filters_skipped"), 1); + assert_eq!(dynamic_metric("dynamic_filter_rows_evaluated"), 1); + assert_eq!(dynamic_metric("dynamic_filter_rows_pruned"), 0); + assert_eq!(pruning_metric(&scan, "row_groups_pruned_statistics"), 0); + assert_eq!(filter.metrics().unwrap().output_rows().unwrap(), 1); + } + } + assert_eq!(outputs, vec![vec![3], vec![3]]); +} + +fn two_batch_build() -> Arc { + let schema = Arc::new(Schema::new(vec![Field::new("key", DataType::Int32, false)])); + memory_exec( + [150, 250] + .into_iter() + .map(|key| { + RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(vec![key]))], + ) + .unwrap() + }) + .collect(), + ) +} + +fn pruning_metric(plan: &Arc, name: &str) -> usize { + let metrics = plan.metrics().unwrap(); + let Some(value) = metrics.sum_by_name(name) else { + return 0; + }; + let MetricValue::PruningMetrics { + pruning_metrics, .. + } = value + else { + panic!("expected pruning metric {name}: {metrics}"); + }; + pruning_metrics.pruned() +} + +async fn run_parquet_join(values: Vec, enabled: bool) -> (usize, usize, usize, usize, usize) { + let mut config = SessionConfig::new() + .with_target_partitions(1) + .with_parquet_page_index_pruning(false); + // Isolate row-group pruning. The residual batch filter remains responsible + // for membership filtering after the reader applies conservative statistics. + config.options_mut().execution.parquet.pushdown_filters = false; + let session = Arc::new(SessionContext::new_with_config(config)); + let (_file, scan) = parquet_probe(values, &session, 100); + let filter = filtered_probe(&scan); + let probe = Arc::clone(&filter) as Arc; + let join = single_key_join_plans(two_batch_build(), probe, PartitionMode::Partitioned); + let plan: Arc = if enabled { + PhysicalPlanner::apply_join_dynamic_filter( + Arc::new(join), + true, + session.copied_config().options(), + ) + .unwrap() + } else { + Arc::new(join) + }; + let output = collect(Arc::clone(&plan), session.task_ctx()) + .await + .unwrap(); + let attached = plan + .metrics() + .and_then(|metrics| metrics.sum_by_name("dynamic_filter_reader_filters_attached")) + .map_or(0, |metric| metric.as_usize()); + ( + row_count(&output), + pruning_metric(&scan, "row_groups_pruned_statistics"), + scan.metrics() + .unwrap() + .sum_by_name("bytes_scanned") + .unwrap() + .as_usize(), + attached, + filter.metrics().unwrap().output_rows().unwrap_or_default(), + ) +} + +#[tokio::test] +async fn broadcast_filter_reaches_parquet_reader_after_complete_build() { + let mut config = SessionConfig::new() + .with_target_partitions(1) + .with_parquet_page_index_pruning(false); + config.options_mut().execution.parquet.pushdown_filters = false; + let session = Arc::new(SessionContext::new_with_config(config)); + let (_file, scan) = parquet_probe((0..400).collect(), &session, 100); + let join = single_key_join_plans( + two_batch_build(), + filtered_probe(&scan), + PartitionMode::Partitioned, + ); + let wrapper = + DynamicFilterJoinExec::new(&join, session.copied_config().options().as_ref().clone()) + .unwrap(); + let runtime = wrapper.build_runtime_join().unwrap(); + assert!(runtime.reader_filter_attached); + let consumer = runtime + .join + .right() + .downcast_ref::() + .unwrap(); + let filter = consumer.children()[0] + .downcast_ref::() + .unwrap(); + let reader = filter.input().downcast_ref::().unwrap(); + let (_, source) = reader.downcast_to_file_source::().unwrap(); + let reader_filter = source.filter().unwrap(); + let reader_filter = find_dynamic_filter(&reader_filter).unwrap(); + let join_filter = produced_join_filter(&runtime.join); + assert_eq!( + reader_filter.expression_id().unwrap(), + join_filter.expression_id().unwrap() + ); + let reader_keys = reader_filter.children(); + assert_eq!(reader_keys.len(), 1); + let remapped_key = reader_keys[0].downcast_ref::().unwrap(); + assert_eq!(remapped_key.name(), "key"); + assert_eq!(remapped_key.index(), 1); + + let (disabled_rows, _, disabled_bytes, _, disabled_filter_rows) = + run_parquet_join((0..400).collect(), false).await; + let (enabled_rows, enabled_pruned, enabled_bytes, attached, enabled_filter_rows) = + run_parquet_join((0..400).collect(), true).await; + assert_eq!(disabled_rows, 2); + assert_eq!(enabled_rows, disabled_rows); + assert_eq!(enabled_pruned, 2); + assert!( + enabled_bytes < disabled_bytes, + "reader filter should avoid data reads: enabled={enabled_bytes}, disabled={disabled_bytes}" + ); + assert_eq!(attached, 1); + assert_eq!(disabled_filter_rows, 400); + assert_eq!(enabled_filter_rows, 200); + + // Every row group's bounds span both build keys, so the same attached + // filter is safe but cannot avoid I/O on this deliberately poor layout. + let mut unclustered = Vec::with_capacity(400); + for group in 0..4 { + unclustered.extend(0..50); + unclustered.extend(350..400); + if group == 0 { + unclustered[group * 100 + 1] = 150; + unclustered[group * 100 + 51] = 250; + } + } + let ( + unclustered_disabled_rows, + _, + unclustered_disabled_bytes, + _, + unclustered_disabled_filter_rows, + ) = run_parquet_join(unclustered.clone(), false).await; + let ( + unclustered_enabled_rows, + unclustered_pruned, + unclustered_enabled_bytes, + attached, + unclustered_enabled_filter_rows, + ) = run_parquet_join(unclustered, true).await; + assert_eq!(unclustered_enabled_rows, 2); + assert_eq!(unclustered_enabled_rows, unclustered_disabled_rows); + assert_eq!(unclustered_pruned, 0); + assert_eq!(unclustered_enabled_bytes, unclustered_disabled_bytes); + assert_eq!(attached, 1); + assert_eq!(unclustered_disabled_filter_rows, 400); + assert_eq!(unclustered_enabled_filter_rows, 400); +} + +fn limited_session(bytes: usize) -> (SessionContext, Arc) { + let pool: Arc = Arc::new(GreedyMemoryPool::new(bytes)); + let runtime = RuntimeEnvBuilder::new() + .with_memory_pool(Arc::clone(&pool)) + .build_arc() + .unwrap(); + ( + SessionContext::new_with_config_rt(SessionConfig::new(), runtime), + pool, + ) +} + +#[tokio::test] +async fn runtime_domains_release_with_streams_while_plans_remain_alive() { + for mode in [PartitionMode::Partitioned, PartitionMode::CollectLeft] { + for finish in [ + "eof", + "cancel_before_poll", + "cancel_after_output", + "build_error", + ] { + let (session, pool) = limited_session(if finish == "build_error" { + 1 + } else { + 8 * 1024 * 1024 + }); + // Sparse keys and more than 150 distinct values require the map + // strategy even without the IN-list override. + let join = single_key_join( + Arc::new(Int64Array::from_iter_values( + (0..4096).map(|i| i * 1_000_003), + )), + Arc::new(Int64Array::from(vec![0, 1])), + mode, + ); + let plan = DynamicFilterJoinExec::new(&join, ConfigOptions::default()).unwrap(); + let producer = plan.build_runtime_join().unwrap(); + let predicate = Arc::downgrade(&produced_join_filter(&producer.join)); + let mut stream = plan + .execute_runtime_join(producer.join, 0, session.task_ctx()) + .unwrap(); + assert!(predicate.upgrade().is_some()); + + match finish { + "cancel_before_poll" => drop(stream), + "build_error" => { + assert!(stream.next().await.unwrap().is_err()); + assert!(stream.next().await.is_none()); + } + _ => { + assert_eq!(stream.next().await.unwrap().unwrap().num_rows(), 1); + assert!(pool.reserved() > 0); + assert!(predicate.upgrade().is_some()); + if finish == "eof" { + assert!(stream.next().await.is_none()); + // Retain the exhausted stream as well as the plan. + assert_eq!(pool.reserved(), 0); + assert!(predicate.upgrade().is_none()); + } else { + drop(stream); + } + let metrics = plan.metrics().unwrap(); + assert_eq!(metrics.output_rows(), Some(1)); + assert_eq!( + metrics + .sum_by_name("dynamic_filter_rows_pruned") + .unwrap() + .as_usize(), + 1 + ); + } + } + // The accumulator also owns this predicate. A zero reservation + // alone would miss the old plan-owned, unaccounted map retention. + assert_eq!(pool.reserved(), 0, "{mode:?}: {finish}"); + assert!(predicate.upgrade().is_none(), "{mode:?}: {finish}"); + assert!(plan.template.dynamic_expressions_produced().is_empty()); + } + } +} + +fn expression_nodes(expr: &Arc) -> usize { + 1 + expr + .children() + .into_iter() + .map(expression_nodes) + .sum::() +} + +#[tokio::test] +async fn duplicate_heavy_builds_do_not_materialize_unreserved_inlists() { + for mode in [PartitionMode::Partitioned, PartitionMode::CollectLeft] { + let (session, pool) = limited_session(1024 * 1024); + let join = single_key_join( + Arc::new(Int8Array::from(vec![1; 65_536])), + Arc::new(Int8Array::from(vec![1, 2])), + mode, + ); + let config = session.copied_config(); + let configured_limit = config + .options() + .optimizer + .hash_join_inlist_pushdown_max_size; + assert!(configured_limit > 0); + let plan = DynamicFilterJoinExec::new(&join, config.options().as_ref().clone()).unwrap(); + let producer = plan.build_runtime_join().unwrap(); + let predicate = Arc::downgrade(&produced_join_filter(&producer.join)); + let mut stream = plan + .execute_runtime_join(producer.join, 0, session.task_ctx()) + .unwrap(); + let mut rows = stream.next().await.unwrap().unwrap().num_rows(); + { + let predicate = predicate.upgrade().unwrap(); + let current = predicate.current().unwrap(); + // Bound the published expression, rather than relying on pool + // reservations, which do not account for IN-list construction. + assert!(expression_nodes(¤t) < 32); + } + while let Some(batch) = stream.next().await { + rows += batch.unwrap().num_rows(); + } + assert_eq!(rows, 65_536, "duplicate build rows must still join"); + assert_eq!(pool.reserved(), 0); + assert!(predicate.upgrade().is_none()); + assert_eq!( + session + .copied_config() + .options() + .optimizer + .hash_join_inlist_pushdown_max_size, + configured_limit, + "the execution override must not change the session" + ); + assert_eq!( + plan.metrics() + .unwrap() + .sum_by_name("dynamic_filter_rows_pruned") + .unwrap() + .as_usize(), + 1 + ); + } +} + +#[tokio::test] +async fn executions_and_resets_have_independent_producers() { + for mode in [PartitionMode::Partitioned, PartitionMode::CollectLeft] { + let (session, pool) = limited_session(1024 * 1024); + let join = single_key_join( + Arc::new(Int64Array::from(vec![5])), + Arc::new(Int64Array::from(vec![5, 90])), + mode, + ); + let plan = Arc::new(DynamicFilterJoinExec::new(&join, ConfigOptions::default()).unwrap()); + let first = plan.build_runtime_join().unwrap(); + let first_filter = Arc::downgrade(&produced_join_filter(&first.join)); + let second = plan.build_runtime_join().unwrap(); + let second_filter = Arc::downgrade(&produced_join_filter(&second.join)); + assert!(!first_filter.ptr_eq(&second_filter)); + let mut first = plan + .execute_runtime_join(first.join, 0, session.task_ctx()) + .unwrap(); + let mut second = plan + .execute_runtime_join(second.join, 0, session.task_ctx()) + .unwrap(); + assert_eq!(first.next().await.unwrap().unwrap().num_rows(), 1); + assert_eq!(second.next().await.unwrap().unwrap().num_rows(), 1); + drop(first); + assert!(first_filter.upgrade().is_none()); + assert!(second_filter.upgrade().is_some()); + assert!(pool.reserved() > 0); + assert!(second.next().await.is_none()); + assert!(second_filter.upgrade().is_none()); + assert_eq!(pool.reserved(), 0); + + let reset = Arc::clone(&plan).reset_state().unwrap(); + assert_eq!( + row_count(&collect(reset, session.task_ctx()).await.unwrap()), + 1 + ); + let replacement = single_key_join( + Arc::new(Int64Array::from(vec![90])), + Arc::new(Int64Array::from(vec![5, 90])), + mode, + ); + let rewritten = Arc::clone(&plan) + .replace_children( + vec![ + Arc::clone(replacement.left()), + Arc::clone(replacement.right()), + ], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + .unwrap(); + let output = collect(rewritten, session.task_ctx()).await.unwrap(); + assert_eq!(row_count(&output), 1); + assert_eq!( + output[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + 90 + ); + assert_eq!(pool.reserved(), 0); + } +} + +#[tokio::test] +async fn child_replacement_rechecks_join_key_types() { + let session = SessionContext::new(); + for mode in [PartitionMode::Partitioned, PartitionMode::CollectLeft] { + let join = plain_join() + .builder() + .with_partition_mode(mode) + .build() + .unwrap(); + let plan = Arc::new( + DynamicFilterJoinExec::try_new(&join, &ConfigOptions::default()) + .unwrap() + .unwrap(), + ); + for (key_type, supported) in [ + (DataType::Int64, true), + (DataType::Float64, false), + (DataType::Utf8, false), + ] { + let rewritten = Arc::clone(&plan) + .replace_children( + vec![ + input(vec![Some(90)], &key_type, 1), + input(vec![Some(5), Some(90)], &key_type, 0), + ], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + .unwrap(); + assert_eq!(rewritten.is::(), supported); + if !supported { + assert!(rewritten + .downcast_ref::() + .unwrap() + .dynamic_expressions_produced() + .is_empty()); + } + let output = collect(Arc::clone(&rewritten), session.task_ctx()) + .await + .unwrap(); + assert_eq!(row_count(&output), 1); + if supported { + assert_eq!(metric(&rewritten, "dynamic_filter_rows_pruned"), 1); + let reset = rewritten.reset_state().unwrap(); + assert!(reset.is::()); + let reset_output = collect(Arc::clone(&reset), session.task_ctx()) + .await + .unwrap(); + assert_eq!(row_count(&reset_output), 1); + assert_eq!(metric(&reset, "dynamic_filter_rows_pruned"), 1); + } + } + } +} + +#[test] +fn child_replacement_rechecks_native_partition_counts() { + for mode in [PartitionMode::Partitioned, PartitionMode::CollectLeft] { + let join = plain_join() + .builder() + .with_partition_mode(mode) + .build() + .unwrap(); + let plan = Arc::new( + DynamicFilterJoinExec::try_new(&join, &ConfigOptions::default()) + .unwrap() + .unwrap(), + ); + for replace_build in [false, true] { + let schema = if replace_build { + join.left().schema() + } else { + join.right().schema() + }; + let two_partitions: Arc = MemorySourceConfig::try_new_exec( + &[ + vec![RecordBatch::new_empty(Arc::clone(&schema))], + vec![RecordBatch::new_empty(Arc::clone(&schema))], + ], + schema, + None, + ) + .unwrap(); + let children = if replace_build { + vec![two_partitions, Arc::clone(join.right())] + } else { + vec![Arc::clone(join.left()), two_partitions] + }; + let rewritten = Arc::clone(&plan) + .replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + .unwrap(); + assert!(rewritten + .downcast_ref::() + .unwrap() + .dynamic_expressions_produced() + .is_empty()); + } + } +} + +#[test] +fn runtime_filter_is_visible_to_expression_visitors() { + let wrapper = DynamicFilterJoinExec::try_new(&plain_join(), &ConfigOptions::default()) + .unwrap() + .unwrap(); + let runtime = wrapper.build_runtime_join().unwrap(); + let predicate = produced_join_filter(&runtime.join); + let consumer = runtime.join.right(); + let mut consumer_expression_ids = Vec::new(); + consumer + .apply_expressions(&mut |expression| { + consumer_expression_ids.push(expression.expression_id()); + Ok(TreeNodeRecursion::Continue) + }) + .unwrap(); + assert_eq!( + consumer_expression_ids, + vec![Some(predicate.expression_id().unwrap())] + ); + assert_eq!( + consumer + .apply_expressions(&mut |_| Ok(TreeNodeRecursion::Stop)) + .unwrap(), + TreeNodeRecursion::Stop + ); + assert!(consumer.dynamic_expressions_produced().is_empty()); + assert!(wrapper.dynamic_expressions_produced().is_empty()); + + let mut join_key_count = 0; + wrapper + .apply_expressions(&mut |expression| { + assert!(expression.expression_id().is_none()); + join_key_count += 1; + Ok(TreeNodeRecursion::Continue) + }) + .unwrap(); + assert_eq!(join_key_count, 2); +} + +#[test] +fn wrapper_preserves_join_statistics_and_distribution() { + for mode in [PartitionMode::Partitioned, PartitionMode::CollectLeft] { + let join = plain_join() + .builder() + .with_partition_mode(mode) + .build() + .unwrap(); + let wrapper = DynamicFilterJoinExec::try_new(&join, &ConfigOptions::default()) + .unwrap() + .unwrap(); + for partition in [None, Some(0)] { + let context = StatisticsContext::new(); + let args = StatisticsArgs::new().with_partition(partition); + let expected = context.compute(&join, &args).unwrap(); + let actual = context.compute(&wrapper, &args).unwrap(); + assert!(expected.num_rows.get_value().is_some()); + assert_eq!(actual, expected, "{mode:?}: {partition:?}"); + } + + let expected = join.input_distribution_requirements(); + let actual = wrapper.input_distribution_requirements(); + assert_eq!( + actual.per_child_distributions().len(), + expected.per_child_distributions().len() + ); + for (actual, expected) in actual + .per_child_distributions() + .zip(expected.per_child_distributions()) + { + match (actual, expected) { + (Distribution::SinglePartition, Distribution::SinglePartition) + | (Distribution::UnspecifiedDistribution, Distribution::UnspecifiedDistribution) => { + } + (Distribution::KeyPartitioned(actual), Distribution::KeyPartitioned(expected)) => { + assert_eq!(actual, expected); + } + _ => panic!("distribution changed: {expected:?} -> {actual:?}"), + } + } + let schema = join.right().schema(); + let two_partitions = MemorySourceConfig::try_new_exec( + &[ + vec![RecordBatch::new_empty(Arc::clone(&schema))], + vec![RecordBatch::new_empty(Arc::clone(&schema))], + ], + schema, + None, + ) + .unwrap(); + let candidate_children = [join.left().as_ref(), two_partitions.as_ref()]; + let expected_unsatisfied = expected + .unsatisfied_co_partitioned_children(join.name(), &candidate_children) + .unwrap(); + let actual_unsatisfied = actual + .unsatisfied_co_partitioned_children(wrapper.name(), &candidate_children) + .unwrap(); + assert_eq!(actual_unsatisfied, expected_unsatisfied); + assert_eq!( + actual_unsatisfied, + if matches!(mode, PartitionMode::Partitioned) { + vec![0, 1] + } else { + vec![] + } + ); + } +} diff --git a/native/core/src/execution/operators/mod.rs b/native/core/src/execution/operators/mod.rs index d09b0b4fb37..ecd8ec73652 100644 --- a/native/core/src/execution/operators/mod.rs +++ b/native/core/src/execution/operators/mod.rs @@ -27,7 +27,8 @@ pub use scan::*; mod aligned_stream_reader; mod copy; mod dynamic_filter; -pub(crate) use dynamic_filter::DynamicFilterJoinExec; +mod dynamic_filter_join; +pub(crate) use dynamic_filter_join::DynamicFilterJoinExec; mod filter; pub(crate) use filter::CometFilterExec; mod expand; From c05e9730d1b76e061f21db0835d052c35ff61745 Mon Sep 17 00:00:00 2001 From: Ping Zhang Date: Tue, 15 Sep 2026 00:09:03 +0000 Subject: [PATCH 2/4] refactor: extract shared Parquet filter attachment --- .../operators/dynamic_filter_join.rs | 99 +------------- .../operators/dynamic_filter_join/tests.rs | 58 +------- native/core/src/execution/operators/mod.rs | 1 + .../operators/parquet_reader_filter.rs | 126 ++++++++++++++++++ .../operators/parquet_reader_filter/tests.rs | 71 ++++++++++ 5 files changed, 205 insertions(+), 150 deletions(-) create mode 100644 native/core/src/execution/operators/parquet_reader_filter.rs create mode 100644 native/core/src/execution/operators/parquet_reader_filter/tests.rs diff --git a/native/core/src/execution/operators/dynamic_filter_join.rs b/native/core/src/execution/operators/dynamic_filter_join.rs index 049739116bf..612dd2e727c 100644 --- a/native/core/src/execution/operators/dynamic_filter_join.rs +++ b/native/core/src/execution/operators/dynamic_filter_join.rs @@ -31,13 +31,8 @@ use arrow::datatypes::DataType; use datafusion::common::config::ConfigOptions; use datafusion::common::tree_node::TreeNodeRecursion; use datafusion::common::{JoinType, NullEquality, Result, Statistics}; -use datafusion::datasource::physical_plan::ParquetSource; -use datafusion::datasource::source::DataSourceExec; use datafusion::execution::TaskContext; -use datafusion::logical_expr::Operator; -use datafusion::physical_expr::expressions::{ - lit, BinaryExpr, Column, DynamicFilterPhysicalExpr, IsNotNullExpr, -}; +use datafusion::physical_expr::expressions::{lit, Column, DynamicFilterPhysicalExpr}; use datafusion::physical_expr::PhysicalExpr; use datafusion::physical_plan::distribution_requirements::InputDistributionRequirements; use datafusion::physical_plan::joins::{HashJoinExec, PartitionMode}; @@ -51,7 +46,7 @@ use datafusion::physical_plan::{ use futures::StreamExt; use super::dynamic_filter::DynamicFilterExec; -use super::CometFilterExec; +use super::parquet_reader_filter::try_attach_parquet_reader_filter; /// A permanent plan must not own a completed join's filter or build accumulator: /// those can retain the hash map after its stream-owned reservation is released. @@ -71,96 +66,6 @@ struct RuntimeDynamicFilterJoin { reader_filter_attached: bool, } -/// Recognize only direct-column null checks joined by AND, without evaluating -/// or changing the predicate. Every accepted leaf is deterministic, infallible, -/// and only discards rows, so reader pruning cannot suppress expression errors -/// or alter stateful evaluation. All other expressions remain a boundary. -fn is_direct_column_null_checks(predicate: &Arc) -> bool { - if let Some(binary) = predicate.downcast_ref::() { - return binary.op() == &Operator::And - && is_direct_column_null_checks(binary.left()) - && is_direct_column_null_checks(binary.right()); - } - predicate - .downcast_ref::() - .is_some_and(|is_not_null| is_not_null.arg().is::()) -} - -fn try_attach_parquet_reader_filter( - input: &Arc, - predicate: Arc, - config: &ConfigOptions, -) -> Result>> { - // Filtering before a fetch can change which rows are selected by its limit. - if input.fetch().is_some() { - log::debug!("Join dynamic filter reader pushdown skipped: probe has a fetch limit"); - return Ok(None); - } - // Spark inserts IS NOT NULL residuals above equijoin inputs, including AND - // chains of inferred null checks. A reader predicate can cross those direct - // checks because both operations only discard rows. Keep every other filter - // as a boundary: reader pruning would change which rows reach stateful - // expressions and can suppress expression errors. - if let Some(filter) = input.downcast_ref::() { - if filter.has_projection() { - log::debug!( - "Join dynamic filter reader pushdown skipped: probe FilterExec has a projection" - ); - return Ok(None); - } - if !is_direct_column_null_checks(filter.predicate()) { - log::debug!( - "Join dynamic filter reader pushdown skipped: probe filter is not direct column IS NOT NULL checks" - ); - return Ok(None); - } - let Some(reader) = - try_attach_parquet_reader_filter(filter.input(), Arc::clone(&predicate), config)? - else { - return Ok(None); - }; - return match filter.with_execution_input(reader) { - Ok(updated) => Ok(Some(updated)), - Err(error) => { - log::debug!( - "Join dynamic filter reader pushdown skipped: probe filter rebuild failed: {error}" - ); - Ok(None) - } - }; - } - let Some(scan) = input.downcast_ref::() else { - log::debug!( - "Join dynamic filter reader pushdown skipped: probe root is {}", - input.name() - ); - return Ok(None); - }; - if scan.downcast_to_file_source::().is_none() { - log::debug!("Join dynamic filter reader pushdown skipped: probe is not Parquet"); - return Ok(None); - } - - let predicate: Arc = predicate; - let propagation = match scan - .data_source() - .try_pushdown_filters(vec![predicate], config) - { - Ok(propagation) => propagation, - Err(error) => { - log::debug!( - "Join dynamic filter reader pushdown skipped: predicate remapping failed: {error}" - ); - return Ok(None); - } - }; - let Some(data_source) = propagation.updated_node else { - log::debug!("Join dynamic filter reader pushdown skipped: Parquet declined the predicate"); - return Ok(None); - }; - Ok(Some(Arc::new(scan.clone().with_data_source(data_source)))) -} - impl DynamicFilterJoinExec { /// Return no wrapper when the join cannot safely use a runtime filter. pub(crate) fn try_new(join: &HashJoinExec, config: &ConfigOptions) -> Result> { diff --git a/native/core/src/execution/operators/dynamic_filter_join/tests.rs b/native/core/src/execution/operators/dynamic_filter_join/tests.rs index 5b3e33f60d0..b9a167463a6 100644 --- a/native/core/src/execution/operators/dynamic_filter_join/tests.rs +++ b/native/core/src/execution/operators/dynamic_filter_join/tests.rs @@ -19,6 +19,7 @@ use super::*; use std::fmt::Display; use std::hash::{Hash, Hasher}; +use crate::execution::operators::CometFilterExec; use crate::execution::planner::PhysicalPlanner; use crate::parquet::parquet_exec::init_datasource_exec; use crate::parquet::parquet_support::ObjectStoreBackend; @@ -29,10 +30,13 @@ use datafusion::common::test_util::batches_to_sort_string; use datafusion::common::tree_node::TreeNodeRecursion; use datafusion::datasource::listing::PartitionedFile; use datafusion::datasource::memory::MemorySourceConfig; +use datafusion::datasource::physical_plan::ParquetSource; +use datafusion::datasource::source::DataSourceExec; use datafusion::execution::memory_pool::{GreedyMemoryPool, MemoryPool}; use datafusion::execution::object_store::ObjectStoreUrl; use datafusion::execution::runtime_env::RuntimeEnvBuilder; -use datafusion::logical_expr::ColumnarValue; +use datafusion::logical_expr::{ColumnarValue, Operator}; +use datafusion::physical_expr::expressions::{BinaryExpr, IsNotNullExpr}; use datafusion::physical_plan::collect; use datafusion::physical_plan::filter::FilterExec; use datafusion::physical_plan::metrics::MetricValue; @@ -684,58 +688,6 @@ fn find_dynamic_filter(expr: &Arc) -> Option<&DynamicFilterPhy .find_map(|child| find_dynamic_filter(child)) } -/// Accept nested null-check conjunctions while retaining all other filter -/// boundaries, including OR and computed or potentially failing expressions. -#[test] -fn reader_filter_crosses_only_direct_column_null_checks() { - let key: Arc = Arc::new(Column::new("key", 0)); - let direct_null_check: Arc = Arc::new(IsNotNullExpr::new(Arc::clone(&key))); - assert!(is_direct_column_null_checks(&direct_null_check)); - let other_null_check: Arc = - Arc::new(IsNotNullExpr::new(Arc::new(Column::new("other", 1)))); - let conjunction: Arc = Arc::new(BinaryExpr::new( - Arc::clone(&direct_null_check), - Operator::And, - Arc::clone(&other_null_check), - )); - let nested: Arc = Arc::new(BinaryExpr::new( - conjunction, - Operator::And, - Arc::clone(&direct_null_check), - )); - assert!(is_direct_column_null_checks(&nested)); - let right_nested: Arc = Arc::new(BinaryExpr::new( - Arc::clone(&direct_null_check), - Operator::And, - Arc::new(BinaryExpr::new( - Arc::clone(&other_null_check), - Operator::And, - Arc::clone(&direct_null_check), - )), - )); - assert!(is_direct_column_null_checks(&right_nested)); - let disjunction: Arc = Arc::new(BinaryExpr::new( - Arc::clone(&direct_null_check), - Operator::Or, - other_null_check, - )); - assert!(!is_direct_column_null_checks(&disjunction)); - - let comparison: Arc = - Arc::new(BinaryExpr::new(Arc::clone(&key), Operator::Gt, lit(0_i32))); - let conjunction: Arc = Arc::new(BinaryExpr::new( - Arc::clone(&direct_null_check), - Operator::And, - comparison, - )); - assert!(!is_direct_column_null_checks(&conjunction)); - - let computed: Arc = - Arc::new(BinaryExpr::new(key, Operator::Plus, lit(1_i32))); - let computed_null_check: Arc = Arc::new(IsNotNullExpr::new(computed)); - assert!(!is_direct_column_null_checks(&computed_null_check)); -} - /// Exercise a real Parquet reader with three distinct nullable columns and /// remapped projection. Both statistics-only and row-filter reads must retain /// the AND residual: build keys 2 and 3 fail its payload checks, leaving key 1. diff --git a/native/core/src/execution/operators/mod.rs b/native/core/src/execution/operators/mod.rs index ecd8ec73652..ac504c2abac 100644 --- a/native/core/src/execution/operators/mod.rs +++ b/native/core/src/execution/operators/mod.rs @@ -28,6 +28,7 @@ mod aligned_stream_reader; mod copy; mod dynamic_filter; mod dynamic_filter_join; +mod parquet_reader_filter; pub(crate) use dynamic_filter_join::DynamicFilterJoinExec; mod filter; pub(crate) use filter::CometFilterExec; diff --git a/native/core/src/execution/operators/parquet_reader_filter.rs b/native/core/src/execution/operators/parquet_reader_filter.rs new file mode 100644 index 00000000000..c705316759e --- /dev/null +++ b/native/core/src/execution/operators/parquet_reader_filter.rs @@ -0,0 +1,126 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Attach a live predicate to a native Parquet reader while preserving input boundaries. + +use std::sync::Arc; + +use datafusion::common::config::ConfigOptions; +use datafusion::common::Result; +use datafusion::datasource::physical_plan::ParquetSource; +use datafusion::datasource::source::DataSourceExec; +use datafusion::logical_expr::Operator; +use datafusion::physical_expr::expressions::{ + BinaryExpr, Column, DynamicFilterPhysicalExpr, IsNotNullExpr, +}; +use datafusion::physical_expr::PhysicalExpr; +use datafusion::physical_plan::ExecutionPlan; + +use super::CometFilterExec; + +/// Recognize only direct-column null checks joined by AND, without evaluating +/// or changing the predicate. Every accepted leaf is deterministic, infallible, +/// and only discards rows, so reader pruning cannot suppress expression errors +/// or alter stateful evaluation. All other expressions remain a boundary. +fn is_direct_column_null_checks(predicate: &Arc) -> bool { + if let Some(binary) = predicate.downcast_ref::() { + return binary.op() == &Operator::And + && is_direct_column_null_checks(binary.left()) + && is_direct_column_null_checks(binary.right()); + } + predicate + .downcast_ref::() + .is_some_and(|is_not_null| is_not_null.arg().is::()) +} + +pub(super) fn try_attach_parquet_reader_filter( + input: &Arc, + predicate: Arc, + config: &ConfigOptions, +) -> Result>> { + // Filtering before a fetch can change which rows are selected by its limit. + if input.fetch().is_some() { + log::debug!("Join dynamic filter reader pushdown skipped: probe has a fetch limit"); + return Ok(None); + } + // Spark inserts IS NOT NULL residuals above equijoin inputs, including AND + // chains of inferred null checks. A reader predicate can cross those direct + // checks because both operations only discard rows. Keep every other filter + // as a boundary: reader pruning would change which rows reach stateful + // expressions and can suppress expression errors. + if let Some(filter) = input.downcast_ref::() { + if filter.has_projection() { + log::debug!( + "Join dynamic filter reader pushdown skipped: probe FilterExec has a projection" + ); + return Ok(None); + } + if !is_direct_column_null_checks(filter.predicate()) { + log::debug!( + "Join dynamic filter reader pushdown skipped: probe filter is not direct column IS NOT NULL checks" + ); + return Ok(None); + } + let Some(reader) = + try_attach_parquet_reader_filter(filter.input(), Arc::clone(&predicate), config)? + else { + return Ok(None); + }; + return match filter.with_execution_input(reader) { + Ok(updated) => Ok(Some(updated)), + Err(error) => { + log::debug!( + "Join dynamic filter reader pushdown skipped: probe filter rebuild failed: {error}" + ); + Ok(None) + } + }; + } + let Some(scan) = input.downcast_ref::() else { + log::debug!( + "Join dynamic filter reader pushdown skipped: probe root is {}", + input.name() + ); + return Ok(None); + }; + if scan.downcast_to_file_source::().is_none() { + log::debug!("Join dynamic filter reader pushdown skipped: probe is not Parquet"); + return Ok(None); + } + + let predicate: Arc = predicate; + let propagation = match scan + .data_source() + .try_pushdown_filters(vec![predicate], config) + { + Ok(propagation) => propagation, + Err(error) => { + log::debug!( + "Join dynamic filter reader pushdown skipped: predicate remapping failed: {error}" + ); + return Ok(None); + } + }; + let Some(data_source) = propagation.updated_node else { + log::debug!("Join dynamic filter reader pushdown skipped: Parquet declined the predicate"); + return Ok(None); + }; + Ok(Some(Arc::new(scan.clone().with_data_source(data_source)))) +} + +#[cfg(test)] +mod tests; diff --git a/native/core/src/execution/operators/parquet_reader_filter/tests.rs b/native/core/src/execution/operators/parquet_reader_filter/tests.rs new file mode 100644 index 00000000000..aacc242731e --- /dev/null +++ b/native/core/src/execution/operators/parquet_reader_filter/tests.rs @@ -0,0 +1,71 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use super::*; +use datafusion::physical_expr::expressions::lit; + +/// Accept nested null-check conjunctions while retaining all other filter +/// boundaries, including OR and computed or potentially failing expressions. +#[test] +fn reader_filter_crosses_only_direct_column_null_checks() { + let key: Arc = Arc::new(Column::new("key", 0)); + let direct_null_check: Arc = Arc::new(IsNotNullExpr::new(Arc::clone(&key))); + assert!(is_direct_column_null_checks(&direct_null_check)); + let other_null_check: Arc = + Arc::new(IsNotNullExpr::new(Arc::new(Column::new("other", 1)))); + let conjunction: Arc = Arc::new(BinaryExpr::new( + Arc::clone(&direct_null_check), + Operator::And, + Arc::clone(&other_null_check), + )); + let nested: Arc = Arc::new(BinaryExpr::new( + conjunction, + Operator::And, + Arc::clone(&direct_null_check), + )); + assert!(is_direct_column_null_checks(&nested)); + let right_nested: Arc = Arc::new(BinaryExpr::new( + Arc::clone(&direct_null_check), + Operator::And, + Arc::new(BinaryExpr::new( + Arc::clone(&other_null_check), + Operator::And, + Arc::clone(&direct_null_check), + )), + )); + assert!(is_direct_column_null_checks(&right_nested)); + let disjunction: Arc = Arc::new(BinaryExpr::new( + Arc::clone(&direct_null_check), + Operator::Or, + other_null_check, + )); + assert!(!is_direct_column_null_checks(&disjunction)); + + let comparison: Arc = + Arc::new(BinaryExpr::new(Arc::clone(&key), Operator::Gt, lit(0_i32))); + let conjunction: Arc = Arc::new(BinaryExpr::new( + Arc::clone(&direct_null_check), + Operator::And, + comparison, + )); + assert!(!is_direct_column_null_checks(&conjunction)); + + let computed: Arc = + Arc::new(BinaryExpr::new(key, Operator::Plus, lit(1_i32))); + let computed_null_check: Arc = Arc::new(IsNotNullExpr::new(computed)); + assert!(!is_direct_column_null_checks(&computed_null_check)); +} From 93b306177abeb5ab72cd10bfb63c0ab35fc4735c Mon Sep 17 00:00:00 2001 From: Ping Zhang Date: Tue, 15 Sep 2026 00:23:14 +0000 Subject: [PATCH 3/4] refactor: use owner prefixes for join filter metrics --- docs/source/user-guide/latest/metrics.md | 16 ++++---- .../src/execution/operators/dynamic_filter.rs | 21 ++++++---- .../operators/dynamic_filter/tests.rs | 11 ++--- .../operators/dynamic_filter_join.rs | 5 ++- .../operators/dynamic_filter_join/tests.rs | 31 +++++++------- native/core/src/execution/planner.rs | 10 +++-- .../spark/sql/comet/CometMetricNode.scala | 12 +++--- .../apache/comet/exec/CometJoinSuite.scala | 40 +++++++++---------- 8 files changed, 77 insertions(+), 69 deletions(-) diff --git a/docs/source/user-guide/latest/metrics.md b/docs/source/user-guide/latest/metrics.md index a400a9fb1c7..2c9f012f321 100644 --- a/docs/source/user-guide/latest/metrics.md +++ b/docs/source/user-guide/latest/metrics.md @@ -35,14 +35,14 @@ With `spark.comet.exec.join.dynamicFilter.enabled=true`, native broadcast and sh report these additional metric keys. See [Join Runtime Filters](tuning.md#join-runtime-filters) for eligibility and reader restrictions. -| Metric | Description | -| ---------------------------------------- | ----------------------------------------------------------------- | -| `dynamic_filter_rows_evaluated` | Probe rows evaluated by the runtime filter. | -| `dynamic_filter_rows_pruned` | Probe rows rejected by that filter before the hash probe. | -| `dynamic_filter_rows_bypassed` | Probe rows passed through while the runtime filter is inactive. | -| `dynamic_filter_eval_time` | Time evaluating the runtime filter. | -| `dynamic_filter_reader_filters_attached` | Executions that attach their runtime filter to a native reader. | -| `dynamic_filter_reader_filters_skipped` | Executions whose probe input is ineligible for reader attachment. | +| Metric | Description | +| -------------------------------------- | ----------------------------------------------------------------- | +| `dynamic_filter_join_rows_evaluated` | Probe rows evaluated by the runtime filter. | +| `dynamic_filter_join_rows_pruned` | Probe rows rejected by that filter before the hash probe. | +| `dynamic_filter_join_rows_bypassed` | Probe rows passed through while the runtime filter is inactive. | +| `dynamic_filter_join_eval_time` | Time evaluating the runtime filter. | +| `dynamic_filter_join_filters_attached` | Executions that attach their runtime filter to a native reader. | +| `dynamic_filter_join_filters_skipped` | Executions whose probe input is ineligible for reader attachment. | The row counters measure residual filtering of decoded probe batches. They exclude rows skipped by the reader. An attached filter does not guarantee that any row groups are pruned: compare the diff --git a/native/core/src/execution/operators/dynamic_filter.rs b/native/core/src/execution/operators/dynamic_filter.rs index 8a294201965..b543747cbeb 100644 --- a/native/core/src/execution/operators/dynamic_filter.rs +++ b/native/core/src/execution/operators/dynamic_filter.rs @@ -43,6 +43,7 @@ pub(crate) struct DynamicFilterExec { input: Arc, predicate: Arc, metrics: ExecutionPlanMetricsSet, + metric_prefix: &'static str, } impl DynamicFilterExec { @@ -50,11 +51,13 @@ impl DynamicFilterExec { input: Arc, predicate: Arc, metrics: ExecutionPlanMetricsSet, + metric_prefix: &'static str, ) -> Self { Self { input, predicate, metrics, + metric_prefix, } } } @@ -116,6 +119,7 @@ impl ExecutionPlan for DynamicFilterExec { children.remove(0), Arc::clone(&self.predicate), ExecutionPlanMetricsSet::new(), + self.metric_prefix, ))) } @@ -131,6 +135,7 @@ impl ExecutionPlan for DynamicFilterExec { Arc::clone(&self.input), predicate, ExecutionPlanMetricsSet::new(), + self.metric_prefix, ))) } @@ -150,16 +155,16 @@ impl ExecutionPlan for DynamicFilterExec { let predicate = Arc::clone(&self.predicate) .with_new_children(vec![Arc::new(Column::new(key.name(), 0))])?; let input = self.input.execute(partition, context)?; - let evaluated = - MetricBuilder::new(&self.metrics).counter("dynamic_filter_rows_evaluated", partition); - let pruned = - MetricBuilder::new(&self.metrics).counter("dynamic_filter_rows_pruned", partition); - let bypassed = - MetricBuilder::new(&self.metrics).counter("dynamic_filter_rows_bypassed", partition); + let evaluated = MetricBuilder::new(&self.metrics) + .counter(format!("{}_rows_evaluated", self.metric_prefix), partition); + let pruned = MetricBuilder::new(&self.metrics) + .counter(format!("{}_rows_pruned", self.metric_prefix), partition); + let bypassed = MetricBuilder::new(&self.metrics) + .counter(format!("{}_rows_bypassed", self.metric_prefix), partition); // Only dedicated metrics: merging this helper into the Spark join must not // add its input/output counts or elapsed time to the join's existing metrics. - let eval_time = - MetricBuilder::new(&self.metrics).subset_time("dynamic_filter_eval_time", partition); + let eval_time = MetricBuilder::new(&self.metrics) + .subset_time(format!("{}_eval_time", self.metric_prefix), partition); let stream = input.map(move |batch| { let batch = batch?; let _timer = eval_time.timer(); diff --git a/native/core/src/execution/operators/dynamic_filter/tests.rs b/native/core/src/execution/operators/dynamic_filter/tests.rs index 9460272c27a..e610f8b4a62 100644 --- a/native/core/src/execution/operators/dynamic_filter/tests.rs +++ b/native/core/src/execution/operators/dynamic_filter/tests.rs @@ -87,6 +87,7 @@ async fn placeholder_updates_and_errors_are_not_hidden() { source, Arc::clone(&predicate), ExecutionPlanMetricsSet::new(), + "test_filter", )); let task = SessionContext::new().task_ctx(); let mut stream = wrapper.execute(0, Arc::clone(&task)).unwrap(); @@ -104,9 +105,9 @@ async fn placeholder_updates_and_errors_are_not_hidden() { while let Some(batch) = stream.next().await { assert_eq!(batch.unwrap().num_rows(), 0); } - assert_eq!(metric(&wrapper, "dynamic_filter_rows_bypassed"), 2); - assert_eq!(metric(&wrapper, "dynamic_filter_rows_pruned"), 8); - assert_eq!(metric(&wrapper, "dynamic_filter_rows_evaluated"), 8); + assert_eq!(metric(&wrapper, "test_filter_rows_bypassed"), 2); + assert_eq!(metric(&wrapper, "test_filter_rows_pruned"), 8); + assert_eq!(metric(&wrapper, "test_filter_rows_evaluated"), 8); // Reset must not preserve an old condition, even while another owner // still holds the previous predicate. @@ -115,8 +116,8 @@ async fn placeholder_updates_and_errors_are_not_hidden() { .await .unwrap(); assert_eq!(row_count(&reset_output), 10); - assert_eq!(metric(&reset, "dynamic_filter_rows_pruned"), 0); - assert_eq!(metric(&reset, "dynamic_filter_rows_bypassed"), 10); + assert_eq!(metric(&reset, "test_filter_rows_pruned"), 0); + assert_eq!(metric(&reset, "test_filter_rows_bypassed"), 10); predicate.update(lit(42i32)).unwrap(); let error = collect(wrapper, task).await.unwrap_err(); diff --git a/native/core/src/execution/operators/dynamic_filter_join.rs b/native/core/src/execution/operators/dynamic_filter_join.rs index 612dd2e727c..cf7d98429e2 100644 --- a/native/core/src/execution/operators/dynamic_filter_join.rs +++ b/native/core/src/execution/operators/dynamic_filter_join.rs @@ -99,6 +99,7 @@ impl DynamicFilterJoinExec { reader.unwrap_or_else(|| Arc::clone(self.template.right())), Arc::clone(&predicate), self.metrics.clone(), + "dynamic_filter_join", )); // In particular, do not share CollectLeft's cached build future with the // template, another execution, or a reset plan. @@ -240,9 +241,9 @@ impl ExecutionPlan for DynamicFilterJoinExec { ) -> Result { let runtime = self.build_runtime_join()?; let attachment_metric = if runtime.reader_filter_attached { - "dynamic_filter_reader_filters_attached" + "dynamic_filter_join_filters_attached" } else { - "dynamic_filter_reader_filters_skipped" + "dynamic_filter_join_filters_skipped" }; MetricBuilder::new(&self.metrics) .counter(attachment_metric, partition) diff --git a/native/core/src/execution/operators/dynamic_filter_join/tests.rs b/native/core/src/execution/operators/dynamic_filter_join/tests.rs index b9a167463a6..aa64656f73d 100644 --- a/native/core/src/execution/operators/dynamic_filter_join/tests.rs +++ b/native/core/src/execution/operators/dynamic_filter_join/tests.rs @@ -288,6 +288,7 @@ async fn completed_filter_evaluates_only_the_shared_probe_key() { probe, predicate, ExecutionPlanMetricsSet::new(), + "dynamic_filter_join", )); let filtered = collect(consumer, session.task_ctx()).await.unwrap(); let selected = BooleanArray::from( @@ -353,9 +354,9 @@ async fn completed_build_filters_both_sides_and_session_inlist_settings() { 4, "duplicates and late build keys must match" ); - assert_eq!(metric(&attached, "dynamic_filter_rows_evaluated"), 202); - assert!(metric(&attached, "dynamic_filter_rows_pruned") >= 199); - assert_eq!(metric(&attached, "dynamic_filter_rows_bypassed"), 0); + assert_eq!(metric(&attached, "dynamic_filter_join_rows_evaluated"), 202); + assert!(metric(&attached, "dynamic_filter_join_rows_pruned") >= 199); + assert_eq!(metric(&attached, "dynamic_filter_join_rows_bypassed"), 0); } else { assert_eq!(row_count(&actual), 0); } @@ -559,7 +560,7 @@ async fn independent_attempts_do_not_share_build_domains() { }) .collect::>(); assert_eq!(keys, vec![build_key]); - assert_eq!(metric(&attached, "dynamic_filter_rows_pruned"), 1); + assert_eq!(metric(&attached, "dynamic_filter_join_rows_pruned"), 1); } } @@ -840,7 +841,7 @@ async fn reader_filter_crosses_null_check_conjunction_and_retains_residual() { assert_eq!( plan.metrics() .unwrap() - .sum_by_name("dynamic_filter_reader_filters_attached") + .sum_by_name("dynamic_filter_join_filters_attached") .unwrap() .as_usize(), 1 @@ -887,7 +888,7 @@ async fn reader_filter_does_not_cross_fetch_limits() { .unwrap(); outputs.push(row_count(&output)); if enabled { - assert_eq!(metric(&plan, "dynamic_filter_reader_filters_skipped"), 1); + assert_eq!(metric(&plan, "dynamic_filter_join_filters_skipped"), 1); } } assert_eq!(outputs, vec![0, 0]); @@ -969,10 +970,10 @@ async fn reader_filter_does_not_cross_seeded_rand_probe_filter() { .and_then(|metrics| metrics.sum_by_name(name)) .map_or(0, |value| value.as_usize()) }; - assert_eq!(dynamic_metric("dynamic_filter_reader_filters_attached"), 0); - assert_eq!(dynamic_metric("dynamic_filter_reader_filters_skipped"), 1); - assert_eq!(dynamic_metric("dynamic_filter_rows_evaluated"), 1); - assert_eq!(dynamic_metric("dynamic_filter_rows_pruned"), 0); + assert_eq!(dynamic_metric("dynamic_filter_join_filters_attached"), 0); + assert_eq!(dynamic_metric("dynamic_filter_join_filters_skipped"), 1); + assert_eq!(dynamic_metric("dynamic_filter_join_rows_evaluated"), 1); + assert_eq!(dynamic_metric("dynamic_filter_join_rows_pruned"), 0); assert_eq!(pruning_metric(&scan, "row_groups_pruned_statistics"), 0); assert_eq!(filter.metrics().unwrap().output_rows().unwrap(), 1); } @@ -1037,7 +1038,7 @@ async fn run_parquet_join(values: Vec, enabled: bool) -> (usize, usize, usi .unwrap(); let attached = plan .metrics() - .and_then(|metrics| metrics.sum_by_name("dynamic_filter_reader_filters_attached")) + .and_then(|metrics| metrics.sum_by_name("dynamic_filter_join_filters_attached")) .map_or(0, |metric| metric.as_usize()); ( row_count(&output), @@ -1207,7 +1208,7 @@ async fn runtime_domains_release_with_streams_while_plans_remain_alive() { assert_eq!(metrics.output_rows(), Some(1)); assert_eq!( metrics - .sum_by_name("dynamic_filter_rows_pruned") + .sum_by_name("dynamic_filter_join_rows_pruned") .unwrap() .as_usize(), 1 @@ -1278,7 +1279,7 @@ async fn duplicate_heavy_builds_do_not_materialize_unreserved_inlists() { assert_eq!( plan.metrics() .unwrap() - .sum_by_name("dynamic_filter_rows_pruned") + .sum_by_name("dynamic_filter_join_rows_pruned") .unwrap() .as_usize(), 1 @@ -1392,14 +1393,14 @@ async fn child_replacement_rechecks_join_key_types() { .unwrap(); assert_eq!(row_count(&output), 1); if supported { - assert_eq!(metric(&rewritten, "dynamic_filter_rows_pruned"), 1); + assert_eq!(metric(&rewritten, "dynamic_filter_join_rows_pruned"), 1); let reset = rewritten.reset_state().unwrap(); assert!(reset.is::()); let reset_output = collect(Arc::clone(&reset), session.task_ctx()) .await .unwrap(); assert_eq!(row_count(&reset_output), 1); - assert_eq!(metric(&reset, "dynamic_filter_rows_pruned"), 1); + assert_eq!(metric(&reset, "dynamic_filter_join_rows_pruned"), 1); } } } diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 60c206617e3..4c174e032dc 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -6266,12 +6266,14 @@ mod tests { assert_eq!(metrics.metrics["build_input_rows"], 4); if enabled { assert_eq!(metrics.metrics["input_rows"], 3); - assert_eq!(metrics.metrics["dynamic_filter_rows_evaluated"], 100); - assert_eq!(metrics.metrics["dynamic_filter_rows_pruned"], 97); - assert_eq!(metrics.metrics["dynamic_filter_rows_bypassed"], 0); + assert_eq!(metrics.metrics["dynamic_filter_join_rows_evaluated"], 100); + assert_eq!(metrics.metrics["dynamic_filter_join_rows_pruned"], 97); + assert_eq!(metrics.metrics["dynamic_filter_join_rows_bypassed"], 0); } else { assert_eq!(metrics.metrics["input_rows"], 100); - assert!(!metrics.metrics.contains_key("dynamic_filter_rows_pruned")); + assert!(!metrics + .metrics + .contains_key("dynamic_filter_join_rows_pruned")); } } } diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometMetricNode.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometMetricNode.scala index 336cbf1a6ee..806c9d00eda 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometMetricNode.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometMetricNode.scala @@ -458,17 +458,17 @@ object CometMetricNode { /** Join runtime filtering and reader attachment. These never replace join or scan metrics. */ def joinDynamicFilterMetrics(sc: SparkContext): Map[String, SQLMetric] = { Map( - "dynamic_filter_rows_evaluated" -> + "dynamic_filter_join_rows_evaluated" -> SQLMetrics.createMetric(sc, "Probe rows evaluated by the join runtime filter"), - "dynamic_filter_rows_pruned" -> + "dynamic_filter_join_rows_pruned" -> SQLMetrics.createMetric(sc, "Probe rows rejected before the hash probe"), - "dynamic_filter_rows_bypassed" -> + "dynamic_filter_join_rows_bypassed" -> SQLMetrics.createMetric(sc, "Probe rows bypassing an inactive join runtime filter"), - "dynamic_filter_eval_time" -> + "dynamic_filter_join_eval_time" -> SQLMetrics.createNanoTimingMetric(sc, "Time evaluating the join runtime filter"), - "dynamic_filter_reader_filters_attached" -> + "dynamic_filter_join_filters_attached" -> SQLMetrics.createMetric(sc, "Join runtime filters attached to native readers"), - "dynamic_filter_reader_filters_skipped" -> + "dynamic_filter_join_filters_skipped" -> SQLMetrics.createMetric(sc, "Join runtime filters not eligible for native readers")) } diff --git a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala index 43f837c9ba1..c18b98aa83b 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala @@ -366,16 +366,16 @@ class CometJoinSuite extends CometTestBase { assert(join.metrics("output_rows").value == 4L) val probeRows = join.metrics("input_rows").value if (enabled) { - val evaluated = join.metrics("dynamic_filter_rows_evaluated").value - val pruned = join.metrics("dynamic_filter_rows_pruned").value - val bypassed = join.metrics("dynamic_filter_rows_bypassed").value + val evaluated = join.metrics("dynamic_filter_join_rows_evaluated").value + val pruned = join.metrics("dynamic_filter_join_rows_pruned").value + val bypassed = join.metrics("dynamic_filter_join_rows_bypassed").value assert(evaluated > 0L && pruned > 0L) assert(probeRows + pruned == evaluated + bypassed) assert(probeRows < unfilteredProbeRows) assert(evaluated + bypassed <= unfilteredProbeRows) - assert(join.metrics("dynamic_filter_eval_time").value > 0L) + assert(join.metrics("dynamic_filter_join_eval_time").value > 0L) } else { - assert(!join.metrics.contains("dynamic_filter_rows_pruned")) + assert(!join.metrics.contains("dynamic_filter_join_rows_pruned")) unfilteredProbeRows = probeRows } } @@ -456,10 +456,8 @@ class CometJoinSuite extends CometTestBase { val bytes = scanMetrics("bytes_scanned").value assert(joins.head.metrics("output_rows").value == 1L) if (enabled) { - assert( - joins.head.metrics("dynamic_filter_reader_filters_attached").value > 0L) - assert( - joins.head.metrics("dynamic_filter_reader_filters_skipped").value == 0L) + assert(joins.head.metrics("dynamic_filter_join_filters_attached").value > 0L) + assert(joins.head.metrics("dynamic_filter_join_filters_skipped").value == 0L) assert( probeFilters.head.metrics("output_rows").value > 0L, "Execution-local reader attachment must preserve probe filter metrics") @@ -545,10 +543,10 @@ class CometJoinSuite extends CometTestBase { assert(probeScans.head.metrics("row_groups_pruned_statistics").value == 0L) if (enabled) { - assert(joins.head.metrics("dynamic_filter_reader_filters_attached").value == 0L) - assert(joins.head.metrics("dynamic_filter_reader_filters_skipped").value == 1L) - assert(joins.head.metrics("dynamic_filter_rows_evaluated").value == 1L) - assert(joins.head.metrics("dynamic_filter_rows_pruned").value == 0L) + assert(joins.head.metrics("dynamic_filter_join_filters_attached").value == 0L) + assert(joins.head.metrics("dynamic_filter_join_filters_skipped").value == 1L) + assert(joins.head.metrics("dynamic_filter_join_rows_evaluated").value == 1L) + assert(joins.head.metrics("dynamic_filter_join_rows_pruned").value == 0L) } } } @@ -579,8 +577,8 @@ class CometJoinSuite extends CometTestBase { val joins = nativeHashJoins(plan) assert(joins.size == 1, s"Expected native hash join:\n$plan") if (build.size == 100) { - assert(joins.head.metrics("dynamic_filter_rows_evaluated").value > 0L) - assert(joins.head.metrics("dynamic_filter_rows_pruned").value == 0L) + assert(joins.head.metrics("dynamic_filter_join_rows_evaluated").value > 0L) + assert(joins.head.metrics("dynamic_filter_join_rows_pruned").value == 0L) } } } @@ -612,8 +610,8 @@ class CometJoinSuite extends CometTestBase { val (_, plan) = checkSparkAnswerAndOperator(sql(query)) val native = nativeHashJoins(plan) assert(native.size == 1, s"Expected native hash join:\n$plan") - assert(native.head.metrics("dynamic_filter_rows_evaluated").value == 0L) - assert(native.head.metrics("dynamic_filter_rows_pruned").value == 0L) + assert(native.head.metrics("dynamic_filter_join_rows_evaluated").value == 0L) + assert(native.head.metrics("dynamic_filter_join_rows_pruned").value == 0L) } // NOT IN must still observe build-side NULLs; never attach a filter here. withSQLConf( @@ -625,7 +623,7 @@ class CometJoinSuite extends CometTestBase { val native = collect(plan) { case join: CometBroadcastHashJoinExec => join } assert(native.size == 1, s"Expected native null-aware anti join:\n$plan") assert(native.head.nativeOp.getHashJoin.getNullAwareAntiJoin) - assert(native.head.metrics("dynamic_filter_rows_evaluated").value == 0L) + assert(native.head.metrics("dynamic_filter_join_rows_evaluated").value == 0L) } } } @@ -650,14 +648,14 @@ class CometJoinSuite extends CometTestBase { assert(join.buildSide == BuildRight) assert(join.nativeOp.getHashJoin.getDynamicFilterEnabled) assert(join.metrics("output_rows").value == 65536L) - assert(join.metrics("dynamic_filter_rows_evaluated").value > 0L) - assert(join.metrics("dynamic_filter_reader_filters_attached").value > 0L) + assert(join.metrics("dynamic_filter_join_rows_evaluated").value > 0L) + assert(join.metrics("dynamic_filter_join_filters_attached").value > 0L) val probeScans = collect(plan) { case scan: CometNativeScanExec if scan.output.size == 1 => scan } assert(probeScans.size == 1, s"Expected one native byte probe scan:\n$plan") val probeRows = probeScans.head.metrics("output_rows").value - val residualPruned = join.metrics("dynamic_filter_rows_pruned").value + val residualPruned = join.metrics("dynamic_filter_join_rows_pruned").value assert( probeRows < 2L || residualPruned > 0L, "Expected the reader or residual filter to prune probe key 2, " + From 7b2f35a7187acab9f7a34ee8bc65b78587d93f55 Mon Sep 17 00:00:00 2001 From: Ping Zhang Date: Tue, 15 Sep 2026 16:54:33 +0000 Subject: [PATCH 4/4] refactor: group runtime filter modules --- .github/labeler.yml | 1 - .../{dynamic_filter_join.rs => dynamic_filter/join.rs} | 4 ++-- .../{dynamic_filter_join => dynamic_filter/join}/tests.rs | 0 .../operators/{dynamic_filter.rs => dynamic_filter/mod.rs} | 7 ++++++- .../parquet_reader.rs} | 2 +- .../parquet_reader}/tests.rs | 0 native/core/src/execution/operators/mod.rs | 4 +--- 7 files changed, 10 insertions(+), 8 deletions(-) rename native/core/src/execution/operators/{dynamic_filter_join.rs => dynamic_filter/join.rs} (99%) rename native/core/src/execution/operators/{dynamic_filter_join => dynamic_filter/join}/tests.rs (100%) rename native/core/src/execution/operators/{dynamic_filter.rs => dynamic_filter/mod.rs} (98%) rename native/core/src/execution/operators/{parquet_reader_filter.rs => dynamic_filter/parquet_reader.rs} (99%) rename native/core/src/execution/operators/{parquet_reader_filter => dynamic_filter/parquet_reader}/tests.rs (100%) diff --git a/.github/labeler.yml b/.github/labeler.yml index 7a0ee982131..5307106f723 100644 --- a/.github/labeler.yml +++ b/.github/labeler.yml @@ -86,7 +86,6 @@ area:expressions: area:joins: - changed-files: - any-glob-to-any-file: - - 'native/core/src/execution/operators/dynamic_filter.rs' - 'native/core/src/execution/operators/dynamic_filter/**' - 'spark/src/main/scala/org/apache/comet/rules/RewriteJoin.scala' - 'spark/src/main/scala/org/apache/comet/rules/CometPlanAdaptiveDynamicPruningFilters.scala' diff --git a/native/core/src/execution/operators/dynamic_filter_join.rs b/native/core/src/execution/operators/dynamic_filter/join.rs similarity index 99% rename from native/core/src/execution/operators/dynamic_filter_join.rs rename to native/core/src/execution/operators/dynamic_filter/join.rs index cf7d98429e2..ba2e2759608 100644 --- a/native/core/src/execution/operators/dynamic_filter_join.rs +++ b/native/core/src/execution/operators/dynamic_filter/join.rs @@ -45,8 +45,8 @@ use datafusion::physical_plan::{ }; use futures::StreamExt; -use super::dynamic_filter::DynamicFilterExec; -use super::parquet_reader_filter::try_attach_parquet_reader_filter; +use super::parquet_reader::try_attach_parquet_reader_filter; +use super::DynamicFilterExec; /// A permanent plan must not own a completed join's filter or build accumulator: /// those can retain the hash map after its stream-owned reservation is released. diff --git a/native/core/src/execution/operators/dynamic_filter_join/tests.rs b/native/core/src/execution/operators/dynamic_filter/join/tests.rs similarity index 100% rename from native/core/src/execution/operators/dynamic_filter_join/tests.rs rename to native/core/src/execution/operators/dynamic_filter/join/tests.rs diff --git a/native/core/src/execution/operators/dynamic_filter.rs b/native/core/src/execution/operators/dynamic_filter/mod.rs similarity index 98% rename from native/core/src/execution/operators/dynamic_filter.rs rename to native/core/src/execution/operators/dynamic_filter/mod.rs index b543747cbeb..7ac920a4ae1 100644 --- a/native/core/src/execution/operators/dynamic_filter.rs +++ b/native/core/src/execution/operators/dynamic_filter/mod.rs @@ -15,7 +15,12 @@ // specific language governing permissions and limitations // under the License. -//! Apply a live runtime predicate to decoded batches. +//! Runtime-filter wiring and shared filtering of decoded batches. + +mod join; +mod parquet_reader; + +pub(crate) use join::DynamicFilterJoinExec; use std::fmt::Formatter; use std::sync::Arc; diff --git a/native/core/src/execution/operators/parquet_reader_filter.rs b/native/core/src/execution/operators/dynamic_filter/parquet_reader.rs similarity index 99% rename from native/core/src/execution/operators/parquet_reader_filter.rs rename to native/core/src/execution/operators/dynamic_filter/parquet_reader.rs index c705316759e..be3eb1c40e9 100644 --- a/native/core/src/execution/operators/parquet_reader_filter.rs +++ b/native/core/src/execution/operators/dynamic_filter/parquet_reader.rs @@ -30,7 +30,7 @@ use datafusion::physical_expr::expressions::{ use datafusion::physical_expr::PhysicalExpr; use datafusion::physical_plan::ExecutionPlan; -use super::CometFilterExec; +use super::super::CometFilterExec; /// Recognize only direct-column null checks joined by AND, without evaluating /// or changing the predicate. Every accepted leaf is deterministic, infallible, diff --git a/native/core/src/execution/operators/parquet_reader_filter/tests.rs b/native/core/src/execution/operators/dynamic_filter/parquet_reader/tests.rs similarity index 100% rename from native/core/src/execution/operators/parquet_reader_filter/tests.rs rename to native/core/src/execution/operators/dynamic_filter/parquet_reader/tests.rs diff --git a/native/core/src/execution/operators/mod.rs b/native/core/src/execution/operators/mod.rs index ac504c2abac..d09b0b4fb37 100644 --- a/native/core/src/execution/operators/mod.rs +++ b/native/core/src/execution/operators/mod.rs @@ -27,9 +27,7 @@ pub use scan::*; mod aligned_stream_reader; mod copy; mod dynamic_filter; -mod dynamic_filter_join; -mod parquet_reader_filter; -pub(crate) use dynamic_filter_join::DynamicFilterJoinExec; +pub(crate) use dynamic_filter::DynamicFilterJoinExec; mod filter; pub(crate) use filter::CometFilterExec; mod expand;