diff --git a/datafusion/expr/src/udaf.rs b/datafusion/expr/src/udaf.rs index 8f7e9cc6cfc2b..b04c70f6c52b5 100644 --- a/datafusion/expr/src/udaf.rs +++ b/datafusion/expr/src/udaf.rs @@ -253,6 +253,16 @@ impl AggregateUDF { self.inner.groups_accumulator_supported(args) } + /// See [`AggregateUDFImpl::groups_accumulator_supported_for_types`] for more details. + pub fn groups_accumulator_supported_for_types( + &self, + arg_types: &[DataType], + is_distinct: bool, + ) -> bool { + self.inner + .groups_accumulator_supported_for_types(arg_types, is_distinct) + } + /// See [`AggregateUDFImpl::create_groups_accumulator`] for more details. pub fn create_groups_accumulator( &self, @@ -617,6 +627,32 @@ pub trait AggregateUDFImpl: Debug + DynEq + DynHash + Send + Sync + Any { false } + /// The same question as [`Self::groups_accumulator_supported`], asked with + /// only the information a logical plan carries. + /// + /// Physical planning has an [`AccumulatorArgs`] to ask with; an optimizer + /// rule does not, and cannot fabricate one, so this is how a logical + /// caller learns whether a call would get a specialized + /// [`GroupsAccumulator`] or fall back to one boxed [`Accumulator`] per + /// group in `GroupsAccumulatorAdapter`. That distinction is worth a rule + /// changing its mind over: the adapter's per-group state can be orders of + /// magnitude larger. + /// + /// The default is `false`, matching the default of + /// [`Self::groups_accumulator_supported`]. An implementation that + /// overrides that one and whose answer is decided by the argument types + /// and `DISTINCT` alone should override this one too, and have the + /// physical method call it so the two cannot disagree. An implementation + /// whose answer needs more than the argument types should leave this at + /// `false`, which claims nothing. + fn groups_accumulator_supported_for_types( + &self, + _arg_types: &[DataType], + _is_distinct: bool, + ) -> bool { + false + } + /// Return a specialized [`GroupsAccumulator`] that manages state /// for all groups. /// @@ -1552,6 +1588,15 @@ impl AggregateUDFImpl for AliasedAggregateUDFImpl { self.inner.groups_accumulator_supported(args) } + fn groups_accumulator_supported_for_types( + &self, + arg_types: &[DataType], + is_distinct: bool, + ) -> bool { + self.inner + .groups_accumulator_supported_for_types(arg_types, is_distinct) + } + fn create_groups_accumulator( &self, args: AccumulatorArgs, diff --git a/datafusion/functions-aggregate/src/count.rs b/datafusion/functions-aggregate/src/count.rs index 4c95ea431809e..ad020cfa8f179 100644 --- a/datafusion/functions-aggregate/src/count.rs +++ b/datafusion/functions-aggregate/src/count.rs @@ -346,14 +346,30 @@ impl AggregateUDFImpl for Count { } fn groups_accumulator_supported(&self, args: AccumulatorArgs) -> bool { - if args.exprs.len() != 1 { + // The answer depends on nothing but the argument types and `DISTINCT`, + // so defer to the logical form and keep one list of supported types. + let arg_types = args + .expr_fields + .iter() + .map(|field| field.data_type().clone()) + .collect::>(); + self.groups_accumulator_supported_for_types(&arg_types, args.is_distinct) + } + + fn groups_accumulator_supported_for_types( + &self, + arg_types: &[DataType], + is_distinct: bool, + ) -> bool { + if arg_types.len() != 1 { return false; } - if !args.is_distinct { + if !is_distinct { return true; } + // Keep in step with `create_distinct_count_groups_accumulator`. matches!( - args.expr_fields[0].data_type(), + arg_types[0], DataType::Int8 | DataType::Int16 | DataType::Int32 diff --git a/datafusion/optimizer/src/single_distinct_to_groupby.rs b/datafusion/optimizer/src/single_distinct_to_groupby.rs index 00c8fab228117..155f6b5a955fb 100644 --- a/datafusion/optimizer/src/single_distinct_to_groupby.rs +++ b/datafusion/optimizer/src/single_distinct_to_groupby.rs @@ -23,14 +23,18 @@ use crate::optimizer::ApplyOrder; use crate::{OptimizerConfig, OptimizerRule}; use datafusion_common::{ - DataFusionError, HashSet, Result, assert_eq_or_internal_err, tree_node::Transformed, + DFSchema, DataFusionError, HashSet, Result, assert_eq_or_internal_err, + tree_node::Transformed, }; use datafusion_expr::builder::project; use datafusion_expr::expr::AggregateFunctionParams; +use datafusion_expr::expr_schema::ExprSchemable; use datafusion_expr::{ - Expr, col, + AggregateUDF, Expr, col, expr::AggregateFunction, + lit, logical_plan::{Aggregate, LogicalPlan}, + when, }; /// single distinct to group by optimizer rule @@ -49,6 +53,38 @@ use datafusion_expr::{ /// ) /// GROUP BY a /// ``` +/// +/// A non-distinct `count` is also allowed alongside the distinct aggregate. It +/// is the one supported function whose outer phase is a *different* function: +/// the inner group-by counts rows per `(group, distinct value)` pair and the +/// outer phase adds those partial counts up with `sum`. +/// +/// ```text +/// Before: +/// SELECT a, count(*), count(DISTINCT b) +/// FROM t +/// GROUP BY a +/// +/// After: +/// SELECT a, +/// CASE WHEN sum(alias2) IS NOT NULL THEN sum(alias2) ELSE 0 END, +/// count(alias1) +/// FROM ( +/// SELECT a, b as alias1, count(*) as alias2 +/// FROM t +/// GROUP BY a, b +/// ) +/// GROUP BY a +/// ``` +/// +/// The `CASE` covers the one input on which the two phases disagree: over an +/// empty input the inner group by produces no rows at all, and a `sum` of no +/// rows is NULL where `count` is 0. +/// +/// That `count` is allowed only when the distinct argument has no specialized +/// `GroupsAccumulator` of its own, so that the rewrite is taking the distinct +/// aggregate off `GroupsAccumulatorAdapter` rather than only adding an inner +/// group by. See `rewrite_pays_for_count` for the measurements behind that. #[derive(Default, Debug)] pub struct SingleDistinctToGroupBy {} @@ -61,10 +97,43 @@ impl SingleDistinctToGroupBy { } } +/// The pair of functions used to compute a non-distinct `count` in two phases: +/// `count` identifies the aggregates that need the treatment, `sum` combines +/// the partial counts the inner group-by produces. +/// +/// Both are resolved from the session's function registry, so a plan built +/// without one keeps the previous behaviour of bailing out on `count`. +struct CountRollup { + count: Arc, + sum: Arc, +} + +impl CountRollup { + fn try_new(config: &dyn OptimizerConfig) -> Option { + let registry = config.function_registry()?; + Some(Self { + count: registry.udaf("count").ok()?, + sum: registry.udaf("sum").ok()?, + }) + } + + /// Whether `func` is the registry's `count`, and so decomposes into + /// `sum` over per-partition counts. + fn is_count(&self, func: &AggregateUDF) -> bool { + self.count.as_ref() == func + } +} + /// Check whether all aggregate exprs are distinct on a single field. -fn is_single_distinct_agg(aggr_expr: &[Expr]) -> Result { +fn is_single_distinct_agg( + aggr_expr: &[Expr], + input_schema: &DFSchema, + count_rollup: Option<&CountRollup>, +) -> Result { let mut fields_set = HashSet::new(); let mut aggregate_count = 0; + let mut distinct_aggs = vec![]; + let mut has_count_rollup = false; for expr in aggr_expr { if let Expr::AggregateFunction(AggregateFunction { func, @@ -86,6 +155,9 @@ fn is_single_distinct_agg(aggr_expr: &[Expr]) -> Result { for e in args { fields_set.insert(e); } + distinct_aggs.push((func, args)); + } else if count_rollup.is_some_and(|rollup| rollup.is_count(func)) { + has_count_rollup = true; } else if func.name() != "sum" && func.name().to_lowercase() != "min" && func.name().to_lowercase() != "max" @@ -96,7 +168,45 @@ fn is_single_distinct_agg(aggr_expr: &[Expr]) -> Result { return Ok(false); } } - Ok(aggregate_count == aggr_expr.len() && fields_set.len() == 1) + if aggregate_count != aggr_expr.len() || fields_set.len() != 1 { + return Ok(false); + } + if has_count_rollup && !rewrite_pays_for_count(&distinct_aggs, input_schema)? { + return Ok(false); + } + Ok(true) +} + +/// Whether the rewrite is worth extending to a plan that only qualifies because +/// of the non-distinct `count`. +/// +/// The rewrite is not free: every other aggregate moves down to the inner group +/// by, which has a row per `(group, distinct value)` pair rather than per group, +/// and each one keeps its state at that finer grain. What pays for it is taking +/// the distinct aggregate off `GroupsAccumulatorAdapter`, whose one boxed +/// accumulator per group is the expensive shape. A distinct aggregate that +/// already has a specialized `GroupsAccumulator` never went near the adapter, so +/// there is nothing to buy and only the inner group by to pay for: ClickBench +/// Q22, whose `count(DISTINCT "UserID")` is over an `Int64`, measured a 132% +/// increase in peak memory pool reservation when the rewrite applied to it. +/// +/// The existing tolerance of a non-distinct `sum`, `min` or `max` predates this +/// and is left alone: narrowing it would change plans that have always been +/// rewritten, which no measurement here calls for. +fn rewrite_pays_for_count( + distinct_aggs: &[(&Arc, &Vec)], + input_schema: &DFSchema, +) -> Result { + for (func, args) in distinct_aggs { + let arg_types = args + .iter() + .map(|arg| arg.get_type(input_schema)) + .collect::>>()?; + if !func.groups_accumulator_supported_for_types(&arg_types, true) { + return Ok(true); + } + } + Ok(false) } /// Check if the first expr is [Expr::GroupingSet]. @@ -120,8 +230,12 @@ impl OptimizerRule for SingleDistinctToGroupBy { fn rewrite( &self, plan: LogicalPlan, - _config: &dyn OptimizerConfig, + config: &dyn OptimizerConfig, ) -> Result, DataFusionError> { + if !matches!(plan, LogicalPlan::Aggregate(_)) { + return Ok(Transformed::no(plan)); + } + let count_rollup = CountRollup::try_new(config); match plan { LogicalPlan::Aggregate(Aggregate { input, @@ -129,8 +243,11 @@ impl OptimizerRule for SingleDistinctToGroupBy { schema, group_expr, .. - }) if is_single_distinct_agg(&aggr_expr)? - && !contains_grouping_set(&group_expr) => + }) if is_single_distinct_agg( + &aggr_expr, + input.schema(), + count_rollup.as_ref(), + )? && !contains_grouping_set(&group_expr) => { let group_size = group_expr.len(); // alias all original group_by exprs @@ -177,7 +294,11 @@ impl OptimizerRule for SingleDistinctToGroupBy { let mut index = 1; let mut group_fields_set = HashSet::new(); let mut inner_aggr_exprs = vec![]; - let outer_aggr_exprs = aggr_expr + // Each aggregate yields the expression the outer `Aggregate` + // computes and the expression the projection above it selects. + // They differ only for `count`, whose projection restores the + // zero that `sum` reports as NULL over an empty input. + let (outer_aggr_exprs, outer_proj_exprs): (Vec, Vec) = aggr_expr .into_iter() .map(|aggr_expr| match aggr_expr { Expr::AggregateFunction(AggregateFunction { @@ -204,18 +325,32 @@ impl OptimizerRule for SingleDistinctToGroupBy { inner_group_exprs .push(arg.alias(SINGLE_DISTINCT_ALIAS)); } - Ok(Expr::AggregateFunction(AggregateFunction::new_udf( - func, - vec![col(SINGLE_DISTINCT_ALIAS)], - false, // intentional to remove distinct here - filter, - order_by, - null_treatment, - ))) + let outer = + Expr::AggregateFunction(AggregateFunction::new_udf( + func, + vec![col(SINGLE_DISTINCT_ALIAS)], + false, // intentional to remove distinct here + filter, + order_by, + null_treatment, + )); + Ok((outer.clone(), outer)) // if the aggregate function is not distinct, we need to rewrite it like two phase aggregation } else { index += 1; let alias_str = format!("alias{index}"); + // `count` is the one function whose two phases + // use different aggregates: the inner group-by + // counts the rows of each `(group, distinct + // value)` partition and the outer phase adds + // those partial counts up. + let rollup = count_rollup + .as_ref() + .filter(|rollup| rollup.is_count(&func)); + let outer_func = match rollup { + Some(rollup) => Arc::clone(&rollup.sum), + None => Arc::clone(&func), + }; inner_aggr_exprs.push( Expr::AggregateFunction(AggregateFunction::new_udf( Arc::clone(&func), @@ -227,19 +362,34 @@ impl OptimizerRule for SingleDistinctToGroupBy { )) .alias(&alias_str), ); - Ok(Expr::AggregateFunction(AggregateFunction::new_udf( - func, - vec![col(&alias_str)], - false, - None, - vec![], - None, - ))) + let outer = + Expr::AggregateFunction(AggregateFunction::new_udf( + outer_func, + vec![col(&alias_str)], + false, + None, + vec![], + None, + )); + if rollup.is_none() { + return Ok((outer.clone(), outer)); + } + // The inner group-by produces no rows at all + // for an empty input, and `sum` reports that as + // NULL where `count` reports 0. Restore the 0, + // which also keeps the column non-nullable as + // `count` had it. + let proj = + when(outer.clone().is_not_null(), outer.clone()) + .otherwise(lit(0_i64))?; + Ok((outer, proj)) } } - _ => Ok(aggr_expr), + _ => Ok((aggr_expr.clone(), aggr_expr)), }) - .collect::>>()?; + .collect::>>()? + .into_iter() + .unzip(); // construct the inner AggrPlan let inner_agg = LogicalPlan::Aggregate(Aggregate::try_new( @@ -265,13 +415,11 @@ impl OptimizerRule for SingleDistinctToGroupBy { } None => group_expr, }) - .chain(outer_aggr_exprs.iter().cloned().enumerate().map( - |(idx, expr)| { - let idx = idx + group_size; - let (qualifier, field) = schema.qualified_field(idx); - expr.alias_qualified(qualifier.cloned(), field.name()) - }, - )) + .chain(outer_proj_exprs.into_iter().enumerate().map(|(idx, expr)| { + let idx = idx + group_size; + let (qualifier, field) = schema.qualified_field(idx); + expr.alias_qualified(qualifier.cloned(), field.name()) + })) .collect(); let outer_aggr = LogicalPlan::Aggregate(Aggregate::try_new( @@ -291,9 +439,18 @@ mod tests { use super::*; use crate::assert_optimized_plan_eq_display_indent_snapshot; use crate::test::*; + use crate::{Optimizer, OptimizerContext}; + use arrow::datatypes::{DataType, Field, Schema}; + use chrono::{DateTime, Utc}; + use datafusion_common::alias::AliasGenerator; + use datafusion_common::config::ConfigOptions; use datafusion_expr::ExprFunctionExt; use datafusion_expr::expr::GroupingSet; - use datafusion_expr::{lit, logical_plan::builder::LogicalPlanBuilder}; + use datafusion_expr::registry::{FunctionRegistry, MemoryFunctionRegistry}; + use datafusion_expr::{ + lit, + logical_plan::builder::{LogicalPlanBuilder, table_scan}, + }; use datafusion_functions_aggregate::count::count_udaf; use datafusion_functions_aggregate::expr_fn::{count, count_distinct, max, min, sum}; use datafusion_functions_aggregate::min_max::max_udaf; @@ -310,17 +467,84 @@ mod tests { )) } + fn count_star() -> Expr { + Expr::AggregateFunction(AggregateFunction::new_udf( + count_udaf(), + vec![lit(1_i64)], + false, + None, + vec![], + None, + )) + } + + /// `test` with a `Utf8` `b`, the column the distinct aggregate reads. + /// + /// `count(DISTINCT b)` over a string has no specialized + /// `GroupsAccumulator`, so it is the case a non-distinct `count` may join. + /// The `UInt32` `b` of [`test_table_scan`] is the case it may not. + fn test_table_scan_utf8_b() -> Result { + let schema = Schema::new(vec![ + Field::new("a", DataType::UInt32, false), + Field::new("b", DataType::Utf8, false), + Field::new("c", DataType::UInt32, false), + ]); + table_scan(Some("test"), &schema, None)?.build() + } + + /// An [`OptimizerConfig`] that resolves functions the way a session does. + /// The rule needs `count` and `sum` from the registry to rewrite a + /// non-distinct `count`. + #[derive(Debug)] + struct RegistryOptimizerContext { + inner: OptimizerContext, + registry: MemoryFunctionRegistry, + } + + impl RegistryOptimizerContext { + fn new() -> Self { + let mut registry = MemoryFunctionRegistry::new(); + registry.register_udaf(count_udaf()).unwrap(); + registry.register_udaf(sum_udaf()).unwrap(); + Self { + inner: OptimizerContext::new(), + registry, + } + } + } + + impl OptimizerConfig for RegistryOptimizerContext { + fn query_execution_start_time(&self) -> Option> { + self.inner.query_execution_start_time() + } + + fn alias_generator(&self) -> &Arc { + self.inner.alias_generator() + } + + fn options(&self) -> Arc { + self.inner.options() + } + + fn function_registry(&self) -> Option<&dyn FunctionRegistry> { + Some(&self.registry) + } + } + macro_rules! assert_optimized_plan_equal { ( $plan:expr, @ $expected:literal $(,)? ) => {{ - let rule: Arc = Arc::new(SingleDistinctToGroupBy::new()); - assert_optimized_plan_eq_display_indent_snapshot!( - rule, - $plan, - @ $expected, - ) + let rule: Arc = + Arc::new(SingleDistinctToGroupBy::new()); + let optimizer = Optimizer::with_rules(vec![rule]); + let optimized_plan = optimizer + .optimize($plan, &RegistryOptimizerContext::new(), |_, _| {}) + .expect("failed to optimize plan"); + insta::assert_snapshot!(optimized_plan.display_indent_schema(), @ $expected); + + Ok::<(), DataFusionError>(()) }}; } @@ -514,6 +738,30 @@ mod tests { #[test] fn distinct_and_common() -> Result<()> { + let table_scan = test_table_scan_utf8_b()?; + + let plan = LogicalPlanBuilder::from(table_scan) + .aggregate( + vec![col("a")], + vec![count_distinct(col("b")), count(col("c"))], + )? + .build()?; + + // Should work: the non-distinct count becomes a sum of the counts the + // inner group-by produces per (test.a, test.b) pair + assert_optimized_plan_equal!( + plan, + @r" + Projection: test.a, count(alias1) AS count(DISTINCT test.b), CASE WHEN sum(alias2) IS NOT NULL THEN sum(alias2) ELSE Int64(0) END AS count(test.c) [a:UInt32, count(DISTINCT test.b):Int64, count(test.c):Int64] + Aggregate: groupBy=[[test.a]], aggr=[[count(alias1), sum(alias2)]] [a:UInt32, count(alias1):Int64, sum(alias2):Int64;N] + Aggregate: groupBy=[[test.a, test.b AS alias1]], aggr=[[count(test.c) AS alias2]] [a:UInt32, alias1:Utf8, alias2:Int64] + TableScan: test [a:UInt32, b:Utf8, c:UInt32] + " + ) + } + + #[test] + fn distinct_and_common_over_a_natively_supported_type() -> Result<()> { let table_scan = test_table_scan()?; let plan = LogicalPlanBuilder::from(table_scan) @@ -523,7 +771,9 @@ mod tests { )? .build()?; - // Do nothing + // Should not work: `count(DISTINCT b)` over a `UInt32` has its own + // `GroupsAccumulator`, so the rewrite would add an inner group-by + // without taking anything off `GroupsAccumulatorAdapter` assert_optimized_plan_equal!( plan, @r" @@ -533,6 +783,31 @@ mod tests { ) } + #[test] + fn distinct_over_a_natively_supported_type_without_a_count() -> Result<()> { + let table_scan = test_table_scan()?; + + let plan = LogicalPlanBuilder::from(table_scan) + .aggregate( + vec![col("a")], + vec![count_distinct(col("b")), sum(col("c"))], + )? + .build()?; + + // Should work: the gate covers only the `count` this change added, so a + // plan that already qualified through `sum`, `min` or `max` is rewritten + // exactly as before + assert_optimized_plan_equal!( + plan, + @r" + Projection: test.a, count(alias1) AS count(DISTINCT test.b), sum(alias2) AS sum(test.c) [a:UInt32, count(DISTINCT test.b):Int64, sum(test.c):UInt64;N] + Aggregate: groupBy=[[test.a]], aggr=[[count(alias1), sum(alias2)]] [a:UInt32, count(alias1):Int64, sum(alias2):UInt64;N] + Aggregate: groupBy=[[test.a, test.b AS alias1]], aggr=[[sum(test.c) AS alias2]] [a:UInt32, alias1:UInt32, alias2:UInt64;N] + TableScan: test [a:UInt32, b:UInt32, c:UInt32] + " + ) + } + #[test] fn group_by_with_expr() -> Result<()> { let table_scan = test_table_scan().unwrap(); @@ -752,4 +1027,128 @@ mod tests { " ) } + + #[test] + fn count_star_and_distinct_without_groupby() -> Result<()> { + let table_scan = test_table_scan_utf8_b()?; + + let plan = LogicalPlanBuilder::from(table_scan) + .aggregate( + Vec::::new(), + vec![count_star(), count_distinct(col("b"))], + )? + .build()?; + + // Should work. Without a group by the outer aggregate sees no rows at + // all for an empty input, so the projection has to turn the NULL that + // `sum` reports there back into the 0 `count` reports. + assert_optimized_plan_equal!( + plan, + @r" + Projection: CASE WHEN sum(alias2) IS NOT NULL THEN sum(alias2) ELSE Int64(0) END AS count(Int64(1)), count(alias1) AS count(DISTINCT test.b) [count(Int64(1)):Int64, count(DISTINCT test.b):Int64] + Aggregate: groupBy=[[]], aggr=[[sum(alias2), count(alias1)]] [sum(alias2):Int64;N, count(alias1):Int64] + Aggregate: groupBy=[[test.b AS alias1]], aggr=[[count(Int64(1)) AS alias2]] [alias1:Utf8, alias2:Int64] + TableScan: test [a:UInt32, b:Utf8, c:UInt32] + " + ) + } + + #[test] + fn count_star_min_max_sum_and_distinct_with_groupby() -> Result<()> { + let table_scan = test_table_scan_utf8_b()?; + + let plan = LogicalPlanBuilder::from(table_scan) + .aggregate( + vec![col("a")], + vec![ + count_star(), + count_distinct(col("b")), + min(col("c")), + max(col("c")), + sum(col("c")), + ], + )? + .build()?; + + // Should work: this is the shape a `count(*)` alongside a + // `count(DISTINCT ..)`, a min, a max and a sum produces + assert_optimized_plan_equal!( + plan, + @r" + Projection: test.a, CASE WHEN sum(alias2) IS NOT NULL THEN sum(alias2) ELSE Int64(0) END AS count(Int64(1)), count(alias1) AS count(DISTINCT test.b), min(alias3) AS min(test.c), max(alias4) AS max(test.c), sum(alias5) AS sum(test.c) [a:UInt32, count(Int64(1)):Int64, count(DISTINCT test.b):Int64, min(test.c):UInt32;N, max(test.c):UInt32;N, sum(test.c):UInt64;N] + Aggregate: groupBy=[[test.a]], aggr=[[sum(alias2), count(alias1), min(alias3), max(alias4), sum(alias5)]] [a:UInt32, sum(alias2):Int64;N, count(alias1):Int64, min(alias3):UInt32;N, max(alias4):UInt32;N, sum(alias5):UInt64;N] + Aggregate: groupBy=[[test.a, test.b AS alias1]], aggr=[[count(Int64(1)) AS alias2, min(test.c) AS alias3, max(test.c) AS alias4, sum(test.c) AS alias5]] [a:UInt32, alias1:Utf8, alias2:Int64, alias3:UInt32;N, alias4:UInt32;N, alias5:UInt64;N] + TableScan: test [a:UInt32, b:Utf8, c:UInt32] + " + ) + } + + #[test] + fn count_with_filter_is_not_rewritten() -> Result<()> { + let table_scan = test_table_scan()?; + + // count(a) FILTER (WHERE a > 5) + let expr = count_udaf() + .call(vec![col("a")]) + .filter(col("a").gt(lit(5))) + .build()?; + let plan = LogicalPlanBuilder::from(table_scan) + .aggregate(vec![col("c")], vec![expr, count_distinct(col("b"))])? + .build()?; + + // Do nothing: the filter would have to be applied per input row, but + // the inner aggregate has already collapsed them + assert_optimized_plan_equal!( + plan, + @r" + Aggregate: groupBy=[[test.c]], aggr=[[count(test.a) FILTER (WHERE test.a > Int32(5)), count(DISTINCT test.b)]] [c:UInt32, count(test.a) FILTER (WHERE test.a > Int32(5)):Int64, count(DISTINCT test.b):Int64] + TableScan: test [a:UInt32, b:UInt32, c:UInt32] + " + ) + } + + #[test] + fn count_with_order_by_is_not_rewritten() -> Result<()> { + let table_scan = test_table_scan()?; + + // count(a ORDER BY a) + let expr = count_udaf() + .call(vec![col("a")]) + .order_by(vec![col("a").sort(true, false)]) + .build()?; + let plan = LogicalPlanBuilder::from(table_scan) + .aggregate(vec![col("c")], vec![expr, count_distinct(col("b"))])? + .build()?; + + // Do nothing + assert_optimized_plan_equal!( + plan, + @r" + Aggregate: groupBy=[[test.c]], aggr=[[count(test.a) ORDER BY [test.a ASC NULLS LAST], count(DISTINCT test.b)]] [c:UInt32, count(test.a) ORDER BY [test.a ASC NULLS LAST]:Int64, count(DISTINCT test.b):Int64] + TableScan: test [a:UInt32, b:UInt32, c:UInt32] + " + ) + } + + #[test] + fn count_without_function_registry_is_not_rewritten() -> Result<()> { + let table_scan = test_table_scan()?; + + let plan = LogicalPlanBuilder::from(table_scan) + .aggregate(vec![col("a")], vec![count_star(), count_distinct(col("b"))])? + .build()?; + + // Do nothing: without a registry the rule cannot resolve the `sum` the + // outer phase needs + let rule: Arc = + Arc::new(SingleDistinctToGroupBy::new()); + assert_optimized_plan_eq_display_indent_snapshot!( + rule, + plan, + @r" + Aggregate: groupBy=[[test.a]], aggr=[[count(Int64(1)), count(DISTINCT test.b)]] [a:UInt32, count(Int64(1)):Int64, count(DISTINCT test.b):Int64] + TableScan: test [a:UInt32, b:UInt32, c:UInt32] + ", + ) + } } diff --git a/datafusion/sqllogictest/test_files/single_distinct_to_groupby.slt b/datafusion/sqllogictest/test_files/single_distinct_to_groupby.slt new file mode 100644 index 0000000000000..ffecec7b4ba23 --- /dev/null +++ b/datafusion/sqllogictest/test_files/single_distinct_to_groupby.slt @@ -0,0 +1,410 @@ +# 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 the `single_distinct_aggregation_to_group_by` rule, which rewrites a +# single `AGG(DISTINCT x)` into a two phase group by. Every result here is +# asserted twice: once with the optimizer disabled +# (`datafusion.optimizer.max_passes = 0`, so the aggregate runs as written) and +# once with it enabled. The two must agree. +# +# A non-distinct `count` may only join the distinct aggregate when the distinct +# argument has no specialized `GroupsAccumulator` of its own. `x` is the +# `VARCHAR` that qualifies, `xi` the `INT` that does not; both hold the same +# values so the two sides can be compared directly. + +statement ok +CREATE TABLE t(g INT, x VARCHAR, xi INT, v INT); + +# g = 1: four rows, two distinct x, one NULL x, one NULL v +# g = 2: two rows, x is NULL throughout +# g = 3: a single row +statement ok +INSERT INTO t VALUES + (1, '10', 10, 100), + (1, '10', 10, NULL), + (1, NULL, NULL, 5), + (1, '20', 20, 7), + (2, NULL, NULL, NULL), + (2, NULL, NULL, 3), + (3, '30', 30, 1); + +statement ok +CREATE TABLE empty_t(g INT, x VARCHAR, xi INT, v INT); + +statement ok +set datafusion.explain.logical_plan_only = true; + +######## +# The rewrite fires for a non-distinct count next to the distinct aggregate +######## + +query TT +EXPLAIN SELECT g, count(*) AS records, count(DISTINCT x) AS distinct_x FROM t GROUP BY g; +---- +logical_plan +01)Projection: t.g, CASE WHEN sum(alias2) IS NOT NULL THEN sum(alias2) ELSE Int64(0) END AS records, count(alias1) AS distinct_x +02)--Aggregate: groupBy=[[t.g]], aggr=[[sum(alias2), count(alias1)]] +03)----Aggregate: groupBy=[[t.g, t.x AS alias1]], aggr=[[count(Int64(1)) AS alias2]] +04)------TableScan: t projection=[g, x] + +# The production shape: a count next to min, max, sum and the distinct count +query TT +EXPLAIN +SELECT g, count(*) AS records, count(DISTINCT x) AS distinct_x, + min(v) AS min_v, max(v) AS max_v, sum(v) AS sum_v +FROM t GROUP BY g; +---- +logical_plan +01)Projection: t.g, CASE WHEN sum(alias2) IS NOT NULL THEN sum(alias2) ELSE Int64(0) END AS records, count(alias1) AS distinct_x, min(alias3) AS min_v, max(alias4) AS max_v, sum(alias5) AS sum_v +02)--Aggregate: groupBy=[[t.g]], aggr=[[sum(alias2), count(alias1), min(alias3), max(alias4), sum(alias5)]] +03)----Aggregate: groupBy=[[t.g, t.x AS alias1]], aggr=[[count(Int64(1)) AS alias2, min(t.v) AS alias3, max(t.v) AS alias4, sum(CAST(t.v AS Int64)) AS alias5]] +04)------TableScan: t projection=[g, x, v] + +# A count with a FILTER still blocks the rewrite: the filter is per input row, +# and the inner aggregate has already collapsed those rows +query TT +EXPLAIN SELECT g, count(*) FILTER (WHERE v > 1) AS records, count(DISTINCT x) AS distinct_x FROM t GROUP BY g; +---- +logical_plan +01)Projection: t.g, count(Int64(1)) FILTER (WHERE t.v > Int64(1)) AS count(*) FILTER (WHERE t.v > Int64(1)) AS records, count(DISTINCT t.x) AS distinct_x +02)--Aggregate: groupBy=[[t.g]], aggr=[[count(Int64(1)) FILTER (WHERE t.v > Int32(1)) AS count(Int64(1)) FILTER (WHERE t.v > Int64(1)), count(DISTINCT t.x)]] +03)----TableScan: t projection=[g, x, v] + +# An unsupported non-distinct aggregate still blocks the rewrite +query TT +EXPLAIN SELECT g, avg(v) AS avg_v, count(DISTINCT x) AS distinct_x FROM t GROUP BY g; +---- +logical_plan +01)Projection: t.g, avg(t.v) AS avg_v, count(DISTINCT t.x) AS distinct_x +02)--Aggregate: groupBy=[[t.g]], aggr=[[avg(CAST(t.v AS Float64)), count(DISTINCT t.x)]] +03)----TableScan: t projection=[g, x, v] + +######## +# The other side of the gate: a distinct argument that has its own +# `GroupsAccumulator` keeps the count out +######## + +# `count(DISTINCT xi)` over an INT is served by +# `PrimitiveDistinctCountGroupsAccumulator`, so the rewrite would add an inner +# group by without taking anything off `GroupsAccumulatorAdapter` +query TT +EXPLAIN SELECT g, count(*) AS records, count(DISTINCT xi) AS distinct_xi FROM t GROUP BY g; +---- +logical_plan +01)Projection: t.g, count(Int64(1)) AS count(*) AS records, count(DISTINCT t.xi) AS distinct_xi +02)--Aggregate: groupBy=[[t.g]], aggr=[[count(Int64(1)), count(DISTINCT t.xi)]] +03)----TableScan: t projection=[g, xi] + +# The gate covers only the count. A plan that already qualified through sum, +# min or max is still rewritten over the same INT column +query TT +EXPLAIN SELECT g, sum(v) AS sum_v, count(DISTINCT xi) AS distinct_xi FROM t GROUP BY g; +---- +logical_plan +01)Projection: t.g, sum(alias2) AS sum_v, count(alias1) AS distinct_xi +02)--Aggregate: groupBy=[[t.g]], aggr=[[sum(alias2), count(alias1)]] +03)----Aggregate: groupBy=[[t.g, t.xi AS alias1]], aggr=[[sum(CAST(t.v AS Int64)) AS alias2]] +04)------TableScan: t projection=[g, xi, v] + +statement ok +set datafusion.explain.logical_plan_only = false; + +######## +# count(*) vs count(col) vs count(1), grouped +######## + +statement ok +set datafusion.optimizer.max_passes = 0; + +query IIIIII +SELECT g, count(*) AS star, count(1) AS one, count(x) AS cnt_x, count(v) AS cnt_v, + count(DISTINCT x) AS distinct_x +FROM t GROUP BY g ORDER BY g; +---- +1 4 4 3 3 2 +2 2 2 0 1 0 +3 1 1 1 1 1 + +statement ok +set datafusion.optimizer.max_passes = 3; + +query IIIIII +SELECT g, count(*) AS star, count(1) AS one, count(x) AS cnt_x, count(v) AS cnt_v, + count(DISTINCT x) AS distinct_x +FROM t GROUP BY g ORDER BY g; +---- +1 4 4 3 3 2 +2 2 2 0 1 0 +3 1 1 1 1 1 + +######## +# count(*) vs count(col) vs count(1), no GROUP BY +######## + +statement ok +set datafusion.optimizer.max_passes = 0; + +query IIIII +SELECT count(*) AS star, count(1) AS one, count(x) AS cnt_x, count(v) AS cnt_v, + count(DISTINCT x) AS distinct_x +FROM t; +---- +7 7 4 5 3 + +statement ok +set datafusion.optimizer.max_passes = 3; + +query IIIII +SELECT count(*) AS star, count(1) AS one, count(x) AS cnt_x, count(v) AS cnt_v, + count(DISTINCT x) AS distinct_x +FROM t; +---- +7 7 4 5 3 + +######## +# The distinct column is NULL for every row of the group +######## + +statement ok +set datafusion.optimizer.max_passes = 0; + +query IIIII +SELECT count(*) AS star, count(x) AS cnt_x, count(v) AS cnt_v, count(DISTINCT x) AS distinct_x, + max(v) AS max_v +FROM t WHERE g = 2; +---- +2 0 1 0 3 + +statement ok +set datafusion.optimizer.max_passes = 3; + +query IIIII +SELECT count(*) AS star, count(x) AS cnt_x, count(v) AS cnt_v, count(DISTINCT x) AS distinct_x, + max(v) AS max_v +FROM t WHERE g = 2; +---- +2 0 1 0 3 + +######## +# Empty input. Without a GROUP BY the aggregate still emits one row, and the +# counts on it must be 0 rather than the NULL an unguarded `sum` would give. +######## + +statement ok +set datafusion.optimizer.max_passes = 0; + +query IIIII +SELECT count(*) AS star, count(1) AS one, count(x) AS cnt_x, count(DISTINCT x) AS distinct_x, + max(v) AS max_v +FROM empty_t; +---- +0 0 0 0 NULL + +query IIIII +SELECT count(*) AS star, count(1) AS one, count(x) AS cnt_x, count(DISTINCT x) AS distinct_x, + max(v) AS max_v +FROM t WHERE g = 99; +---- +0 0 0 0 NULL + +query IIII +SELECT g, count(*) AS star, count(x) AS cnt_x, count(DISTINCT x) AS distinct_x +FROM empty_t GROUP BY g ORDER BY g; +---- + +statement ok +set datafusion.optimizer.max_passes = 3; + +query IIIII +SELECT count(*) AS star, count(1) AS one, count(x) AS cnt_x, count(DISTINCT x) AS distinct_x, + max(v) AS max_v +FROM empty_t; +---- +0 0 0 0 NULL + +query IIIII +SELECT count(*) AS star, count(1) AS one, count(x) AS cnt_x, count(DISTINCT x) AS distinct_x, + max(v) AS max_v +FROM t WHERE g = 99; +---- +0 0 0 0 NULL + +query IIII +SELECT g, count(*) AS star, count(x) AS cnt_x, count(DISTINCT x) AS distinct_x +FROM empty_t GROUP BY g ORDER BY g; +---- + +######## +# HAVING and ORDER BY read the rewritten count by name from the node above +######## + +statement ok +set datafusion.explain.logical_plan_only = true; + +query TT +EXPLAIN +SELECT g, count(*) AS records, count(DISTINCT x) AS distinct_x +FROM t GROUP BY g HAVING count(*) > 1 AND count(DISTINCT x) > 0 +ORDER BY count(*) DESC, g; +---- +logical_plan +01)Sort: records DESC NULLS FIRST, t.g ASC NULLS LAST +02)--Projection: t.g, CASE WHEN sum(alias2) IS NOT NULL THEN sum(alias2) ELSE Int64(0) END AS records, count(alias1) AS distinct_x +03)----Filter: CASE WHEN sum(alias2) IS NOT NULL THEN sum(alias2) ELSE Int64(0) END > Int64(1) AND count(alias1) > Int64(0) +04)------Aggregate: groupBy=[[t.g]], aggr=[[sum(alias2), count(alias1)]] +05)--------Aggregate: groupBy=[[t.g, t.x AS alias1]], aggr=[[count(Int64(1)) AS alias2]] +06)----------TableScan: t projection=[g, x] + +statement ok +set datafusion.explain.logical_plan_only = false; + +statement ok +set datafusion.optimizer.max_passes = 0; + +query III +SELECT g, count(*) AS records, count(DISTINCT x) AS distinct_x +FROM t GROUP BY g HAVING count(*) > 1 AND count(DISTINCT x) > 0 +ORDER BY count(*) DESC, g; +---- +1 4 2 + +statement ok +set datafusion.optimizer.max_passes = 3; + +query III +SELECT g, count(*) AS records, count(DISTINCT x) AS distinct_x +FROM t GROUP BY g HAVING count(*) > 1 AND count(DISTINCT x) > 0 +ORDER BY count(*) DESC, g; +---- +1 4 2 + +######## +# The whole production shape, aggregating over a join +######## + +statement ok +CREATE TABLE failed(x VARCHAR); + +statement ok +INSERT INTO failed VALUES ('10'), ('30'); + +statement ok +set datafusion.optimizer.max_passes = 0; + +query IIIIII +SELECT t.g, count(*) AS records, count(DISTINCT t.x) AS distinct_x, + min(t.v) AS min_v, max(t.v) AS max_v, sum(t.v) AS sum_v +FROM t JOIN failed f ON t.x = f.x +GROUP BY t.g ORDER BY t.g; +---- +1 2 1 100 100 100 +3 1 1 1 1 1 + +statement ok +set datafusion.optimizer.max_passes = 3; + +query IIIIII +SELECT t.g, count(*) AS records, count(DISTINCT t.x) AS distinct_x, + min(t.v) AS min_v, max(t.v) AS max_v, sum(t.v) AS sum_v +FROM t JOIN failed f ON t.x = f.x +GROUP BY t.g ORDER BY t.g; +---- +1 2 1 100 100 100 +3 1 1 1 1 1 + +######## +# Both sides of the gate produce the same results as the unoptimized plan +######## + +statement ok +set datafusion.optimizer.max_passes = 0; + +# gated: the count keeps the INT distinct unrewritten +query IIIII +SELECT g, count(*) AS records, count(DISTINCT xi) AS distinct_xi, sum(v) AS sum_v, max(v) AS max_v +FROM t GROUP BY g ORDER BY g; +---- +1 4 2 112 100 +2 2 0 3 3 +3 1 1 1 1 + +# rewritten: the same shape over the VARCHAR distinct +query IIIII +SELECT g, count(*) AS records, count(DISTINCT x) AS distinct_x, sum(v) AS sum_v, max(v) AS max_v +FROM t GROUP BY g ORDER BY g; +---- +1 4 2 112 100 +2 2 0 3 3 +3 1 1 1 1 + +# rewritten through sum alone, over the INT distinct the gate excludes +query III +SELECT g, count(DISTINCT xi) AS distinct_xi, sum(v) AS sum_v +FROM t GROUP BY g ORDER BY g; +---- +1 2 112 +2 0 3 +3 1 1 + +query II +SELECT count(*) AS records, count(DISTINCT xi) AS distinct_xi FROM empty_t; +---- +0 0 + +statement ok +set datafusion.optimizer.max_passes = 3; + +# gated: the count keeps the INT distinct unrewritten +query IIIII +SELECT g, count(*) AS records, count(DISTINCT xi) AS distinct_xi, sum(v) AS sum_v, max(v) AS max_v +FROM t GROUP BY g ORDER BY g; +---- +1 4 2 112 100 +2 2 0 3 3 +3 1 1 1 1 + +# rewritten: the same shape over the VARCHAR distinct +query IIIII +SELECT g, count(*) AS records, count(DISTINCT x) AS distinct_x, sum(v) AS sum_v, max(v) AS max_v +FROM t GROUP BY g ORDER BY g; +---- +1 4 2 112 100 +2 2 0 3 3 +3 1 1 1 1 + +# rewritten through sum alone, over the INT distinct the gate excludes +query III +SELECT g, count(DISTINCT xi) AS distinct_xi, sum(v) AS sum_v +FROM t GROUP BY g ORDER BY g; +---- +1 2 112 +2 0 3 +3 1 1 + +query II +SELECT count(*) AS records, count(DISTINCT xi) AS distinct_xi FROM empty_t; +---- +0 0 + +statement ok +DROP TABLE failed; + +statement ok +DROP TABLE empty_t; + +statement ok +DROP TABLE t; diff --git a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs b/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs index c9f874dd9b095..6c5cd182ac8a8 100644 --- a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs +++ b/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs @@ -333,8 +333,13 @@ async fn simple_aggregate() -> Result<()> { #[tokio::test] async fn aggregate_distinct_with_having() -> Result<()> { - roundtrip("SELECT a, count(distinct b) FROM data GROUP BY a, c HAVING count(b) > 100") - .await + let ctx = create_context_without_single_distinct_to_group_by().await?; + roundtrip_with_ctx( + "SELECT a, count(distinct b) FROM data GROUP BY a, c HAVING count(b) > 100", + ctx, + ) + .await?; + Ok(()) } #[tokio::test] @@ -2886,6 +2891,31 @@ async fn create_context() -> Result { create_context_with_dialect(None).await } +/// [`create_context`] with `single_distinct_aggregation_to_group_by` removed from +/// the optimizer. +/// +/// That rule rewrites a single `AGG(DISTINCT x)` into a two phase group by whose +/// inner aggregate aliases its grouping and measure expressions (`alias1`, +/// `alias2`). Substrait carries no names for those expressions, so the consumer +/// derives them from the expressions themselves and a rewritten plan does not +/// round trip to an identical plan. That holds for every output of the rule, not +/// only for the query under test. +async fn create_context_without_single_distinct_to_group_by() -> Result { + let ctx = create_context().await?; + let rules = ctx + .state() + .optimizer() + .rules + .iter() + .filter(|rule| rule.name() != "single_distinct_aggregation_to_group_by") + .cloned() + .collect(); + let state = SessionStateBuilder::new_from_existing(ctx.state()) + .with_optimizer_rules(rules) + .build(); + Ok(SessionContext::new_with_state(state)) +} + async fn create_context_with_dialect(dialect: Option) -> Result { let mut session_config = SessionConfig::default();