adriangb commented on code in PR #25376: URL: https://github.com/apache/datafusion/pull/25376#discussion_r4032095498
########## datafusion/physical-plan/src/aggregates/builder.rs: ########## @@ -0,0 +1,1130 @@ +// 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; + } Review Comment: Confirmed. `AggregateFunctionExpr`'s `PartialEq` ignores `DISTINCT`, `IGNORE NULLS` and `ORDER BY`, so `count(b)` and `count(DISTINCT b)` with the same alias compare equal and the short-circuit dropped the replacement. The PR has since been pared back to a pure refactor, and `with_aggr_exprs` now applies the replacement unconditionally, so the comparison is gone. Recorded in https://github.com/apache/datafusion/issues/25394 so that a future short-circuit compares by allocation instead. ########## datafusion/physical-plan/src/aggregates/builder.rs: ########## @@ -0,0 +1,1130 @@ +// 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 Review Comment: Confirmed, including the fix you suggest: carry over only the inherited output schema and derive the rest from the new expressions, handling the dynamic filter separately. `array_agg(b)` -> `array_agg(b ORDER BY b)` gets past the schema check and loses the requirement, and `SanityCheckPlan` runs after `OptimizeAggregateOrder`, so nothing downstream catches it. This PR has since been pared back to a pure refactor, so it keeps the behaviour `with_new_aggr_exprs` has on `main` rather than fixing it here. Filed as https://github.com/apache/datafusion/issues/25394 with the reproducer. ########## docs/source/library-user-guide/upgrading/56.0.0.md: ########## @@ -394,3 +394,62 @@ let description = ChildFilterDescription::from_child_with_column_mapping( &child, )?; ``` + +### `AggregateExec` is built and rewritten with `AggregateExecBuilder` + +`AggregateExec::try_new` takes six positional arguments, and the fields that +optimizer rules change afterwards (the limit hint, the aggregate expressions) +were set with `with_*` methods that copied the remaining fields by hand. This +made it easy to build an aggregate that only fails when it is executed, for +example by pushing a limit into an aggregate that cannot execute one. + +The new `AggregateExecBuilder` names every argument, defaults the `FILTER` +expressions to "no filter", and validates the node once in `build()`: + +```rust,ignore +let exec = AggregateExec::builder(AggregateMode::Single, input) + .with_group_by(group_by) + .with_aggr_exprs(aggr_exprs) + .with_limit_options(LimitOptions::new(10)) + .build()?; +``` + +An existing node is rewritten with `AggregateExec::to_builder`, which carries +over the output schema and plan properties of the original node so a rewrite +cannot rename output fields: + +```rust,ignore +let with_limit = exec + .to_builder() + .with_limit_options(LimitOptions::new_with_order(10, true)) + .build()?; +``` + +Building and rewriting an `AggregateExec` is how DataFusion's own physical +optimizer rules work, not a public API. The builder, `AggregateExec::builder`, +`AggregateExec::to_builder` and the methods it replaces are therefore all +`#[doc(hidden)]`, and may change without notice. + +**Who is affected:** + +- Callers of `AggregateExec::with_limit_options`, + `AggregateExec::with_new_limit_options` and + `AggregateExec::with_new_aggr_exprs`, which are deprecated and hidden. +- Anyone building an `AggregateExec` with a limit that the aggregate cannot + execute, or decoding such a plan from protobuf: this is now an error at plan + time instead of an internal error, a panic, or silently ignored `FILTER` + expressions at execution time. A limit is only executable on an aggregate + with no aggregate expressions (a `SELECT DISTINCT`-style aggregate), where the + groups themselves are the result, or on one with a single `MIN`/`MAX` + expression and a single group by expression, where the limit either carries no + ordering direction or carries the one that aggregate already implies. Review Comment: Right on both counts. The validation this described has been carved out of the PR entirely, and the upgrade note no longer describes any, so the section is gone. The three accepted shapes and the full set of top-k conditions are written up in https://github.com/apache/datafusion/issues/25393, and will be documented with the check that enforces them. -- 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]
