Copilot commented on code in PR #25376: URL: https://github.com/apache/datafusion/pull/25376#discussion_r4031623348
########## datafusion/physical-plan/src/aggregates/builder.rs: ########## @@ -0,0 +1,1045 @@ +// 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(()); + } + + // Everything below is reached only by `GroupedTopKAggregateStream`, which + // keeps a bounded priority queue of `(group key, min/max value)` pairs. + // + // That queue needs an ordering direction, which comes either from a MIN/MAX + // aggregate or from the limit itself. Without one `execute_typed` runs the + // regular grouped streams, which treat the limit as a soft hint, so there is + // nothing here to check. Keep this condition in step with `execute_typed`: + // rejecting a limit it would have executed turns a working query into a + // planning error. + let minmax_desc = exec.get_minmax_desc(); + if minmax_desc.is_none() && limit_options.descending().is_none() { + return Ok(()); + } Review Comment: A direction-less limit is not safe for aggregates such as `COUNT`/`AVG`. Both hash implementations stop consuming input once the group threshold is reached, so later rows for retained groups are never accumulated and the aggregate values can be wrong. Only the group-by-only case can safely take this early return; non-MIN/MAX aggregate expressions should be rejected even when `descending()` is `None`. This issue also appears on line 835 of the same file. ########## datafusion/proto/tests/cases/plans/aggregates.rs: ########## @@ -205,23 +207,61 @@ fn roundtrip_aggregate_with_limit() -> Result<()> { let groups: Vec<(Arc<dyn PhysicalExpr>, String)> = vec![(col("a", &schema)?, "unused".to_string())]; + // a limit is only valid on an aggregate that can execute it: a single + // MIN/MAX aggregate, or a `SELECT DISTINCT`-style aggregate let aggregates = vec![ - AggregateExprBuilder::new(avg_udaf(), vec![col("b", &schema)?]) + AggregateExprBuilder::new(min_udaf(), vec![col("b", &schema)?]) .schema(Arc::clone(&schema)) - .alias("AVG(b)") + .alias("MIN(b)") .build() .map(Arc::new)?, ]; - let agg = AggregateExec::try_new( + let agg = AggregateExec::builder( AggregateMode::Final, - PhysicalGroupBy::new_single(groups.clone()), - aggregates, - vec![None], Arc::new(EmptyExec::new(schema.clone())), - schema, - )?; - let agg = agg.with_limit_options(Some(LimitOptions::new_with_order(12, false))); + ) + .with_group_by(PhysicalGroupBy::new_single(groups.clone())) + .with_aggr_exprs(aggregates) + .with_input_schema(schema) + .with_limit_options(LimitOptions::new_with_order(12, false)) + .build()?; + roundtrip_test(Arc::new(agg)) +} + +#[test] +fn roundtrip_aggregate_with_soft_limit() -> Result<()> { + let field_a = Field::new("a", DataType::Int64, false); + let field_b = Field::new("b", DataType::Int64, false); + let schema = Arc::new(Schema::new(vec![field_a, field_b])); + + // A direction-less limit is a soft hint that the hash streams honour, so it + // is executable whatever the aggregate looks like. Decoding has to + // reproduce what was encoded rather than re-deciding that: these are more + // group keys than the top-k stream takes, and the aggregate expression + // stops it being a `SELECT DISTINCT`-style one, so a decoder that checked Review Comment: This roundtrip now blesses a plan that can return partial `COUNT` values: the hash aggregate stops reading once its soft group limit is reached, even though later rows may still contribute to retained groups. Keep the soft-limit roundtrip group-by-only (or expect decoding to reject this COUNT shape) rather than asserting that a direction-less limit is valid for every aggregate. -- 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]
