From 9acbb37f837774dbada0a6e91947ecf3a4c98224 Mon Sep 17 00:00:00 2001 From: wudidapaopao <664920313@qq.com> Date: Sun, 20 Sep 2026 17:44:56 +0800 Subject: [PATCH] feat: decompose shared float64 avg aggregates --- datafusion/expr/src/udaf.rs | 78 ++++- datafusion/expr/src/utils.rs | 229 ++++++++++++++ datafusion/functions-aggregate/src/average.rs | 71 ++++- .../aggregate_decomposition.rs | 291 ++++++++++++++++++ .../optimizer/src/simplify_expressions/mod.rs | 1 + .../simplify_expressions/simplify_exprs.rs | 72 ++++- .../sqllogictest/test_files/aggregate.slt | 14 +- .../test_files/aggregates_simplify.slt | 21 ++ .../test_files/avg_to_sum_count.slt | 250 +++++++++++++++ .../optimizer_group_by_constant.slt | 7 +- .../consumer/rel/aggregate_rel.rs | 3 +- .../src/logical_plan/consumer/rel/mod.rs | 2 +- .../logical_plan/consumer/rel/project_rel.rs | 3 +- .../src/logical_plan/consumer/utils.rs | 223 +------------- 14 files changed, 1021 insertions(+), 244 deletions(-) create mode 100644 datafusion/optimizer/src/simplify_expressions/aggregate_decomposition.rs create mode 100644 datafusion/sqllogictest/test_files/avg_to_sum_count.slt diff --git a/datafusion/expr/src/udaf.rs b/datafusion/expr/src/udaf.rs index 458b39c969b6f..d3ee973164814 100644 --- a/datafusion/expr/src/udaf.rs +++ b/datafusion/expr/src/udaf.rs @@ -40,6 +40,7 @@ use crate::function::{ AccumulatorArgs, AggregateFunctionSimplification, StateFieldsArgs, }; use crate::groups_accumulator::GroupsAccumulator; +use crate::simplify::SimplifyContext; use crate::udf_eq::UdfEq; use crate::utils::AggregateOrderSensitivity; use crate::utils::format_state_name; @@ -312,6 +313,17 @@ impl AggregateUDF { self.inner.simplify() } + /// Returns this aggregate function's candidate decomposition, if any. + /// + /// See [`AggregateUDFImpl::decompose`] for more details. + pub fn decompose( + &self, + aggregate_function: &AggregateFunction, + info: &SimplifyContext, + ) -> Result> { + self.inner.decompose(aggregate_function, info) + } + /// Rewrite aggregate to have simpler arguments /// /// See [`AggregateUDFImpl::simplify_expr_op_literal`] for more details @@ -751,6 +763,25 @@ pub trait AggregateUDFImpl: Debug + DynEq + DynHash + Send + Sync + Any { None } + /// Returns an optional candidate decomposition into simpler aggregates. + /// + /// Unlike [`Self::simplify`], the optimizer only applies this rewrite when + /// at least one aggregate in the returned expression can be shared with + /// another aggregate in the same plan node. This makes the hook suitable + /// for rewrites such as `AVG(x)` into `SUM(x) / COUNT(x)`, which may be + /// slower when neither component can be reused. + /// + /// A returned candidate expression must have the same data type and + /// nullability as the original aggregate expression. Return `None` when + /// this aggregate cannot be decomposed. + fn decompose( + &self, + _aggregate_function: &AggregateFunction, + _info: &SimplifyContext, + ) -> Result> { + Ok(None) + } + /// Rewrite the aggregate to have simpler arguments /// /// This query pattern is not common in most real workloads, and most @@ -1635,6 +1666,14 @@ impl AggregateUDFImpl for AliasedAggregateUDFImpl { self.inner.simplify() } + fn decompose( + &self, + aggregate_function: &AggregateFunction, + info: &SimplifyContext, + ) -> Result> { + self.inner.decompose(aggregate_function, info) + } + fn simplify_expr_op_literal( &self, agg_function: &AggregateFunction, @@ -1715,7 +1754,9 @@ pub enum SetMonotonicity { #[cfg(test)] mod test { - use crate::{AggregateUDF, AggregateUDFImpl}; + use crate::expr::AggregateFunction; + use crate::simplify::SimplifyContext; + use crate::{AggregateUDF, AggregateUDFImpl, Expr, col}; use arrow::datatypes::{DataType, FieldRef}; use datafusion_common::Result; use datafusion_expr_common::accumulator::Accumulator; @@ -1762,6 +1803,13 @@ mod test { fn state_fields(&self, _args: StateFieldsArgs) -> Result> { unimplemented!() } + fn decompose( + &self, + _aggregate_function: &AggregateFunction, + _info: &SimplifyContext, + ) -> Result> { + Ok(Some(col("decomposed"))) + } } #[derive(Debug, Clone, PartialEq, Eq, Hash)] @@ -1824,6 +1872,34 @@ mod test { assert!(!(a1 == b1)); } + #[test] + fn test_decompose_forwarded_through_aliases() -> Result<()> { + let udf = AggregateUDF::from(AMeanUdf::new()).with_aliases(["alias"]); + let Expr::AggregateFunction(aggregate_function) = udf.call(vec![col("a")]) else { + panic!("expected aggregate function") + }; + + assert_eq!( + udf.decompose(&aggregate_function, &SimplifyContext::default())?, + Some(col("decomposed")) + ); + Ok(()) + } + + #[test] + fn test_default_decompose_returns_none() -> Result<()> { + let udf = AggregateUDF::from(BMeanUdf::new()); + let Expr::AggregateFunction(aggregate_function) = udf.call(vec![col("a")]) else { + panic!("expected aggregate function") + }; + + assert_eq!( + udf.decompose(&aggregate_function, &SimplifyContext::default())?, + None + ); + Ok(()) + } + fn hash(value: T) -> u64 { let hasher = &mut DefaultHasher::new(); value.hash(hasher); diff --git a/datafusion/expr/src/utils.rs b/datafusion/expr/src/utils.rs index 4ca12c5e0339e..84b155547da7c 100644 --- a/datafusion/expr/src/utils.rs +++ b/datafusion/expr/src/utils.rs @@ -50,6 +50,110 @@ pub use datafusion_functions_aggregate_common::order::AggregateOrderSensitivity; /// `COUNT()` expressions pub use datafusion_common::utils::expr::COUNT_STAR_EXPANSION; +/// Tracks expression names and generates aliases that do not conflict with +/// names already in an output schema. +/// +/// In addition to duplicate schema names, this detects ambiguity between a +/// qualified field such as `t.a` and an unqualified field named `a`. +#[derive(Default)] +pub struct NameTracker { + /// Tracks seen schema names (from expr.schema_name()). + /// Used to detect duplicates that would fail validate_unique_names. + seen_schema_names: HashSet, + /// Tracks column names that have been seen with a qualifier. + /// Used to detect ambiguous references (qualified + unqualified with same name). + qualified_names: HashSet, + /// Tracks column names that have been seen without a qualifier. + /// Used to detect ambiguous references. + unqualified_names: HashSet, +} + +impl NameTracker { + pub fn new() -> Self { + Self::default() + } + + /// Reserve the names of `exprs` without changing the expressions. + /// + /// This is useful when existing output expressions must retain their names + /// and subsequently generated expressions need to avoid them. + pub fn reserve(&mut self, exprs: &[Expr]) { + for expr in exprs { + self.insert(expr); + } + } + + fn would_conflict(&self, expr: &Expr) -> bool { + let (qualifier, name) = expr.qualified_name(); + let schema_name = expr.schema_name().to_string(); + self.would_conflict_inner((qualifier, &name), &schema_name) + } + + fn would_conflict_inner( + &self, + qualified_name: (Option, &str), + schema_name: &str, + ) -> bool { + // Check for duplicate schema_name (would fail validate_unique_names) + if self.seen_schema_names.contains(schema_name) { + return true; + } + + // Check for ambiguous reference (would fail DFSchema::check_names) + // This happens when a qualified field and unqualified field have the same name + let (qualifier, name) = qualified_name; + match qualifier { + Some(_) => { + // Adding a qualified name - conflicts if unqualified version exists + self.unqualified_names.contains(name) + } + None => { + // Adding an unqualified name - conflicts if qualified version exists + self.qualified_names.contains(name) + } + } + } + + fn insert(&mut self, expr: &Expr) { + let schema_name = expr.schema_name().to_string(); + self.seen_schema_names.insert(schema_name); + + let (qualifier, name) = expr.qualified_name(); + match qualifier { + Some(_) => { + self.qualified_names.insert(name); + } + None => { + self.unqualified_names.insert(name); + } + } + } + + /// Return `expr` unchanged if its name is available, or alias it to a + /// unique name otherwise. + pub fn get_uniquely_named_expr(&mut self, expr: Expr) -> Result { + if !self.would_conflict(&expr) { + self.insert(&expr); + return Ok(expr); + } + + // Name collision - need to generate a unique alias + let schema_name = expr.schema_name().to_string(); + let mut counter = 0; + let candidate_name = loop { + let candidate_name = format!("{schema_name}__temp__{counter}"); + // .alias always produces an unqualified name so check for conflicts accordingly. + if !self.would_conflict_inner((None, &candidate_name), &candidate_name) { + break candidate_name; + } + counter += 1; + }; + let candidate_expr = expr.alias(&candidate_name); + self.insert(&candidate_expr); + Ok(candidate_expr) + } +} + /// Count the number of distinct exprs in a list of group by expressions. If the /// first element is a `GroupingSet` expression then it must be the only expr. pub fn grouping_set_expr_count(group_expr: &[Expr]) -> Result { @@ -1515,6 +1619,131 @@ mod tests { use arrow::datatypes::{UnionFields, UnionMode}; use datafusion_expr_common::signature::Volatility; + #[test] + fn name_tracker_unique_names_pass_through() -> Result<()> { + let mut tracker = NameTracker::new(); + + // First expression should pass through unchanged + let expr1 = col("a"); + let result1 = tracker.get_uniquely_named_expr(expr1.clone())?; + assert_eq!(result1, col("a")); + + // Different name should also pass through unchanged + let expr2 = col("b"); + let result2 = tracker.get_uniquely_named_expr(expr2)?; + assert_eq!(result2, col("b")); + + Ok(()) + } + + #[test] + fn name_tracker_duplicate_schema_name_gets_alias() -> Result<()> { + let mut tracker = NameTracker::new(); + + // First expression with name "a" + let expr1 = col("a"); + let result1 = tracker.get_uniquely_named_expr(expr1)?; + assert_eq!(result1, col("a")); + + // Second expression with same name "a" should get aliased + let expr2 = col("a"); + let result2 = tracker.get_uniquely_named_expr(expr2)?; + assert_eq!(result2, col("a").alias("a__temp__0")); + + // Third expression with same name "a" should get a different alias + let expr3 = col("a"); + let result3 = tracker.get_uniquely_named_expr(expr3)?; + assert_eq!(result3, col("a").alias("a__temp__1")); + + Ok(()) + } + + #[test] + fn name_tracker_qualified_then_unqualified_conflicts() -> Result<()> { + let mut tracker = NameTracker::new(); + + // First: qualified column "table.a" + let qualified_col = Expr::Column(Column::new(Some("table"), "a")); + let result1 = tracker.get_uniquely_named_expr(qualified_col)?; + assert_eq!(result1, Expr::Column(Column::new(Some("table"), "a"))); + + // Second: unqualified column "a" - should conflict (ambiguous reference) + let unqualified_col = col("a"); + let result2 = tracker.get_uniquely_named_expr(unqualified_col)?; + // Should be aliased to avoid ambiguous reference + assert_eq!(result2, col("a").alias("a__temp__0")); + + Ok(()) + } + + #[test] + fn name_tracker_unqualified_then_qualified_conflicts() -> Result<()> { + let mut tracker = NameTracker::new(); + + // First: unqualified column "a" + let unqualified_col = col("a"); + let result1 = tracker.get_uniquely_named_expr(unqualified_col)?; + assert_eq!(result1, col("a")); + + // Second: qualified column "table.a" - should conflict (ambiguous reference) + let qualified_col = Expr::Column(Column::new(Some("table"), "a")); + let result2 = tracker.get_uniquely_named_expr(qualified_col)?; + // Should be aliased to avoid ambiguous reference + assert_eq!( + result2, + Expr::Column(Column::new(Some("table"), "a")).alias("table.a__temp__0") + ); + + Ok(()) + } + + #[test] + fn name_tracker_different_qualifiers_no_conflict() -> Result<()> { + let mut tracker = NameTracker::new(); + + // First: qualified column "table1.a" + let col1 = Expr::Column(Column::new(Some("table1"), "a")); + let result1 = tracker.get_uniquely_named_expr(col1.clone())?; + assert_eq!(result1, col1); + + // Second: qualified column "table2.a" - different qualifier, different schema_name + // so should NOT conflict + let col2 = Expr::Column(Column::new(Some("table2"), "a")); + let result2 = tracker.get_uniquely_named_expr(col2.clone())?; + assert_eq!(result2, col2); + + Ok(()) + } + + #[test] + fn name_tracker_aliased_expressions() -> Result<()> { + let mut tracker = NameTracker::new(); + + // First: col("x").alias("result") + let expr1 = col("x").alias("result"); + let result1 = tracker.get_uniquely_named_expr(expr1.clone())?; + assert_eq!(result1, col("x").alias("result")); + + // Second: col("y").alias("result") - same alias name, should conflict + let expr2 = col("y").alias("result"); + let result2 = tracker.get_uniquely_named_expr(expr2)?; + assert_eq!(result2, col("y").alias("result").alias("result__temp__0")); + + Ok(()) + } + + #[test] + fn name_tracker_avoids_reserved_qualified_name() -> Result<()> { + let mut tracker = NameTracker::new(); + tracker.reserve(&[Expr::Column(Column::new(Some("t"), "a"))]); + + assert_eq!( + tracker.get_uniquely_named_expr(col("a"))?, + col("a").alias("a__temp__0") + ); + Ok(()) + } + #[test] fn test_group_window_expr_by_sort_keys_empty_case() -> Result<()> { let result = group_window_expr_by_sort_keys(vec![])?; diff --git a/datafusion/functions-aggregate/src/average.rs b/datafusion/functions-aggregate/src/average.rs index f59ac816a930a..d935073340dd8 100644 --- a/datafusion/functions-aggregate/src/average.rs +++ b/datafusion/functions-aggregate/src/average.rs @@ -32,10 +32,14 @@ use arrow::datatypes::{ DurationSecondType, Field, FieldRef, Float64Type, TimeUnit, UInt64Type, }; use datafusion_common::types::{NativeType, logical_float64}; +use datafusion_common::utils::expr::COUNT_STAR_EXPANSION; use datafusion_common::{ Result, ScalarValue, exec_datafusion_err, exec_err, internal_err, not_impl_err, }; +use datafusion_expr::expr::AggregateFunction; +use datafusion_expr::expr_fn::cast; use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs}; +use datafusion_expr::simplify::SimplifyContext; use datafusion_expr::utils::format_state_name; use datafusion_expr::{ Accumulator, AggregateUDFImpl, Coercion, Documentation, EmitTo, Expr, GroupSelection, @@ -66,7 +70,7 @@ make_udaf_expr_and_func!( ); pub fn avg_distinct(expr: Expr) -> Expr { - Expr::AggregateFunction(datafusion_expr::expr::AggregateFunction::new_udf( + Expr::AggregateFunction(AggregateFunction::new_udf( avg_udaf(), vec![expr], true, @@ -492,11 +496,76 @@ impl AggregateUDFImpl for Avg { ReversedUDAF::Identical } + fn decompose( + &self, + aggregate_function: &AggregateFunction, + info: &SimplifyContext, + ) -> Result> { + decompose_avg_aggregate(aggregate_function, info) + } + fn documentation(&self) -> Option<&Documentation> { self.doc() } } +/// Returns a candidate rewrite of `AVG(x)` to +/// `SUM(x) / CAST(COUNT(x) AS Float64)`. +/// +/// The optimizer applies the candidate only when at least one resulting +/// aggregate can be shared without increasing the number of distinct aggregate +/// expressions. Duplicate components are removed when the optimizer rebuilds +/// the inner aggregate. +/// +/// Only the Float64 path is rewritten: decimal and duration AVGs use widened +/// intermediate sum types and different division semantics. +fn decompose_avg_aggregate( + aggregate_function: &AggregateFunction, + info: &SimplifyContext, +) -> Result> { + let params = &aggregate_function.params; + + // DISTINCT / FILTER / ORDER BY / null treatment would have to be + // replicated onto both SUM and COUNT; keep those AVGs as-is. + if params.distinct + || params.filter.is_some() + || !params.order_by.is_empty() + || params.null_treatment.is_some() + { + return Ok(None); + } + + let [arg] = params.args.as_slice() else { + return Ok(None); + }; + + // Numeric inputs are coerced to Float64 by AVG's signature, so this single + // check selects the Float64 path and skips decimal/duration inputs. + if info.get_data_type(arg)? != DataType::Float64 { + return Ok(None); + } + + // After the split the argument is evaluated once per aggregate, which is + // not equivalent for volatile expressions. + if arg.is_volatile() { + return Ok(None); + } + + // For non-nullable arguments count rows so the component can be shared + // with an existing COUNT(*). + let count_arg = if info.nullable(arg)? { + arg.clone() + } else { + Expr::Literal(COUNT_STAR_EXPANSION, None) + }; + + // SUM keeps AVG's Float64 argument so the arithmetic matches AVG exactly. + let sum_expr = crate::sum::sum(arg.clone()); + let count_expr = crate::count::count(count_arg); + + Ok(Some(sum_expr / cast(count_expr, DataType::Float64))) +} + /// The precision and scale of a decimal `DataType` fn decimal_parts(data_type: &DataType) -> Result<(u8, i8)> { match data_type { diff --git a/datafusion/optimizer/src/simplify_expressions/aggregate_decomposition.rs b/datafusion/optimizer/src/simplify_expressions/aggregate_decomposition.rs new file mode 100644 index 0000000000000..0f95fa746b78a --- /dev/null +++ b/datafusion/optimizer/src/simplify_expressions/aggregate_decomposition.rs @@ -0,0 +1,291 @@ +// 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. + +//! Applies aggregate decompositions whose components can be shared. + +use datafusion_common::Result; +use datafusion_expr::Expr; +use datafusion_expr::simplify::SimplifyContext; +use datafusion_expr::utils::find_aggregate_exprs; + +#[derive(Debug)] +struct AggregateDecomposition { + /// Original aggregate expression, without its output alias. + original: Expr, + /// Full replacement expression, with the original output name preserved. + expression: Expr, + /// Distinct aggregate expressions needed to evaluate the replacement. + components: Vec, +} + +/// Applies aggregate decompositions only when at least one resulting aggregate +/// expression occurs elsewhere in the same aggregate node. +/// +/// Candidates connected by shared components are considered together. A group +/// is applied only when deduplication will not increase the number of distinct +/// aggregate expressions. +pub(super) fn rewrite_shared_aggregate_components( + aggr_expr: &mut [Expr], + info: &SimplifyContext, +) -> Result { + let mut candidates = Vec::with_capacity(aggr_expr.len()); + let mut existing_components = vec![]; + + // Phase 1: ask each aggregate for a decomposition candidate. + for expr in aggr_expr.iter() { + let unaliased_expr = expr.clone().unalias_nested().data; + let candidate = + if let Expr::AggregateFunction(aggregate_function) = unaliased_expr { + if let Some(decomposed) = aggregate_function + .func + .decompose(&aggregate_function, info)? + { + let components = find_aggregate_exprs([&decomposed]); + let original_name = expr.name_for_alias()?; + let expression = decomposed.alias_if_changed(original_name)?; + Some(AggregateDecomposition { + original: Expr::AggregateFunction(aggregate_function), + expression, + components, + }) + } else { + None + } + } else { + None + }; + + if candidate.is_none() { + existing_components.extend(find_aggregate_exprs([expr])); + } + candidates.push(candidate); + } + + // Phase 2: decide which connected groups of candidates are profitable. + let apply_candidate = + select_profitable_decompositions(&candidates, &existing_components); + + // Phase 3: replace only the candidates selected by the profitability check. + let mut rewritten = false; + for ((expr, candidate), apply) in + aggr_expr.iter_mut().zip(candidates).zip(apply_candidate) + { + if apply { + let candidate = + candidate.expect("an applied decomposition candidate must exist"); + debug_assert!(!candidate.components.is_empty()); + *expr = candidate.expression; + rewritten = true; + } + } + + Ok(rewritten) +} + +/// Selects which aggregate decomposition candidates should be applied. +/// +/// It favors groups that reuse existing or duplicate components and rejects a +/// group when decomposition would increase the number of distinct aggregates. +/// For AVG decompositions into SUM and COUNT, this usually selects a beneficial +/// set of rewrites. +fn select_profitable_decompositions( + candidates: &[Option], + existing_components: &[Expr], +) -> Vec { + let mut apply_candidate = vec![false; candidates.len()]; + let mut visited = vec![false; candidates.len()]; + + for start in 0..candidates.len() { + if visited[start] || candidates[start].is_none() { + continue; + } + + visited[start] = true; + // Find the connected group of candidates linked by shared components. + let mut group = vec![start]; + let mut pending = vec![start]; + while let Some(current) = pending.pop() { + let current_components = &candidates[current] + .as_ref() + .expect("a pending decomposition candidate must exist") + .components; + for (other, candidate) in candidates.iter().enumerate() { + if visited[other] { + continue; + } + let Some(candidate) = candidate else { + continue; + }; + if current_components + .iter() + .any(|component| candidate.components.contains(component)) + { + visited[other] = true; + group.push(other); + pending.push(other); + } + } + } + + // Existing components are free. Count only distinct components that + // this group would add after deduplication. + let mut shared_component = false; + let mut new_components = vec![]; + let mut original_aggregates = vec![]; + for &index in &group { + let candidate = candidates[index] + .as_ref() + .expect("a grouped decomposition candidate must exist"); + if !original_aggregates.contains(&candidate.original) { + original_aggregates.push(candidate.original.clone()); + } + for component in &candidate.components { + if existing_components.contains(component) + || new_components.contains(component) + { + shared_component = true; + } else { + new_components.push(component.clone()); + } + } + } + + // Replacing N distinct aggregates is worthwhile only if there is actual + // sharing and the group adds at most N distinct aggregates. + if shared_component && new_components.len() <= original_aggregates.len() { + for index in group { + apply_candidate[index] = true; + } + } + } + + apply_candidate +} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion_expr::col; + + fn candidate(name: &str, components: &[&str]) -> Option { + Some(AggregateDecomposition { + original: col(name), + expression: col(name), + components: components.iter().map(|name| col(*name)).collect(), + }) + } + + fn aliased_candidate( + name: &str, + original: &str, + components: &[&str], + ) -> Option { + Some(AggregateDecomposition { + original: col(original), + expression: col(name), + components: components.iter().map(|name| col(*name)).collect(), + }) + } + + #[test] + fn applies_when_component_already_exists() { + let candidates = vec![candidate("avg", &["sum", "count"])]; + assert_eq!( + select_profitable_decompositions(&candidates, &[col("sum")]), + vec![true] + ); + } + + #[test] + fn rejects_without_shared_component() { + let candidates = vec![candidate("avg", &["sum", "count"])]; + assert_eq!( + select_profitable_decompositions(&candidates, &[]), + vec![false] + ); + } + + #[test] + fn rejects_when_candidate_sharing_increases_aggregate_count() { + let candidates = vec![ + candidate("avg_a", &["sum_a", "count"]), + candidate("avg_b", &["sum_b", "count"]), + ]; + assert_eq!( + select_profitable_decompositions(&candidates, &[]), + vec![false, false] + ); + } + + #[test] + fn applies_candidate_sharing_without_increasing_aggregate_count() { + let candidates = vec![ + candidate("aggregate_a", &["sum", "count"]), + candidate("aggregate_b", &["sum", "count"]), + ]; + assert_eq!( + select_profitable_decompositions(&candidates, &[]), + vec![true, true] + ); + } + + #[test] + fn selects_disconnected_groups_independently() { + let candidates = vec![ + candidate("profitable", &["existing", "new"]), + candidate("unprofitable", &["other_a", "other_b"]), + ]; + assert_eq!( + select_profitable_decompositions(&candidates, &[col("existing")]), + vec![true, false] + ); + } + + #[test] + fn rejects_duplicate_original_aggregates_that_expand() { + let candidates = vec![ + aliased_candidate("avg_a", "avg", &["sum", "count"]), + aliased_candidate("avg_b", "avg", &["sum", "count"]), + ]; + assert_eq!( + select_profitable_decompositions(&candidates, &[]), + vec![false, false] + ); + } + + #[test] + fn applies_transitively_connected_profitable_candidates() { + let candidates = vec![ + candidate("a", &["existing", "b"]), + candidate("b", &["b", "c"]), + candidate("c", &["c", "d"]), + ]; + assert_eq!( + select_profitable_decompositions(&candidates, &[col("existing")]), + vec![true, true, true] + ); + } + + #[test] + fn rejects_empty_decomposition() { + let candidates = vec![candidate("empty", &[])]; + assert_eq!( + select_profitable_decompositions(&candidates, &[]), + vec![false] + ); + } +} diff --git a/datafusion/optimizer/src/simplify_expressions/mod.rs b/datafusion/optimizer/src/simplify_expressions/mod.rs index e0b53b79d468c..ed5c28453775a 100644 --- a/datafusion/optimizer/src/simplify_expressions/mod.rs +++ b/datafusion/optimizer/src/simplify_expressions/mod.rs @@ -18,6 +18,7 @@ //! [`SimplifyExpressions`] simplifies expressions in the logical plan, //! [`ExprSimplifier`] simplifies individual `Expr`s. +mod aggregate_decomposition; pub mod expr_simplifier; mod inlist_simplifier; mod linear_aggregates; diff --git a/datafusion/optimizer/src/simplify_expressions/simplify_exprs.rs b/datafusion/optimizer/src/simplify_expressions/simplify_exprs.rs index 1c5a4a1869ddb..33034c95c7adf 100644 --- a/datafusion/optimizer/src/simplify_expressions/simplify_exprs.rs +++ b/datafusion/optimizer/src/simplify_expressions/simplify_exprs.rs @@ -19,17 +19,19 @@ use std::sync::Arc; -use datafusion_common::tree_node::{Transformed, TreeNode}; +use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRecursion}; use datafusion_common::{Column, DFSchema, DFSchemaRef, DataFusionError, Result}; use datafusion_expr::logical_plan::{Aggregate, LogicalPlan, Projection}; use datafusion_expr::simplify::SimplifyContext; use datafusion_expr::utils::{ - columnize_expr, find_aggregate_exprs, grouping_set_to_exprlist, merge_schema, + NameTracker, columnize_expr, find_aggregate_exprs, grouping_set_to_exprlist, + merge_schema, }; use datafusion_expr::{DmlStatement, Expr, WriteOp}; use super::ExprSimplifier; use crate::optimizer::ApplyOrder; +use crate::simplify_expressions::aggregate_decomposition::rewrite_shared_aggregate_components; use crate::simplify_expressions::linear_aggregates::rewrite_multiple_linear_aggregates; use crate::utils::NamePreserver; use crate::{OptimizerConfig, OptimizerRule}; @@ -120,7 +122,7 @@ impl SimplifyExpressions { // Inputs have already been rewritten (due to bottom-up traversal handled by Optimizer) // Just need to rewrite our own expressions - let simplifier = ExprSimplifier::new(info); + let simplifier = ExprSimplifier::new(info.clone()); // The left and right expressions in a Join on clause are not // commutative, for reasons that are not entirely clear. Thus, do not @@ -154,7 +156,7 @@ impl SimplifyExpressions { rewrite_expr(expr) } })? - .transform_data(rewrite_aggregate_non_aggregate_aggr_expr) + .transform_data(|plan| rewrite_aggregate_non_aggregate_aggr_expr(plan, &info)) } } @@ -168,7 +170,8 @@ impl SimplifyExpressions { /// Ensures that `LogicalPlan::Aggregate` is well formed after rewrites /// by potentially introducing an extra `Projection`. /// -/// Also applies the [`rewrite_multiple_linear_aggregates`] special case +/// Also applies shared aggregate decompositions and the +/// [`rewrite_multiple_linear_aggregates`] special case. /// /// # Rationale: /// @@ -186,6 +189,7 @@ impl SimplifyExpressions { /// * ` Aggregate(group_expr, aggr_expr=[agg(exp2) AS _X])` fn rewrite_aggregate_non_aggregate_aggr_expr( plan: LogicalPlan, + info: &SimplifyContext, ) -> Result> { let LogicalPlan::Aggregate(Aggregate { input, @@ -198,7 +202,10 @@ fn rewrite_aggregate_non_aggregate_aggr_expr( return Ok(Transformed::no(plan)); }; - let rewrote_aggs = rewrite_multiple_linear_aggregates(&mut aggr_expr)?; + let rewrote_linear = rewrite_multiple_linear_aggregates(&mut aggr_expr)?; + let rewrote_decompositions = + rewrite_shared_aggregate_components(&mut aggr_expr, info)?; + let rewrote_aggs = rewrote_decompositions || rewrote_linear; // Ensure that all Aggregate arguments are AggregateExpr if aggr_expr.iter().all(is_top_level_aggregate_expr) { @@ -215,7 +222,25 @@ fn rewrite_aggregate_non_aggregate_aggr_expr( // Otherwise we need to add a Projection above Aggregate to calculate // the final output expressions. - let inner_aggr_expr = find_aggregate_exprs(aggr_expr.iter()); + // Decomposing `SELECT SUM(x), AVG(x), COUNT(*)` for an integer column `x` + // produces separate Int64 and Float64 sums, both named "sum(x)". + // A GROUP BY column named "sum(x)" can also conflict, so reserve group names. + let mut projection_exprs = aggregate_output_exprs(&group_expr)?; + let mut name_tracker = NameTracker::new(); + name_tracker.reserve(&projection_exprs); + + let mut renames = Vec::new(); + let inner_aggr_expr = find_aggregate_exprs(aggr_expr.iter()) + .into_iter() + .map(|agg| { + let named_agg = name_tracker.get_uniquely_named_expr(agg.clone())?; + if named_agg != agg { + let (relation, name) = named_agg.qualified_name(); + renames.push((agg, Column::new(relation, name))); + } + Ok(named_agg) + }) + .collect::>>()?; let inner_aggregate = LogicalPlan::Aggregate(Aggregate::try_new( Arc::clone(&input), group_expr.clone(), @@ -223,11 +248,21 @@ fn rewrite_aggregate_non_aggregate_aggr_expr( )?); let inner_aggregate = Arc::new(inner_aggregate); - let mut projection_exprs = aggregate_output_exprs(&group_expr)?; projection_exprs.extend(aggr_expr); let projection_exprs = projection_exprs .into_iter() - .map(|expr| columnize_expr(expr, inner_aggregate.as_ref())) + .map(|expr| { + let original_name = expr.schema_name().to_string(); + let replaced = replace_renamed_aggregates(expr, &renames)?; + let expr = columnize_expr(replaced.data, inner_aggregate.as_ref())?; + // Replacing an aggregate with its disambiguation alias changes the + // schema name; restore it so references above keep resolving. + if replaced.transformed && expr.schema_name().to_string() != original_name { + Ok(expr.alias(original_name)) + } else { + Ok(expr) + } + }) .collect::>>()?; Ok(Transformed::yes(LogicalPlan::Projection( @@ -235,6 +270,25 @@ fn rewrite_aggregate_non_aggregate_aggr_expr( ))) } +/// Replaces aggregates that were aliased for disambiguation with a reference +/// to the aliased column, as `columnize_expr` only matches unaliased exprs. +fn replace_renamed_aggregates( + expr: Expr, + renames: &[(Expr, Column)], +) -> Result> { + if renames.is_empty() { + return Ok(Transformed::no(expr)); + } + expr.transform_down(|e| match renames.iter().find(|(agg, _)| agg == &e) { + Some((_, column)) => Ok(Transformed::new( + Expr::Column(column.clone()), + true, + TreeNodeRecursion::Jump, + )), + None => Ok(Transformed::no(e)), + }) +} + fn is_top_level_aggregate_expr(expr: &Expr) -> bool { matches!( expr.clone().unalias_nested().data, diff --git a/datafusion/sqllogictest/test_files/aggregate.slt b/datafusion/sqllogictest/test_files/aggregate.slt index 3eabd5c4ff7bb..7bffaae93f77f 100644 --- a/datafusion/sqllogictest/test_files/aggregate.slt +++ b/datafusion/sqllogictest/test_files/aggregate.slt @@ -9462,14 +9462,16 @@ ORDER BY g; logical_plan 01)Sort: stream_test.g ASC NULLS LAST 02)--Projection: stream_test.g, count(Int64(1)) AS count(*), sum(stream_test.x), avg(stream_test.x), avg(stream_test.x) AS mean(stream_test.x), min(stream_test.x), max(stream_test.y), bit_and(stream_test.i), bit_or(stream_test.i), bit_xor(stream_test.i), bool_and(stream_test.b), bool_or(stream_test.b), median(stream_test.x), Int32(0) AS grouping(stream_test.g), var(stream_test.x), var(stream_test.x) AS var_samp(stream_test.x), var_pop(stream_test.x), var(stream_test.x) AS var_sample(stream_test.x), var_pop(stream_test.x) AS var_population(stream_test.x), stddev(stream_test.x), stddev(stream_test.x) AS stddev_samp(stream_test.x), stddev_pop(stream_test.x) -03)----Aggregate: groupBy=[[stream_test.g]], aggr=[[count(Int64(1)), sum(stream_test.x), avg(stream_test.x), min(stream_test.x), max(stream_test.y), bit_and(stream_test.i), bit_or(stream_test.i), bit_xor(stream_test.i), bool_and(stream_test.b), bool_or(stream_test.b), median(stream_test.x), var(stream_test.x), var_pop(stream_test.x), stddev(stream_test.x), stddev_pop(stream_test.x)]] -04)------Sort: stream_test.g ASC NULLS LAST, fetch=10000 -05)--------TableScan: stream_test projection=[g, x, y, i, b] +03)----Projection: stream_test.g, count(Int64(1)), sum(stream_test.x), sum(stream_test.x) / CAST(count(stream_test.x) AS Float64) AS avg(stream_test.x), min(stream_test.x), max(stream_test.y), bit_and(stream_test.i), bit_or(stream_test.i), bit_xor(stream_test.i), bool_and(stream_test.b), bool_or(stream_test.b), median(stream_test.x), var(stream_test.x), var_pop(stream_test.x), stddev(stream_test.x), stddev_pop(stream_test.x) +04)------Aggregate: groupBy=[[stream_test.g]], aggr=[[count(Int64(1)), sum(stream_test.x), count(stream_test.x), min(stream_test.x), max(stream_test.y), bit_and(stream_test.i), bit_or(stream_test.i), bit_xor(stream_test.i), bool_and(stream_test.b), bool_or(stream_test.b), median(stream_test.x), var(stream_test.x), var_pop(stream_test.x), stddev(stream_test.x), stddev_pop(stream_test.x)]] +05)--------Sort: stream_test.g ASC NULLS LAST, fetch=10000 +06)----------TableScan: stream_test projection=[g, x, y, i, b] physical_plan 01)ProjectionExec: expr=[g@0 as g, count(Int64(1))@1 as count(*), sum(stream_test.x)@2 as sum(stream_test.x), avg(stream_test.x)@3 as avg(stream_test.x), avg(stream_test.x)@3 as mean(stream_test.x), min(stream_test.x)@4 as min(stream_test.x), max(stream_test.y)@5 as max(stream_test.y), bit_and(stream_test.i)@6 as bit_and(stream_test.i), bit_or(stream_test.i)@7 as bit_or(stream_test.i), bit_xor(stream_test.i)@8 as bit_xor(stream_test.i), bool_and(stream_test.b)@9 as bool_and(stream_test.b), bool_or(stream_test.b)@10 as bool_or(stream_test.b), median(stream_test.x)@11 as median(stream_test.x), 0 as grouping(stream_test.g), var(stream_test.x)@12 as var(stream_test.x), var(stream_test.x)@12 as var_samp(stream_test.x), var_pop(stream_test.x)@13 as var_pop(stream_test.x), var(stream_test.x)@12 as var_sample(stream_test.x), var_pop(stream_test.x)@13 as var_population(stream_test.x), stddev(stream_test.x)@14 as stddev(stream_test.x), stddev(stream_test.x)@14 as stddev_samp(stream_test.x), stddev_pop(stream_test.x)@15 as stddev_pop(stream_test.x)] -02)--AggregateExec: mode=Single, gby=[g@0 as g], aggr=[count(Int64(1)), sum(stream_test.x), avg(stream_test.x), min(stream_test.x), max(stream_test.y), bit_and(stream_test.i), bit_or(stream_test.i), bit_xor(stream_test.i), bool_and(stream_test.b), bool_or(stream_test.b), median(stream_test.x), var(stream_test.x), var_pop(stream_test.x), stddev(stream_test.x), stddev_pop(stream_test.x)], ordering_mode=Sorted -03)----SortExec: TopK(fetch=10000), expr=[g@0 ASC NULLS LAST], preserve_partitioning=[false] -04)------DataSourceExec: partitions=1, partition_sizes=[1] +02)--ProjectionExec: expr=[g@0 as g, count(Int64(1))@1 as count(Int64(1)), sum(stream_test.x)@2 as sum(stream_test.x), sum(stream_test.x)@2 / CAST(count(stream_test.x)@3 AS Float64) as avg(stream_test.x), min(stream_test.x)@4 as min(stream_test.x), max(stream_test.y)@5 as max(stream_test.y), bit_and(stream_test.i)@6 as bit_and(stream_test.i), bit_or(stream_test.i)@7 as bit_or(stream_test.i), bit_xor(stream_test.i)@8 as bit_xor(stream_test.i), bool_and(stream_test.b)@9 as bool_and(stream_test.b), bool_or(stream_test.b)@10 as bool_or(stream_test.b), median(stream_test.x)@11 as median(stream_test.x), var(stream_test.x)@12 as var(stream_test.x), var_pop(stream_test.x)@13 as var_pop(stream_test.x), stddev(stream_test.x)@14 as stddev(stream_test.x), stddev_pop(stream_test.x)@15 as stddev_pop(stream_test.x)] +03)----AggregateExec: mode=Single, gby=[g@0 as g], aggr=[count(Int64(1)), sum(stream_test.x), count(stream_test.x), min(stream_test.x), max(stream_test.y), bit_and(stream_test.i), bit_or(stream_test.i), bit_xor(stream_test.i), bool_and(stream_test.b), bool_or(stream_test.b), median(stream_test.x), var(stream_test.x), var_pop(stream_test.x), stddev(stream_test.x), stddev_pop(stream_test.x)], ordering_mode=Sorted +04)------SortExec: TopK(fetch=10000), expr=[g@0 ASC NULLS LAST], preserve_partitioning=[false] +05)--------DataSourceExec: partitions=1, partition_sizes=[1] query IIRRRRRIIIBBRIRRRRRRRR SELECT diff --git a/datafusion/sqllogictest/test_files/aggregates_simplify.slt b/datafusion/sqllogictest/test_files/aggregates_simplify.slt index c4055d17396c5..b767318d287de 100644 --- a/datafusion/sqllogictest/test_files/aggregates_simplify.slt +++ b/datafusion/sqllogictest/test_files/aggregates_simplify.slt @@ -351,8 +351,29 @@ physical_plan 03)----AggregateExec: mode=Single, gby=[], aggr=[sum(tbl.val), count(tbl.val)] 04)------DataSourceExec: partitions=1, partition_sizes=[2] +# Extracted aggregate names must not conflict with group output names. +statement ok +CREATE TABLE linear_name_collision_t( + a INT NOT NULL, + "sum(linear_name_collision_t.a)" INT NOT NULL +) AS VALUES (1, 10), (2, 20); + +query III rowsort +SELECT + "sum(linear_name_collision_t.a)", + SUM(a + 1), + SUM(a + 2) +FROM linear_name_collision_t +GROUP BY "sum(linear_name_collision_t.a)"; +---- +10 2 3 +20 3 4 + statement ok DROP TABLE IF EXISTS tbl; statement ok DROP TABLE sum_simplify_t; + +statement ok +DROP TABLE linear_name_collision_t; diff --git a/datafusion/sqllogictest/test_files/avg_to_sum_count.slt b/datafusion/sqllogictest/test_files/avg_to_sum_count.slt new file mode 100644 index 0000000000000..c2d77ccd06eee --- /dev/null +++ b/datafusion/sqllogictest/test_files/avg_to_sum_count.slt @@ -0,0 +1,250 @@ +# 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. + +# Tests for decomposing Float64 AVG(x) into SUM(x) / COUNT(x) when at least one +# component can be shared with another aggregate. + +statement ok +CREATE TABLE t(a INT NOT NULL, b INT, c DOUBLE) AS VALUES + (1, 10, 1.5), (2, NULL, 2.5), (3, 30, NULL), (4, 40, 4.5); + +# AVG over a non-nullable column decomposes into sum / count(*), and the count +# is shared with an explicit count(*). +query TT +EXPLAIN SELECT count(*), avg(a) FROM t; +---- +logical_plan +01)Projection: count(Int64(1)) AS count(*), sum(t.a) / CAST(count(Int64(1)) AS Float64) AS avg(t.a) +02)--Aggregate: groupBy=[[]], aggr=[[count(Int64(1)), sum(CAST(t.a AS Float64))]] +03)----TableScan: t projection=[a] +physical_plan +01)ProjectionExec: expr=[count(Int64(1))@0 as count(*), sum(t.a)@1 / CAST(count(Int64(1))@0 AS Float64) as avg(t.a)] +02)--AggregateExec: mode=Single, gby=[], aggr=[count(Int64(1)), sum(t.a)] +03)----DataSourceExec: partitions=1, partition_sizes=[1] + +# AVG without a shared SUM or COUNT keeps its combined accumulator. +query TT +EXPLAIN SELECT avg(b) FROM t; +---- +logical_plan +01)Aggregate: groupBy=[[]], aggr=[[avg(CAST(t.b AS Float64))]] +02)--TableScan: t projection=[b] +physical_plan +01)AggregateExec: mode=Single, gby=[], aggr=[avg(t.b)] +02)--DataSourceExec: partitions=1, partition_sizes=[1] + +# Decomposition does not normalize COUNT(CAST(expr)) to COUNT(expr); that +# equivalence belongs in a separate optimizer rewrite. +query TT +EXPLAIN SELECT count(b), avg(b) FROM t; +---- +logical_plan +01)Aggregate: groupBy=[[]], aggr=[[count(t.b), avg(CAST(t.b AS Float64))]] +02)--TableScan: t projection=[b] +physical_plan +01)AggregateExec: mode=Single, gby=[], aggr=[count(t.b), avg(t.b)] +02)--DataSourceExec: partitions=1, partition_sizes=[1] + +# Sharing only between decomposition candidates would still increase the +# number of distinct aggregates (two AVGs to two SUMs plus one COUNT), so both +# AVGs keep their combined accumulators. +query TT +EXPLAIN SELECT avg(a), avg(a + 1) FROM t; +---- +logical_plan +01)Aggregate: groupBy=[[]], aggr=[[avg(CAST(t.a AS Float64)), avg(CAST(CAST(t.a AS Int64) + Int64(1) AS Float64))]] +02)--TableScan: t projection=[a] +physical_plan +01)AggregateExec: mode=Single, gby=[], aggr=[avg(t.a), avg(t.a + Int64(1))] +02)--DataSourceExec: partitions=1, partition_sizes=[1] + +# A nullable AVG decomposes when its count can be shared. +query TT +EXPLAIN SELECT count(c), avg(c) FROM t; +---- +logical_plan +01)Projection: count(t.c), sum(t.c) / CAST(count(t.c) AS Float64) AS avg(t.c) +02)--Aggregate: groupBy=[[]], aggr=[[count(t.c), sum(t.c)]] +03)----TableScan: t projection=[c] +physical_plan +01)ProjectionExec: expr=[count(t.c)@0 as count(t.c), sum(t.c)@1 / CAST(count(t.c)@0 AS Float64) as avg(t.c)] +02)--AggregateExec: mode=Single, gby=[], aggr=[count(t.c), sum(t.c)] +03)----DataSourceExec: partitions=1, partition_sizes=[1] + +# AVG decomposes when its SUM can be shared. +query TT +EXPLAIN SELECT sum(c), avg(c) FROM t; +---- +logical_plan +01)Projection: sum(t.c), sum(t.c) / CAST(count(t.c) AS Float64) AS avg(t.c) +02)--Aggregate: groupBy=[[]], aggr=[[sum(t.c), count(t.c)]] +03)----TableScan: t projection=[c] +physical_plan +01)ProjectionExec: expr=[sum(t.c)@0 as sum(t.c), sum(t.c)@0 / CAST(count(t.c)@1 AS Float64) as avg(t.c)] +02)--AggregateExec: mode=Single, gby=[], aggr=[sum(t.c), count(t.c)] +03)----DataSourceExec: partitions=1, partition_sizes=[1] + +# AVG uses the same SUM and COUNT when both are already requested. +query TT +EXPLAIN SELECT sum(c), count(c), avg(c) FROM t; +---- +logical_plan +01)Projection: sum(t.c), count(t.c), sum(t.c) / CAST(count(t.c) AS Float64) AS avg(t.c) +02)--Aggregate: groupBy=[[]], aggr=[[sum(t.c), count(t.c)]] +03)----TableScan: t projection=[c] +physical_plan +01)ProjectionExec: expr=[sum(t.c)@0 as sum(t.c), count(t.c)@1 as count(t.c), sum(t.c)@0 / CAST(count(t.c)@1 AS Float64) as avg(t.c)] +02)--AggregateExec: mode=Single, gby=[], aggr=[sum(t.c), count(t.c)] +03)----DataSourceExec: partitions=1, partition_sizes=[1] + +# The decomposed sum(CAST(a AS Float64)) collides with the user-written +# sum(a) on schema name; the inner aggregate is disambiguated with an alias. +query TT +EXPLAIN SELECT count(*), avg(a), sum(a) FROM t; +---- +logical_plan +01)Projection: count(Int64(1)) AS count(*), sum(t.a) / CAST(count(Int64(1)) AS Float64) AS avg(t.a), sum(t.a)__temp__0 AS sum(t.a) +02)--Aggregate: groupBy=[[]], aggr=[[count(Int64(1)), sum(CAST(t.a AS Float64)), sum(CAST(t.a AS Int64)) AS sum(t.a)__temp__0]] +03)----TableScan: t projection=[a] +physical_plan +01)ProjectionExec: expr=[count(Int64(1))@0 as count(*), sum(t.a)@1 / CAST(count(Int64(1))@0 AS Float64) as avg(t.a), sum(t.a)__temp__0@2 as sum(t.a)] +02)--AggregateExec: mode=Single, gby=[], aggr=[count(Int64(1)), sum(t.a), sum(t.a) as sum(t.a)__temp__0] +03)----DataSourceExec: partitions=1, partition_sizes=[1] + +query IRI +SELECT count(*), avg(a), sum(a) FROM t; +---- +4 2.5 10 + +# Results with NULLs in both nullable columns. +query RRR +SELECT avg(a), avg(b), avg(c) FROM t; +---- +2.5 26.666666666667 2.833333333333 + +# Decomposition preserves AVG semantics for empty input... +query RI +SELECT avg(a), count(*) FROM t WHERE a > 100; +---- +NULL 0 + +# ...and for all-NULL input. +query RI +SELECT avg(x), count(x) +FROM (VALUES (CAST(NULL AS DOUBLE)), (NULL)) v(x); +---- +NULL 0 + +# Grouped decomposition. +query IRI rowsort +SELECT b % 20, avg(a), count(*) FROM t GROUP BY b % 20; +---- +0 4 1 +10 2 2 +NULL 2 1 + +# Grouping sets decomposition preserves expanded grouping columns and the +# internal grouping id needed by GROUPING(). +query IIRI +SELECT b, grouping(b), avg(a), count(*) +FROM t +GROUP BY GROUPING SETS ((b), ()) +ORDER BY grouping(b), b NULLS FIRST; +---- +NULL 0 2 1 +10 0 1 1 +30 0 3 1 +40 0 4 1 +NULL 1 2.5 4 + +# DISTINCT / FILTER AVGs are not decomposed. +query TT +EXPLAIN SELECT avg(DISTINCT a), avg(a) FILTER (WHERE a > 1) FROM t; +---- +logical_plan +01)Aggregate: groupBy=[[]], aggr=[[avg(DISTINCT __common_expr_1 AS t.a), avg(__common_expr_1) FILTER (WHERE t.a > Int32(1)) AS avg(t.a) FILTER (WHERE t.a > Int64(1))]] +02)--Projection: CAST(t.a AS Float64) AS __common_expr_1, t.a +03)----TableScan: t projection=[a] +physical_plan +01)AggregateExec: mode=Single, gby=[], aggr=[avg(DISTINCT t.a), avg(__common_expr_1) FILTER (WHERE t.a > Int32(1)) as avg(t.a) FILTER (WHERE t.a > Int64(1))] +02)--ProjectionExec: expr=[CAST(a@0 AS Float64) as __common_expr_1, a@0 as a] +03)----DataSourceExec: partitions=1, partition_sizes=[1] + +# Ordered and volatile AVG arguments are not decomposed. +query TT +EXPLAIN SELECT sum(c), avg(c ORDER BY a), sum(random()), avg(random()) FROM t; +---- +logical_plan +01)Aggregate: groupBy=[[]], aggr=[[sum(t.c), avg(t.c) ORDER BY [t.a ASC NULLS LAST], sum(random()), avg(random())]] +02)--TableScan: t projection=[a, c] +physical_plan +01)AggregateExec: mode=Single, gby=[], aggr=[sum(t.c), avg(t.c) ORDER BY [t.a ASC NULLS LAST], sum(random()), avg(random())] +02)--SortExec: expr=[a@0 ASC NULLS LAST], preserve_partitioning=[false] +03)----DataSourceExec: partitions=1, partition_sizes=[1] + +# Decimal AVG is not decomposed: its widened sum type and division semantics +# differ from SUM / COUNT. +statement ok +CREATE TABLE td(d DECIMAL(10,2) NOT NULL) AS VALUES (1.10), (2.20); + +query TT +EXPLAIN SELECT sum(d), avg(d) FROM td; +---- +logical_plan +01)Aggregate: groupBy=[[]], aggr=[[sum(td.d), avg(td.d)]] +02)--TableScan: td projection=[d] +physical_plan +01)AggregateExec: mode=Single, gby=[], aggr=[sum(td.d), avg(td.d)] +02)--DataSourceExec: partitions=1, partition_sizes=[1] + +query RR +SELECT sum(d), avg(d) FROM td; +---- +3.30 1.650000 + +# Internal aggregate aliases must not conflict with group output names, +# including names that look like generated aliases. +statement ok +CREATE TABLE name_collision_t( + a INT NOT NULL, + "sum(name_collision_t.a):1" INT NOT NULL, + "sum(name_collision_t.a)__temp__0" INT NOT NULL +) AS VALUES (1, 10, 100), (2, 20, 200); + +query IIRII rowsort +SELECT + "sum(name_collision_t.a):1", + "sum(name_collision_t.a)__temp__0", + avg(a), + sum(a), + count(*) +FROM name_collision_t +GROUP BY + "sum(name_collision_t.a):1", + "sum(name_collision_t.a)__temp__0"; +---- +10 100 1 1 1 +20 200 2 2 1 + +statement ok +DROP TABLE t; + +statement ok +DROP TABLE td; + +statement ok +DROP TABLE name_collision_t; diff --git a/datafusion/sqllogictest/test_files/optimizer_group_by_constant.slt b/datafusion/sqllogictest/test_files/optimizer_group_by_constant.slt index 8c0556547496a..2f4e60f63794a 100644 --- a/datafusion/sqllogictest/test_files/optimizer_group_by_constant.slt +++ b/datafusion/sqllogictest/test_files/optimizer_group_by_constant.slt @@ -60,9 +60,10 @@ FROM test_table t group by 1, 2, 3 ---- logical_plan -01)Aggregate: groupBy=[[Int64(123), Int64(456), Int64(789)]], aggr=[[count(Int64(1)), avg(t.c12)]] -02)--SubqueryAlias: t -03)----TableScan: test_table projection=[c12] +01)Projection: Int64(123), Int64(456), Int64(789), count(Int64(1)), sum(t.c12) / CAST(count(Int64(1)) AS Float64) AS avg(t.c12) +02)--Aggregate: groupBy=[[Int64(123), Int64(456), Int64(789)]], aggr=[[count(Int64(1)), sum(t.c12)]] +03)----SubqueryAlias: t +04)------TableScan: test_table projection=[c12] query TT EXPLAIN diff --git a/datafusion/substrait/src/logical_plan/consumer/rel/aggregate_rel.rs b/datafusion/substrait/src/logical_plan/consumer/rel/aggregate_rel.rs index 982a87d6d5e83..ff8b4f0a9a388 100644 --- a/datafusion/substrait/src/logical_plan/consumer/rel/aggregate_rel.rs +++ b/datafusion/substrait/src/logical_plan/consumer/rel/aggregate_rel.rs @@ -15,10 +15,11 @@ // specific language governing permissions and limitations // under the License. -use crate::logical_plan::consumer::{NameTracker, SubstraitConsumer}; +use crate::logical_plan::consumer::SubstraitConsumer; use crate::logical_plan::consumer::{from_substrait_agg_func, from_substrait_sorts}; use datafusion::common::{Column, DFSchemaRef, internal_err, not_impl_err}; use datafusion::logical_expr::builder::project; +use datafusion::logical_expr::utils::NameTracker; use datafusion::logical_expr::{ Aggregate, Expr, GroupingSet, LogicalPlan, LogicalPlanBuilder, }; diff --git a/datafusion/substrait/src/logical_plan/consumer/rel/mod.rs b/datafusion/substrait/src/logical_plan/consumer/rel/mod.rs index 038ada115b9d8..1bf0deb585648 100644 --- a/datafusion/substrait/src/logical_plan/consumer/rel/mod.rs +++ b/datafusion/substrait/src/logical_plan/consumer/rel/mod.rs @@ -38,10 +38,10 @@ pub use set_rel::*; pub use sort_rel::*; use crate::logical_plan::consumer::SubstraitConsumer; -use crate::logical_plan::consumer::utils::NameTracker; use async_recursion::async_recursion; use datafusion::common::{Column, not_impl_err, substrait_datafusion_err, substrait_err}; use datafusion::logical_expr::builder::project; +use datafusion::logical_expr::utils::NameTracker; use datafusion::logical_expr::{Expr, LogicalPlan, Projection}; use std::sync::Arc; use substrait::proto::rel::RelType; diff --git a/datafusion/substrait/src/logical_plan/consumer/rel/project_rel.rs b/datafusion/substrait/src/logical_plan/consumer/rel/project_rel.rs index 5aea6c809b701..dc1762b925823 100644 --- a/datafusion/substrait/src/logical_plan/consumer/rel/project_rel.rs +++ b/datafusion/substrait/src/logical_plan/consumer/rel/project_rel.rs @@ -16,11 +16,10 @@ // under the License. use crate::logical_plan::consumer::SubstraitConsumer; -use crate::logical_plan::consumer::utils::NameTracker; use async_recursion::async_recursion; use datafusion::common::{Column, not_impl_err}; use datafusion::logical_expr::builder::project; -use datafusion::logical_expr::utils::find_window_exprs; +use datafusion::logical_expr::utils::{NameTracker, find_window_exprs}; use datafusion::logical_expr::{Expr, LogicalPlan, LogicalPlanBuilder}; use std::collections::HashSet; use std::sync::Arc; diff --git a/datafusion/substrait/src/logical_plan/consumer/utils.rs b/datafusion/substrait/src/logical_plan/consumer/utils.rs index 824c79452d86e..5b079557b40a6 100644 --- a/datafusion/substrait/src/logical_plan/consumer/utils.rs +++ b/datafusion/substrait/src/logical_plan/consumer/utils.rs @@ -18,12 +18,11 @@ use crate::logical_plan::consumer::SubstraitConsumer; use datafusion::arrow::datatypes::{DataType, Field, Schema, TimeUnit, UnionFields}; use datafusion::common::{ - DFSchema, DFSchemaRef, TableReference, exec_err, not_impl_err, - substrait_datafusion_err, substrait_err, + DFSchema, DFSchemaRef, exec_err, not_impl_err, substrait_datafusion_err, + substrait_err, }; use datafusion::logical_expr::expr::Sort; use datafusion::logical_expr::{Cast, Expr, ExprSchemable}; -use std::collections::HashSet; use std::sync::Arc; use substrait::proto::SortField; use substrait::proto::sort_field::SortDirection; @@ -403,102 +402,6 @@ fn compatible_nullabilities( || (!datafusion_nullability && substrait_nullability) } -pub(super) struct NameTracker { - /// Tracks seen schema names (from expr.schema_name()). - /// Used to detect duplicates that would fail validate_unique_names. - seen_schema_names: HashSet, - /// Tracks column names that have been seen with a qualifier. - /// Used to detect ambiguous references (qualified + unqualified with same name). - qualified_names: HashSet, - /// Tracks column names that have been seen without a qualifier. - /// Used to detect ambiguous references. - unqualified_names: HashSet, -} - -impl NameTracker { - pub(super) fn new() -> Self { - NameTracker { - seen_schema_names: HashSet::default(), - qualified_names: HashSet::default(), - unqualified_names: HashSet::default(), - } - } - - /// Check if the expression would cause a conflict either in: - /// 1. validate_unique_names (duplicate schema_name) - /// 2. DFSchema::check_names (ambiguous reference) - fn would_conflict(&self, expr: &Expr) -> bool { - let (qualifier, name) = expr.qualified_name(); - let schema_name = expr.schema_name().to_string(); - self.would_conflict_inner((qualifier, &name), &schema_name) - } - - fn would_conflict_inner( - &self, - qualified_name: (Option, &str), - schema_name: &str, - ) -> bool { - // Check for duplicate schema_name (would fail validate_unique_names) - if self.seen_schema_names.contains(schema_name) { - return true; - } - - // Check for ambiguous reference (would fail DFSchema::check_names) - // This happens when a qualified field and unqualified field have the same name - let (qualifier, name) = qualified_name; - match qualifier { - Some(_) => { - // Adding a qualified name - conflicts if unqualified version exists - self.unqualified_names.contains(name) - } - None => { - // Adding an unqualified name - conflicts if qualified version exists - self.qualified_names.contains(name) - } - } - } - - fn insert(&mut self, expr: &Expr) { - let schema_name = expr.schema_name().to_string(); - self.seen_schema_names.insert(schema_name); - - let (qualifier, name) = expr.qualified_name(); - match qualifier { - Some(_) => { - self.qualified_names.insert(name); - } - None => { - self.unqualified_names.insert(name); - } - } - } - - pub(super) fn get_uniquely_named_expr( - &mut self, - expr: Expr, - ) -> datafusion::common::Result { - if !self.would_conflict(&expr) { - self.insert(&expr); - return Ok(expr); - } - - // Name collision - need to generate a unique alias - let schema_name = expr.schema_name().to_string(); - let mut counter = 0; - let candidate_name = loop { - let candidate_name = format!("{schema_name}__temp__{counter}"); - // .alias always produces an unqualified name so check for conflicts accordingly. - if !self.would_conflict_inner((None, &candidate_name), &candidate_name) { - break candidate_name; - } - counter += 1; - }; - let candidate_expr = expr.alias(&candidate_name); - self.insert(&candidate_expr); - Ok(candidate_expr) - } -} - /// Convert Substrait Sorts to DataFusion Exprs pub async fn from_substrait_sorts( consumer: &impl SubstraitConsumer, @@ -565,14 +468,13 @@ pub(crate) fn from_substrait_precision( #[cfg(test)] pub(crate) mod tests { - use super::{NameTracker, ensure_schema_compatibility, make_renamed_schema}; + use super::{ensure_schema_compatibility, make_renamed_schema}; use crate::extensions::Extensions; use crate::logical_plan::consumer::DefaultSubstraitConsumer; use datafusion::arrow::datatypes::{DataType, Field, Fields, Schema}; use datafusion::common::{DFSchema, TableReference}; use datafusion::error::Result; use datafusion::execution::SessionState; - use datafusion::logical_expr::{Expr, col}; use datafusion::prelude::SessionContext; use std::collections::HashMap; use std::sync::{Arc, LazyLock}; @@ -738,125 +640,6 @@ pub(crate) mod tests { Ok(()) } - #[test] - fn name_tracker_unique_names_pass_through() -> Result<()> { - let mut tracker = NameTracker::new(); - - // First expression should pass through unchanged - let expr1 = col("a"); - let result1 = tracker.get_uniquely_named_expr(expr1.clone())?; - assert_eq!(result1, col("a")); - - // Different name should also pass through unchanged - let expr2 = col("b"); - let result2 = tracker.get_uniquely_named_expr(expr2)?; - assert_eq!(result2, col("b")); - - Ok(()) - } - - #[test] - fn name_tracker_duplicate_schema_name_gets_alias() -> Result<()> { - let mut tracker = NameTracker::new(); - - // First expression with name "a" - let expr1 = col("a"); - let result1 = tracker.get_uniquely_named_expr(expr1)?; - assert_eq!(result1, col("a")); - - // Second expression with same name "a" should get aliased - let expr2 = col("a"); - let result2 = tracker.get_uniquely_named_expr(expr2)?; - assert_eq!(result2, col("a").alias("a__temp__0")); - - // Third expression with same name "a" should get a different alias - let expr3 = col("a"); - let result3 = tracker.get_uniquely_named_expr(expr3)?; - assert_eq!(result3, col("a").alias("a__temp__1")); - - Ok(()) - } - - #[test] - fn name_tracker_qualified_then_unqualified_conflicts() -> Result<()> { - let mut tracker = NameTracker::new(); - - // First: qualified column "table.a" - let qualified_col = - Expr::Column(datafusion::common::Column::new(Some("table"), "a")); - let result1 = tracker.get_uniquely_named_expr(qualified_col)?; - assert_eq!( - result1, - Expr::Column(datafusion::common::Column::new(Some("table"), "a")) - ); - - // Second: unqualified column "a" - should conflict (ambiguous reference) - let unqualified_col = col("a"); - let result2 = tracker.get_uniquely_named_expr(unqualified_col)?; - // Should be aliased to avoid ambiguous reference - assert_eq!(result2, col("a").alias("a__temp__0")); - - Ok(()) - } - - #[test] - fn name_tracker_unqualified_then_qualified_conflicts() -> Result<()> { - let mut tracker = NameTracker::new(); - - // First: unqualified column "a" - let unqualified_col = col("a"); - let result1 = tracker.get_uniquely_named_expr(unqualified_col)?; - assert_eq!(result1, col("a")); - - // Second: qualified column "table.a" - should conflict (ambiguous reference) - let qualified_col = - Expr::Column(datafusion::common::Column::new(Some("table"), "a")); - let result2 = tracker.get_uniquely_named_expr(qualified_col)?; - // Should be aliased to avoid ambiguous reference - assert_eq!( - result2, - Expr::Column(datafusion::common::Column::new(Some("table"), "a")) - .alias("table.a__temp__0") - ); - - Ok(()) - } - - #[test] - fn name_tracker_different_qualifiers_no_conflict() -> Result<()> { - let mut tracker = NameTracker::new(); - - // First: qualified column "table1.a" - let col1 = Expr::Column(datafusion::common::Column::new(Some("table1"), "a")); - let result1 = tracker.get_uniquely_named_expr(col1.clone())?; - assert_eq!(result1, col1); - - // Second: qualified column "table2.a" - different qualifier, different schema_name - // so should NOT conflict - let col2 = Expr::Column(datafusion::common::Column::new(Some("table2"), "a")); - let result2 = tracker.get_uniquely_named_expr(col2.clone())?; - assert_eq!(result2, col2); - - Ok(()) - } - - #[test] - fn name_tracker_aliased_expressions() -> Result<()> { - let mut tracker = NameTracker::new(); - - // First: col("x").alias("result") - let expr1 = col("x").alias("result"); - let result1 = tracker.get_uniquely_named_expr(expr1.clone())?; - assert_eq!(result1, col("x").alias("result")); - - // Second: col("y").alias("result") - same alias name, should conflict - let expr2 = col("y").alias("result"); - let result2 = tracker.get_uniquely_named_expr(expr2)?; - assert_eq!(result2, col("y").alias("result").alias("result__temp__0")); - - Ok(()) - } - fn schema_with_struct_inner(inner_nullable: bool) -> DFSchema { let inner = Field::new("inner", DataType::Int32, inner_nullable); let outer = Field::new("s", DataType::Struct(Fields::from(vec![inner])), false);