Copilot commented on code in PR #25376: URL: https://github.com/apache/datafusion/pull/25376#discussion_r4031868442
########## datafusion/physical-plan/src/aggregates/builder.rs: ########## @@ -0,0 +1,1063 @@ +// 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. + +//! [`AggregateExecBuilder`]: build and rewrite [`AggregateExec`] nodes + +use std::sync::Arc; + +use super::{ + AggrDynFilter, AggregateExec, AggregateMode, LimitOptions, PhysicalGroupBy, + create_schema, get_finer_aggregate_exprs_requirement, topk_types_supported, +}; +use crate::metrics::ExecutionPlanMetricsSet; +use crate::{ExecutionPlan, ExecutionPlanProperties, InputOrderMode, PlanProperties}; + +use arrow::datatypes::SchemaRef; +use datafusion_common::{Result, assert_eq_or_internal_err, internal_err, plan_err}; +use datafusion_physical_expr::aggregate::AggregateFunctionExpr; +use datafusion_physical_expr::equivalence::ProjectionMapping; +use datafusion_physical_expr_common::physical_expr::PhysicalExpr; +use datafusion_physical_expr_common::sort_expr::{ + LexRequirement, OrderingRequirements, PhysicalSortRequirement, +}; + +/// The `FILTER` expression of each aggregate expression, `None` where an +/// aggregate has no filter. +type FilterExprs = Arc<[Option<Arc<dyn PhysicalExpr>>]>; + +/// Builds an [`AggregateExec`], and is the single place one is constructed and +/// validated. +/// +/// Reached through [`AggregateExec::builder`] for a new node and +/// [`AggregateExec::to_builder`] for a rewrite of an existing one; a rewrite +/// keeps the derived state (output schema, plan properties, ordering +/// requirements, dynamic filter) of the node it came from unless a field it is +/// computed from changes. +/// +/// Public for internal use only: this is how DataFusion's own physical +/// optimizer rules build and rewrite aggregates, and it may change without +/// notice. +/// +/// ``` +/// # use std::sync::Arc; +/// # use arrow::datatypes::{DataType, Field, Schema}; +/// # use datafusion_physical_plan::aggregates::{ +/// # AggregateExec, AggregateMode, LimitOptions, PhysicalGroupBy, +/// # }; +/// # use datafusion_physical_plan::{ExecutionPlan, empty::EmptyExec}; +/// # use datafusion_physical_expr::expressions::col; +/// # fn main() -> datafusion_common::Result<()> { +/// # let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])); +/// # let input = Arc::new(EmptyExec::new(Arc::clone(&schema))); +/// # let group_by = +/// # PhysicalGroupBy::new_single(vec![(col("a", &schema)?, "a".to_string())]); +/// let exec = AggregateExec::builder(AggregateMode::Single, input) +/// .with_group_by(group_by) +/// .build()?; +/// +/// // push a limit into it, keeping its schema and plan properties +/// let limited = exec.to_builder().with_limit_options(LimitOptions::new(10)).build()?; +/// assert_eq!(limited.schema(), exec.schema()); +/// # Ok(()) +/// # } +/// ``` +#[doc(hidden)] +#[derive(Debug, Clone)] +pub struct AggregateExecBuilder { + mode: AggregateMode, + group_by: Arc<PhysicalGroupBy>, + aggr_expr: Arc<[Arc<AggregateFunctionExpr>]>, + /// `None` means "no filter for any aggregate expression" + filter_expr: Option<FilterExprs>, + input: Arc<dyn ExecutionPlan>, + /// `None` means "the schema of `input`" + input_schema: Option<SchemaRef>, + limit_options: Option<LimitOptions>, + /// Output schema explicitly supplied by the caller, see + /// [`AggregateExecBuilder::with_output_schema`]. Always honored. + output_schema: Option<SchemaRef>, + /// State carried over from the [`AggregateExec`] this builder was derived + /// from, dropped as soon as a field it is computed from changes. + derived: Option<DerivedState>, + /// Whether the aggregate expressions were replaced, which is what makes the + /// inherited output schema and dynamic filter worth re-checking. + aggr_expr_replaced: bool, +} + +/// State of an [`AggregateExec`] that is computed from its inputs, and which is +/// preserved verbatim when a node is rewritten without touching what it is +/// computed from. +#[derive(Debug, Clone)] +struct DerivedState { + schema: SchemaRef, + cache: Arc<PlanProperties>, + required_input_ordering: Option<OrderingRequirements>, + input_order_mode: InputOrderMode, + dynamic_filter: Option<Arc<AggrDynFilter>>, +} + +impl AggregateExecBuilder { + /// Create a builder for an aggregate over `input`. + /// + /// Unless overridden the aggregate has no group by expressions, no + /// aggregate expressions, no filters, no limit, and uses the schema of + /// `input` as its [input schema](AggregateExec::input_schema). + pub fn new(mode: AggregateMode, input: Arc<dyn ExecutionPlan>) -> Self { + Self { + mode, + group_by: Arc::new(PhysicalGroupBy::default()), + aggr_expr: Arc::from([]), + filter_expr: None, + input, + input_schema: None, + limit_options: None, + output_schema: None, + derived: None, + aggr_expr_replaced: false, + } + } + + /// Create a builder pre-populated from `exec`. + /// + /// Takes `&AggregateExec` rather than ownership because every caller holds + /// a borrow from `downcast_ref` on an `Arc<dyn ExecutionPlan>`; nothing is + /// deep-copied, the fields are `Arc`s. + pub(crate) fn from_exec(exec: &AggregateExec) -> Self { + Self { + mode: exec.mode, + group_by: Arc::clone(&exec.group_by), + aggr_expr: Arc::clone(&exec.aggr_expr), + filter_expr: Some(Arc::clone(&exec.filter_expr)), + input: Arc::clone(&exec.input), + input_schema: Some(Arc::clone(&exec.input_schema)), + limit_options: exec.limit_options, + output_schema: None, + derived: Some(DerivedState { + schema: Arc::clone(&exec.schema), + cache: Arc::clone(&exec.cache), + required_input_ordering: exec.required_input_ordering.clone(), + input_order_mode: exec.input_order_mode.clone(), + dynamic_filter: exec.dynamic_filter.clone(), + }), + aggr_expr_replaced: false, + } + } + + /// Set the [`AggregateMode`]. + pub fn with_mode(mut self, mode: AggregateMode) -> Self { + if mode == self.mode { + return self; + } + self.mode = mode; + self.invalidate_derived() + } + + /// Set the group by expressions. + pub fn with_group_by(mut self, group_by: impl Into<Arc<PhysicalGroupBy>>) -> Self { + let group_by = group_by.into(); + if Arc::ptr_eq(&self.group_by, &group_by) || *self.group_by == *group_by { + return self; + } + self.group_by = group_by; + self.invalidate_derived() + } + + /// Set the aggregate expressions. + /// + /// A builder derived from an existing node keeps that node's output schema, + /// so rewriting the aggregate expressions (for example reversing them in + /// `OptimizeAggregateOrder`) cannot change output field names. `build` + /// checks the new expressions against that schema, and against the dynamic + /// filter the node carried. + pub fn with_aggr_exprs( + mut self, + aggr_expr: impl Into<Arc<[Arc<AggregateFunctionExpr>]>>, + ) -> Self { + let aggr_expr = aggr_expr.into(); + if aggr_expr == self.aggr_expr { + return self; + } + self.aggr_expr = aggr_expr; + self.aggr_expr_replaced = true; + self + } + + /// Set the `FILTER` expression of each aggregate expression. + /// + /// Must have the same length as the aggregate expressions; `build` returns + /// an error otherwise. If never called, no aggregate is filtered. + pub fn with_filter_exprs(mut self, filter_expr: impl Into<FilterExprs>) -> Self { + let filter_expr = filter_expr.into(); + if self.filter_expr.as_ref().is_some_and(|existing| { + Arc::ptr_eq(existing, &filter_expr) || **existing == *filter_expr + }) { + return self; + } + self.filter_expr = Some(filter_expr); + self.invalidate_derived() + } + + /// Set the input plan. + pub fn with_input(mut self, input: Arc<dyn ExecutionPlan>) -> Self { + if Arc::ptr_eq(&self.input, &input) { + return self; + } + self.input = input; + self.invalidate_derived() + } + + /// Set the [input schema](AggregateExec::input_schema): the schema of the + /// data *before* any aggregation is applied. + /// + /// For `Partial` and `Single` aggregates this is the schema of the input + /// plan (the default). For `Final` and `FinalPartitioned` aggregates it is + /// the input schema of the matching partial aggregate, which is *not* the + /// schema of the input plan. + pub fn with_input_schema(mut self, input_schema: SchemaRef) -> Self { + self.input_schema = Some(input_schema); + self + } + + /// Set the limit pushed down into this aggregate, or `None` to remove it. + /// + /// The limit is a hint: operators above the aggregate still enforce it. + /// `build` rejects limits this aggregate cannot execute, rather than + /// letting them fail (or be silently ignored) at execution time. Accepts + /// both `LimitOptions` and `Option<LimitOptions>`. + pub fn with_limit_options( + mut self, + limit_options: impl Into<Option<LimitOptions>>, + ) -> Self { + self.limit_options = limit_options.into(); + self + } + + /// Use `schema` as the output schema instead of computing it. + /// + /// For callers that must preserve a schema exactly, such as decoding a + /// serialized plan. The caller owns the schema being correct. + pub(crate) fn with_output_schema(mut self, schema: SchemaRef) -> Self { + self.output_schema = Some(schema); + self + } + + /// Drop state derived from the node this builder came from, because a field + /// it is computed from was replaced. + fn invalidate_derived(mut self) -> Self { + self.derived = None; + self + } + + /// Build the [`AggregateExec`], validating it. + pub fn build(self) -> Result<AggregateExec> { + let Self { + mode, + group_by, + aggr_expr, + filter_expr, + input, + input_schema, + limit_options, + output_schema, + derived, + aggr_expr_replaced, + } = self; + + let input_schema = input_schema.unwrap_or_else(|| input.schema()); + let filter_expr = filter_expr + .unwrap_or_else(|| std::iter::repeat_n(None, aggr_expr.len()).collect()); + + assert_eq_or_internal_err!( + aggr_expr.len(), + filter_expr.len(), + "Inconsistent aggregate expr: {:?} and filter expr: {:?} for AggregateExec, their size should match", + aggr_expr, + filter_expr + ); + + let mut exec = match derived { + // Nothing the derived state is computed from changed: clone the + // node this builder came from with the new values rather than + // recomputing. In particular its output schema is kept, so a + // rewrite of the aggregate expressions cannot rename output fields. + Some(derived) if output_schema.is_none() => { + let mut dynamic_filter = derived.dynamic_filter; + if aggr_expr_replaced { + check_schema_compatible( + &derived.schema, + &input, + &group_by, + &aggr_expr, + mode, + )?; + dynamic_filter = dynamic_filter.as_ref().and_then(|existing| { + rederive_dynamic_filter(existing, &aggr_expr) + }); + } + AggregateExec { + mode, + group_by, + aggr_expr, + filter_expr, + input, + schema: derived.schema, + input_schema, + metrics: ExecutionPlanMetricsSet::new(), + required_input_ordering: derived.required_input_ordering, + input_order_mode: derived.input_order_mode, + cache: derived.cache, + limit_options: None, + dynamic_filter, + } + } + _ => build_from_scratch( + mode, + group_by, + &aggr_expr, + filter_expr, + input, + input_schema, + output_schema, + )?, + }; + + exec.limit_options = limit_options; + validate_limit_options(&exec)?; + Ok(exec) + } +} + +/// Compute every derived part of an [`AggregateExec`] from its inputs: the +/// output schema (unless `output_schema` supplies one), the ordering the input +/// must have, how the input is ordered relative to the group by, the plan +/// properties, and the dynamic filter. +fn build_from_scratch( + mode: AggregateMode, + group_by: Arc<PhysicalGroupBy>, + aggr_expr: &[Arc<AggregateFunctionExpr>], + filter_expr: FilterExprs, + input: Arc<dyn ExecutionPlan>, + input_schema: SchemaRef, + output_schema: Option<SchemaRef>, +) -> Result<AggregateExec> { + // `get_finer_aggregate_exprs_requirement` may rewrite the aggregate + // expressions (e.g. reverse them), so it needs them owned. + let mut aggr_expr = aggr_expr.to_vec(); + + // The output schema is computed from the aggregate expressions *as given*, + // before the requirement analysis below may rewrite them: output field + // names come from those expressions and must not change as a side effect + // of a rewrite. This is why an explicitly supplied schema exists at all. + let schema = match output_schema { + Some(schema) => schema, + None => Arc::new(create_schema(&input.schema(), &group_by, &aggr_expr, mode)?), + }; + + let input_eq_properties = input.equivalence_properties(); + // Get GROUP BY expressions: + let groupby_exprs = group_by.input_exprs(); + // If existing ordering satisfies a prefix of the GROUP BY expressions, + // prefix requirements with this section. In this case, aggregation will + // work more efficiently. + // Copy the `PhysicalSortExpr`s to retain the sort options. + let (new_sort_exprs, indices) = + input_eq_properties.find_longest_permutation(&groupby_exprs)?; + + let mut new_requirements = new_sort_exprs + .into_iter() + .map(PhysicalSortRequirement::from) + .collect::<Vec<_>>(); + + let req = get_finer_aggregate_exprs_requirement( + &mut aggr_expr, + &group_by, + input_eq_properties, + &mode, + )?; + new_requirements.extend(req); + + let required_input_ordering = + LexRequirement::new(new_requirements).map(OrderingRequirements::new_soft); + + // Constant expressions never change, so they cannot mark a completed group. + // Exclude them from both the ordering indices and the group expression count. + // If our aggregation has grouping sets then our base grouping exprs will + // be expanded based on the flags in `group_by.groups` where for each + // group we swap the grouping expr for `null` if the flag is `true` + // That means that each index in `indices` is valid if and only if + // it is not null in every group + let indices: Vec<usize> = indices + .into_iter() + .filter(|idx| group_by.groups.iter().all(|group| !group[*idx])) + .filter(|idx| { + input_eq_properties + .is_expr_constant(&groupby_exprs[*idx]) + .is_none() + }) + .collect(); + + let num_non_constant_groupby_exprs = groupby_exprs + .iter() + .filter(|expr| input_eq_properties.is_expr_constant(expr).is_none()) + .count(); + let mut input_order_mode = if indices.len() == num_non_constant_groupby_exprs + && !indices.is_empty() + && group_by.groups.len() == 1 + { + InputOrderMode::Sorted + } else if !indices.is_empty() { + InputOrderMode::PartiallySorted(indices) + } else { + InputOrderMode::Linear + }; + + // Input order mode is also used to advertise plan output ordering, grouping + // sets handling, and partial reduce aggregation can't promise that. + if group_by.has_grouping_set() || mode == AggregateMode::PartialReduce { + input_order_mode = InputOrderMode::Linear; + } + + // construct a map from the input expression to the output expression of the Aggregation group by + let group_expr_mapping = + ProjectionMapping::try_new(group_by.expr.clone(), &input.schema())?; + + let cache = if group_by.has_grouping_set() { + AggregateExec::compute_grouping_set_properties(&input, Arc::clone(&schema)) + } else { + AggregateExec::compute_properties( + &input, + Arc::clone(&schema), + &group_expr_mapping, + group_by.is_true_no_grouping(), + &mode, + &input_order_mode, + aggr_expr.as_ref(), + )? + }; + + let mut exec = AggregateExec { + mode, + group_by, + aggr_expr: aggr_expr.into(), + filter_expr, + input, + schema, + input_schema, + metrics: ExecutionPlanMetricsSet::new(), + required_input_ordering, + limit_options: None, + input_order_mode, + cache: Arc::new(cache), + dynamic_filter: None, + }; + + exec.init_dynamic_filter(); + + Ok(exec) +} + +/// Verify that `schema`, carried over from the node a builder was derived from, +/// still describes the output of the (replaced) aggregate expressions. +/// +/// Field *names* are allowed to differ: preserving them is the whole point of +/// carrying the schema over. Anything else means the rewrite produced a +/// different aggregate and the schema must not be reused. +fn check_schema_compatible( + schema: &SchemaRef, + input: &Arc<dyn ExecutionPlan>, + group_by: &PhysicalGroupBy, + aggr_expr: &[Arc<AggregateFunctionExpr>], + mode: AggregateMode, +) -> Result<()> { + let computed = create_schema(&input.schema(), group_by, aggr_expr, mode)?; + + let incompatible = || { + internal_err!( + "New aggregate expressions are not compatible with the output schema of the \ + aggregate they replace.\nExpected: {schema}\nGot: {computed}" + ) + }; + + let field_count_matches = computed.fields().len() == schema.fields().len(); + if !field_count_matches { + return incompatible(); + } + + for (computed_field, existing_field) in computed.fields().iter().zip(schema.fields()) + { + let same_type = computed_field.data_type() == existing_field.data_type(); + let same_nullability = + computed_field.is_nullable() == existing_field.is_nullable(); + if !same_type || !same_nullability { + return incompatible(); + } + } + + Ok(()) +} + +/// Keep `existing`, a dynamic filter carried over from the node a builder was +/// derived from, only while it still describes the (replaced) aggregate +/// expressions. +/// +/// The filter records which aggregate expressions are `MIN`/`MAX`, at which +/// index, and over which column, and the pushed-down predicate is built from +/// that: a `MIN` produces `col < bound`, a `MAX` produces `col > bound`. +/// Keeping the old state across a rewrite that turns a `MIN` into a `MAX` (or +/// that changes the column) would push down the wrong predicate and prune rows +/// the aggregate needs, so in that case the filter is rebuilt from the new +/// expressions. +/// +/// Rebuilding loses the link to whichever child accepted the old filter during +/// pushdown, which costs the optimization but cannot give wrong results: an +/// abandoned filter is never narrowed, so the child keeps reading every row. +/// The common case — a rewrite that only reorders or reverses the aggregate +/// expressions — leaves the state unchanged and keeps the original filter. +fn rederive_dynamic_filter( + existing: &Arc<AggrDynFilter>, + aggr_expr: &[Arc<AggregateFunctionExpr>], +) -> Option<Arc<AggrDynFilter>> { + // The new expressions support no dynamic filter at all: drop it rather than + // leave one that describes expressions this node no longer has. + let rederived = AggregateExec::derive_dynamic_filter(aggr_expr)?; + if describes_same_aggregates(existing, &rederived) { + Some(Arc::clone(existing)) + } else { + Some(rederived) + } +} + +/// Whether two dynamic filter states describe the same aggregate expressions: +/// the same `MIN`/`MAX` kinds at the same indices, over the same columns. +fn describes_same_aggregates(a: &AggrDynFilter, b: &AggrDynFilter) -> bool { + let same_accumulators = a.accumulator_dyn_filter_info.len() + == b.accumulator_dyn_filter_info.len() + && a.accumulator_dyn_filter_info + .iter() + .zip(b.accumulator_dyn_filter_info.iter()) + .all(|(a, b)| a.aggr_type == b.aggr_type && a.aggr_index == b.aggr_index); + + let same_columns = a.filter.children().len() == b.filter.children().len() + && a.filter + .children() + .iter() + .zip(b.filter.children()) + .all(|(a, b)| a.eq(&b)); + + same_accumulators && same_columns +} + +/// Reject [`LimitOptions`] that `exec` cannot execute. +/// +/// A limit is only pushed into an aggregate by the optimizer, but nothing stops +/// another rule (or an external consumer) from copying one onto an aggregate +/// with a different shape. Without this check such a plan builds fine and then +/// panics, errors, or silently drops `FILTER` clauses when it is executed, so +/// the conditions the limited execution paths rely on are checked here, once. +fn validate_limit_options(exec: &AggregateExec) -> Result<()> { + let Some(limit_options) = exec.limit_options else { + return Ok(()); + }; + + // Aggregates without a group by produce a single row: the limit is a no-op + // and is ignored by `AggregateStream`. + if exec.group_by.is_true_no_grouping() { + return Ok(()); + } + + // A soft limit on a distinct-style aggregate: the hash streams stop + // accumulating new groups once they have enough, no further requirements. + if exec.is_unordered_unfiltered_group_by_distinct() { + return Ok(()); + } + + // Without a single MIN/MAX aggregate there is nothing for a top-k queue to + // order by, so the limit has to be executed by the regular grouped streams. + // Those honor it by stopping as soon as they hold enough groups, which is + // only correct when the groups themselves are the answer: with aggregate + // expressions, later batches falling into groups already accumulated are + // never read, so a `COUNT` comes back short. A `SELECT DISTINCT`-style + // aggregate has no aggregate expressions and is accepted above. + let minmax_desc = exec.get_minmax_desc(); Review Comment: A `MIN`/`MAX` aggregate is accepted even when the explicit limit direction is the opposite of `get_minmax_desc()`. `GroupedTopKAggregateStream::new` then ignores `LimitOptions::descending` for this case and uses the aggregate's direction, so it retains the wrong K groups (for example, `MIN(...)` with a descending limit retains the smallest minima). Reject an explicit mismatched direction here, as `TopKAggregation::transform_agg` already does. -- 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]
