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/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/join.rs similarity index 51% rename from native/core/src/execution/operators/dynamic_filter.rs rename to native/core/src/execution/operators/dynamic_filter/join.rs index 27e650a2b78..ba2e2759608 100644 --- a/native/core/src/execution/operators/dynamic_filter.rs +++ b/native/core/src/execution/operators/dynamic_filter/join.rs @@ -27,188 +27,26 @@ 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::{JoinType, NullEquality, Result, Statistics}; use datafusion::execution::TaskContext; -use datafusion::logical_expr::{ColumnarValue, 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::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, + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, + PlanProperties, ReplaceChildrenOptions, SendableRecordBatchStream, }; use futures::StreamExt; -use super::CometFilterExec; - -/// A task-local consumer of DataFusion's build-side runtime filter. -#[derive(Debug)] -pub(crate) struct DynamicFilterExec { - input: Arc, - predicate: Arc, - metrics: ExecutionPlanMetricsSet, -} - -impl DynamicFilterExec { - fn new(input: Arc, predicate: Arc) -> Self { - Self { - input, - predicate, - metrics: ExecutionPlanMetricsSet::new(), - } - } -} - -impl DisplayAs for DynamicFilterExec { - fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { - write!(f, "CometDynamicFilterExec") - } -} - -impl ExecutionPlan for DynamicFilterExec { - fn name(&self) -> &str { - "CometDynamicFilterExec" - } - - fn properties(&self) -> &Arc { - // Removing rows preserves the input's schema, ordering and partitioning. - self.input.properties() - } - - fn children(&self) -> Vec<&Arc> { - vec![&self.input] - } - - fn apply_expressions( - &self, - f: &mut dyn FnMut(&Arc) -> Result, - ) -> Result { - apply_expression_roots([Arc::clone(&self.predicate) as Arc], f) - } - - fn maintains_input_order(&self) -> Vec { - vec![true] - } - - fn cardinality_effect(&self) -> CardinalityEffect { - CardinalityEffect::LowerEqual - } - - fn with_new_children( - self: Arc, - children: Vec>, - ) -> Result> { - self.replace_children( - children, - ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), - ) - } - - fn replace_children( - self: Arc, - mut children: Vec>, - _options: ReplaceChildrenOptions, - ) -> Result> { - if children.len() != 1 { - return internal_err!("CometDynamicFilterExec requires one child"); - } - Ok(Arc::new(Self::new( - children.remove(0), - Arc::clone(&self.predicate), - ))) - } - - fn reset_state(self: Arc) -> Result> { - // HashJoinExec resets its producer on reexecution. Never retain a previous - // build's domain in the consumer. A reset plan safely bypasses filtering; - // ordinary Spark task attempts each construct a fresh, connected plan. - let predicate = Arc::new(DynamicFilterPhysicalExpr::new( - self.predicate.children().into_iter().cloned().collect(), - lit(true), - )); - Ok(Arc::new(Self::new(Arc::clone(&self.input), predicate))) - } - - fn execute( - &self, - partition: usize, - context: Arc, - ) -> Result { - let children = self.predicate.children(); - let [key] = children.as_slice() else { - return internal_err!("CometDynamicFilterExec requires one join-key column"); - }; - let Some(key) = key.downcast_ref::() else { - return internal_err!("CometDynamicFilterExec requires a direct join-key column"); - }; - let key_index = key.index(); - 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); - // 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 stream = input.map(move |batch| { - let batch = batch?; - let _timer = eval_time.timer(); - // AND may prefilter its input before evaluating hash membership. A - // zero-copy key projection keeps payload columns out of that temporary - // batch. The remapped expression still observes live producer updates. - let key_batch = batch.project(&[key_index])?; - match predicate.evaluate(&key_batch)? { - // DataFusion leaves this placeholder unchanged until the complete - // build is available, or if it declines to populate the filter. - ColumnarValue::Scalar(ScalarValue::Boolean(Some(true))) => { - bypassed.add(batch.num_rows()); - Ok(batch) - } - ColumnarValue::Scalar(ScalarValue::Boolean(Some(false) | None)) => { - evaluated.add(batch.num_rows()); - pruned.add(batch.num_rows()); - Ok(batch.slice(0, 0)) - } - ColumnarValue::Array(mask) => { - let filtered = filter_record_batch(&batch, as_boolean_array(&mask)?)?; - evaluated.add(batch.num_rows()); - pruned.add(batch.num_rows() - filtered.num_rows()); - Ok(filtered) - } - _ => internal_err!("Join dynamic filter must evaluate to a Boolean"), - } - }); - // Return even empty batches. Each poll consumes at most one input batch, - // so a selective filter cannot drain a ready input in an unbounded loop. - Ok(Box::pin(RecordBatchStreamAdapter::new( - self.schema(), - stream, - ))) - } - - fn metrics(&self) -> Option { - Some(self.metrics.clone_inner()) - } -} +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. @@ -228,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> { @@ -347,11 +95,12 @@ impl DynamicFilterJoinExec { &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(), - }); + let consumer = Arc::new(DynamicFilterExec::new( + 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. let join = self @@ -492,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 new file mode 100644 index 00000000000..aa64656f73d --- /dev/null +++ b/native/core/src/execution/operators/dynamic_filter/join/tests.rs @@ -0,0 +1,1563 @@ +// 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::operators::CometFilterExec; +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::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, 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; +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(), + "dynamic_filter_join", + )); + 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_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); + } + } + } + } + } +} + +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_join_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)) +} + +/// 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_join_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_join_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_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); + } + } + 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_join_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_join_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_join_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_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_join_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/mod.rs b/native/core/src/execution/operators/dynamic_filter/mod.rs new file mode 100644 index 00000000000..7ac920a4ae1 --- /dev/null +++ b/native/core/src/execution/operators/dynamic_filter/mod.rs @@ -0,0 +1,215 @@ +// 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. + +//! 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; + +use arrow::compute::filter_record_batch; +use datafusion::common::cast::as_boolean_array; +use datafusion::common::tree_node::TreeNodeRecursion; +use datafusion::common::{internal_err, Result, ScalarValue}; +use datafusion::execution::TaskContext; +use datafusion::logical_expr::ColumnarValue; +use datafusion::physical_expr::expressions::{lit, Column, DynamicFilterPhysicalExpr}; +use datafusion::physical_expr::PhysicalExpr; +use datafusion::physical_plan::execution_plan::CardinalityEffect; +use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, MetricBuilder, MetricsSet}; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; +use datafusion::physical_plan::{ + apply_expression_roots, ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, + PlanProperties, ReplaceChildrenOptions, SendableRecordBatchStream, +}; +use futures::StreamExt; + +/// A task-local consumer of a live runtime predicate. +#[derive(Debug)] +pub(crate) struct DynamicFilterExec { + input: Arc, + predicate: Arc, + metrics: ExecutionPlanMetricsSet, + metric_prefix: &'static str, +} + +impl DynamicFilterExec { + pub(super) fn new( + input: Arc, + predicate: Arc, + metrics: ExecutionPlanMetricsSet, + metric_prefix: &'static str, + ) -> Self { + Self { + input, + predicate, + metrics, + metric_prefix, + } + } +} + +impl DisplayAs for DynamicFilterExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { + write!(f, "CometDynamicFilterExec") + } +} + +impl ExecutionPlan for DynamicFilterExec { + fn name(&self) -> &str { + "CometDynamicFilterExec" + } + + fn properties(&self) -> &Arc { + // Removing rows preserves the input's schema, ordering and partitioning. + self.input.properties() + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + apply_expression_roots([Arc::clone(&self.predicate) as Arc], f) + } + + fn maintains_input_order(&self) -> Vec { + vec![true] + } + + fn cardinality_effect(&self) -> CardinalityEffect { + CardinalityEffect::LowerEqual + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn replace_children( + self: Arc, + mut children: Vec>, + _options: ReplaceChildrenOptions, + ) -> Result> { + if children.len() != 1 { + return internal_err!("CometDynamicFilterExec requires one child"); + } + Ok(Arc::new(Self::new( + children.remove(0), + Arc::clone(&self.predicate), + ExecutionPlanMetricsSet::new(), + self.metric_prefix, + ))) + } + + fn reset_state(self: Arc) -> Result> { + // HashJoinExec resets its producer on reexecution. Never retain a previous + // build's domain in the consumer. A reset plan safely bypasses filtering; + // ordinary Spark task attempts each construct a fresh, connected plan. + let predicate = Arc::new(DynamicFilterPhysicalExpr::new( + self.predicate.children().into_iter().cloned().collect(), + lit(true), + )); + Ok(Arc::new(Self::new( + Arc::clone(&self.input), + predicate, + ExecutionPlanMetricsSet::new(), + self.metric_prefix, + ))) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + let children = self.predicate.children(); + let [key] = children.as_slice() else { + return internal_err!("CometDynamicFilterExec requires one join-key column"); + }; + let Some(key) = key.downcast_ref::() else { + return internal_err!("CometDynamicFilterExec requires a direct join-key column"); + }; + let key_index = key.index(); + 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(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(format!("{}_eval_time", self.metric_prefix), partition); + let stream = input.map(move |batch| { + let batch = batch?; + let _timer = eval_time.timer(); + // AND may prefilter its input before evaluating hash membership. A + // zero-copy key projection keeps payload columns out of that temporary + // batch. The remapped expression still observes live producer updates. + let key_batch = batch.project(&[key_index])?; + match predicate.evaluate(&key_batch)? { + // DataFusion leaves this placeholder unchanged until the complete + // build is available, or if it declines to populate the filter. + ColumnarValue::Scalar(ScalarValue::Boolean(Some(true))) => { + bypassed.add(batch.num_rows()); + Ok(batch) + } + ColumnarValue::Scalar(ScalarValue::Boolean(Some(false) | None)) => { + evaluated.add(batch.num_rows()); + pruned.add(batch.num_rows()); + Ok(batch.slice(0, 0)) + } + ColumnarValue::Array(mask) => { + let filtered = filter_record_batch(&batch, as_boolean_array(&mask)?)?; + evaluated.add(batch.num_rows()); + pruned.add(batch.num_rows() - filtered.num_rows()); + Ok(filtered) + } + _ => internal_err!("Join dynamic filter must evaluate to a Boolean"), + } + }); + // Return even empty batches. Each poll consumes at most one input batch, + // so a selective filter cannot drain a ready input in an unbounded loop. + Ok(Box::pin(RecordBatchStreamAdapter::new( + self.schema(), + stream, + ))) + } + + fn metrics(&self) -> Option { + Some(self.metrics.clone_inner()) + } +} + +#[cfg(test)] +mod tests; diff --git a/native/core/src/execution/operators/dynamic_filter/parquet_reader.rs b/native/core/src/execution/operators/dynamic_filter/parquet_reader.rs new file mode 100644 index 00000000000..be3eb1c40e9 --- /dev/null +++ b/native/core/src/execution/operators/dynamic_filter/parquet_reader.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::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/dynamic_filter/parquet_reader/tests.rs b/native/core/src/execution/operators/dynamic_filter/parquet_reader/tests.rs new file mode 100644 index 00000000000..aacc242731e --- /dev/null +++ b/native/core/src/execution/operators/dynamic_filter/parquet_reader/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)); +} diff --git a/native/core/src/execution/operators/dynamic_filter/tests.rs b/native/core/src/execution/operators/dynamic_filter/tests.rs index fe6b0d06539..e610f8b4a62 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,12 @@ 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(), + "test_filter", + )); let task = SessionContext::new().task_ctx(); let mut stream = wrapper.execute(0, Arc::clone(&task)).unwrap(); let first = stream.next().await.unwrap().unwrap(); @@ -381,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. @@ -392,1254 +116,10 @@ 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(); 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/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, " +