diff --git a/datafusion/core/src/optimizer_rule_reference.md b/datafusion/core/src/optimizer_rule_reference.md index 1367ed0843c59..298a01edd64dc 100644 --- a/datafusion/core/src/optimizer_rule_reference.md +++ b/datafusion/core/src/optimizer_rule_reference.md @@ -56,12 +56,13 @@ Rule order matters. The default pipeline may change between releases. | 17 | `eliminate_outer_join` | Rewrites outer joins to inner joins when later filters reject the NULL-extended rows. | | 18 | `push_down_limit` | Moves literal limits closer to scans and unions and merges adjacent limits. | | 19 | `push_down_filter` | Moves filters as early as possible through filter-commutative operators. | -| 20 | `single_distinct_aggregation_to_group_by` | Rewrites single-column `DISTINCT` aggregations into two-stage `GROUP BY` plans. | -| 21 | `eliminate_group_by_constant` | Removes constant or functionally redundant expressions from `GROUP BY`. | -| 22 | `common_sub_expression_eliminate` | Computes repeated subexpressions once and reuses the result. | -| 23 | `extract_leaf_expressions` | Pulls cheap leaf expressions closer to data sources so later pruning and filter rules can act earlier. | -| 24 | `push_down_leaf_projections` | Pushes the helper projections created by leaf extraction toward leaf inputs. | -| 25 | `optimize_projections` | Prunes unused columns and removes unnecessary logical projections. | +| 20 | `eliminate_aggregate_distinct` | Drops the `DISTINCT` modifier from aggregates whose result cannot change, such as `min`, `max` and `bit_or`. | +| 21 | `single_distinct_aggregation_to_group_by` | Rewrites single-column `DISTINCT` aggregations into two-stage `GROUP BY` plans. | +| 22 | `eliminate_group_by_constant` | Removes constant or functionally redundant expressions from `GROUP BY`. | +| 23 | `common_sub_expression_eliminate` | Computes repeated subexpressions once and reuses the result. | +| 24 | `extract_leaf_expressions` | Pulls cheap leaf expressions closer to data sources so later pruning and filter rules can act earlier. | +| 25 | `push_down_leaf_projections` | Pushes the helper projections created by leaf extraction toward leaf inputs. | +| 26 | `optimize_projections` | Prunes unused columns and removes unnecessary logical projections. | ### Physical Optimizer Rules diff --git a/datafusion/expr/src/lib.rs b/datafusion/expr/src/lib.rs index a904422989942..649f3de2520ac 100644 --- a/datafusion/expr/src/lib.rs +++ b/datafusion/expr/src/lib.rs @@ -137,10 +137,10 @@ pub use partition_evaluator::PartitionEvaluator; pub use sqlparser; pub use table_source::{TableProviderFilterPushDown, TableSource, TableType}; pub use udaf::{ - AggregateUDF, AggregateUDFImpl, ReversedUDAF, SetMonotonicity, StatisticsArgs, - UdafDisplayNameBuilder, UdafHumanDisplayBuilder, UdafSchemaNameBuilder, - UdafWindowFunctionDisplayNameBuilder, UdafWindowFunctionSchemaNameBuilder, - udaf_default_return_field, + AggregateUDF, AggregateUDFImpl, DistinctHandling, ReversedUDAF, SetMonotonicity, + StatisticsArgs, UdafDisplayNameBuilder, UdafHumanDisplayBuilder, + UdafSchemaNameBuilder, UdafWindowFunctionDisplayNameBuilder, + UdafWindowFunctionSchemaNameBuilder, udaf_default_return_field, }; #[expect(deprecated)] pub use udaf::{ diff --git a/datafusion/expr/src/udaf.rs b/datafusion/expr/src/udaf.rs index 458b39c969b6f..d28e697d309c9 100644 --- a/datafusion/expr/src/udaf.rs +++ b/datafusion/expr/src/udaf.rs @@ -360,6 +360,11 @@ impl AggregateUDF { self.inner.supports_within_group_clause() } + /// See [`AggregateUDFImpl::distinct_handling`] for more details. + pub fn distinct_handling(&self) -> DistinctHandling { + self.inner.distinct_handling() + } + /// Returns the documentation for this Aggregate UDF. /// /// Documentation can be accessed programmatically as well as @@ -940,6 +945,20 @@ pub trait AggregateUDFImpl: Debug + DynEq + DynHash + Send + Sync + Any { false } + /// How this function treats the `DISTINCT` modifier. + /// + /// Return [`DistinctHandling::Ignored`] for duplicate-insensitive + /// functions so that `f(DISTINCT x)` is planned as `f(x)`. + /// + /// Return [`DistinctHandling::Unsupported`] if the accumulator does not + /// implement `DISTINCT`, that is, it does not read `is_distinct`, or it + /// rejects `DISTINCT` with an error. The planner then has to deduplicate + /// the input or reject the query. Nothing reads this variant yet: + /// rejecting such queries at planning time is a follow-up change. + fn distinct_handling(&self) -> DistinctHandling { + DistinctHandling::Honored + } + /// Returns the documentation for this Aggregate UDF. /// /// Documentation can be accessed programmatically as well as @@ -1687,6 +1706,10 @@ impl AggregateUDFImpl for AliasedAggregateUDFImpl { self.inner.set_monotonicity(data_type) } + fn distinct_handling(&self) -> DistinctHandling { + self.inner.distinct_handling() + } + fn documentation(&self) -> Option<&Documentation> { self.inner.documentation() } @@ -1713,6 +1736,29 @@ pub enum SetMonotonicity { NotMonotonic, } +/// How an aggregate function treats the `DISTINCT` modifier. +/// +/// Mathematically, `Ignored` means the function's merge operation is +/// idempotent (its state forms a semilattice): f(S ⊎ S) = f(S), so +/// removing duplicates from the input cannot change the result. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[non_exhaustive] +pub enum DistinctHandling { + /// The result is the same with or without `DISTINCT`, so the planner + /// is free to drop it. `min`, `max`, `bool_and`, `bit_or`, ... + Ignored, + /// The accumulator reads `AccumulatorArgs::is_distinct` and deduplicates + /// its input, so the planner must leave the flag alone. `count`, `sum`, + /// `avg`, `var_samp`, `array_agg`, ... This is the default. + Honored, + /// The accumulator does not implement `DISTINCT`: it does not read + /// `is_distinct`, or it rejects `DISTINCT` with an error. The planner has + /// to deduplicate the input first (today `SingleDistinctToGroupBy` does + /// that for single-argument functions) or reject the query. `stddev`, + /// `approx_median`, `corr`, `regr_*`, `nth_value`, ... + Unsupported, +} + #[cfg(test)] mod test { use crate::{AggregateUDF, AggregateUDFImpl}; diff --git a/datafusion/functions-aggregate/src/any_value.rs b/datafusion/functions-aggregate/src/any_value.rs index dc3bd23d806fc..a7bbdc7359c0e 100644 --- a/datafusion/functions-aggregate/src/any_value.rs +++ b/datafusion/functions-aggregate/src/any_value.rs @@ -122,4 +122,9 @@ impl AggregateUDFImpl for AnyValue { fn documentation(&self) -> Option<&Documentation> { self.doc() } + + // TODO: this is arguably `DistinctHandling::Ignored` — the accumulator + // ignores `is_distinct` and returns an unspecified input value either + // way. Grouped with `first_value`/`last_value` and left at the default + // `Honored` until that family is settled together. } diff --git a/datafusion/functions-aggregate/src/approx_distinct.rs b/datafusion/functions-aggregate/src/approx_distinct.rs index 672150ee9b67e..e29f2f3679d4a 100644 --- a/datafusion/functions-aggregate/src/approx_distinct.rs +++ b/datafusion/functions-aggregate/src/approx_distinct.rs @@ -38,6 +38,7 @@ use datafusion_common::{ DataFusionError, Result, downcast_value, internal_datafusion_err, internal_err, not_impl_err, }; +use datafusion_expr::DistinctHandling; use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs}; use datafusion_expr::utils::format_state_name; use datafusion_expr::{ @@ -871,6 +872,11 @@ impl AggregateUDFImpl for ApproxDistinct { fn documentation(&self) -> Option<&Documentation> { self.doc() } + + fn distinct_handling(&self) -> DistinctHandling { + // Updating an HLL register with a value already seen is a no-op. + DistinctHandling::Ignored + } } fn is_fixed_domain_type(data_type: &DataType) -> bool { diff --git a/datafusion/functions-aggregate/src/approx_median.rs b/datafusion/functions-aggregate/src/approx_median.rs index 162dc224f2ccb..ddaa883a3fd68 100644 --- a/datafusion/functions-aggregate/src/approx_median.rs +++ b/datafusion/functions-aggregate/src/approx_median.rs @@ -147,4 +147,10 @@ impl AggregateUDFImpl for ApproxMedian { fn documentation(&self) -> Option<&Documentation> { self.doc() } + + fn distinct_handling(&self) -> datafusion_expr::DistinctHandling { + // The accumulator rejects `DISTINCT` with `not_impl_err!`, so the + // planner has to deduplicate the input first. + datafusion_expr::DistinctHandling::Unsupported + } } diff --git a/datafusion/functions-aggregate/src/approx_percentile_cont.rs b/datafusion/functions-aggregate/src/approx_percentile_cont.rs index 4af3574d8bd74..5c8441aec092b 100644 --- a/datafusion/functions-aggregate/src/approx_percentile_cont.rs +++ b/datafusion/functions-aggregate/src/approx_percentile_cont.rs @@ -31,6 +31,7 @@ use datafusion_common::{ DataFusionError, Result, ScalarValue, downcast_value, internal_err, not_impl_err, plan_err, }; +use datafusion_expr::DistinctHandling; use datafusion_expr::expr::{AggregateFunction, Sort}; use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs}; use datafusion_expr::utils::format_state_name; @@ -324,6 +325,13 @@ impl AggregateUDFImpl for ApproxPercentileCont { fn documentation(&self) -> Option<&Documentation> { self.doc() } + + fn distinct_handling(&self) -> DistinctHandling { + // Duplicate-sensitive, but the accumulator does not read + // `is_distinct` and today silently returns the non-distinct answer. + // The tag records the intent; enforcement is a follow-up change. + DistinctHandling::Unsupported + } } #[derive(Debug)] diff --git a/datafusion/functions-aggregate/src/approx_percentile_cont_with_weight.rs b/datafusion/functions-aggregate/src/approx_percentile_cont_with_weight.rs index 90b24dc3c678d..1a810a3c36685 100644 --- a/datafusion/functions-aggregate/src/approx_percentile_cont_with_weight.rs +++ b/datafusion/functions-aggregate/src/approx_percentile_cont_with_weight.rs @@ -26,6 +26,7 @@ use arrow::{array::ArrayRef, datatypes::DataType}; use datafusion_common::ScalarValue; use datafusion_common::types::{NativeType, logical_float64}; use datafusion_common::{Result, not_impl_err, plan_err}; +use datafusion_expr::DistinctHandling; use datafusion_expr::expr::{AggregateFunction, Sort}; use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs}; use datafusion_expr::{ @@ -281,6 +282,11 @@ impl AggregateUDFImpl for ApproxPercentileContWithWeight { fn documentation(&self) -> Option<&Documentation> { self.doc() } + + fn distinct_handling(&self) -> DistinctHandling { + // The accumulator rejects `DISTINCT` with `not_impl_err!`. + DistinctHandling::Unsupported + } } #[derive(Debug)] diff --git a/datafusion/functions-aggregate/src/bit_and_or_xor.rs b/datafusion/functions-aggregate/src/bit_and_or_xor.rs index 92212b328ee7d..8c6b6fe1b8e38 100644 --- a/datafusion/functions-aggregate/src/bit_and_or_xor.rs +++ b/datafusion/functions-aggregate/src/bit_and_or_xor.rs @@ -31,6 +31,7 @@ use datafusion_common::hash_utils::RandomState; use datafusion_common::cast::as_list_array; use datafusion_common::{Result, ScalarValue, not_impl_err}; +use datafusion_expr::DistinctHandling; use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs}; use datafusion_expr::utils::format_state_name; use datafusion_expr::{ @@ -318,6 +319,18 @@ impl AggregateUDFImpl for BitwiseOperation { fn documentation(&self) -> Option<&Documentation> { Some(self.documentation) } + + fn distinct_handling(&self) -> DistinctHandling { + match self.operation { + // Bitwise AND/OR are idempotent: duplicates cannot change the + // result. Only XOR has a distinct accumulator. + BitwiseOperationType::And | BitwiseOperationType::Or => { + DistinctHandling::Ignored + } + // XOR cancels duplicate pairs, so `DISTINCT` is meaningful. + BitwiseOperationType::Xor => DistinctHandling::Honored, + } + } } struct BitAndAccumulator { diff --git a/datafusion/functions-aggregate/src/bool_and_or.rs b/datafusion/functions-aggregate/src/bool_and_or.rs index 33c9880e2786b..34449d63ef2cb 100644 --- a/datafusion/functions-aggregate/src/bool_and_or.rs +++ b/datafusion/functions-aggregate/src/bool_and_or.rs @@ -29,6 +29,7 @@ use arrow::datatypes::{DataType, FieldRef}; use datafusion_common::internal_err; use datafusion_common::{Result, ScalarValue}; use datafusion_common::{downcast_value, not_impl_err}; +use datafusion_expr::DistinctHandling; use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs}; use datafusion_expr::utils::{AggregateOrderSensitivity, format_state_name}; use datafusion_expr::{ @@ -183,6 +184,11 @@ impl AggregateUDFImpl for BoolAnd { fn documentation(&self) -> Option<&Documentation> { self.doc() } + + fn distinct_handling(&self) -> DistinctHandling { + // Boolean AND/OR are idempotent: duplicates cannot change the result. + DistinctHandling::Ignored + } } #[derive(Debug, Default)] @@ -313,6 +319,11 @@ impl AggregateUDFImpl for BoolOr { fn documentation(&self) -> Option<&Documentation> { self.doc() } + + fn distinct_handling(&self) -> DistinctHandling { + // Boolean AND/OR are idempotent: duplicates cannot change the result. + DistinctHandling::Ignored + } } #[derive(Debug, Default)] diff --git a/datafusion/functions-aggregate/src/correlation.rs b/datafusion/functions-aggregate/src/correlation.rs index 525f6801d2fb7..c3832f8de61cd 100644 --- a/datafusion/functions-aggregate/src/correlation.rs +++ b/datafusion/functions-aggregate/src/correlation.rs @@ -31,6 +31,7 @@ use arrow::{ array::ArrayRef, datatypes::{DataType, Field}, }; +use datafusion_expr::DistinctHandling; use datafusion_expr::{EmitTo, GroupSelection, GroupsAccumulator}; use datafusion_functions_aggregate_common::aggregate::groups_accumulator::accumulate::accumulate_multiple; use log::debug; @@ -145,6 +146,13 @@ impl AggregateUDFImpl for Correlation { debug!("GroupsAccumulator is created for aggregate function `corr(c1, c2)`"); Ok(Box::new(CorrelationGroupsAccumulator::new())) } + + fn distinct_handling(&self) -> DistinctHandling { + // Duplicate-sensitive, but the accumulator does not read + // `is_distinct` and today silently returns the non-distinct answer. + // The tag records the intent; enforcement is a follow-up change. + DistinctHandling::Unsupported + } } /// An accumulator to compute correlation diff --git a/datafusion/functions-aggregate/src/covariance.rs b/datafusion/functions-aggregate/src/covariance.rs index 454d56f8ea577..46abc825175f1 100644 --- a/datafusion/functions-aggregate/src/covariance.rs +++ b/datafusion/functions-aggregate/src/covariance.rs @@ -21,6 +21,7 @@ use arrow::array::ArrayRef; use arrow::datatypes::{DataType, Field, FieldRef}; use datafusion_common::cast::{as_float64_array, as_uint64_array}; use datafusion_common::{Result, ScalarValue}; +use datafusion_expr::DistinctHandling; use datafusion_expr::{ Accumulator, AggregateUDFImpl, Documentation, Signature, Volatility, function::{AccumulatorArgs, StateFieldsArgs}, @@ -128,6 +129,13 @@ impl AggregateUDFImpl for CovarianceSample { fn documentation(&self) -> Option<&Documentation> { self.doc() } + + fn distinct_handling(&self) -> DistinctHandling { + // Duplicate-sensitive, but the accumulator does not read + // `is_distinct` and today silently returns the non-distinct answer. + // The tag records the intent; enforcement is a follow-up change. + DistinctHandling::Unsupported + } } #[user_doc( @@ -206,6 +214,13 @@ impl AggregateUDFImpl for CovariancePopulation { fn documentation(&self) -> Option<&Documentation> { self.doc() } + + fn distinct_handling(&self) -> DistinctHandling { + // Duplicate-sensitive, but the accumulator does not read + // `is_distinct` and today silently returns the non-distinct answer. + // The tag records the intent; enforcement is a follow-up change. + DistinctHandling::Unsupported + } } /// An accumulator to compute covariance diff --git a/datafusion/functions-aggregate/src/first_last.rs b/datafusion/functions-aggregate/src/first_last.rs index f2d6af5bcf4c1..c2992d276f80f 100644 --- a/datafusion/functions-aggregate/src/first_last.rs +++ b/datafusion/functions-aggregate/src/first_last.rs @@ -380,6 +380,12 @@ impl AggregateUDFImpl for FirstValue { fn documentation(&self) -> Option<&Documentation> { self.doc() } + + // TODO: whether this is `DistinctHandling::Ignored` depends on `ORDER BY`. + // `first_value(DISTINCT x ORDER BY y)` deduplicates `x` and leaves the `y` + // ordering meaningless, while `first_value(DISTINCT x ORDER BY x)` is just + // `min(x)`. Left at the default `Honored` until that is settled, even + // though the accumulator ignores `is_distinct` today. } struct FirstLastGroupsAccumulator { @@ -1294,6 +1300,12 @@ impl AggregateUDFImpl for LastValue { ) -> Result> { create_groups_accumulator(&args, false, self.is_input_pre_ordered, self.name()) } + + // TODO: whether this is `DistinctHandling::Ignored` depends on `ORDER BY`. + // `last_value(DISTINCT x ORDER BY y)` deduplicates `x` and leaves the `y` + // ordering meaningless, while `last_value(DISTINCT x ORDER BY x)` is + // `max(x)` when `x` has no NULL. Left at the default `Honored` until that is settled, even + // though the accumulator ignores `is_distinct` today. } /// This accumulator is used when there is no ordering specified for the diff --git a/datafusion/functions-aggregate/src/grouping.rs b/datafusion/functions-aggregate/src/grouping.rs index 720a63aab6884..1c35b2cd04a28 100644 --- a/datafusion/functions-aggregate/src/grouping.rs +++ b/datafusion/functions-aggregate/src/grouping.rs @@ -20,6 +20,7 @@ use arrow::datatypes::Field; use arrow::datatypes::{DataType, FieldRef}; use datafusion_common::{Result, not_impl_err}; +use datafusion_expr::DistinctHandling; use datafusion_expr::function::AccumulatorArgs; use datafusion_expr::function::StateFieldsArgs; use datafusion_expr::utils::format_state_name; @@ -110,4 +111,13 @@ impl AggregateUDFImpl for Grouping { fn documentation(&self) -> Option<&Documentation> { self.doc() } + + fn distinct_handling(&self) -> DistinctHandling { + // The result depends only on which grouping set a row belongs to, not + // on how many rows share a value, so duplicates cannot change it. + // `ResolveGroupingFunction` replaces the call before the optimizer + // runs, so this tag is not reachable from SQL and the accumulator + // above is never built. + DistinctHandling::Ignored + } } diff --git a/datafusion/functions-aggregate/src/min_max.rs b/datafusion/functions-aggregate/src/min_max.rs index 89a1e8114f5e4..5e11981bb8db3 100644 --- a/datafusion/functions-aggregate/src/min_max.rs +++ b/datafusion/functions-aggregate/src/min_max.rs @@ -45,6 +45,7 @@ use arrow::datatypes::{ use crate::min_max::min_max_bytes::MinMaxBytesAccumulator; use crate::min_max::min_max_struct::MinMaxStructAccumulator; use datafusion_common::ScalarValue; +use datafusion_expr::DistinctHandling; use datafusion_expr::{ Accumulator, AggregateUDFImpl, Documentation, SetMonotonicity, Signature, Volatility, function::AccumulatorArgs, @@ -399,6 +400,11 @@ impl AggregateUDFImpl for Max { // the same as new values are seen. SetMonotonicity::Increasing } + + fn distinct_handling(&self) -> DistinctHandling { + // `MAX` is idempotent: duplicates cannot change the maximum. + DistinctHandling::Ignored + } } #[derive(Debug)] @@ -694,6 +700,11 @@ impl AggregateUDFImpl for Min { // the same as new values are seen. SetMonotonicity::Decreasing } + + fn distinct_handling(&self) -> DistinctHandling { + // `MIN` is idempotent: duplicates cannot change the minimum. + DistinctHandling::Ignored + } } #[derive(Debug)] diff --git a/datafusion/functions-aggregate/src/nth_value.rs b/datafusion/functions-aggregate/src/nth_value.rs index 5e7f9c6c3186b..d3571bc925c24 100644 --- a/datafusion/functions-aggregate/src/nth_value.rs +++ b/datafusion/functions-aggregate/src/nth_value.rs @@ -30,6 +30,7 @@ use datafusion_common::utils::{SingleRowListArrayBuilder, get_row_at_idx}; use datafusion_common::{ Result, ScalarValue, assert_or_internal_err, exec_err, not_impl_err, }; +use datafusion_expr::DistinctHandling; use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs}; use datafusion_expr::utils::format_state_name; use datafusion_expr::{ @@ -189,6 +190,13 @@ impl AggregateUDFImpl for NthValueAgg { fn documentation(&self) -> Option<&Documentation> { self.doc() } + + fn distinct_handling(&self) -> DistinctHandling { + // Duplicate-sensitive, but the accumulator does not read + // `is_distinct` and today silently returns the non-distinct answer. + // The tag records the intent; enforcement is a follow-up change. + DistinctHandling::Unsupported + } } #[derive(Debug)] diff --git a/datafusion/functions-aggregate/src/regr.rs b/datafusion/functions-aggregate/src/regr.rs index 5b5a144fd2322..6d679aabea368 100644 --- a/datafusion/functions-aggregate/src/regr.rs +++ b/datafusion/functions-aggregate/src/regr.rs @@ -22,6 +22,7 @@ use arrow::{array::ArrayRef, datatypes::DataType, datatypes::Field}; use datafusion_common::cast::{as_float64_array, as_uint64_array}; use datafusion_common::{HashMap, Result, ScalarValue}; use datafusion_doc::aggregate_doc_sections::DOC_SECTION_STATISTICAL; +use datafusion_expr::DistinctHandling; use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs}; use datafusion_expr::utils::format_state_name; use datafusion_expr::{ @@ -513,6 +514,13 @@ impl AggregateUDFImpl for Regr { fn documentation(&self) -> Option<&Documentation> { self.regr_type.documentation() } + + fn distinct_handling(&self) -> DistinctHandling { + // Duplicate-sensitive, but the accumulator does not read + // `is_distinct` and today silently returns the non-distinct answer. + // The tag records the intent; enforcement is a follow-up change. + DistinctHandling::Unsupported + } } /// `RegrAccumulator` is used to compute linear regression aggregate functions diff --git a/datafusion/functions-aggregate/src/stddev.rs b/datafusion/functions-aggregate/src/stddev.rs index af7e849d9939c..14a28fa809979 100644 --- a/datafusion/functions-aggregate/src/stddev.rs +++ b/datafusion/functions-aggregate/src/stddev.rs @@ -140,6 +140,12 @@ impl AggregateUDFImpl for Stddev { fn documentation(&self) -> Option<&Documentation> { self.doc() } + + fn distinct_handling(&self) -> datafusion_expr::DistinctHandling { + // The accumulator rejects `DISTINCT` with `not_impl_err!`, so the + // planner has to deduplicate the input first. + datafusion_expr::DistinctHandling::Unsupported + } } make_udaf_expr_and_func!( @@ -240,6 +246,12 @@ impl AggregateUDFImpl for StddevPop { fn documentation(&self) -> Option<&Documentation> { self.doc() } + + fn distinct_handling(&self) -> datafusion_expr::DistinctHandling { + // The accumulator rejects `DISTINCT` with `not_impl_err!`, so the + // planner has to deduplicate the input first. + datafusion_expr::DistinctHandling::Unsupported + } } /// An accumulator to compute the average diff --git a/datafusion/functions-aggregate/src/variance.rs b/datafusion/functions-aggregate/src/variance.rs index 072d064f76bd4..3c4003f099c15 100644 --- a/datafusion/functions-aggregate/src/variance.rs +++ b/datafusion/functions-aggregate/src/variance.rs @@ -150,6 +150,9 @@ impl AggregateUDFImpl for VarianceSample { fn documentation(&self) -> Option<&Documentation> { self.doc() } + + // Left at the default `Honored`: `DistinctVarianceAccumulator` + // deduplicates the input when `is_distinct` is set. } #[user_doc( @@ -252,6 +255,9 @@ impl AggregateUDFImpl for VariancePopulation { fn documentation(&self) -> Option<&Documentation> { self.doc() } + + // Left at the default `Honored`: `DistinctVarianceAccumulator` + // deduplicates the input when `is_distinct` is set. } /// An accumulator to compute variance diff --git a/datafusion/optimizer/src/eliminate_aggregate_distinct.rs b/datafusion/optimizer/src/eliminate_aggregate_distinct.rs new file mode 100644 index 0000000000000..2af8bbfa7df12 --- /dev/null +++ b/datafusion/optimizer/src/eliminate_aggregate_distinct.rs @@ -0,0 +1,380 @@ +// 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. + +//! [`EliminateAggregateDistinct`] drops the `DISTINCT` modifier from aggregate +//! functions that report [`DistinctHandling::Ignored`] + +use crate::optimizer::ApplyOrder; +use crate::{OptimizerConfig, OptimizerRule}; + +use datafusion_common::Result; +use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRecursion}; +use datafusion_expr::expr::AggregateFunction; +use datafusion_expr::expr_rewriter::NamePreserver; +use datafusion_expr::{DistinctHandling, Expr, LogicalPlan}; + +/// Optimizer rule that removes a `DISTINCT` modifier that cannot change the +/// result of the aggregate it is attached to. +/// +/// `min`, `max`, `bool_and`, `bit_or` and friends have an idempotent merge, so +/// `min(DISTINCT x)` and `min(x)` return the same value. Removing the flag here +/// keeps [`crate::single_distinct_to_groupby::SingleDistinctToGroupBy`] from +/// rewriting the plan into an inner group by that only exists to deduplicate. +/// +/// An aggregate states how it treats duplicates through +/// [`datafusion_expr::AggregateUDFImpl::distinct_handling`]. +/// +/// ```text +/// Aggregate: groupBy=[[g]], aggr=[[min(DISTINCT x)]] +/// ``` +/// +/// becomes +/// +/// ```text +/// Aggregate: groupBy=[[g]], aggr=[[min(x) AS "min(DISTINCT x)"]] +/// ``` +/// +/// The alias keeps the output schema unchanged so the parent projection still +/// resolves. +#[derive(Default, Debug)] +pub struct EliminateAggregateDistinct {} + +impl EliminateAggregateDistinct { + pub fn new() -> Self { + Self {} + } +} + +impl OptimizerRule for EliminateAggregateDistinct { + fn name(&self) -> &str { + "eliminate_aggregate_distinct" + } + + fn apply_order(&self) -> Option { + Some(ApplyOrder::BottomUp) + } + + fn supports_rewrite(&self) -> bool { + true + } + + fn rewrite( + &self, + plan: LogicalPlan, + _config: &dyn OptimizerConfig, + ) -> Result> { + // Aggregate expressions only appear on Aggregate nodes, so every other + // node is a cheap no-op. Window functions carry their own `distinct` + // flag and are out of scope. + let LogicalPlan::Aggregate(aggregate) = &plan else { + return Ok(Transformed::no(plan)); + }; + if !can_strip_every_distinct(&aggregate.aggr_expr)? { + return Ok(Transformed::no(plan)); + } + + // Dropping `DISTINCT` changes `Expr::schema_name`, and with it the + // output schema of the Aggregate, so restore the original name. The + // aggregate may sit under an alias that type coercion added, so walk + // the expression rather than matching only its root. + let name_preserver = NamePreserver::new(&plan); + plan.map_expressions(|expr| { + let saved_name = name_preserver.save(&expr); + expr.transform_down(strip_ignored_distinct) + .map(|t| t.update_data(|e| saved_name.restore(e))) + }) + } +} + +/// Whether the node has at least one `DISTINCT` and every one is `Ignored`. +/// +/// If only some of them were stripped, an `Honored` `DISTINCT` could stay +/// beside a stripped aggregate. A stripped aggregate carries an alias, and +/// [`crate::single_distinct_to_groupby::SingleDistinctToGroupBy`] does not +/// rewrite a node whose `aggr_expr` contains an alias. So `min(DISTINCT x), +/// count(DISTINCT x)` would lose the rewrite that `count` needs. When every +/// `DISTINCT` is `Ignored`, none is left after stripping, so that rule has +/// nothing to rewrite. This is conservative: `min(DISTINCT x), +/// count(DISTINCT y)` keeps the `min` flag. That flag costs nothing at run +/// time, because `min` ignores `is_distinct` when it selects its accumulator. +fn can_strip_every_distinct(aggr_expr: &[Expr]) -> Result { + let mut found_distinct = false; + let mut all_ignored = true; + for expr in aggr_expr { + expr.apply(|e| { + if let Expr::AggregateFunction(AggregateFunction { func, params }) = e + && params.distinct + { + found_distinct = true; + if func.distinct_handling() != DistinctHandling::Ignored { + all_ignored = false; + return Ok(TreeNodeRecursion::Stop); + } + } + Ok(TreeNodeRecursion::Continue) + })?; + if !all_ignored { + break; + } + } + Ok(found_distinct && all_ignored) +} + +/// Drops `DISTINCT` from `expr` if it is an aggregate that ignores duplicates. +/// +/// The handling is checked again here rather than trusted to +/// [`can_strip_every_distinct`], which only inspects `aggr_expr`, while +/// `map_expressions` also visits the group expressions. +/// +/// An idempotent merge is not always commutative: `first_value` is +/// idempotent but order-sensitive. The rule does not need commutativity. +/// `Ignored` means that the result does not change when duplicates are +/// removed, and stripping `DISTINCT` only stops that removal. `order_by` and +/// `filter` are carried over untouched, so the function sees the same rows in +/// the same order, plus the duplicates that it ignores. +fn strip_ignored_distinct(expr: Expr) -> Result> { + Ok(match expr { + Expr::AggregateFunction(mut agg) + if agg.params.distinct + && agg.func.distinct_handling() == DistinctHandling::Ignored => + { + agg.params.distinct = false; + Transformed::yes(Expr::AggregateFunction(agg)) + } + _ => Transformed::no(expr), + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::OptimizerContext; + use crate::assert_optimized_plan_eq_snapshot; + use crate::test::*; + + use crate::single_distinct_to_groupby::SingleDistinctToGroupBy; + use datafusion_expr::{ExprFunctionExt, LogicalPlanBuilder, col, lit}; + use datafusion_functions_aggregate::expr_fn::{bit_xor, max, min, sum}; + + use std::sync::Arc; + + macro_rules! assert_optimized_plan_equal { + ( + $plan:expr, + @ $expected:literal $(,)? + ) => {{ + let optimizer_ctx = OptimizerContext::new().with_max_passes(1); + let rules: Vec> = + vec![Arc::new(EliminateAggregateDistinct::new())]; + assert_optimized_plan_eq_snapshot!( + optimizer_ctx, + rules, + $plan, + @ $expected, + ) + }}; + } + + /// `min(DISTINCT b)` loses the flag but keeps its column name. + #[test] + fn eliminate_distinct_from_min() -> Result<()> { + let table_scan = test_table_scan()?; + let plan = LogicalPlanBuilder::from(table_scan) + .aggregate(vec![col("a")], vec![min(col("b")).distinct().build()?])? + .build()?; + + assert_optimized_plan_equal!(plan, @r" + Aggregate: groupBy=[[test.a]], aggr=[[min(test.b) AS min(DISTINCT test.b)]] + TableScan: test + ") + } + + /// `sum` deduplicates for real, so the flag stays. + #[test] + fn keep_distinct_on_sum() -> Result<()> { + let table_scan = test_table_scan()?; + let plan = LogicalPlanBuilder::from(table_scan) + .aggregate(vec![col("a")], vec![sum(col("b")).distinct().build()?])? + .build()?; + + assert_optimized_plan_equal!(plan, @r" + Aggregate: groupBy=[[test.a]], aggr=[[sum(DISTINCT test.b)]] + TableScan: test + ") + } + + /// XOR cancels duplicate pairs, unlike its `bit_and`/`bit_or` siblings. + #[test] + fn keep_distinct_on_bit_xor() -> Result<()> { + let table_scan = test_table_scan()?; + let plan = LogicalPlanBuilder::from(table_scan) + .aggregate(vec![col("a")], vec![bit_xor(col("b")).distinct().build()?])? + .build()?; + + assert_optimized_plan_equal!(plan, @r" + Aggregate: groupBy=[[test.a]], aggr=[[bit_xor(DISTINCT test.b)]] + TableScan: test + ") + } + + /// A node keeps every flag when one of them has to stay, so that this rule + /// cannot change which plans `SingleDistinctToGroupBy` rewrites. + #[test] + fn mixed_node_is_left_alone() -> Result<()> { + let table_scan = test_table_scan()?; + let plan = LogicalPlanBuilder::from(table_scan) + .aggregate( + vec![col("a")], + vec![ + min(col("b")).distinct().build()?, + sum(col("c")).distinct().build()?, + ], + )? + .build()?; + + assert_optimized_plan_equal!(plan, @r" + Aggregate: groupBy=[[test.a]], aggr=[[min(DISTINCT test.b), sum(DISTINCT test.c)]] + TableScan: test + ") + } + + /// Several duplicate-insensitive aggregates all lose the flag together. + #[test] + fn eliminate_distinct_from_every_ignored_aggregate() -> Result<()> { + let table_scan = test_table_scan()?; + let plan = LogicalPlanBuilder::from(table_scan) + .aggregate( + vec![col("a")], + vec![ + min(col("b")).distinct().build()?, + max(col("c")).distinct().build()?, + ], + )? + .build()?; + + assert_optimized_plan_equal!(plan, @r" + Aggregate: groupBy=[[test.a]], aggr=[[min(test.b) AS min(DISTINCT test.b), max(test.c) AS max(DISTINCT test.c)]] + TableScan: test + ") + } + + /// A non-distinct aggregate alongside is no obstacle. + #[test] + fn eliminate_distinct_beside_non_distinct_aggregate() -> Result<()> { + let table_scan = test_table_scan()?; + let plan = LogicalPlanBuilder::from(table_scan) + .aggregate( + vec![col("a")], + vec![min(col("b")).distinct().build()?, sum(col("c"))], + )? + .build()?; + + assert_optimized_plan_equal!(plan, @r" + Aggregate: groupBy=[[test.a]], aggr=[[min(test.b) AS min(DISTINCT test.b), sum(test.c)]] + TableScan: test + ") + } + + /// The gate only inspects `aggr_expr`, so a distinct aggregate that honors + /// the flag in the group expressions must keep it. + #[test] + fn keep_honored_distinct_in_group_expr() -> Result<()> { + let table_scan = test_table_scan()?; + let plan = LogicalPlanBuilder::from(table_scan) + .aggregate( + vec![sum(col("c")).distinct().build()?], + vec![min(col("b")).distinct().build()?], + )? + .build()?; + + assert_optimized_plan_equal!(plan, @r" + Aggregate: groupBy=[[sum(DISTINCT test.c)]], aggr=[[min(test.b) AS min(DISTINCT test.b)]] + TableScan: test + ") + } + + /// `FILTER` is applied before deduplication, so it rides along untouched. + #[test] + fn eliminate_distinct_keeps_filter() -> Result<()> { + let table_scan = test_table_scan()?; + let plan = LogicalPlanBuilder::from(table_scan) + .aggregate( + vec![col("a")], + vec![ + min(col("b")) + .distinct() + .filter(col("c").gt(lit(0u32))) + .build()?, + ], + )? + .build()?; + + assert_optimized_plan_equal!(plan, @r" + Aggregate: groupBy=[[test.a]], aggr=[[min(test.b) FILTER (WHERE test.c > UInt32(0)) AS min(DISTINCT test.b) FILTER (WHERE test.c > UInt32(0))]] + TableScan: test + ") + } + + /// The conservative gate exists so `SingleDistinctToGroupBy` still sees the + /// plan exactly as it did before this rule was added. + #[test] + fn mixed_node_still_reaches_single_distinct_to_groupby() -> Result<()> { + let table_scan = test_table_scan()?; + let plan = LogicalPlanBuilder::from(table_scan) + .aggregate( + vec![col("a")], + vec![ + min(col("b")).distinct().build()?, + sum(col("b")).distinct().build()?, + ], + )? + .build()?; + + let optimizer_ctx = OptimizerContext::new().with_max_passes(1); + let rules: Vec> = vec![ + Arc::new(EliminateAggregateDistinct::new()), + Arc::new(SingleDistinctToGroupBy::new()), + ]; + assert_optimized_plan_eq_snapshot!( + optimizer_ctx, + rules, + plan, + @r" + Projection: test.a, min(alias1) AS min(DISTINCT test.b), sum(alias1) AS sum(DISTINCT test.b) + Aggregate: groupBy=[[test.a]], aggr=[[min(alias1), sum(alias1)]] + Aggregate: groupBy=[[test.a, test.b AS alias1]], aggr=[[]] + TableScan: test + ", + ) + } + + /// A plan with no Aggregate takes the no-op path. + #[test] + fn non_aggregate_plan_is_unchanged() -> Result<()> { + let table_scan = test_table_scan()?; + let plan = LogicalPlanBuilder::from(table_scan) + .filter(col("b").gt(lit(1u32)))? + .project(vec![col("a")])? + .build()?; + + assert_optimized_plan_equal!(plan, @r" + Projection: test.a + Filter: test.b > UInt32(1) + TableScan: test + ") + } +} diff --git a/datafusion/optimizer/src/lib.rs b/datafusion/optimizer/src/lib.rs index fbe7ad2f4d327..1f5fa191d60ca 100644 --- a/datafusion/optimizer/src/lib.rs +++ b/datafusion/optimizer/src/lib.rs @@ -43,6 +43,7 @@ pub mod common_subexpr_eliminate; pub mod decorrelate; pub mod decorrelate_lateral_join; pub mod decorrelate_predicate_subquery; +pub mod eliminate_aggregate_distinct; pub mod eliminate_cross_join; pub mod eliminate_duplicated_expr; pub mod eliminate_filter; diff --git a/datafusion/optimizer/src/optimizer.rs b/datafusion/optimizer/src/optimizer.rs index ca4a6688b50c5..49c59014fc110 100644 --- a/datafusion/optimizer/src/optimizer.rs +++ b/datafusion/optimizer/src/optimizer.rs @@ -44,6 +44,7 @@ use datafusion_expr::{ use crate::common_subexpr_eliminate::CommonSubexprEliminate; use crate::decorrelate_lateral_join::DecorrelateLateralJoin; use crate::decorrelate_predicate_subquery::DecorrelatePredicateSubquery; +use crate::eliminate_aggregate_distinct::EliminateAggregateDistinct; use crate::eliminate_cross_join::EliminateCrossJoin; use crate::eliminate_duplicated_expr::EliminateDuplicatedExpr; use crate::eliminate_filter::EliminateFilter; @@ -308,6 +309,7 @@ impl Optimizer { // Filters can't be pushed down past Limits, we should do PushDownFilter after PushDownLimit Arc::new(PushDownLimit::new()), Arc::new(PushDownFilter::new()), + Arc::new(EliminateAggregateDistinct::new()), Arc::new(SingleDistinctToGroupBy::new()), // The previous optimizations added expressions and projections, // that might benefit from the following rules diff --git a/datafusion/sqllogictest/test_files/aggregates_simplify.slt b/datafusion/sqllogictest/test_files/aggregates_simplify.slt index c4055d17396c5..a97135470caad 100644 --- a/datafusion/sqllogictest/test_files/aggregates_simplify.slt +++ b/datafusion/sqllogictest/test_files/aggregates_simplify.slt @@ -356,3 +356,246 @@ DROP TABLE IF EXISTS tbl; statement ok DROP TABLE sum_simplify_t; + +####### +# EliminateAggregateDistinct: DISTINCT is dropped from duplicate-insensitive +# aggregates, so no inner group by is planned to deduplicate the input. +####### + +statement ok +CREATE TABLE distinct_simplify_t (g INT, v INT, b BOOLEAN) AS VALUES + (1, 3, true), + (1, 3, true), + (1, 5, false), + (2, 7, true), + (2, 7, true), + (2, NULL, NULL); + +# min: DISTINCT cannot change the minimum +query II +SELECT g, min(DISTINCT v) FROM distinct_simplify_t GROUP BY g ORDER BY g; +---- +1 3 +2 7 + +query TT +EXPLAIN SELECT g, min(DISTINCT v) FROM distinct_simplify_t GROUP BY g; +---- +logical_plan +01)Aggregate: groupBy=[[distinct_simplify_t.g]], aggr=[[min(distinct_simplify_t.v) AS min(DISTINCT distinct_simplify_t.v)]] +02)--TableScan: distinct_simplify_t projection=[g, v] +physical_plan +01)AggregateExec: mode=FinalPartitioned, gby=[g@0 as g], aggr=[min(distinct_simplify_t.v) as min(DISTINCT distinct_simplify_t.v)] +02)--RepartitionExec: partitioning=Hash([g@0], 4), input_partitions=1 +03)----AggregateExec: mode=Partial, gby=[g@0 as g], aggr=[min(distinct_simplify_t.v) as min(DISTINCT distinct_simplify_t.v)] +04)------DataSourceExec: partitions=1, partition_sizes=[1] + +# max +query II +SELECT g, max(DISTINCT v) FROM distinct_simplify_t GROUP BY g ORDER BY g; +---- +1 5 +2 7 + +query TT +EXPLAIN SELECT g, max(DISTINCT v) FROM distinct_simplify_t GROUP BY g; +---- +logical_plan +01)Aggregate: groupBy=[[distinct_simplify_t.g]], aggr=[[max(distinct_simplify_t.v) AS max(DISTINCT distinct_simplify_t.v)]] +02)--TableScan: distinct_simplify_t projection=[g, v] +physical_plan +01)AggregateExec: mode=FinalPartitioned, gby=[g@0 as g], aggr=[max(distinct_simplify_t.v) as max(DISTINCT distinct_simplify_t.v)] +02)--RepartitionExec: partitioning=Hash([g@0], 4), input_partitions=1 +03)----AggregateExec: mode=Partial, gby=[g@0 as g], aggr=[max(distinct_simplify_t.v) as max(DISTINCT distinct_simplify_t.v)] +04)------DataSourceExec: partitions=1, partition_sizes=[1] + +# bool_and +query IB +SELECT g, bool_and(DISTINCT b) FROM distinct_simplify_t GROUP BY g ORDER BY g; +---- +1 false +2 true + +query TT +EXPLAIN SELECT g, bool_and(DISTINCT b) FROM distinct_simplify_t GROUP BY g; +---- +logical_plan +01)Aggregate: groupBy=[[distinct_simplify_t.g]], aggr=[[bool_and(distinct_simplify_t.b) AS bool_and(DISTINCT distinct_simplify_t.b)]] +02)--TableScan: distinct_simplify_t projection=[g, b] +physical_plan +01)AggregateExec: mode=FinalPartitioned, gby=[g@0 as g], aggr=[bool_and(distinct_simplify_t.b) as bool_and(DISTINCT distinct_simplify_t.b)] +02)--RepartitionExec: partitioning=Hash([g@0], 4), input_partitions=1 +03)----AggregateExec: mode=Partial, gby=[g@0 as g], aggr=[bool_and(distinct_simplify_t.b) as bool_and(DISTINCT distinct_simplify_t.b)] +04)------DataSourceExec: partitions=1, partition_sizes=[1] + +# bool_or +query IB +SELECT g, bool_or(DISTINCT b) FROM distinct_simplify_t GROUP BY g ORDER BY g; +---- +1 true +2 true + +query TT +EXPLAIN SELECT g, bool_or(DISTINCT b) FROM distinct_simplify_t GROUP BY g; +---- +logical_plan +01)Aggregate: groupBy=[[distinct_simplify_t.g]], aggr=[[bool_or(distinct_simplify_t.b) AS bool_or(DISTINCT distinct_simplify_t.b)]] +02)--TableScan: distinct_simplify_t projection=[g, b] +physical_plan +01)AggregateExec: mode=FinalPartitioned, gby=[g@0 as g], aggr=[bool_or(distinct_simplify_t.b) as bool_or(DISTINCT distinct_simplify_t.b)] +02)--RepartitionExec: partitioning=Hash([g@0], 4), input_partitions=1 +03)----AggregateExec: mode=Partial, gby=[g@0 as g], aggr=[bool_or(distinct_simplify_t.b) as bool_or(DISTINCT distinct_simplify_t.b)] +04)------DataSourceExec: partitions=1, partition_sizes=[1] + +# bit_and +query II +SELECT g, bit_and(DISTINCT v) FROM distinct_simplify_t GROUP BY g ORDER BY g; +---- +1 1 +2 7 + +query TT +EXPLAIN SELECT g, bit_and(DISTINCT v) FROM distinct_simplify_t GROUP BY g; +---- +logical_plan +01)Aggregate: groupBy=[[distinct_simplify_t.g]], aggr=[[bit_and(distinct_simplify_t.v) AS bit_and(DISTINCT distinct_simplify_t.v)]] +02)--TableScan: distinct_simplify_t projection=[g, v] +physical_plan +01)AggregateExec: mode=FinalPartitioned, gby=[g@0 as g], aggr=[bit_and(distinct_simplify_t.v) as bit_and(DISTINCT distinct_simplify_t.v)] +02)--RepartitionExec: partitioning=Hash([g@0], 4), input_partitions=1 +03)----AggregateExec: mode=Partial, gby=[g@0 as g], aggr=[bit_and(distinct_simplify_t.v) as bit_and(DISTINCT distinct_simplify_t.v)] +04)------DataSourceExec: partitions=1, partition_sizes=[1] + +# bit_or +query II +SELECT g, bit_or(DISTINCT v) FROM distinct_simplify_t GROUP BY g ORDER BY g; +---- +1 7 +2 7 + +query TT +EXPLAIN SELECT g, bit_or(DISTINCT v) FROM distinct_simplify_t GROUP BY g; +---- +logical_plan +01)Aggregate: groupBy=[[distinct_simplify_t.g]], aggr=[[bit_or(distinct_simplify_t.v) AS bit_or(DISTINCT distinct_simplify_t.v)]] +02)--TableScan: distinct_simplify_t projection=[g, v] +physical_plan +01)AggregateExec: mode=FinalPartitioned, gby=[g@0 as g], aggr=[bit_or(distinct_simplify_t.v) as bit_or(DISTINCT distinct_simplify_t.v)] +02)--RepartitionExec: partitioning=Hash([g@0], 4), input_partitions=1 +03)----AggregateExec: mode=Partial, gby=[g@0 as g], aggr=[bit_or(distinct_simplify_t.v) as bit_or(DISTINCT distinct_simplify_t.v)] +04)------DataSourceExec: partitions=1, partition_sizes=[1] + +# Negative case: bit_xor cancels duplicate pairs, so DISTINCT is kept and +# SingleDistinctToGroupBy still rewrites the plan. +query II +SELECT g, bit_xor(DISTINCT v) FROM distinct_simplify_t GROUP BY g ORDER BY g; +---- +1 6 +2 7 + +query TT +EXPLAIN SELECT g, bit_xor(DISTINCT v) FROM distinct_simplify_t GROUP BY g; +---- +logical_plan +01)Projection: distinct_simplify_t.g, bit_xor(alias1) AS bit_xor(DISTINCT distinct_simplify_t.v) +02)--Aggregate: groupBy=[[distinct_simplify_t.g]], aggr=[[bit_xor(alias1)]] +03)----Aggregate: groupBy=[[distinct_simplify_t.g, distinct_simplify_t.v AS alias1]], aggr=[[]] +04)------TableScan: distinct_simplify_t projection=[g, v] +physical_plan +01)ProjectionExec: expr=[g@0 as g, bit_xor(alias1)@1 as bit_xor(DISTINCT distinct_simplify_t.v)] +02)--AggregateExec: mode=FinalPartitioned, gby=[g@0 as g], aggr=[bit_xor(alias1)] +03)----RepartitionExec: partitioning=Hash([g@0], 4), input_partitions=4 +04)------AggregateExec: mode=Partial, gby=[g@0 as g], aggr=[bit_xor(alias1)] +05)--------AggregateExec: mode=FinalPartitioned, gby=[g@0 as g, alias1@1 as alias1], aggr=[] +06)----------RepartitionExec: partitioning=Hash([g@0, alias1@1], 4), input_partitions=1 +07)------------AggregateExec: mode=Partial, gby=[g@0 as g, v@1 as alias1], aggr=[] +08)--------------DataSourceExec: partitions=1, partition_sizes=[1] + +# Negative case: count deduplicates for real +query II +SELECT g, count(DISTINCT v) FROM distinct_simplify_t GROUP BY g ORDER BY g; +---- +1 2 +2 1 + +query TT +EXPLAIN SELECT g, count(DISTINCT v) FROM distinct_simplify_t GROUP BY g; +---- +logical_plan +01)Projection: distinct_simplify_t.g, count(alias1) AS count(DISTINCT distinct_simplify_t.v) +02)--Aggregate: groupBy=[[distinct_simplify_t.g]], aggr=[[count(alias1)]] +03)----Aggregate: groupBy=[[distinct_simplify_t.g, distinct_simplify_t.v AS alias1]], aggr=[[]] +04)------TableScan: distinct_simplify_t projection=[g, v] +physical_plan +01)ProjectionExec: expr=[g@0 as g, count(alias1)@1 as count(DISTINCT distinct_simplify_t.v)] +02)--AggregateExec: mode=FinalPartitioned, gby=[g@0 as g], aggr=[count(alias1)] +03)----RepartitionExec: partitioning=Hash([g@0], 4), input_partitions=4 +04)------AggregateExec: mode=Partial, gby=[g@0 as g], aggr=[count(alias1)] +05)--------AggregateExec: mode=FinalPartitioned, gby=[g@0 as g, alias1@1 as alias1], aggr=[] +06)----------RepartitionExec: partitioning=Hash([g@0, alias1@1], 4), input_partitions=1 +07)------------AggregateExec: mode=Partial, gby=[g@0 as g, v@1 as alias1], aggr=[] +08)--------------DataSourceExec: partitions=1, partition_sizes=[1] + +# Mixed: a node that still needs one DISTINCT keeps all of them, so this rule +# cannot change which plans SingleDistinctToGroupBy rewrites. `count` needs the +# rewrite, and the plan below shows that it still happens. +query III +SELECT g, min(DISTINCT v), count(DISTINCT v) FROM distinct_simplify_t GROUP BY g ORDER BY g; +---- +1 3 2 +2 7 1 + +query TT +EXPLAIN SELECT g, min(DISTINCT v), count(DISTINCT v) FROM distinct_simplify_t GROUP BY g; +---- +logical_plan +01)Projection: distinct_simplify_t.g, min(alias1) AS min(DISTINCT distinct_simplify_t.v), count(alias1) AS count(DISTINCT distinct_simplify_t.v) +02)--Aggregate: groupBy=[[distinct_simplify_t.g]], aggr=[[min(alias1), count(alias1)]] +03)----Aggregate: groupBy=[[distinct_simplify_t.g, distinct_simplify_t.v AS alias1]], aggr=[[]] +04)------TableScan: distinct_simplify_t projection=[g, v] +physical_plan +01)ProjectionExec: expr=[g@0 as g, min(alias1)@1 as min(DISTINCT distinct_simplify_t.v), count(alias1)@2 as count(DISTINCT distinct_simplify_t.v)] +02)--AggregateExec: mode=FinalPartitioned, gby=[g@0 as g], aggr=[min(alias1), count(alias1)] +03)----RepartitionExec: partitioning=Hash([g@0], 4), input_partitions=4 +04)------AggregateExec: mode=Partial, gby=[g@0 as g], aggr=[min(alias1), count(alias1)] +05)--------AggregateExec: mode=FinalPartitioned, gby=[g@0 as g, alias1@1 as alias1], aggr=[] +06)----------RepartitionExec: partitioning=Hash([g@0, alias1@1], 4), input_partitions=1 +07)------------AggregateExec: mode=Partial, gby=[g@0 as g, v@1 as alias1], aggr=[] +08)--------------DataSourceExec: partitions=1, partition_sizes=[1] + +# FILTER is applied before deduplication, so it survives the rewrite +query II +SELECT g, min(DISTINCT v) FILTER (WHERE v > 3) FROM distinct_simplify_t GROUP BY g ORDER BY g; +---- +1 5 +2 7 + +query TT +EXPLAIN SELECT g, min(DISTINCT v) FILTER (WHERE v > 3) FROM distinct_simplify_t GROUP BY g; +---- +logical_plan +01)Aggregate: groupBy=[[distinct_simplify_t.g]], aggr=[[min(distinct_simplify_t.v) FILTER (WHERE distinct_simplify_t.v > Int32(3)) AS min(DISTINCT distinct_simplify_t.v) FILTER (WHERE distinct_simplify_t.v > Int64(3))]] +02)--TableScan: distinct_simplify_t projection=[g, v] +physical_plan +01)AggregateExec: mode=FinalPartitioned, gby=[g@0 as g], aggr=[min(distinct_simplify_t.v) FILTER (WHERE distinct_simplify_t.v > Int32(3)) as min(DISTINCT distinct_simplify_t.v) FILTER (WHERE distinct_simplify_t.v > Int64(3))] +02)--RepartitionExec: partitioning=Hash([g@0], 4), input_partitions=1 +03)----AggregateExec: mode=Partial, gby=[g@0 as g], aggr=[min(distinct_simplify_t.v) FILTER (WHERE distinct_simplify_t.v > Int32(3)) as min(DISTINCT distinct_simplify_t.v) FILTER (WHERE distinct_simplify_t.v > Int64(3))] +04)------DataSourceExec: partitions=1, partition_sizes=[1] + +# min/max(DISTINCT) over floats now sees the raw values. Before this rule, the +# inner group by of SingleDistinctToGroupBy turned -0.0 into 0.0 in its hash +# key, so max(DISTINCT v) printed 0.0 here. Now it agrees with max(v). +# The R column type prints -0.0 as 0, so compare the text form. +statement ok +CREATE TABLE distinct_simplify_float_t (v DOUBLE) AS VALUES (-1.0), (-0.0), (-0.0); + +query TTTT +SELECT cast(min(DISTINCT v) AS VARCHAR), cast(max(DISTINCT v) AS VARCHAR), cast(min(v) AS VARCHAR), cast(max(v) AS VARCHAR) FROM distinct_simplify_float_t; +---- +-1.0 -0.0 -1.0 -0.0 + +statement ok +DROP TABLE distinct_simplify_float_t; + +statement ok +DROP TABLE distinct_simplify_t; diff --git a/datafusion/sqllogictest/test_files/explain.slt b/datafusion/sqllogictest/test_files/explain.slt index b6837002086ad..e3490644068e1 100644 --- a/datafusion/sqllogictest/test_files/explain.slt +++ b/datafusion/sqllogictest/test_files/explain.slt @@ -195,6 +195,7 @@ logical_plan after filter_null_join_keys SAME TEXT AS ABOVE logical_plan after eliminate_outer_join SAME TEXT AS ABOVE logical_plan after push_down_limit SAME TEXT AS ABOVE logical_plan after push_down_filter SAME TEXT AS ABOVE +logical_plan after eliminate_aggregate_distinct SAME TEXT AS ABOVE logical_plan after single_distinct_aggregation_to_group_by SAME TEXT AS ABOVE logical_plan after eliminate_group_by_constant SAME TEXT AS ABOVE logical_plan after common_sub_expression_eliminate SAME TEXT AS ABOVE @@ -220,6 +221,7 @@ logical_plan after filter_null_join_keys SAME TEXT AS ABOVE logical_plan after eliminate_outer_join SAME TEXT AS ABOVE logical_plan after push_down_limit SAME TEXT AS ABOVE logical_plan after push_down_filter SAME TEXT AS ABOVE +logical_plan after eliminate_aggregate_distinct SAME TEXT AS ABOVE logical_plan after single_distinct_aggregation_to_group_by SAME TEXT AS ABOVE logical_plan after eliminate_group_by_constant SAME TEXT AS ABOVE logical_plan after common_sub_expression_eliminate SAME TEXT AS ABOVE @@ -574,6 +576,7 @@ logical_plan after filter_null_join_keys SAME TEXT AS ABOVE logical_plan after eliminate_outer_join SAME TEXT AS ABOVE logical_plan after push_down_limit SAME TEXT AS ABOVE logical_plan after push_down_filter SAME TEXT AS ABOVE +logical_plan after eliminate_aggregate_distinct SAME TEXT AS ABOVE logical_plan after single_distinct_aggregation_to_group_by SAME TEXT AS ABOVE logical_plan after eliminate_group_by_constant SAME TEXT AS ABOVE logical_plan after common_sub_expression_eliminate SAME TEXT AS ABOVE @@ -599,6 +602,7 @@ logical_plan after filter_null_join_keys SAME TEXT AS ABOVE logical_plan after eliminate_outer_join SAME TEXT AS ABOVE logical_plan after push_down_limit SAME TEXT AS ABOVE logical_plan after push_down_filter SAME TEXT AS ABOVE +logical_plan after eliminate_aggregate_distinct SAME TEXT AS ABOVE logical_plan after single_distinct_aggregation_to_group_by SAME TEXT AS ABOVE logical_plan after eliminate_group_by_constant SAME TEXT AS ABOVE logical_plan after common_sub_expression_eliminate SAME TEXT AS ABOVE diff --git a/datafusion/sqllogictest/test_files/single_distinct_to_groupby.slt b/datafusion/sqllogictest/test_files/single_distinct_to_groupby.slt index 8b032536fd420..fddfc661b8cc5 100644 --- a/datafusion/sqllogictest/test_files/single_distinct_to_groupby.slt +++ b/datafusion/sqllogictest/test_files/single_distinct_to_groupby.slt @@ -124,14 +124,16 @@ logical_plan 02)--Aggregate: groupBy=[[t.g]], aggr=[[count(Int64(1)), sum(DISTINCT CAST(t.v AS Int64))]] 03)----TableScan: t projection=[g, v] -# `min(DISTINCT v)` is the same value as `min(v)`, so the unrewritten plan keeps -# one scalar per group while the rewrite would build a row per distinct pair +# `min(DISTINCT v)` is the same value as `min(v)`. EliminateAggregateDistinct +# drops the flag before this rule runs, so there is no distinct aggregate left +# to rewrite and the plan keeps one scalar per group instead of a row per +# distinct pair query TT EXPLAIN SELECT g, count(*) AS records, min(DISTINCT v) AS distinct_min_v FROM t GROUP BY g; ---- logical_plan 01)Projection: t.g, count(Int64(1)) AS count(*) AS records, min(DISTINCT t.v) AS distinct_min_v -02)--Aggregate: groupBy=[[t.g]], aggr=[[count(Int64(1)), min(DISTINCT t.v)]] +02)--Aggregate: groupBy=[[t.g]], aggr=[[count(Int64(1)), min(t.v) AS min(DISTINCT t.v)]] 03)----TableScan: t projection=[g, v] # The gate covers only the count. A plan that already qualified through sum, diff --git a/docs/source/library-user-guide/functions/adding-udfs.md b/docs/source/library-user-guide/functions/adding-udfs.md index c3a40557a006d..c330d3db9612d 100644 --- a/docs/source/library-user-guide/functions/adding-udfs.md +++ b/docs/source/library-user-guide/functions/adding-udfs.md @@ -1116,6 +1116,24 @@ impl Accumulator for GeometricMean { } ``` +### Declaring how an Aggregate UDF treats `DISTINCT` + +By default DataFusion assumes an aggregate honors the `DISTINCT` modifier, which means the accumulator is expected to +read `AccumulatorArgs::is_distinct` and deduplicate its input. Override +[`AggregateUDFImpl::distinct_handling`] when that is not what your function does: + +- Return `DistinctHandling::Ignored` when duplicates cannot change the result, that is, when merging a value the + accumulator has already seen is a no-op. `min`, `max`, `bool_and` and `bit_or` are all in this group. The optimizer + then plans `f(DISTINCT x)` as `f(x)`, which skips both the per-group hash set and the extra grouping stage that + `SingleDistinctToGroupBy` would otherwise introduce. +- Return `DistinctHandling::Unsupported` when the accumulator does not implement `DISTINCT`: it does not read + `is_distinct`, or it rejects `DISTINCT` with an error. The planner must then deduplicate the input first or reject + the query. Today this is a declaration only; rejecting such queries at planning time is a follow-up change. +- Leave the default `DistinctHandling::Honored` when the accumulator reads `AccumulatorArgs::is_distinct` and + deduplicates its input itself. + +Getting this wrong changes query results, so only claim `Ignored` if your merge is genuinely idempotent. + ### Registering an Aggregate UDF To register a Aggregate UDF, you need to wrap the function implementation in a [`AggregateUDF`] struct and then register @@ -1370,6 +1388,7 @@ async fn main() -> Result<()> { [`aggregateudf`]: https://docs.rs/datafusion/latest/datafusion/logical_expr/struct.AggregateUDF.html [`create_udaf`]: https://docs.rs/datafusion/latest/datafusion/logical_expr/fn.create_udaf.html +[`aggregateudfimpl::distinct_handling`]: https://docs.rs/datafusion/latest/datafusion/logical_expr/trait.AggregateUDFImpl.html#method.distinct_handling [`advanced_udaf.rs`]: https://github.com/apache/datafusion/blob/main/datafusion-examples/examples/udf/advanced_udaf.rs ## Adding a Table UDF