sunchao commented on code in PR #4810: URL: https://github.com/apache/datafusion-comet/pull/4810#discussion_r4103394742
########## native/core/src/execution/operators/dynamic_filter.rs: ########## @@ -0,0 +1,495 @@ +// 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 dynamic filter for hash join probe sides. +//! +//! [`DynamicFilterExec`] evaluates a join's [`DynamicFilterPhysicalExpr`] against +//! probe-side batches before they reach the hash probe. The expression starts as a +//! `lit(true)` placeholder and is populated by DataFusion's `HashJoinExec` build phase +//! (min/max bounds plus `InList` or hash-table-lookup membership). Until then — or if +//! the join never populates it — batches pass through untouched, so correctness never +//! depends on population. +//! +//! [`attach_join_dynamic_filter`] rewires an eligible `HashJoinExec` so that the join +//! and a new `DynamicFilterExec` wrapping its probe child share the same filter. + +use std::fmt::Formatter; +use std::pin::Pin; +use std::sync::Arc; +use std::task::{Context, Poll}; + +use arrow::array::RecordBatch; +use arrow::compute::filter_record_batch; +use arrow::datatypes::SchemaRef; +use datafusion::common::cast::as_boolean_array; +use datafusion::common::config::ConfigOptions; +use datafusion::common::{DataFusionError, Result as DataFusionResult, ScalarValue}; +use datafusion::execution::TaskContext; +use datafusion::logical_expr::ColumnarValue; +use datafusion::physical_expr::expressions::{lit, DynamicFilterPhysicalExpr}; +use datafusion::physical_expr::PhysicalExpr; +use datafusion::physical_plan::joins::{HashJoinExec, PartitionMode}; +use datafusion::physical_plan::metrics::{ + BaselineMetrics, Count, ExecutionPlanMetricsSet, MetricBuilder, MetricsSet, Time, +}; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, RecordBatchStream, + SendableRecordBatchStream, +}; +use futures::{Stream, StreamExt}; + +/// Stop evaluating the filter for the remainder of a partition stream when, after at +/// least [`GUARD_MIN_ROWS`] filtered rows, it keeps more than this fraction of them. +const GUARD_MAX_SELECTIVITY: f64 = 0.95; +/// Minimum number of rows to observe before the selectivity guard may disable the +/// filter, so a few unrepresentative leading batches don't make the decision. +const GUARD_MIN_ROWS: usize = 65_536; + +/// Filters probe-side batches with a join's shared [`DynamicFilterPhysicalExpr`]. +/// +/// Distinct from a generic `FilterExec` in three ways: a pass-through fast path while +/// the filter is still the constant-`true` placeholder, a selectivity guard that +/// disables evaluation on non-selective streams, and dedicated metrics +/// (`dynamic_filter_rows_pruned`). +#[derive(Debug)] +pub struct DynamicFilterExec { + input: Arc<dyn ExecutionPlan>, + predicate: Arc<DynamicFilterPhysicalExpr>, + metrics: ExecutionPlanMetricsSet, + cache: Arc<PlanProperties>, +} + +impl DynamicFilterExec { + pub fn new(input: Arc<dyn ExecutionPlan>, predicate: Arc<DynamicFilterPhysicalExpr>) -> Self { + // Filtering preserves schema, ordering, and partitioning. + let cache = Arc::clone(input.properties()); + Self { + input, + predicate, + metrics: ExecutionPlanMetricsSet::new(), + cache, + } + } + + pub fn predicate(&self) -> &Arc<DynamicFilterPhysicalExpr> { + &self.predicate + } +} + +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<PlanProperties> { + &self.cache + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + vec![&self.input] + } + + fn maintains_input_order(&self) -> Vec<bool> { + vec![true] + } + + fn with_new_children( + self: Arc<Self>, + mut children: Vec<Arc<dyn ExecutionPlan>>, + ) -> DataFusionResult<Arc<dyn ExecutionPlan>> { + Ok(Arc::new(DynamicFilterExec::new( + children.swap_remove(0), + Arc::clone(&self.predicate), + ))) + } + + fn execute( + &self, + partition: usize, + context: Arc<TaskContext>, + ) -> DataFusionResult<SendableRecordBatchStream> { + let input = self.input.execute(partition, context)?; + let baseline_metrics = BaselineMetrics::new(&self.metrics, partition); + let rows_pruned = + MetricBuilder::new(&self.metrics).counter("dynamic_filter_rows_pruned", partition); + // A dedicated timer rather than the baseline elapsed_compute: this operator is + // registered as an additional native plan on the join's SparkPlan node for + // metrics collection, and the metric merge sums same-named metrics — timing + // via elapsed_compute would inflate the join's reported compute time. + let eval_time = + MetricBuilder::new(&self.metrics).subset_time("dynamic_filter_eval_time", partition); + Ok(Box::pin(DynamicFilterStream { + schema: self.input.schema(), + input, + predicate: Arc::clone(&self.predicate), + baseline_metrics, + rows_pruned, + eval_time, + rows_evaluated: 0, + rows_kept: 0, + guard_disabled: false, + })) + } + + fn metrics(&self) -> Option<MetricsSet> { + Some(self.metrics.clone_inner()) + } +} + +struct DynamicFilterStream { + schema: SchemaRef, + input: SendableRecordBatchStream, + predicate: Arc<DynamicFilterPhysicalExpr>, + baseline_metrics: BaselineMetrics, + rows_pruned: Count, + eval_time: Time, + /// Rows seen since the filter became a real (non-placeholder) predicate. + rows_evaluated: usize, + rows_kept: usize, + /// Set once the selectivity guard decides the filter is not worth evaluating. + guard_disabled: bool, +} + +impl DynamicFilterStream { + fn filter_batch(&mut self, batch: RecordBatch) -> DataFusionResult<Option<RecordBatch>> { + let _timer = self.eval_time.timer(); + match self.predicate.evaluate(&batch)? { + // Placeholder (or degenerate all-true) predicate: pass through untouched. + // Not counted toward the selectivity guard — the real filter may not have + // arrived yet. + ColumnarValue::Scalar(ScalarValue::Boolean(Some(true))) => Ok(Some(batch)), + // Constant false/null (e.g. empty build side): the whole batch is pruned. + ColumnarValue::Scalar(_) => { + self.rows_pruned.add(batch.num_rows()); + self.rows_evaluated += batch.num_rows(); + Ok(None) + } + ColumnarValue::Array(array) => { + let mask = as_boolean_array(&array)?; + let filtered = filter_record_batch(&batch, mask)?; + let kept = filtered.num_rows(); + self.rows_pruned.add(batch.num_rows() - kept); + self.rows_evaluated += batch.num_rows(); + self.rows_kept += kept; + if self.rows_evaluated >= GUARD_MIN_ROWS + && (self.rows_kept as f64 / self.rows_evaluated as f64) > GUARD_MAX_SELECTIVITY + { + self.guard_disabled = true; + } + if kept == 0 { + Ok(None) + } else { + Ok(Some(filtered)) + } + } + } + } +} + +impl Stream for DynamicFilterStream { + type Item = DataFusionResult<RecordBatch>; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> { + loop { + match self.input.poll_next_unpin(cx) { + Poll::Ready(Some(Ok(batch))) => { + if self.guard_disabled { + self.baseline_metrics.record_output(batch.num_rows()); + return Poll::Ready(Some(Ok(batch))); + } + match self.filter_batch(batch) { + Ok(Some(filtered)) => { + self.baseline_metrics.record_output(filtered.num_rows()); + return Poll::Ready(Some(Ok(filtered))); + } + // Entire batch pruned: keep polling the input. + Ok(None) => continue, + Err(e) => return Poll::Ready(Some(Err(e))), + } + } + other => return other, + } + } + } +} + +impl RecordBatchStream for DynamicFilterStream { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +} + +/// Attaches a runtime dynamic filter to an eligible [`HashJoinExec`], wrapping its +/// probe (right) child in a [`DynamicFilterExec`] that shares the same filter. +/// +/// Accepts either a `HashJoinExec` or a `ProjectionExec` directly above one (the shape +/// `HashJoinExec::swap_inputs` produces). Returns the (possibly rewritten) plan plus +/// the installed wrapper, so the planner can register the wrapper for metrics +/// collection (`SparkPlan::new_with_additional`); when the join is not eligible the +/// input plan is returned unchanged with `None`. +/// +/// Eligibility mirrors DataFusion's own `allow_join_dynamic_filter_pushdown` gate: +/// - the session option `optimizer.enable_join_dynamic_filter_pushdown` must be on, +/// - `optimizer.preserve_file_partitions` with `PartitionMode::Partitioned` is +/// excluded (file-group partitions are not hash-distributed by the join keys), +/// - the probe side must be preserved under the ON clause +/// (`JoinType::on_lr_is_preserved().1`), which admits Inner, Left, LeftSemi, +/// RightSemi, LeftAnti, and LeftMark joins — a probe row removed by the filter +/// could not have matched any build row, so results are unchanged. +/// +/// DataFusion re-checks the same gate at execute time (an ineligible join never +/// populates the filter, leaving the wrapper a harmless pass-through); checking here +/// as well avoids installing a wrapper that can never engage. +/// +/// Callers must not pass null-aware anti joins (Spark NOT IN semantics); that gate +/// lives at the call site where the flag is known. +pub fn attach_join_dynamic_filter( + plan: Arc<dyn ExecutionPlan>, + config: &ConfigOptions, +) -> DataFusionResult<PlanWithDynamicFilter> { + // swap_inputs may have inserted a projection above the join to restore column order. + if plan.is::<datafusion::physical_plan::projection::ProjectionExec>() { + let child = Arc::clone(plan.children()[0]); + let (new_child, wrapper) = attach_join_dynamic_filter(child, config)?; + return Ok((plan.with_new_children(vec![new_child])?, wrapper)); + } + + let Some(hash_join) = plan.downcast_ref::<HashJoinExec>() else { + return Ok((plan, None)); + }; + if !config.optimizer.enable_join_dynamic_filter_pushdown { + return Ok((plan, None)); + } + if config.optimizer.preserve_file_partitions > 0 + && matches!(hash_join.partition_mode(), PartitionMode::Partitioned) + { + return Ok((plan, None)); + } + if !hash_join.join_type().on_lr_is_preserved().1 { + return Ok((plan, None)); + } + let probe_keys: Vec<Arc<dyn PhysicalExpr>> = hash_join + .on() + .iter() + .map(|(_, right)| Arc::clone(right)) + .collect(); + if probe_keys.is_empty() { + return Ok((plan, None)); + } + + let dynamic_filter = Arc::new(DynamicFilterPhysicalExpr::new(probe_keys, lit(true))); Review Comment: [P1] Gate nested join keys before attaching this filter. With `spark.comet.exec.join.dynamicFilter.enabled=true`, an inner hash join whose two inputs contain struct keys `{x:1}`, `{x:null}`, and `{x:3}` returns only two matches instead of three. DataFusion 54 generates bounds `{x:1} <= k <= {x:3}` that discard the valid nested-null match. Array keys `[1]`, `[null]`, and `[3]` also reproduce incorrect pruning, losing all three matches. Spark and the unfiltered native join preserve these matches, so enabling this optimization silently changes query results. Skip affected struct/array key types until upstream bounds preserve their comparison semantics, and add differential regression coverage. Evidence: An isolated Rust harness imports this exact head’s unchanged `dynamic_filter.rs` and constructs single-partition `HashJoinExec` plans with `PartitionMode::Partitioned`, `NullEquality::NullEqualsNothing`, and column keys. Comparing `collect(plain)` with `collect(attach_join_dynamic_filter(plain, ConfigOptions::default()))` produced struct rows 3→2 and array rows 3→0. Commands: `cargo test --manifest-path /tmp/comet-4810-test/Cargo.toml --test key_types struct_keys -- --nocapture` and the same command with `list_null_element`; both fail parity assertions. Spark 4.1.3 returned three matches for both datasets using `SELECT /*+ BROADCAST(b) */ p.k, b.k FROM keys p JOIN keys b ON p.k=b.k`, and likewise with `SHUFFLE_HASH(b)`. Logs: `/tmp/comet-4810-struct.log`, `/tmp/comet-4810-list.log`, and `/tmp/comet-4810-spark-reference.log`. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
