adriangb commented on code in PR #25376:
URL: https://github.com/apache/datafusion/pull/25376#discussion_r4026414992
##########
datafusion/physical-optimizer/src/update_aggr_exprs.rs:
##########
@@ -89,43 +89,29 @@ impl PhysicalOptimizerRule for OptimizeAggregateOrder {
let input = aggr_exec.input();
let mut aggr_exprs = aggr_exec.aggr_expr().to_vec();
+ let groupby_exprs = aggr_exec.group_expr().input_exprs();
// If the existing ordering satisfies a prefix of the GROUP BY
// expressions, prefix requirements with this section. In this
// case, aggregation will work more efficiently.
- //
- // Only do this for a single grouping set. With grouping sets
- // (e.g. ROLLUP), the stream feeds the same rows once per
- // grouping set, and within a coarser set's group the rows are
- // ordered by the *full* group-by prefix, not by the
- // aggregate's own ORDER BY. Proving the prefixed requirement
- // there would mark first/last aggregates as pre-ordered when
- // their groups are not, producing wrong results in positional
- // paths. Without the prefix, a globally satisfied ORDER BY
- // still holds within every group of every grouping set.
- let requirement = if aggr_exec.group_expr().is_single() {
- let groupby_exprs = aggr_exec.group_expr().input_exprs();
- let indices =
- get_ordered_partition_by_indices(&groupby_exprs,
input)?;
- indices
- .iter()
- .map(|&idx| {
- PhysicalSortRequirement::new(
- Arc::clone(&groupby_exprs[idx]),
- None,
- )
- })
- .collect::<Vec<_>>()
- } else {
- vec![]
- };
+ let indices = get_ordered_partition_by_indices(&groupby_exprs,
input)?;
+ let requirement = indices
+ .iter()
+ .map(|&idx| {
+ PhysicalSortRequirement::new(
+ Arc::clone(&groupby_exprs[idx]),
+ None,
+ )
+ })
+ .collect::<Vec<_>>();
aggr_exprs = try_convert_aggregate_if_better(
aggr_exprs,
&requirement,
input.equivalence_properties(),
)?;
- let aggr_exec = aggr_exec.with_new_aggr_exprs(aggr_exprs);
+ let aggr_exec =
+
aggr_exec.to_builder().with_aggr_exprs(aggr_exprs).build()?;
Review Comment:
Please lets check when `to_builder()` is used (and discriminate tests vs.
production code) and if it should be `into_builder()` and take ownership
instead (or maybe we need both). It'd be nice to avoid clones if we can.
##########
docs/source/library-user-guide/upgrading/56.0.0.md:
##########
@@ -262,98 +97,65 @@ The output type of the `floor` and `ceil` UDFs has been
changed from the exact i
Change the expected type or wrap the expression in `CAST`. It's recommended to
avoid relying on decimal's exact precision and scale.
-### `map_extract` / `element_at` return an empty list for absent keys
+[#24703]: https://github.com/apache/datafusion/pull/24703
-`map_extract` (and its alias `element_at`) previously returned a single-element
-list containing `NULL` when the key was not present in the map. It now returns
-an empty list, matching the documented behavior and DuckDB. Two related cases
-changed at the same time, also matching DuckDB:
+### `AggregateExec` is built and rewritten with `AggregateExecBuilder`
-- A `NULL` map input now yields `NULL` instead of `[NULL]`.
-- A `NULL` lookup key now yields `[]` instead of `[NULL]`.
+`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.
-A key that is present with a `NULL` value still returns `[NULL]`, so absent
-keys and `NULL` values are now distinguishable.
+The new `AggregateExecBuilder` names every argument, defaults the `FILTER`
+expressions to "no filter", and validates the node once in `build()`:
-**Migration guide:**
-
-```sql
--- Before
-SELECT map_extract(MAP {'a': 1}, 'missing'); -- [NULL]
-
--- After
-SELECT map_extract(MAP {'a': 1}, 'missing'); -- []
+```rust
+# /* comment to avoid running as a doctest
+let exec = AggregateExec::builder(AggregateMode::Single, input)
Review Comment:
Use the `rust,ignore` directive instead of a comment
##########
datafusion/physical-plan/src/aggregates/builder.rs:
##########
@@ -0,0 +1,880 @@
+// 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, topk_types_supported,
+};
+use crate::metrics::ExecutionPlanMetricsSet;
+use crate::{ExecutionPlan, 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_common::physical_expr::PhysicalExpr;
+use datafusion_physical_expr_common::sort_expr::OrderingRequirements;
+
+/// Builder for [`AggregateExec`].
+///
+/// This is the recommended way to create an [`AggregateExec`], and the only
+/// supported way to derive a new [`AggregateExec`] from an existing one (see
+/// [`AggregateExec::to_builder`]).
+///
+/// Like the methods it replaces, this is public for internal use only and is
+/// not part of the public API: it is how DataFusion's own physical optimizer
+/// rules build and rewrite aggregates, and it may change without notice. It is
+/// `#[doc(hidden)]` for that reason, not because it is unfinished.
+///
+/// Compared to calling [`AggregateExec::try_new`] and then mutating individual
+/// fields, the builder:
+///
+/// 1. Names every argument, so `input` / `input_schema` and `aggr_expr` /
+/// `filter_expr` can't be transposed by accident.
+/// 2. Defaults `filter_expr` to "no filter for each aggregate", which is what
+/// the vast majority of callers want and removes a common source of
+/// length-mismatch panics.
+/// 3. Validates the plan once, at the end, so combinations that would panic or
+/// return an internal error during execution (for example a limit pushed
+/// into an aggregate that cannot execute it) are rejected up front.
+/// 4. Keeps derived state (output schema, plan properties, ordering
+/// requirements, dynamic filter) consistent when rewriting an existing
+/// node, instead of asking every caller to copy the fields by hand.
+///
+/// # Example: creating a new aggregate
Review Comment:
This is way too much docstring
##########
datafusion/physical-plan/src/aggregates/builder.rs:
##########
@@ -0,0 +1,880 @@
+// 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, topk_types_supported,
+};
+use crate::metrics::ExecutionPlanMetricsSet;
+use crate::{ExecutionPlan, 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_common::physical_expr::PhysicalExpr;
+use datafusion_physical_expr_common::sort_expr::OrderingRequirements;
+
+/// Builder for [`AggregateExec`].
+///
+/// This is the recommended way to create an [`AggregateExec`], and the only
+/// supported way to derive a new [`AggregateExec`] from an existing one (see
+/// [`AggregateExec::to_builder`]).
+///
+/// Like the methods it replaces, this is public for internal use only and is
+/// not part of the public API: it is how DataFusion's own physical optimizer
+/// rules build and rewrite aggregates, and it may change without notice. It is
+/// `#[doc(hidden)]` for that reason, not because it is unfinished.
+///
+/// Compared to calling [`AggregateExec::try_new`] and then mutating individual
+/// fields, the builder:
+///
+/// 1. Names every argument, so `input` / `input_schema` and `aggr_expr` /
+/// `filter_expr` can't be transposed by accident.
+/// 2. Defaults `filter_expr` to "no filter for each aggregate", which is what
+/// the vast majority of callers want and removes a common source of
+/// length-mismatch panics.
+/// 3. Validates the plan once, at the end, so combinations that would panic or
+/// return an internal error during execution (for example a limit pushed
+/// into an aggregate that cannot execute it) are rejected up front.
+/// 4. Keeps derived state (output schema, plan properties, ordering
+/// requirements, dynamic filter) consistent when rewriting an existing
+/// node, instead of asking every caller to copy the fields by hand.
+///
+/// # Example: creating a new aggregate
+/// ```
+/// # use std::sync::Arc;
+/// # use arrow::datatypes::{DataType, Field, Schema};
+/// # use datafusion_physical_plan::aggregates::{
+/// # AggregateExec, AggregateMode, PhysicalGroupBy,
+/// # };
+/// # use datafusion_physical_plan::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()?;
+/// assert_eq!(exec.mode(), &AggregateMode::Single);
+/// # Ok(())
+/// # }
+/// ```
+///
+/// # Example: rewriting an existing aggregate
+/// ```
+/// # 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 an existing `SELECT DISTINCT a`-style aggregate
+/// let limited = exec
+/// .to_builder()
+/// .with_limit_options(LimitOptions::new(10))
+/// .build()?;
+/// assert_eq!(limited.limit_options().map(|o| o.limit()), Some(10));
+/// // the rewritten node keeps the original output schema
+/// assert_eq!(limited.schema(), exec.schema());
+/// # Ok(())
+/// # }
+/// ```
+#[doc(hidden)]
+#[derive(Debug, Clone)]
+pub struct AggregateExecBuilder {
+ mode: AggregateMode,
+ group_by: Arc<PhysicalGroupBy>,
+ aggr_expr: Vec<Arc<AggregateFunctionExpr>>,
+ /// `None` means "no filter for any aggregate expression"
+ filter_expr: Option<Vec<Option<Arc<dyn PhysicalExpr>>>>,
+ 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. Always honored, see
+ /// [`AggregateExecBuilder::with_output_schema`].
+ output_schema: Option<SchemaRef>,
+ /// State carried over from the [`AggregateExec`] this builder was derived
+ /// from. Dropped as soon as a field it depends on changes, see
+ /// [`AggregateExecBuilder::invalidate_derived`].
+ derived: Option<DerivedState>,
+ /// Whether the aggregate expressions were replaced. Only used to decide
+ /// whether a derived output schema still has to be checked.
+ 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 the inputs it
+/// depends on.
+#[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: vec![],
+ 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`.
+ ///
+ /// The derived state of `exec` (output schema, plan properties, ordering
+ /// requirements and dynamic filter) is reused unless a field it depends on
+ /// is replaced. Execution metrics are always reset, since `build` returns
a
+ /// new plan node.
+ pub(crate) fn from_exec(exec: &AggregateExec) -> Self {
+ Self {
+ mode: exec.mode,
+ group_by: Arc::clone(&exec.group_by),
+ aggr_expr: exec.aggr_expr.to_vec(),
+ filter_expr: Some(exec.filter_expr.to_vec()),
+ 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 {
+ 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 {
+ self.group_by = group_by.into();
+ self.invalidate_derived()
+ }
+
+ /// Set the aggregate expressions.
+ ///
+ /// When the builder was created with [`AggregateExec::to_builder`], the
+ /// output schema of the original node is kept, so that rewriting the
+ /// aggregate expressions (for example reversing them in
+ /// `OptimizeAggregateOrder`) cannot change output field names. `build`
+ /// verifies that the new expressions still produce a compatible schema.
+ pub fn with_aggr_exprs(
+ mut self,
+ aggr_expr: impl IntoIterator<Item = Arc<AggregateFunctionExpr>>,
+ ) -> Self {
+ self.aggr_expr = aggr_expr.into_iter().collect();
+ 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 IntoIterator<Item = Option<Arc<dyn PhysicalExpr>>>,
+ ) -> Self {
+ self.filter_expr = Some(filter_expr.into_iter().collect());
+ self.invalidate_derived()
+ }
+
+ /// Set the input plan.
+ pub fn with_input(mut self, input: Arc<dyn ExecutionPlan>) -> 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.
+ ///
+ /// Used when decoding a serialized plan, where the output schema is part
of
+ /// the message and must be preserved exactly. The caller is responsible
for
+ /// the schema matching the aggregate; prefer letting `build` compute it.
+ #[cfg_attr(not(feature = "proto"), allow(dead_code))]
+ 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());
+
+ let mut exec = match (output_schema, derived) {
+ // An explicitly supplied output schema is honored as is,
everything
+ // else is computed from the inputs.
+ (Some(schema), _) => AggregateExec::try_new_with_schema(
+ mode,
+ group_by,
+ aggr_expr,
+ filter_expr,
+ input,
+ input_schema,
+ schema,
+ )?,
+ // Nothing the derived state depends on changed: clone the node
this
+ // builder came from with the new values instead of recomputing. In
+ // particular its output schema is kept, so a rewrite of the
+ // aggregate expressions cannot rename output fields.
+ (None, Some(derived)) => {
+ 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 dynamic_filter = derived.dynamic_filter;
+ if aggr_expr_replaced {
+ check_schema_compatible(
+ &derived.schema,
+ &input,
+ &group_by,
+ &aggr_expr,
+ mode,
+ )?;
+ dynamic_filter = rederive_dynamic_filter(dynamic_filter,
&aggr_expr);
+ }
+ AggregateExec {
+ mode,
+ group_by,
+ aggr_expr: aggr_expr.into(),
+ filter_expr: filter_expr.into(),
+ 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,
+ }
+ }
+ (None, None) => AggregateExec::try_new(
+ mode,
+ group_by,
+ aggr_expr,
+ filter_expr,
+ input,
+ input_schema,
+ )?,
+ };
+
+ exec.limit_options = limit_options;
+ validate_limit_options(&exec)?;
+ 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 compatible =
+ computed.fields().len() == schema.fields().len()
+ && computed.fields().iter().zip(schema.fields()).all(
+ |(computed, existing)| {
+ computed.data_type() == existing.data_type()
+ && computed.is_nullable() == existing.is_nullable()
+ },
+ );
Review Comment:
This comparison is impossible for a human to reason through. Can we break it
down into steps w/ named variables? E.g.
```rust
let field_count_matches: bool = computed.fields().len() ==
schema.fields().len();
if !field_count_matches { ... }
...
```
##########
datafusion/physical-plan/src/aggregates/mod.rs:
##########
@@ -905,9 +902,68 @@ pub struct AggregateExec {
}
impl AggregateExec {
+ /// Create a builder for a new [`AggregateExec`] over `input`.
+ ///
+ /// See [`AggregateExecBuilder`] for details and examples.
+ ///
+ /// This is public for internal use only and is not part of the public API:
+ /// it is how DataFusion's own physical optimizer rules build and rewrite
+ /// aggregates, and it may change without notice.
+ #[doc(hidden)]
+ pub fn builder(
+ mode: AggregateMode,
+ input: Arc<dyn ExecutionPlan>,
+ ) -> AggregateExecBuilder {
+ AggregateExecBuilder::new(mode, input)
+ }
+
+ /// Create a builder pre-populated with the fields of this
+ /// [`AggregateExec`], to derive a new node from it.
+ ///
+ /// This is the supported way to rewrite an existing aggregate: the derived
+ /// output schema and plan properties are carried over (so a rewrite cannot
+ /// rename output fields), and the result is validated.
+ ///
+ /// This is public for internal use only and is not part of the public API:
+ /// it 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};
Review Comment:
Also way too much docstring for an internal / non public method.
##########
datafusion/physical-plan/src/aggregates/mod.rs:
##########
Review Comment:
I wonder if we could minimize code duplication by having
`try_new_with_schema` / other existing constructors delegate to the builder.
That way there is only one code block that builds and verifies.
##########
datafusion/physical-plan/src/aggregates/builder.rs:
##########
@@ -0,0 +1,880 @@
+// 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, topk_types_supported,
+};
+use crate::metrics::ExecutionPlanMetricsSet;
+use crate::{ExecutionPlan, 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_common::physical_expr::PhysicalExpr;
+use datafusion_physical_expr_common::sort_expr::OrderingRequirements;
+
+/// Builder for [`AggregateExec`].
+///
+/// This is the recommended way to create an [`AggregateExec`], and the only
+/// supported way to derive a new [`AggregateExec`] from an existing one (see
+/// [`AggregateExec::to_builder`]).
+///
+/// Like the methods it replaces, this is public for internal use only and is
+/// not part of the public API: it is how DataFusion's own physical optimizer
+/// rules build and rewrite aggregates, and it may change without notice. It is
+/// `#[doc(hidden)]` for that reason, not because it is unfinished.
+///
+/// Compared to calling [`AggregateExec::try_new`] and then mutating individual
+/// fields, the builder:
+///
+/// 1. Names every argument, so `input` / `input_schema` and `aggr_expr` /
+/// `filter_expr` can't be transposed by accident.
+/// 2. Defaults `filter_expr` to "no filter for each aggregate", which is what
+/// the vast majority of callers want and removes a common source of
+/// length-mismatch panics.
+/// 3. Validates the plan once, at the end, so combinations that would panic or
+/// return an internal error during execution (for example a limit pushed
+/// into an aggregate that cannot execute it) are rejected up front.
+/// 4. Keeps derived state (output schema, plan properties, ordering
+/// requirements, dynamic filter) consistent when rewriting an existing
+/// node, instead of asking every caller to copy the fields by hand.
+///
+/// # Example: creating a new aggregate
+/// ```
+/// # use std::sync::Arc;
+/// # use arrow::datatypes::{DataType, Field, Schema};
+/// # use datafusion_physical_plan::aggregates::{
+/// # AggregateExec, AggregateMode, PhysicalGroupBy,
+/// # };
+/// # use datafusion_physical_plan::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()?;
+/// assert_eq!(exec.mode(), &AggregateMode::Single);
+/// # Ok(())
+/// # }
+/// ```
+///
+/// # Example: rewriting an existing aggregate
+/// ```
+/// # 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 an existing `SELECT DISTINCT a`-style aggregate
+/// let limited = exec
+/// .to_builder()
+/// .with_limit_options(LimitOptions::new(10))
+/// .build()?;
+/// assert_eq!(limited.limit_options().map(|o| o.limit()), Some(10));
+/// // the rewritten node keeps the original output schema
+/// assert_eq!(limited.schema(), exec.schema());
+/// # Ok(())
+/// # }
+/// ```
+#[doc(hidden)]
+#[derive(Debug, Clone)]
+pub struct AggregateExecBuilder {
+ mode: AggregateMode,
+ group_by: Arc<PhysicalGroupBy>,
+ aggr_expr: Vec<Arc<AggregateFunctionExpr>>,
+ /// `None` means "no filter for any aggregate expression"
+ filter_expr: Option<Vec<Option<Arc<dyn PhysicalExpr>>>>,
+ 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. Always honored, see
+ /// [`AggregateExecBuilder::with_output_schema`].
+ output_schema: Option<SchemaRef>,
+ /// State carried over from the [`AggregateExec`] this builder was derived
+ /// from. Dropped as soon as a field it depends on changes, see
+ /// [`AggregateExecBuilder::invalidate_derived`].
+ derived: Option<DerivedState>,
+ /// Whether the aggregate expressions were replaced. Only used to decide
+ /// whether a derived output schema still has to be checked.
+ 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 the inputs it
+/// depends on.
+#[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: vec![],
+ 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`.
+ ///
+ /// The derived state of `exec` (output schema, plan properties, ordering
+ /// requirements and dynamic filter) is reused unless a field it depends on
+ /// is replaced. Execution metrics are always reset, since `build` returns
a
+ /// new plan node.
+ pub(crate) fn from_exec(exec: &AggregateExec) -> Self {
+ Self {
+ mode: exec.mode,
+ group_by: Arc::clone(&exec.group_by),
+ aggr_expr: exec.aggr_expr.to_vec(),
+ filter_expr: Some(exec.filter_expr.to_vec()),
+ 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 {
+ self.mode = mode;
+ self.invalidate_derived()
Review Comment:
Should these methods check `if mode = self.mode { return self }` or
something to avoid invalidating the derived state if not needed?
##########
datafusion/physical-plan/src/aggregates/builder.rs:
##########
@@ -0,0 +1,880 @@
+// 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, topk_types_supported,
+};
+use crate::metrics::ExecutionPlanMetricsSet;
+use crate::{ExecutionPlan, 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_common::physical_expr::PhysicalExpr;
+use datafusion_physical_expr_common::sort_expr::OrderingRequirements;
+
+/// Builder for [`AggregateExec`].
+///
+/// This is the recommended way to create an [`AggregateExec`], and the only
+/// supported way to derive a new [`AggregateExec`] from an existing one (see
+/// [`AggregateExec::to_builder`]).
+///
+/// Like the methods it replaces, this is public for internal use only and is
+/// not part of the public API: it is how DataFusion's own physical optimizer
+/// rules build and rewrite aggregates, and it may change without notice. It is
+/// `#[doc(hidden)]` for that reason, not because it is unfinished.
+///
+/// Compared to calling [`AggregateExec::try_new`] and then mutating individual
+/// fields, the builder:
+///
+/// 1. Names every argument, so `input` / `input_schema` and `aggr_expr` /
+/// `filter_expr` can't be transposed by accident.
+/// 2. Defaults `filter_expr` to "no filter for each aggregate", which is what
+/// the vast majority of callers want and removes a common source of
+/// length-mismatch panics.
+/// 3. Validates the plan once, at the end, so combinations that would panic or
+/// return an internal error during execution (for example a limit pushed
+/// into an aggregate that cannot execute it) are rejected up front.
+/// 4. Keeps derived state (output schema, plan properties, ordering
+/// requirements, dynamic filter) consistent when rewriting an existing
+/// node, instead of asking every caller to copy the fields by hand.
+///
+/// # Example: creating a new aggregate
+/// ```
+/// # use std::sync::Arc;
+/// # use arrow::datatypes::{DataType, Field, Schema};
+/// # use datafusion_physical_plan::aggregates::{
+/// # AggregateExec, AggregateMode, PhysicalGroupBy,
+/// # };
+/// # use datafusion_physical_plan::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()?;
+/// assert_eq!(exec.mode(), &AggregateMode::Single);
+/// # Ok(())
+/// # }
+/// ```
+///
+/// # Example: rewriting an existing aggregate
+/// ```
+/// # 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 an existing `SELECT DISTINCT a`-style aggregate
+/// let limited = exec
+/// .to_builder()
+/// .with_limit_options(LimitOptions::new(10))
+/// .build()?;
+/// assert_eq!(limited.limit_options().map(|o| o.limit()), Some(10));
+/// // the rewritten node keeps the original output schema
+/// assert_eq!(limited.schema(), exec.schema());
+/// # Ok(())
+/// # }
+/// ```
+#[doc(hidden)]
+#[derive(Debug, Clone)]
+pub struct AggregateExecBuilder {
+ mode: AggregateMode,
+ group_by: Arc<PhysicalGroupBy>,
+ aggr_expr: Vec<Arc<AggregateFunctionExpr>>,
+ /// `None` means "no filter for any aggregate expression"
+ filter_expr: Option<Vec<Option<Arc<dyn PhysicalExpr>>>>,
+ 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. Always honored, see
+ /// [`AggregateExecBuilder::with_output_schema`].
+ output_schema: Option<SchemaRef>,
+ /// State carried over from the [`AggregateExec`] this builder was derived
+ /// from. Dropped as soon as a field it depends on changes, see
+ /// [`AggregateExecBuilder::invalidate_derived`].
+ derived: Option<DerivedState>,
+ /// Whether the aggregate expressions were replaced. Only used to decide
+ /// whether a derived output schema still has to be checked.
+ 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 the inputs it
+/// depends on.
+#[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: vec![],
+ 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`.
+ ///
+ /// The derived state of `exec` (output schema, plan properties, ordering
+ /// requirements and dynamic filter) is reused unless a field it depends on
+ /// is replaced. Execution metrics are always reset, since `build` returns
a
+ /// new plan node.
+ pub(crate) fn from_exec(exec: &AggregateExec) -> Self {
+ Self {
+ mode: exec.mode,
+ group_by: Arc::clone(&exec.group_by),
+ aggr_expr: exec.aggr_expr.to_vec(),
+ filter_expr: Some(exec.filter_expr.to_vec()),
+ 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 {
+ 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 {
+ self.group_by = group_by.into();
+ self.invalidate_derived()
+ }
+
+ /// Set the aggregate expressions.
+ ///
+ /// When the builder was created with [`AggregateExec::to_builder`], the
+ /// output schema of the original node is kept, so that rewriting the
+ /// aggregate expressions (for example reversing them in
+ /// `OptimizeAggregateOrder`) cannot change output field names. `build`
+ /// verifies that the new expressions still produce a compatible schema.
+ pub fn with_aggr_exprs(
+ mut self,
+ aggr_expr: impl IntoIterator<Item = Arc<AggregateFunctionExpr>>,
+ ) -> Self {
+ self.aggr_expr = aggr_expr.into_iter().collect();
+ 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 IntoIterator<Item = Option<Arc<dyn PhysicalExpr>>>,
+ ) -> Self {
+ self.filter_expr = Some(filter_expr.into_iter().collect());
+ self.invalidate_derived()
+ }
+
+ /// Set the input plan.
+ pub fn with_input(mut self, input: Arc<dyn ExecutionPlan>) -> 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.
+ ///
+ /// Used when decoding a serialized plan, where the output schema is part
of
+ /// the message and must be preserved exactly. The caller is responsible
for
+ /// the schema matching the aggregate; prefer letting `build` compute it.
+ #[cfg_attr(not(feature = "proto"), allow(dead_code))]
+ 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());
+
+ let mut exec = match (output_schema, derived) {
+ // An explicitly supplied output schema is honored as is,
everything
+ // else is computed from the inputs.
+ (Some(schema), _) => AggregateExec::try_new_with_schema(
+ mode,
+ group_by,
+ aggr_expr,
+ filter_expr,
+ input,
+ input_schema,
+ schema,
+ )?,
+ // Nothing the derived state depends on changed: clone the node
this
+ // builder came from with the new values instead of recomputing. In
+ // particular its output schema is kept, so a rewrite of the
+ // aggregate expressions cannot rename output fields.
+ (None, Some(derived)) => {
+ 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 dynamic_filter = derived.dynamic_filter;
+ if aggr_expr_replaced {
+ check_schema_compatible(
+ &derived.schema,
+ &input,
+ &group_by,
+ &aggr_expr,
+ mode,
+ )?;
+ dynamic_filter = rederive_dynamic_filter(dynamic_filter,
&aggr_expr);
+ }
+ AggregateExec {
+ mode,
+ group_by,
+ aggr_expr: aggr_expr.into(),
+ filter_expr: filter_expr.into(),
+ 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,
+ }
+ }
+ (None, None) => AggregateExec::try_new(
+ mode,
+ group_by,
+ aggr_expr,
+ filter_expr,
+ input,
+ input_schema,
+ )?,
+ };
+
+ exec.limit_options = limit_options;
+ validate_limit_options(&exec)?;
+ 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 compatible =
+ computed.fields().len() == schema.fields().len()
+ && computed.fields().iter().zip(schema.fields()).all(
+ |(computed, existing)| {
+ computed.data_type() == existing.data_type()
+ && computed.is_nullable() == existing.is_nullable()
+ },
+ );
+ if !compatible {
+ return internal_err!(
+ "New aggregate expressions are not compatible with the output
schema of the \
+ aggregate they replace.\nExpected: {schema}\nGot: {computed}"
+ );
+ }
+ Ok(())
+}
+
+/// Keep 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`.
+/// Carrying the old state over 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: Option<Arc<AggrDynFilter>>,
+ aggr_expr: &[Arc<AggregateFunctionExpr>],
+) -> Option<Arc<AggrDynFilter>> {
+ let existing = existing?;
Review Comment:
This seems a bit superfluous. Maybe the caller should do this and we accept
`&Arc<AggrDynFilter>`?
--
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]
