sunchao commented on code in PR #7598: URL: https://github.com/apache/datafusion-comet/pull/7598#discussion_r4532043323
########## native/core/src/execution/operators/dynamic_filter/min_max.rs: ########## @@ -0,0 +1,888 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Connect one partial MIN/MAX aggregate's improving bound to its Parquet reader. + +use std::fmt::Formatter; +use std::sync::Arc; + +use arrow::datatypes::DataType; +use datafusion::common::config::ConfigOptions; +use datafusion::common::tree_node::TreeNodeRecursion; +use datafusion::common::{internal_err, Result, Statistics}; +use datafusion::datasource::source::DataSourceExec; +use datafusion::execution::TaskContext; +use datafusion::functions_aggregate::min_max::{Max, Min}; +use datafusion::physical_expr::expressions::{lit, Column, DynamicFilterPhysicalExpr}; +use datafusion::physical_expr::{OrderingRequirements, PhysicalExpr}; +use datafusion::physical_plan::aggregates::{AggregateExec, AggregateMode}; +use datafusion::physical_plan::distribution_requirements::InputDistributionRequirements; +use datafusion::physical_plan::execution_plan::CardinalityEffect; +use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, MetricBuilder, MetricsSet}; +use datafusion::physical_plan::statistics::{ChildStats, StatisticsArgs}; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; +use datafusion::physical_plan::{ + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, + PlanProperties, ReplaceChildrenOptions, SendableRecordBatchStream, +}; +use futures::StreamExt; + +use super::parquet_reader::{parquet_file_source, try_attach_parquet_reader_filter}; + +/// Own an unexecuted aggregate template and metric handles. Each execution owns +/// a fresh aggregate, accumulator bounds, and reader predicate; none of that live +/// state is retained by the template or shared with another Spark task. +#[derive(Debug)] +pub(crate) struct MinMaxReaderFilterExec { + template: AggregateExec, + config: ConfigOptions, + metrics: ExecutionPlanMetricsSet, +} + +struct RuntimeAggregate { + aggregate: AggregateExec, + reader_filter_attached: bool, +} + +impl MinMaxReaderFilterExec { + pub(crate) fn try_new( + aggregate: &AggregateExec, + config: &ConfigOptions, + ) -> Result<Option<Self>> { + if !config.optimizer.enable_dynamic_filter_pushdown + || !config.optimizer.enable_aggregate_dynamic_filter_pushdown + || aggregate.mode() != &AggregateMode::Partial + || !aggregate.group_expr().is_true_no_grouping() + || aggregate.limit_options().is_some() + || aggregate.filter_expr().iter().any(Option::is_some) + || aggregate.input().output_partitioning().partition_count() != 1 + || aggregate.input().fetch().is_some() + { + return Ok(None); + } + let [expression] = aggregate.aggr_expr() else { + return Ok(None); + }; + let function = expression.fun().inner(); + if !(function.is::<Min>() || function.is::<Max>()) + || expression.is_distinct() + || !expression.order_bys().is_empty() + { + return Ok(None); + } + let arguments = expression.expressions(); + let [key] = arguments.as_slice() else { + return Ok(None); + }; + if !key.is::<Column>() + || !matches!( + key.data_type(aggregate.input().schema().as_ref())?, + DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 + ) + { + return Ok(None); + } + // Keep expressions, residuals, limits, and native/JVM boundaries out of + // this first path. The shared helper still handles reader column mapping + // and per-file schema guards once this direct scan is accepted. + let Some(scan) = aggregate.input().downcast_ref::<DataSourceExec>() else { + return Ok(None); + }; + let Some((_, source)) = parquet_file_source(scan) else { + return Ok(None); + }; + if source.filter().is_some() { + return Ok(None); + } + Ok(Some(Self::new(aggregate, config.clone())?)) + } + + fn new(aggregate: &AggregateExec, config: ConfigOptions) -> Result<Self> { + Ok(Self { + template: Self::fresh_aggregate(aggregate, Arc::clone(aggregate.input()))?, + config, + metrics: ExecutionPlanMetricsSet::new(), + }) + } + + fn fresh_aggregate( + template: &AggregateExec, + input: Arc<dyn ExecutionPlan>, + ) -> Result<AggregateExec> { + // AggregateExec::clone/replace_children keep shared accumulator bounds. + // Replacing just its dynamic expression would leave those bounds stale. + AggregateExec::try_new( + *template.mode(), + template.group_expr().clone(), + template.aggr_expr().to_vec(), + template.filter_expr().to_vec(), + input, + template.input_schema(), + ) + } + + fn build_runtime_aggregate(&self) -> Result<RuntimeAggregate> { + let predicate = Arc::new(DynamicFilterPhysicalExpr::new( + self.template.aggr_expr()[0].expressions(), + lit(true), + )); + let reader = try_attach_parquet_reader_filter( + self.template.input(), + Arc::clone(&predicate), + &self.config, + )?; + let reader_filter_attached = reader.is_some(); + let input = reader.unwrap_or_else(|| Arc::clone(self.template.input())); + let aggregate = + Self::fresh_aggregate(&self.template, input)?.with_dynamic_filter_expr(predicate)?; + Ok(RuntimeAggregate { + aggregate, + reader_filter_attached, + }) + } + + fn execute_runtime_aggregate( + &self, + aggregate: AggregateExec, + partition: usize, + context: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + let result = aggregate.execute(partition, context); + // The ungrouped stream registers its counters in execute(). Retain only + // their handles, including on error, so the runtime state can be dropped. + for metric in aggregate.metrics().unwrap_or_default().iter() { + self.metrics.register(Arc::clone(metric)); + } + drop(aggregate); + let input = result?; + // Also release state at EOF/error if the caller retains the spent stream. + let stream = futures::stream::unfold(Some(input), |input| async move { + let mut input = input?; + let batch = input.next().await?; + let remaining = if batch.is_ok() { Some(input) } else { None }; + Some((batch, remaining)) + }); + Ok(Box::pin(RecordBatchStreamAdapter::new( + self.schema(), + stream, + ))) + } +} + +impl DisplayAs for MinMaxReaderFilterExec { + fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { + write!(f, "CometMinMaxReaderFilterExec: ")?; + self.template.fmt_as(t, f) + } +} + +impl ExecutionPlan for MinMaxReaderFilterExec { + fn name(&self) -> &str { + "CometMinMaxReaderFilterExec" + } + + fn properties(&self) -> &Arc<PlanProperties> { + self.template.properties() + } + + fn input_distribution_requirements(&self) -> InputDistributionRequirements { + self.template.input_distribution_requirements() + } + + fn required_input_ordering(&self) -> Vec<Option<OrderingRequirements>> { + self.template.required_input_ordering() + } + + fn maintains_input_order(&self) -> Vec<bool> { + self.template.maintains_input_order() + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + self.template.children() + } + + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + // The template's unused predicate is not a permanent-plan consumer. + datafusion::physical_plan::apply_expression_roots( + self.template.aggr_expr()[0].expressions(), + f, + ) + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn replace_children( + self: Arc<Self>, + mut children: Vec<Arc<dyn ExecutionPlan>>, + _options: ReplaceChildrenOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + if children.len() != 1 { + return internal_err!("CometMinMaxReaderFilterExec requires one child"); + } + let aggregate = Self::fresh_aggregate(&self.template, children.remove(0))?; + match Self::try_new(&aggregate, &self.config)? { + Some(wrapper) => Ok(Arc::new(wrapper)), + None => Ok(Arc::new(aggregate)), + } + } + + fn reset_state(self: Arc<Self>) -> Result<Arc<dyn ExecutionPlan>> { + Ok(Arc::new(Self::new(&self.template, self.config.clone())?)) + } + + fn execute( + &self, + partition: usize, + context: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + let runtime = self.build_runtime_aggregate()?; + let attachment_metric = if runtime.reader_filter_attached { + "dynamic_filter_minmax_filters_attached" + } else { + "dynamic_filter_minmax_filters_skipped" + }; + MetricBuilder::new(&self.metrics) + .counter(attachment_metric, partition) + .add(1); + self.execute_runtime_aggregate(runtime.aggregate, partition, context) + } + + fn metrics(&self) -> Option<MetricsSet> { + Some(self.metrics.clone_inner()) + } + + fn cardinality_effect(&self) -> CardinalityEffect { + self.template.cardinality_effect() + } + + fn child_stats_requests(&self, partition: Option<usize>) -> Vec<ChildStats> { + self.template.child_stats_requests(partition) + } + + fn statistics_from_inputs( + &self, + input_stats: &[Arc<Statistics>], + args: &StatisticsArgs, + ) -> Result<Arc<Statistics>> { + self.template.statistics_from_inputs(input_stats, args) + } +} + +#[cfg(test)] +mod tests { + use super::super::parquet_reader::concrete_parquet_source; + use super::*; + use crate::parquet::parquet_exec::init_datasource_exec; + use crate::parquet::parquet_support::ObjectStoreBackend; + use arrow::array::{ArrayRef, Int32Array}; + use arrow::compute::cast; + use arrow::datatypes::{Field, Schema}; + use arrow::record_batch::RecordBatch; + use datafusion::datasource::listing::PartitionedFile; + use datafusion::datasource::memory::MemorySourceConfig; + use datafusion::execution::object_store::ObjectStoreUrl; + use datafusion::functions_aggregate::min_max::{max_udaf, min_udaf}; + use datafusion::functions_aggregate::sum::sum_udaf; + use datafusion::physical_expr::aggregate::AggregateExprBuilder; + use datafusion::physical_expr::expressions::{CastExpr, IsNotNullExpr}; + use datafusion::physical_plan::aggregates::PhysicalGroupBy; + use datafusion::physical_plan::collect; + use datafusion::physical_plan::filter::FilterExec; + use datafusion::physical_plan::projection::ProjectionExec; + use datafusion::prelude::{SessionConfig, SessionContext}; + use parquet::arrow::ArrowWriter; + use parquet::file::properties::{EnabledStatistics, WriterProperties}; + + fn session() -> Arc<SessionContext> { + let mut config = SessionConfig::new() + .with_target_partitions(1) + .with_batch_size(100) + .with_parquet_page_index_pruning(false); + // Measure row-group pruning independently of decoder filtering and page indexes. + config.options_mut().execution.parquet.pushdown_filters = false; + Arc::new(SessionContext::new_with_config(config)) + } + + fn aggregate(input: Arc<dyn ExecutionPlan>, is_max: bool) -> AggregateExec { + let expr = AggregateExprBuilder::new( + if is_max { max_udaf() } else { min_udaf() }, + vec![Arc::new(Column::new("key", 0))], + ) + .schema(input.schema()) + .alias("extremum") + .build() + .unwrap(); + AggregateExec::try_new( + AggregateMode::Partial, + PhysicalGroupBy::default(), + vec![Arc::new(expr)], + vec![None], + Arc::clone(&input), + input.schema(), + ) + .unwrap() + } + + fn wrapper(aggregate: &AggregateExec) -> MinMaxReaderFilterExec { + MinMaxReaderFilterExec::try_new(aggregate, &ConfigOptions::default()) + .unwrap() + .expect("eligible MIN/MAX") + } + + fn result(batches: &[RecordBatch]) -> Option<i32> { + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 1); + let batch = batches.iter().find(|batch| batch.num_rows() == 1).unwrap(); + let column = cast(batch.column(0), &DataType::Int32).unwrap(); + column + .as_any() + .downcast_ref::<Int32Array>() + .unwrap() + .iter() + .next() + .unwrap() + } + + fn metric(plan: &dyn ExecutionPlan, name: &str) -> usize { + plan.metrics() + .and_then(|metrics| metrics.sum_by_name(name)) + .map_or(0, |value| value.as_usize()) + } + + fn produced_filter(aggregate: &AggregateExec) -> Arc<DynamicFilterPhysicalExpr> { + let mut expressions = aggregate.dynamic_expressions_produced(); + assert_eq!(expressions.len(), 1); + (expressions.pop().unwrap() as Arc<dyn std::any::Any + Send + Sync>) + .downcast::<DynamicFilterPhysicalExpr>() + .unwrap() + } + + #[tokio::test] + async fn signed_integer_extrema_preserve_empty_null_and_duplicate_results() { + let session = session(); + for key_type in [ + DataType::Int8, + DataType::Int16, + DataType::Int32, + DataType::Int64, + ] { + for values in [ + vec![], + vec![None; 4], + vec![None, Some(-8), Some(4), Some(-8), None], + ] { + for is_max in [false, true] { + let expected = if is_max { + values.iter().flatten().max() + } else { + values.iter().flatten().min() + } + .copied(); + let (_file, scan) = parquet_input(values.clone(), &key_type, &session, 2); + let plain = aggregate(scan, is_max); + let filtered = Arc::new(wrapper(&plain)); + let actual = collect(Arc::clone(&filtered) as _, session.task_ctx()) + .await + .unwrap(); + let baseline = collect(Arc::new(plain), session.task_ctx()).await.unwrap(); + assert_eq!(result(&actual), expected, "{key_type:?}, max={is_max}"); + assert_eq!(result(&actual), result(&baseline)); + assert_eq!( + metric(filtered.as_ref(), "dynamic_filter_minmax_filters_attached"), + 1 + ); + } + } + } + } + + #[tokio::test] + async fn live_bounds_prune_remapped_row_groups_and_keep_scan_metrics_on_the_scan() { + let session = session(); + for is_max in [false, true] { + for favorable in [false, true] { + // MIN improves 200 -> 100 -> 0; MAX uses the negated layout. + let mut starts = if favorable { + vec![200, 100, 150, 0, 50, 300] + } else { + vec![500, 400, 300, 200, 100, 0] + }; + let sign = if is_max { -1 } else { 1 }; + let values: Vec<_> = starts + .drain(..) + .flat_map(|start| start..start + 100) + .map(|value| Some(sign * value)) + .collect(); + let (_file, scan) = parquet_input(values.clone(), &DataType::Int32, &session, 100); + let filtered = Arc::new(wrapper(&aggregate(Arc::clone(&scan) as _, is_max))); + let actual = collect(Arc::clone(&filtered) as _, session.task_ctx()) + .await + .unwrap(); + let (_baseline_file, baseline_scan) = + parquet_input(values, &DataType::Int32, &session, 100); + let baseline = collect( + Arc::new(aggregate(Arc::clone(&baseline_scan) as _, is_max)), + session.task_ctx(), + ) + .await + .unwrap(); + assert_eq!(result(&actual), result(&baseline)); + assert_eq!(metric(baseline_scan.as_ref(), "output_rows"), 600); + assert_eq!( + metric(baseline_scan.as_ref(), "row_groups_pruned_dynamic_filter"), + 0 + ); + assert_eq!(metric(filtered.as_ref(), "output_rows"), 1); + assert_eq!( + metric(filtered.as_ref(), "dynamic_filter_minmax_filters_attached"), + 1 + ); + assert!(filtered + .metrics() + .unwrap() + .sum_by_name("bytes_scanned") + .is_none()); + if favorable { + assert!(metric(scan.as_ref(), "row_groups_pruned_dynamic_filter") >= 3); + assert!(metric(scan.as_ref(), "output_rows") <= 300); + assert!( + metric(scan.as_ref(), "bytes_scanned") + < metric(baseline_scan.as_ref(), "bytes_scanned") + ); + } else { + assert_eq!(metric(scan.as_ref(), "row_groups_pruned_dynamic_filter"), 0); + assert_eq!(metric(scan.as_ref(), "output_rows"), 600); + } + } + } + } + + #[tokio::test] + async fn each_execution_resets_both_the_predicate_and_accumulator_bounds() { + let session = session(); + let (_file, scan) = parquet_input( + vec![Some(4), Some(3), Some(2), Some(1)], + &DataType::Int32, + &session, + 2, + ); + let plan = Arc::new(wrapper(&aggregate(scan, false))); + let first = plan.build_runtime_aggregate().unwrap(); + let second = plan.build_runtime_aggregate().unwrap(); + let first_filter = produced_filter(&first.aggregate); + let second_filter = produced_filter(&second.aggregate); + assert_ne!(first_filter.expression_id(), second_filter.expression_id()); + for (runtime, predicate) in [(first, first_filter), (second, second_filter)] { + assert_eq!(predicate.snapshot_generation(), 1); + let rows = collect(Arc::new(runtime.aggregate), session.task_ctx()) + .await + .unwrap(); + assert_eq!(result(&rows), Some(1)); + // A fresh predicate with stale shared bounds would never publish this bound. + assert!(predicate.snapshot_generation() > 1); + } + for execution in [ + Arc::clone(&plan) as Arc<dyn ExecutionPlan>, + Arc::clone(&plan) as _, + Arc::clone(&plan).reset_state().unwrap(), + ] { + assert_eq!( + result(&collect(execution, session.task_ctx()).await.unwrap()), + Some(1) + ); + } + let (_replacement_file, replacement) = parquet_input( + vec![Some(104), Some(103), Some(102), Some(101)], + &DataType::Int32, + &session, + 2, + ); + let replaced = plan.with_new_children(vec![replacement]).unwrap(); Review Comment: [P2] Use the non-deprecated child-replacement API in this test. The exact-head fork Rust job runs Clippy with warnings denied and fails here because `ExecutionPlan::with_new_children` is deprecated, before the new tests execute. The failure is visible in [the fork CI job](https://github.com/pingzh/datafusion-comet/actions/runs/42924255591/job/130208797598). Updating this call to the supported replacement helper will unblock that check. -- 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]
