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, and fixed in 5832f0f835. `AggregateFunctionExpr`'s `PartialEq` ignores `DISTINCT`, `IGNORE NULLS` and `ORDER BY`, so `with_aggr_exprs` now treats only the same allocations as a no-op. New test `an_equal_aggr_expr_is_still_a_replacement` builds `count(b)`, asserts the `count(DISTINCT b)` replacement compares equal to it, and asserts the built node has the distinct expression. ########## 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, and fixed in 5832f0f835 along the lines you suggest: when the aggregate expressions are replaced the rewrite path now calls `build_from_scratch` and passes the inherited output schema as the only carried-over state, so the required input ordering, the input order mode and the plan properties are all derived from the new expressions. The dynamic filter is handled separately afterwards, kept only while it still describes them. New test `replacing_aggr_exprs_rederives_the_required_input_ordering` replaces `array_agg(b)` with `array_agg(b ORDER BY b)` (same output field, so it gets past the schema check) and asserts the requirement appears. ########## 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: Fixed in fb79b5de94. The note now lists the three shapes that accept a limit separately: no group by (the limit is ignored), a direction-less `SELECT DISTINCT` limit, and a top-k limit with its full set of conditions (nonzero limit, exactly one group by expression and no grouping sets, no `FILTER`, supported key and value types, and a direction that matches the `MIN`/`MAX` when both are present). -- 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]
