diff --git a/datafusion/common/src/config.rs b/datafusion/common/src/config.rs index 302cb0b25548c..c2d3ea02748e9 100644 --- a/datafusion/common/src/config.rs +++ b/datafusion/common/src/config.rs @@ -1915,6 +1915,23 @@ config_namespace! { /// branches share the same source and compatible wrapper nodes such as identical /// projections or aliases. pub enable_unions_to_filter: bool, default = false + + /// When `true`, the `SingleDistinctToGroupBy` optimizer rule applies its cost guard + /// universally — including to lone-distinct queries that have no companion non-distinct + /// aggregates. The rewrite is skipped whenever the distinct function already has a + /// native `GroupsAccumulator`, reducing peak memory by avoiding the larger intermediate + /// `(group, distinct value)` result set that the two-phase plan produces. + /// + /// The trade-off is reduced parallelism: skipping the rewrite causes the engine to fall + /// back to serial per-group accumulation, which is typically 2–3× slower for queries + /// with high-cardinality distinct values. + /// + /// When `false` (the default), lone-distinct queries are always rewritten into the + /// two-phase `GROUP BY` so they can fan out across partitions and achieve maximum + /// throughput. + /// + /// See + pub single_distinct_aggregation_prefer_memory: bool, default = false } } diff --git a/datafusion/core/tests/sql/unparser.rs b/datafusion/core/tests/sql/unparser.rs index e3674d8c387e1..30af9e14f7b26 100644 --- a/datafusion/core/tests/sql/unparser.rs +++ b/datafusion/core/tests/sql/unparser.rs @@ -542,36 +542,9 @@ async fn assert_issue_23317_unparsed_sql_plans( Ok(()) } -#[tokio::test] -async fn optimized_duckdb_unparse_preserves_nested_aggregate_scope() -> Result<()> { - let ctx = issue_23317_context()?; - let plan = ctx.sql(ISSUE_23317_QUERY).await?.into_optimized_plan()?; - let dialect = DuckDBDialect::new(); - let unparser = Unparser::new(&dialect); - let sql = unparser.plan_to_sql(&plan)?.to_string(); - - assert_issue_23317_unparsed_sql_plans(&ctx, &sql).await?; - - assert!( - sql.contains(concat!( - r#"FROM (SELECT sum("total_revenue") AS "alias2", "#, - r#"date_part('year', "signup_date") AS "group_alias_0", "#, - r#""customer_id" AS "alias1" "# - )), - "inner aggregate should define the aliases before the outer aggregate uses them: {sql}", - ); - assert!( - !sql.contains(r#"date_part('year', "c"."signup_date") AS "group_alias_0""#), - "inner aggregate must not reference out-of-scope alias c: {sql}", - ); - - Ok(()) -} - #[tokio::test] async fn optimized_duckdb_unparse_unqualifies_aggregate_input_projection() -> Result<()> { let ctx = issue_23317_context()?; - assert!(ctx.remove_optimizer_rule(SingleDistinctToGroupBy::new().name())); let plan = ctx.sql(ISSUE_23317_QUERY).await?.into_optimized_plan()?; let dialect = DuckDBDialect::new(); diff --git a/datafusion/optimizer/src/eliminate_aggregate_distinct.rs b/datafusion/optimizer/src/eliminate_aggregate_distinct.rs index d2b68853a2327..4c11e167b2275 100644 --- a/datafusion/optimizer/src/eliminate_aggregate_distinct.rs +++ b/datafusion/optimizer/src/eliminate_aggregate_distinct.rs @@ -346,7 +346,7 @@ mod tests { optimizer_ctx, rules, plan, - @r" + @" Projection: test.a, min(alias1) AS min(DISTINCT test.b), sum(alias1) AS sum(DISTINCT test.b) Aggregate: groupBy=[[test.a]], aggr=[[min(alias1), sum(alias1)]] Aggregate: groupBy=[[test.a, test.b AS alias1]], aggr=[[]] diff --git a/datafusion/optimizer/src/single_distinct_to_groupby.rs b/datafusion/optimizer/src/single_distinct_to_groupby.rs index 8b176e904e9b9..f90056730ad10 100644 --- a/datafusion/optimizer/src/single_distinct_to_groupby.rs +++ b/datafusion/optimizer/src/single_distinct_to_groupby.rs @@ -107,11 +107,28 @@ impl CountRollup { } } -/// Check whether all aggregate exprs are distinct on a single field. +/// Check whether all aggregate exprs are distinct on a single field and whether +/// rewriting the aggregation into a two-phase GROUP BY is advantageous. +/// +/// The decision logic follows this cost and parallelism model: +/// +/// Does the query have a GROUP BY? +/// | +/// +- No -> Is there more than 1 partition, and is it not count(DISTINCT tinyint/smallint)? +/// | +- Yes -> REWRITE (2.6-5.5x faster: the only way the distinct runs in parallel) +/// | +- No -> SKIP (faster and uses less memory) +/// | +/// +- Yes -> Does the distinct aggregate have its own fast grouped accumulator? +/// +- Yes (count(DISTINCT ``)) -> SKIP (faster or a tie, less memory) +/// +- No (sum/avg DISTINCT, floats, decimals, strings, user-defined) -> REWRITE (1.4-20x faster) +/// +- config option to skip it when memory matters more (few groups: 2.4-6x more memory) fn is_single_distinct_agg( aggr_expr: &[Expr], + group_expr: &[Expr], input_schema: &DFSchema, count_rollup: Option<&CountRollup>, + target_partitions: usize, + prefer_memory: bool, ) -> Result { let mut fields_set = HashSet::new(); let mut aggregate_count = 0; @@ -141,10 +158,12 @@ fn is_single_distinct_agg( distinct_aggs.push((func, args.as_slice())); } 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" + } else if func.name() == "sum" + || func.name().eq_ignore_ascii_case("min") + || func.name().eq_ignore_ascii_case("max") { + // Companion non-distinct aggregates + } else { return Ok(false); } } else { @@ -154,47 +173,80 @@ fn is_single_distinct_agg( 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); + + // Invariant: Distinct aggregates with DistinctHandling::Unsupported (e.g. approx_median) + // cannot execute without being rewritten into a GROUP BY first. + if distinct_aggs.iter().any(|(func, _)| { + func.distinct_handling() == datafusion_expr::DistinctHandling::Unsupported + }) { + return Ok(true); } - 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 keeps its state at that finer grain. What pays for that is taking the -/// distinct aggregate off `GroupsAccumulatorAdapter` and its one boxed -/// accumulator per group. An 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. -/// -/// `Some(false)` is therefore the only answer that permits the rewrite. `None`, -/// the default for almost every function, is not evidence that it pays. -/// -/// This is a proxy: what actually decides the outcome is the cost per distinct -/// value on each side, which this rule cannot see. It is deliberately -/// conservative in the direction that leaves a plan alone. The pre-existing -/// tolerance of a non-distinct `sum`, `min` or `max` is not gated, since -/// narrowing it would change plans that have always been rewritten. -/// -/// Measured in . -fn rewrite_pays_for_count( - distinct_aggs: &[(&Arc, &[Expr])], - input_schema: &DFSchema, -) -> Result { - for (func, args) in distinct_aggs { - let arg_types = args + let has_fast_grouped_accumulator = distinct_aggs.iter().all(|(func, args)| { + let arg_types: Vec<_> = args .iter() - .map(|arg| arg.get_type(input_schema)) - .collect::>>()?; - if func.groups_accumulator_supported_for_types(&arg_types, true) == Some(false) { - return Ok(true); + .map(|arg| { + arg.get_type(input_schema) + .unwrap_or(arrow::datatypes::DataType::Null) + }) + .collect(); + func.groups_accumulator_supported_for_types(&arg_types, true) == Some(true) + }); + + // If there is a companion non-distinct count rollup alongside a distinct aggregate + // that already has a fast native grouped accumulator (e.g. count(DISTINCT )), + // rolling up the count into sum(count) adds extra intermediate grouping state without + // taking anything off GroupsAccumulatorAdapter. + if has_count_rollup && has_fast_grouped_accumulator { + return Ok(false); + } + + let has_group_by = !group_expr.is_empty(); + + if !has_group_by { + // Query has NO GROUP BY: + // Is there more than 1 partition, and is it not count(DISTINCT tinyint/smallint)? + // Yes -> REWRITE (2.6-5.5x faster: the only way the distinct runs in parallel) + // No -> SKIP (faster and uses less memory) + let is_multi_partition = target_partitions > 1; + let is_count_distinct_small_int = distinct_aggs.iter().all(|(func, args)| { + count_rollup.is_some_and(|rollup| rollup.is_count(func)) + && args.len() == 1 + && matches!( + args[0] + .get_type(input_schema) + .unwrap_or(arrow::datatypes::DataType::Null), + arrow::datatypes::DataType::Int8 + | arrow::datatypes::DataType::UInt8 + | arrow::datatypes::DataType::Int16 + | arrow::datatypes::DataType::UInt16 + ) + }); + + if is_multi_partition && !is_count_distinct_small_int { + Ok(true) + } else { + Ok(false) + } + } else { + // Query HAS a GROUP BY: + // Does the distinct aggregate have its own fast grouped accumulator? + // Yes (count(DISTINCT )) -> SKIP (faster or a tie, less memory) + // No (sum/avg DISTINCT, floats, decimals, strings, user-defined) -> REWRITE (1.4-20x faster) + // +- config option to skip it when memory matters more (few groups: 2.4-6x more memory) + if has_fast_grouped_accumulator { + // Fast native grouped accumulator exists (e.g. count(DISTINCT )). + // Single-phase grouping is faster or equivalent, and avoids two-phase memory overhead. + Ok(false) + } else if prefer_memory { + // No fast grouped accumulator, but caller opted in to prefer lower memory over throughput. + Ok(false) + } else { + // No fast grouped accumulator (e.g. sum/avg DISTINCT, Utf8 strings). + // Rewriting into two-phase GROUP BY is 1.4-20x faster by avoiding GroupsAccumulatorAdapter. + Ok(true) } } - Ok(false) } /// Check if the first expr is [Expr::GroupingSet]. @@ -219,6 +271,11 @@ impl OptimizerRule for SingleDistinctToGroupBy { if !matches!(plan, LogicalPlan::Aggregate(_)) { return Ok(Transformed::no(plan)); } + let prefer_memory = config + .options() + .optimizer + .single_distinct_aggregation_prefer_memory; + let target_partitions = config.options().execution.target_partitions; let count_rollup = CountRollup::try_new(config); match plan { LogicalPlan::Aggregate(Aggregate { @@ -227,11 +284,15 @@ impl OptimizerRule for SingleDistinctToGroupBy { schema, group_expr, .. - }) if is_single_distinct_agg( - &aggr_expr, - input.schema(), - count_rollup.as_ref(), - )? && !contains_grouping_set(&group_expr) => + }) if !contains_grouping_set(&group_expr) + && is_single_distinct_agg( + &aggr_expr, + &group_expr, + input.schema(), + count_rollup.as_ref(), + target_partitions, + prefer_memory, + )? => { let group_size = group_expr.len(); // alias all original group_by exprs @@ -448,7 +509,9 @@ mod tests { }; use datafusion_functions_aggregate::correlation::corr_udaf; use datafusion_functions_aggregate::count::count_udaf; - use datafusion_functions_aggregate::expr_fn::{count, count_distinct, max, min, sum}; + use datafusion_functions_aggregate::expr_fn::{ + avg, count, count_distinct, max, min, sum, + }; use datafusion_functions_aggregate::min_max::max_udaf; use datafusion_functions_aggregate::sum::sum_udaf; @@ -558,6 +621,24 @@ mod tests { table_scan(Some("test"), &schema, None)?.build() } + fn test_table_scan_int8_b() -> Result { + let schema = Schema::new(vec![ + Field::new("a", DataType::UInt32, false), + Field::new("b", DataType::Int8, false), + Field::new("c", DataType::UInt32, false), + ]); + table_scan(Some("test"), &schema, None).unwrap().build() + } + + fn test_table_scan_uint8_b() -> Result { + let schema = Schema::new(vec![ + Field::new("a", DataType::UInt32, false), + Field::new("b", DataType::UInt8, false), + Field::new("c", DataType::UInt32, false), + ]); + table_scan(Some("test"), &schema, None).unwrap().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`. @@ -569,11 +650,27 @@ mod tests { impl RegistryOptimizerContext { fn new() -> Self { + Self::with_prefer_memory(false) + } + + /// Build a context that exposes a function registry and optionally turns on + /// `single_distinct_aggregation_prefer_memory` so tests can verify both sides + /// of the memory / throughput trade-off. + fn with_prefer_memory(prefer_memory: bool) -> Self { + Self::with_options(prefer_memory, 4) + } + + fn with_options(prefer_memory: bool, target_partitions: usize) -> Self { let mut registry = MemoryFunctionRegistry::new(); registry.register_udaf(count_udaf()).unwrap(); registry.register_udaf(sum_udaf()).unwrap(); + + let mut opts = ConfigOptions::new(); + opts.optimizer.single_distinct_aggregation_prefer_memory = prefer_memory; + opts.execution.target_partitions = target_partitions; + Self { - inner: OptimizerContext::new(), + inner: OptimizerContext::new_with_config_options(Arc::new(opts)), registry, } } @@ -614,6 +711,24 @@ mod tests { }}; } + macro_rules! assert_optimized_plan_equal_with_context { + ( + $ctx:expr, + $plan:expr, + @ $expected:literal $(,)? + ) => {{ + let rule: Arc = + Arc::new(SingleDistinctToGroupBy::new()); + let optimizer = Optimizer::with_rules(vec![rule]); + let optimized_plan = optimizer + .optimize($plan, &$ctx, |_, _| {}) + .expect("failed to optimize plan"); + insta::assert_snapshot!(optimized_plan.display_indent_schema(), @ $expected); + + Ok::<(), DataFusionError>(()) + }}; + } + #[test] fn not_exist_distinct() -> Result<()> { let table_scan = test_table_scan()?; @@ -643,7 +758,7 @@ mod tests { // Should work assert_optimized_plan_equal!( plan, - @r" + @" Projection: count(alias1) AS count(DISTINCT test.b) [count(DISTINCT test.b):Int64] Aggregate: groupBy=[[]], aggr=[[count(alias1)]] [count(alias1):Int64] Aggregate: groupBy=[[test.b AS alias1]], aggr=[[]] [alias1:UInt32] @@ -729,7 +844,7 @@ mod tests { assert_optimized_plan_equal!( plan, - @r" + @" Projection: count(alias1) AS count(DISTINCT Int32(2) * test.b) [count(DISTINCT Int32(2) * test.b):Int64] Aggregate: groupBy=[[]], aggr=[[count(alias1)]] [count(alias1):Int64] Aggregate: groupBy=[[Int32(2) * test.b AS alias1]], aggr=[[]] [alias1:Int64] @@ -746,14 +861,36 @@ mod tests { .aggregate(vec![col("a")], vec![count_distinct(col("b"))])? .build()?; - // Should work + // `count(DISTINCT )` has a specialized fast grouped accumulator: + // single-phase grouping is faster or equivalent and avoids two-phase memory overhead. assert_optimized_plan_equal!( plan, @r" + Aggregate: groupBy=[[test.a]], aggr=[[count(DISTINCT test.b)]] [a:UInt32, count(DISTINCT test.b):Int64] + TableScan: test [a:UInt32, b:UInt32, c:UInt32] + " + ) + } + + #[test] + fn single_distinct_and_groupby_without_fast_accumulator() -> Result<()> { + let table_scan = test_table_scan_utf8_b().unwrap(); + + let plan = LogicalPlanBuilder::from(table_scan) + .aggregate(vec![col("a")], vec![count_distinct(col("b"))]) + .unwrap() + .build() + .unwrap(); + + // `count(DISTINCT )` lacks a specialized fast grouped accumulator: + // rewrites into two-phase GROUP BY to avoid GroupsAccumulatorAdapter. + assert_optimized_plan_equal!( + plan, + @" Projection: test.a, count(alias1) AS count(DISTINCT test.b) [a:UInt32, count(DISTINCT test.b):Int64] Aggregate: groupBy=[[test.a]], aggr=[[count(alias1)]] [a:UInt32, count(alias1):Int64] - Aggregate: groupBy=[[test.a, test.b AS alias1]], aggr=[[]] [a:UInt32, alias1:UInt32] - TableScan: test [a:UInt32, b:UInt32, c:UInt32] + Aggregate: groupBy=[[test.a, test.b AS alias1]], aggr=[[]] [a:UInt32, alias1:Utf8] + TableScan: test [a:UInt32, b:Utf8, c:UInt32] " ) } @@ -833,7 +970,7 @@ mod tests { // Should work assert_optimized_plan_equal!( plan, - @r" + @" Projection: test.a, count(alias1) AS count(DISTINCT test.b), max(alias1) AS max(DISTINCT test.b) [a:UInt32, count(DISTINCT test.b):Int64, max(DISTINCT test.b):UInt32;N] Aggregate: groupBy=[[test.a]], aggr=[[count(alias1), max(alias1)]] [a:UInt32, count(alias1):Int64, max(alias1):UInt32;N] Aggregate: groupBy=[[test.a, test.b AS alias1]], aggr=[[]] [a:UInt32, alias1:UInt32] @@ -905,32 +1042,55 @@ mod tests { // 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] + @" + Aggregate: groupBy=[[test.a]], aggr=[[count(DISTINCT test.b), sum(test.c)]] [a:UInt32, count(DISTINCT test.b):Int64, sum(test.c):UInt64;N] + TableScan: test [a:UInt32, b:UInt32, c:UInt32] " ) } #[test] - fn distinct_over_a_function_with_no_opinion_is_not_rewritten() -> Result<()> { - let table_scan = test_table_scan()?; + fn distinct_over_a_function_with_no_opinion_rewrites_by_default() -> Result<()> { + let table_scan = test_table_scan().unwrap(); let plan = LogicalPlanBuilder::from(table_scan) .aggregate( vec![col("a")], vec![silent_distinct(col("b")), count(col("c"))], - )? - .build()?; + ) + .unwrap() + .build() + .unwrap(); - // Should not work: `silent` leaves - // `groups_accumulator_supported_for_types` at its default, so it - // reports nothing about `GroupsAccumulatorAdapter`. Silence is not - // evidence that the rewrite pays, and the rule must not read it as - // permission + // User-defined / no-opinion aggregate lacks a fast grouped accumulator: + // by default it rewrites into two-phase GROUP BY to avoid GroupsAccumulatorAdapter (1.4-20x faster) assert_optimized_plan_equal!( + plan, + @" + Projection: test.a, silent(alias1) AS silent(DISTINCT test.b), CASE WHEN sum(alias2) IS NOT NULL THEN sum(alias2) ELSE Int64(0) END AS count(test.c) [a:UInt32, silent(DISTINCT test.b):UInt32;N, count(test.c):Int64] + Aggregate: groupBy=[[test.a]], aggr=[[silent(alias1), sum(alias2)]] [a:UInt32, silent(alias1):UInt32;N, sum(alias2):Int64;N] + Aggregate: groupBy=[[test.a, test.b AS alias1]], aggr=[[count(test.c) AS alias2]] [a:UInt32, alias1:UInt32, alias2:Int64] + TableScan: test [a:UInt32, b:UInt32, c:UInt32] + " + ) + } + + #[test] + fn distinct_over_a_function_with_no_opinion_prefer_memory_skips() -> Result<()> { + let table_scan = test_table_scan().unwrap(); + + let plan = LogicalPlanBuilder::from(table_scan) + .aggregate( + vec![col("a")], + vec![silent_distinct(col("b")), count(col("c"))], + ) + .unwrap() + .build() + .unwrap(); + + let ctx = RegistryOptimizerContext::with_prefer_memory(true); + assert_optimized_plan_equal_with_context!( + ctx, plan, @r" Aggregate: groupBy=[[test.a]], aggr=[[silent(DISTINCT test.b), count(test.c)]] [a:UInt32, silent(DISTINCT test.b):UInt32;N, count(test.c):Int64] @@ -940,19 +1100,47 @@ mod tests { } #[test] - fn distinct_sum_and_count_is_not_rewritten() -> Result<()> { - let table_scan = test_table_scan()?; + fn distinct_sum_and_count_rewrites_by_default() -> Result<()> { + let table_scan = test_table_scan().unwrap(); let plan = LogicalPlanBuilder::from(table_scan) .aggregate( vec![col("a")], vec![sum_distinct(col("b")), count(col("c"))], - )? - .build()?; + ) + .unwrap() + .build() + .unwrap(); - // Should not work: `sum` reports nothing about - // `GroupsAccumulatorAdapter` either, so the rewrite is not known to pay + // `sum(DISTINCT)` lacks a fast grouped accumulator: + // by default it rewrites into two-phase GROUP BY to avoid GroupsAccumulatorAdapter (1.4-20x faster) assert_optimized_plan_equal!( + plan, + @" + Projection: test.a, sum(alias1) AS sum(DISTINCT test.b), CASE WHEN sum(alias2) IS NOT NULL THEN sum(alias2) ELSE Int64(0) END AS count(test.c) [a:UInt32, sum(DISTINCT test.b):UInt64;N, count(test.c):Int64] + Aggregate: groupBy=[[test.a]], aggr=[[sum(alias1), sum(alias2)]] [a:UInt32, sum(alias1):UInt64;N, sum(alias2):Int64;N] + Aggregate: groupBy=[[test.a, test.b AS alias1]], aggr=[[count(test.c) AS alias2]] [a:UInt32, alias1:UInt32, alias2:Int64] + TableScan: test [a:UInt32, b:UInt32, c:UInt32] + " + ) + } + + #[test] + fn distinct_sum_and_count_prefer_memory_skips() -> Result<()> { + let table_scan = test_table_scan().unwrap(); + + let plan = LogicalPlanBuilder::from(table_scan) + .aggregate( + vec![col("a")], + vec![sum_distinct(col("b")), count(col("c"))], + ) + .unwrap() + .build() + .unwrap(); + + let ctx = RegistryOptimizerContext::with_prefer_memory(true); + assert_optimized_plan_equal_with_context!( + ctx, plan, @r" Aggregate: groupBy=[[test.a]], aggr=[[sum(DISTINCT test.b), count(test.c)]] [a:UInt32, sum(DISTINCT test.b):UInt64;N, count(test.c):Int64] @@ -963,20 +1151,22 @@ mod tests { #[test] fn group_by_with_expr() -> Result<()> { - let table_scan = test_table_scan().unwrap(); + let table_scan = test_table_scan_utf8_b().unwrap(); let plan = LogicalPlanBuilder::from(table_scan) - .aggregate(vec![col("a") + lit(1)], vec![count_distinct(col("c"))])? - .build()?; + .aggregate(vec![col("a") + lit(1)], vec![count_distinct(col("b"))]) + .unwrap() + .build() + .unwrap(); // Should work assert_optimized_plan_equal!( plan, - @r" - Projection: group_alias_0 AS test.a + Int32(1), count(alias1) AS count(DISTINCT test.c) [test.a + Int32(1):Int64, count(DISTINCT test.c):Int64] + @" + Projection: group_alias_0 AS test.a + Int32(1), count(alias1) AS count(DISTINCT test.b) [test.a + Int32(1):Int64, count(DISTINCT test.b):Int64] Aggregate: groupBy=[[group_alias_0]], aggr=[[count(alias1)]] [group_alias_0:Int64, count(alias1):Int64] - Aggregate: groupBy=[[test.a + Int32(1) AS group_alias_0, test.c AS alias1]], aggr=[[]] [group_alias_0:Int64, alias1:UInt32] - TableScan: test [a:UInt32, b:UInt32, c:UInt32] + Aggregate: groupBy=[[test.a + Int32(1) AS group_alias_0, test.b AS alias1]], aggr=[[]] [group_alias_0:Int64, alias1:Utf8] + TableScan: test [a:UInt32, b:Utf8, c:UInt32] " ) } @@ -996,10 +1186,10 @@ mod tests { )? .build()?; - // Should work + // `max(DISTINCT b)` lacks a fast grouped accumulator: rewrites into two-phase GROUP BY assert_optimized_plan_equal!( plan, - @r" + @" Projection: test.a, sum(alias2) AS sum(test.c), count(alias1) AS count(DISTINCT test.b), max(alias1) AS max(DISTINCT test.b) [a:UInt32, sum(test.c):UInt64;N, count(DISTINCT test.b):Int64, max(DISTINCT test.b):UInt32;N] Aggregate: groupBy=[[test.a]], aggr=[[sum(alias2), count(alias1), max(alias1)]] [a:UInt32, sum(alias2):UInt64;N, count(alias1):Int64, max(alias1):UInt32;N] Aggregate: groupBy=[[test.a, test.b AS alias1]], aggr=[[sum(test.c) AS alias2]] [a:UInt32, alias1:UInt32, alias2:UInt64;N] @@ -1008,6 +1198,261 @@ mod tests { ) } + /// Ensures the `UInt16` branch of the small-int guard is hit. + /// `count(DISTINCT uint16)` without GROUP BY should SKIP the rewrite even + /// on 4 partitions because the cardinality is bounded and the accumulator is cheap. + #[test] + fn no_groupby_uint16_small_int_skips() -> Result<()> { + let schema = Schema::new(vec![ + Field::new("a", DataType::UInt32, false), + Field::new("b", DataType::UInt16, false), + Field::new("c", DataType::UInt32, false), + ]); + let table_scan = table_scan(Some("test"), &schema, None) + .unwrap() + .build() + .unwrap(); + + let plan = LogicalPlanBuilder::from(table_scan) + .aggregate(Vec::::new(), vec![count_distinct(col("b"))]) + .unwrap() + .build() + .unwrap(); + + let ctx = RegistryOptimizerContext::with_options(false, 4); + assert_optimized_plan_equal_with_context!( + ctx, + plan, + @r" + Aggregate: groupBy=[[]], aggr=[[count(DISTINCT test.b)]] [count(DISTINCT test.b):Int64] + TableScan: test [a:UInt32, b:UInt16, c:UInt32] + " + ) + } + + /// Ensures the `Int16` branch of the small-int guard is hit. + #[test] + fn no_groupby_int16_small_int_skips() -> Result<()> { + let schema = Schema::new(vec![ + Field::new("a", DataType::UInt32, false), + Field::new("b", DataType::Int16, false), + Field::new("c", DataType::UInt32, false), + ]); + let table_scan = table_scan(Some("test"), &schema, None) + .unwrap() + .build() + .unwrap(); + + let plan = LogicalPlanBuilder::from(table_scan) + .aggregate(Vec::::new(), vec![count_distinct(col("b"))]) + .unwrap() + .build() + .unwrap(); + + let ctx = RegistryOptimizerContext::with_options(false, 4); + assert_optimized_plan_equal_with_context!( + ctx, + plan, + @r" + Aggregate: groupBy=[[]], aggr=[[count(DISTINCT test.b)]] [count(DISTINCT test.b):Int64] + TableScan: test [a:UInt32, b:Int16, c:UInt32] + " + ) + } + + /// Ensures the `UInt8` branch of the small-int guard is hit. + #[test] + fn no_groupby_uint8_small_int_skips() -> Result<()> { + let table_scan = test_table_scan_uint8_b().unwrap(); + + let plan = LogicalPlanBuilder::from(table_scan) + .aggregate(Vec::::new(), vec![count_distinct(col("b"))]) + .unwrap() + .build() + .unwrap(); + + let ctx = RegistryOptimizerContext::with_options(false, 4); + assert_optimized_plan_equal_with_context!( + ctx, + plan, + @r" + Aggregate: groupBy=[[]], aggr=[[count(DISTINCT test.b)]] [count(DISTINCT test.b):Int64] + TableScan: test [a:UInt32, b:UInt8, c:UInt32] + " + ) + } + + /// Verifies the `is_count_distinct_small_int` check works correctly when + /// the function is NOT count (not from the `count_rollup` registry). + /// `sum(DISTINCT int8)` without GROUP BY on multi-partitions must still + /// rewrite because it is not the small-int cardinality exception. + #[test] + fn no_groupby_non_count_small_int_still_rewrites() -> Result<()> { + let table_scan = test_table_scan_int8_b().unwrap(); + + let plan = LogicalPlanBuilder::from(table_scan) + .aggregate(Vec::::new(), vec![sum_distinct(col("b"))]) + .unwrap() + .build() + .unwrap(); + + // sum(DISTINCT int8) is not the small-int cardinality exception: + // it must rewrite on multi-partition to enable parallel deduplication. + let ctx = RegistryOptimizerContext::with_options(false, 4); + assert_optimized_plan_equal_with_context!( + ctx, + plan, + @" + Projection: sum(alias1) AS sum(DISTINCT test.b) [sum(DISTINCT test.b):Int64;N] + Aggregate: groupBy=[[]], aggr=[[sum(alias1)]] [sum(alias1):Int64;N] + Aggregate: groupBy=[[test.b AS alias1]], aggr=[[]] [alias1:Int8] + TableScan: test [a:UInt32, b:Int8, c:UInt32] + " + ) + } + + /// Verifies `min` as the sole non-distinct companion forces a skip when + /// the distinct aggregate already has a fast grouped accumulator. + /// This exercises the `func.name().eq_ignore_ascii_case("min")` branch. + #[test] + fn single_distinct_and_groupby_with_min_companion_and_fast_accumulator_skips() + -> Result<()> { + let table_scan = test_table_scan().unwrap(); + + let plan = LogicalPlanBuilder::from(table_scan) + .aggregate( + vec![col("a")], + vec![count_distinct(col("b")), min(col("c"))], + ) + .unwrap() + .build() + .unwrap(); + + // count(DISTINCT UInt32) has a fast grouped accumulator: skip rewrite. + assert_optimized_plan_equal!( + plan, + @r" + Aggregate: groupBy=[[test.a]], aggr=[[count(DISTINCT test.b), min(test.c)]] [a:UInt32, count(DISTINCT test.b):Int64, min(test.c):UInt32;N] + TableScan: test [a:UInt32, b:UInt32, c:UInt32] + " + ) + } + + /// Verifies `avg` as companion aggregate skips rewrite because `avg` + /// cannot be rolled up in the two-phase GROUP BY (not count, sum, min, max). + /// This directly exercises the companion aggregate fallback `else { return Ok(false); }`. + #[test] + fn single_distinct_and_groupby_with_avg_companion_skips() -> Result<()> { + let table_scan = test_table_scan().unwrap(); + + let plan = LogicalPlanBuilder::from(table_scan) + .aggregate( + vec![col("a")], + vec![count_distinct(col("b")), avg(col("c"))], + ) + .unwrap() + .build() + .unwrap(); + + assert_optimized_plan_equal!( + plan, + @r" + Aggregate: groupBy=[[test.a]], aggr=[[count(DISTINCT test.b), avg(test.c)]] [a:UInt32, count(DISTINCT test.b):Int64, avg(test.c):Float64;N] + TableScan: test [a:UInt32, b:UInt32, c:UInt32] + " + ) + } + + /// With GROUP BY + `prefer_memory = true`, if the distinct aggregate has a + /// fast grouped accumulator, it still skips (the fast-accumulator branch + /// takes priority and `prefer_memory` never even needs to be checked). + #[test] + fn prefer_memory_true_with_fast_accumulator_skips_not_prefer_memory_branch() + -> Result<()> { + let table_scan = test_table_scan().unwrap(); + + let plan = LogicalPlanBuilder::from(table_scan) + .aggregate(vec![col("a")], vec![count_distinct(col("b"))]) + .unwrap() + .build() + .unwrap(); + + let ctx = RegistryOptimizerContext::with_options(false, 4); + assert_optimized_plan_equal_with_context!( + ctx, + plan, + @r" + Aggregate: groupBy=[[test.a]], aggr=[[count(DISTINCT test.b)]] [a:UInt32, count(DISTINCT test.b):Int64] + TableScan: test [a:UInt32, b:UInt32, c:UInt32] + " + ) + } + + #[test] + fn no_groupby_multi_partition_rewrites() -> Result<()> { + let table_scan = test_table_scan().unwrap(); + + let plan = LogicalPlanBuilder::from(table_scan) + .aggregate(Vec::::new(), vec![count_distinct(col("b"))]) + .unwrap() + .build() + .unwrap(); + + let ctx = RegistryOptimizerContext::with_options(false, 4); + assert_optimized_plan_equal_with_context!( + ctx, + plan, + @" + Projection: count(alias1) AS count(DISTINCT test.b) [count(DISTINCT test.b):Int64] + Aggregate: groupBy=[[]], aggr=[[count(alias1)]] [count(alias1):Int64] + Aggregate: groupBy=[[test.b AS alias1]], aggr=[[]] [alias1:UInt32] + TableScan: test [a:UInt32, b:UInt32, c:UInt32] + " + ) + } + + #[test] + fn no_groupby_single_partition_skips() -> Result<()> { + let table_scan = test_table_scan().unwrap(); + + let plan = LogicalPlanBuilder::from(table_scan) + .aggregate(Vec::::new(), vec![count_distinct(col("b"))]) + .unwrap() + .build() + .unwrap(); + + let ctx = RegistryOptimizerContext::with_options(false, 1); + assert_optimized_plan_equal_with_context!( + ctx, + plan, + @r" + Aggregate: groupBy=[[]], aggr=[[count(DISTINCT test.b)]] [count(DISTINCT test.b):Int64] + TableScan: test [a:UInt32, b:UInt32, c:UInt32] + " + ) + } + + #[test] + fn no_groupby_small_int_skips() -> Result<()> { + let table_scan = test_table_scan_int8_b().unwrap(); + + let plan = LogicalPlanBuilder::from(table_scan) + .aggregate(Vec::::new(), vec![count_distinct(col("b"))]) + .unwrap() + .build() + .unwrap(); + + let ctx = RegistryOptimizerContext::with_options(false, 4); + assert_optimized_plan_equal_with_context!( + ctx, + plan, + @r" + Aggregate: groupBy=[[]], aggr=[[count(DISTINCT test.b)]] [count(DISTINCT test.b):Int64] + TableScan: test [a:UInt32, b:Int8, c:UInt32] + " + ) + } + #[test] fn one_distinct_and_two_common() -> Result<()> { let table_scan = test_table_scan()?; @@ -1022,11 +1467,9 @@ mod tests { // Should work assert_optimized_plan_equal!( plan, - @r" - Projection: test.a, sum(alias2) AS sum(test.c), max(alias3) AS max(test.c), count(alias1) AS count(DISTINCT test.b) [a:UInt32, sum(test.c):UInt64;N, max(test.c):UInt32;N, count(DISTINCT test.b):Int64] - Aggregate: groupBy=[[test.a]], aggr=[[sum(alias2), max(alias3), count(alias1)]] [a:UInt32, sum(alias2):UInt64;N, max(alias3):UInt32;N, count(alias1):Int64] - Aggregate: groupBy=[[test.a, test.b AS alias1]], aggr=[[sum(test.c) AS alias2, max(test.c) AS alias3]] [a:UInt32, alias1:UInt32, alias2:UInt64;N, alias3:UInt32;N] - TableScan: test [a:UInt32, b:UInt32, c:UInt32] + @" + Aggregate: groupBy=[[test.a]], aggr=[[sum(test.c), max(test.c), count(DISTINCT test.b)]] [a:UInt32, sum(test.c):UInt64;N, max(test.c):UInt32;N, count(DISTINCT test.b):Int64] + TableScan: test [a:UInt32, b:UInt32, c:UInt32] " ) } @@ -1045,11 +1488,9 @@ mod tests { // Should work assert_optimized_plan_equal!( plan, - @r" - Projection: test.c, min(alias2) AS min(test.a), count(alias1) AS count(DISTINCT test.b) [c:UInt32, min(test.a):UInt32;N, count(DISTINCT test.b):Int64] - Aggregate: groupBy=[[test.c]], aggr=[[min(alias2), count(alias1)]] [c:UInt32, min(alias2):UInt32;N, count(alias1):Int64] - Aggregate: groupBy=[[test.c, test.b AS alias1]], aggr=[[min(test.a) AS alias2]] [c:UInt32, alias1:UInt32, alias2:UInt32;N] - TableScan: test [a:UInt32, b:UInt32, c:UInt32] + @" + Aggregate: groupBy=[[test.c]], aggr=[[min(test.a), count(DISTINCT test.b)]] [c:UInt32, min(test.a):UInt32;N, count(DISTINCT test.b):Int64] + TableScan: test [a:UInt32, b:UInt32, c:UInt32] " ) } @@ -1305,6 +1746,122 @@ mod tests { ) } + #[test] + fn test_unsupported_distinct_handling_is_rewritten() { + let table_scan = test_table_scan().unwrap(); + + // approx_median returns DistinctHandling::Unsupported + let aggr_expr = Expr::AggregateFunction(AggregateFunction::new_udf( + datafusion_functions_aggregate::approx_median::approx_median_udaf(), + vec![col("b")], + true, + None, + vec![], + None, + )); + + let plan = LogicalPlanBuilder::from(table_scan) + .aggregate(vec![col("a")], vec![aggr_expr, sum(col("c"))]) + .unwrap() + .build() + .unwrap(); + + // The rule should rewrite the plan because approx_median doesn't support DISTINCT natively, + // so it requires SingleDistinctToGroupBy to deduplicate the input even when accompanied + // by non-distinct aggregates. + let rule: Arc = + Arc::new(SingleDistinctToGroupBy::new()); + assert_optimized_plan_eq_display_indent_snapshot!( + rule, + plan, + @r" + Projection: test.a, approx_median(alias1) AS approx_median(DISTINCT test.b), sum(alias2) AS sum(test.c) [a:UInt32, approx_median(DISTINCT test.b):Float64;N, sum(test.c):UInt64;N] + Aggregate: groupBy=[[test.a]], aggr=[[approx_median(alias1), sum(alias2)]] [a:UInt32, approx_median(alias1):Float64;N, 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] + ", + ).unwrap(); + } + + #[test] + fn single_distinct_aggregation_prefer_memory_default_rewrites_lone_distinct() + -> Result<()> { + let table_scan = test_table_scan().unwrap(); + + let plan = LogicalPlanBuilder::from(table_scan) + .aggregate(vec![col("a")], vec![sum_distinct(col("b"))]) + .unwrap() + .build() + .unwrap(); + + let ctx = RegistryOptimizerContext::with_prefer_memory(false); + assert_optimized_plan_equal_with_context!( + ctx, + plan, + @r" + Projection: test.a, sum(alias1) AS sum(DISTINCT test.b) [a:UInt32, sum(DISTINCT test.b):UInt64;N] + Aggregate: groupBy=[[test.a]], aggr=[[sum(alias1)]] [a:UInt32, sum(alias1):UInt64;N] + Aggregate: groupBy=[[test.a, test.b AS alias1]], aggr=[[]] [a:UInt32, alias1:UInt32] + TableScan: test [a:UInt32, b:UInt32, c:UInt32] + " + ) + } + + #[test] + fn single_distinct_aggregation_prefer_memory_true_skips_rewrite_for_groups_accumulator() + -> Result<()> { + let table_scan = test_table_scan().unwrap(); + + let plan = LogicalPlanBuilder::from(table_scan) + .aggregate(vec![col("a")], vec![sum_distinct(col("b"))]) + .unwrap() + .build() + .unwrap(); + + let ctx = RegistryOptimizerContext::with_prefer_memory(true); + assert_optimized_plan_equal_with_context!( + ctx, + plan, + @r" + Aggregate: groupBy=[[test.a]], aggr=[[sum(DISTINCT test.b)]] [a:UInt32, sum(DISTINCT test.b):UInt64;N] + TableScan: test [a:UInt32, b:UInt32, c:UInt32] + " + ) + } + + #[test] + fn single_distinct_aggregation_prefer_memory_true_still_rewrites_unsupported() + -> Result<()> { + let table_scan = test_table_scan().unwrap(); + + let aggr_expr = Expr::AggregateFunction(AggregateFunction::new_udf( + datafusion_functions_aggregate::approx_median::approx_median_udaf(), + vec![col("b")], + true, + None, + vec![], + None, + )); + + let plan = LogicalPlanBuilder::from(table_scan) + .aggregate(vec![col("a")], vec![aggr_expr]) + .unwrap() + .build() + .unwrap(); + + let ctx = RegistryOptimizerContext::with_prefer_memory(true); + assert_optimized_plan_equal_with_context!( + ctx, + plan, + @r" + Projection: test.a, approx_median(alias1) AS approx_median(DISTINCT test.b) [a:UInt32, approx_median(DISTINCT test.b):Float64;N] + Aggregate: groupBy=[[test.a]], aggr=[[approx_median(alias1)]] [a:UInt32, approx_median(alias1):Float64;N] + Aggregate: groupBy=[[test.a, test.b AS alias1]], aggr=[[]] [a:UInt32, alias1:UInt32] + TableScan: test [a:UInt32, b:UInt32, c:UInt32] + " + ) + } + #[test] fn sum_decimal_with_distinct_preserves_declared_type() -> Result<()> { let schema = Schema::new(vec![ diff --git a/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt b/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt index dc6af9b1e4e47..0da9ec2d40e19 100644 --- a/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt +++ b/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt @@ -92,6 +92,8 @@ FROM ( 100000 100000 # Assert spill happened (same reasoning as Case A). +# count(DISTINCT v) with GROUP BY stays single-phase: Int64 has a fast native +# grouped accumulator, so the optimizer skips the two-phase GROUP BY rewrite. query TT EXPLAIN ANALYZE SELECT count(*), sum(d) @@ -102,7 +104,7 @@ FROM ( ) ---- -04)------AggregateExec: mode=Single, gby=[group_alias_0@0 as group_alias_0], aggr=[count(alias1)], metrics=[spill_count=, spilled_bytes=, spilled_rows=100.0 K,] +04)------AggregateExec: mode=Single, gby=[v@0 * 7 % 100000 as t.v * Int64(7) % Int64(100000)], aggr=[count(DISTINCT t.v)], metrics=[spill_count=, spilled_bytes=, spilled_rows=100.0 K,] # --- Case D: multiple aggregates (sum/min/max) under memory limit --- diff --git a/datafusion/sqllogictest/test_files/aggregates_simplify.slt b/datafusion/sqllogictest/test_files/aggregates_simplify.slt index b9c4ba147bb2d..94cb9954d90e9 100644 --- a/datafusion/sqllogictest/test_files/aggregates_simplify.slt +++ b/datafusion/sqllogictest/test_files/aggregates_simplify.slt @@ -183,6 +183,9 @@ SELECT SUM(DISTINCT column1 + 1), SUM(column1 + 1) FROM sum_simplify_t; ---- 5 7 +# On a multi-partition runner the distinct sum has no fast grouped accumulator, +# so SingleDistinctToGroupBy rewrites to a two-phase plan (the only way to +# run the distinct in parallel when there is no GROUP BY). query TT EXPLAIN SELECT SUM(DISTINCT column1 + 1), SUM(column1 + 1) FROM sum_simplify_t; ---- @@ -568,19 +571,13 @@ query TT EXPLAIN SELECT g, count(DISTINCT v) FROM distinct_simplify_t GROUP BY g; ---- logical_plan -01)Projection: distinct_simplify_t.g, count(alias1) AS count(DISTINCT distinct_simplify_t.v) -02)--Aggregate: groupBy=[[distinct_simplify_t.g]], aggr=[[count(alias1)]] -03)----Aggregate: groupBy=[[distinct_simplify_t.g, distinct_simplify_t.v AS alias1]], aggr=[[]] -04)------TableScan: distinct_simplify_t projection=[g, v] +01)Aggregate: groupBy=[[distinct_simplify_t.g]], aggr=[[count(DISTINCT distinct_simplify_t.v)]] +02)--TableScan: distinct_simplify_t projection=[g, v] physical_plan -01)ProjectionExec: expr=[g@0 as g, count(alias1)@1 as count(DISTINCT distinct_simplify_t.v)] -02)--AggregateExec: mode=FinalPartitioned, gby=[g@0 as g], aggr=[count(alias1)] -03)----RepartitionExec: partitioning=Hash([g@0], 4), input_partitions=4 -04)------AggregateExec: mode=Partial, gby=[g@0 as g], aggr=[count(alias1)] -05)--------AggregateExec: mode=FinalPartitioned, gby=[g@0 as g, alias1@1 as alias1], aggr=[] -06)----------RepartitionExec: partitioning=Hash([g@0, alias1@1], 4), input_partitions=1 -07)------------AggregateExec: mode=Partial, gby=[g@0 as g, v@1 as alias1], aggr=[] -08)--------------DataSourceExec: partitions=1, partition_sizes=[1] +01)AggregateExec: mode=FinalPartitioned, gby=[g@0 as g], aggr=[count(DISTINCT distinct_simplify_t.v)] +02)--RepartitionExec: partitioning=Hash([g@0], 4), input_partitions=1 +03)----AggregateExec: mode=Partial, gby=[g@0 as g], aggr=[count(DISTINCT distinct_simplify_t.v)] +04)------DataSourceExec: partitions=1, partition_sizes=[1] # Mixed: a node that still needs one DISTINCT keeps all of them, so this rule # cannot change which plans SingleDistinctToGroupBy rewrites. `count` needs the diff --git a/datafusion/sqllogictest/test_files/clickbench.slt b/datafusion/sqllogictest/test_files/clickbench.slt index d89880ca33789..5cac67ffeee07 100644 --- a/datafusion/sqllogictest/test_files/clickbench.slt +++ b/datafusion/sqllogictest/test_files/clickbench.slt @@ -219,27 +219,25 @@ SELECT "AdvEngineID", COUNT(*) FROM hits WHERE "AdvEngineID" <> 0 GROUP BY "AdvE ---- ## Q8 +# count(DISTINCT UserID) stays single-phase: UserID is UInt64 (integer with fast +# grouped accumulator), so the optimizer skips the two-phase GROUP BY rewrite. query TT EXPLAIN SELECT "RegionID", COUNT(DISTINCT "UserID") AS u FROM hits GROUP BY "RegionID" ORDER BY u DESC LIMIT 10; ---- logical_plan 01)Sort: u DESC NULLS FIRST, fetch=10 -02)--Projection: hits.RegionID, count(alias1) AS u -03)----Aggregate: groupBy=[[hits.RegionID]], aggr=[[count(alias1)]] -04)------Aggregate: groupBy=[[hits.RegionID, hits.UserID AS alias1]], aggr=[[]] -05)--------SubqueryAlias: hits -06)----------TableScan: hits_raw projection=[RegionID, UserID] +02)--Projection: hits.RegionID, count(DISTINCT hits.UserID) AS u +03)----Aggregate: groupBy=[[hits.RegionID]], aggr=[[count(DISTINCT hits.UserID)]] +04)------SubqueryAlias: hits +05)--------TableScan: hits_raw projection=[RegionID, UserID] physical_plan 01)SortPreservingMergeExec: [u@1 DESC], fetch=10 -02)--ProjectionExec: expr=[RegionID@0 as RegionID, count(alias1)@1 as u] -03)----SortExec: TopK(fetch=10), expr=[count(alias1)@1 DESC], preserve_partitioning=[true] -04)------AggregateExec: mode=FinalPartitioned, gby=[RegionID@0 as RegionID], aggr=[count(alias1)] -05)--------RepartitionExec: partitioning=Hash([RegionID@0], 4), input_partitions=4 -06)----------AggregateExec: mode=Partial, gby=[RegionID@0 as RegionID], aggr=[count(alias1)] -07)------------AggregateExec: mode=FinalPartitioned, gby=[RegionID@0 as RegionID, alias1@1 as alias1], aggr=[] -08)--------------RepartitionExec: partitioning=Hash([RegionID@0, alias1@1], 4), input_partitions=1 -09)----------------AggregateExec: mode=Partial, gby=[RegionID@0 as RegionID, UserID@1 as alias1], aggr=[] -10)------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/clickbench_hits_10.parquet]]}, projection=[RegionID, UserID], file_type=parquet +02)--ProjectionExec: expr=[RegionID@0 as RegionID, count(DISTINCT hits.UserID)@1 as u] +03)----SortExec: TopK(fetch=10), expr=[count(DISTINCT hits.UserID)@1 DESC], preserve_partitioning=[true] +04)------AggregateExec: mode=FinalPartitioned, gby=[RegionID@0 as RegionID], aggr=[count(DISTINCT hits.UserID)] +05)--------RepartitionExec: partitioning=Hash([RegionID@0], 4), input_partitions=1 +06)----------AggregateExec: mode=Partial, gby=[RegionID@0 as RegionID], aggr=[count(DISTINCT hits.UserID)] +07)------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/clickbench_hits_10.parquet]]}, projection=[RegionID, UserID], file_type=parquet query II rowsort SELECT "RegionID", COUNT(DISTINCT "UserID") AS u FROM hits GROUP BY "RegionID" ORDER BY u DESC LIMIT 10; @@ -282,25 +280,21 @@ EXPLAIN SELECT "MobilePhoneModel", COUNT(DISTINCT "UserID") AS u FROM hits WHERE ---- logical_plan 01)Sort: u DESC NULLS FIRST, fetch=10 -02)--Projection: hits.MobilePhoneModel, count(alias1) AS u -03)----Aggregate: groupBy=[[hits.MobilePhoneModel]], aggr=[[count(alias1)]] -04)------Aggregate: groupBy=[[hits.MobilePhoneModel, hits.UserID AS alias1]], aggr=[[]] -05)--------SubqueryAlias: hits -06)----------Filter: hits_raw.MobilePhoneModel != Utf8View("") -07)------------TableScan: hits_raw projection=[UserID, MobilePhoneModel], partial_filters=[hits_raw.MobilePhoneModel != Utf8View("")] +02)--Projection: hits.MobilePhoneModel, count(DISTINCT hits.UserID) AS u +03)----Aggregate: groupBy=[[hits.MobilePhoneModel]], aggr=[[count(DISTINCT hits.UserID)]] +04)------SubqueryAlias: hits +05)--------Filter: hits_raw.MobilePhoneModel != Utf8View("") +06)----------TableScan: hits_raw projection=[UserID, MobilePhoneModel], partial_filters=[hits_raw.MobilePhoneModel != Utf8View("")] physical_plan 01)SortPreservingMergeExec: [u@1 DESC], fetch=10 -02)--ProjectionExec: expr=[MobilePhoneModel@0 as MobilePhoneModel, count(alias1)@1 as u] -03)----SortExec: TopK(fetch=10), expr=[count(alias1)@1 DESC], preserve_partitioning=[true] -04)------AggregateExec: mode=FinalPartitioned, gby=[MobilePhoneModel@0 as MobilePhoneModel], aggr=[count(alias1)] +02)--ProjectionExec: expr=[MobilePhoneModel@0 as MobilePhoneModel, count(DISTINCT hits.UserID)@1 as u] +03)----SortExec: TopK(fetch=10), expr=[count(DISTINCT hits.UserID)@1 DESC], preserve_partitioning=[true] +04)------AggregateExec: mode=FinalPartitioned, gby=[MobilePhoneModel@0 as MobilePhoneModel], aggr=[count(DISTINCT hits.UserID)] 05)--------RepartitionExec: partitioning=Hash([MobilePhoneModel@0], 4), input_partitions=4 -06)----------AggregateExec: mode=Partial, gby=[MobilePhoneModel@0 as MobilePhoneModel], aggr=[count(alias1)] -07)------------AggregateExec: mode=FinalPartitioned, gby=[MobilePhoneModel@0 as MobilePhoneModel, alias1@1 as alias1], aggr=[] -08)--------------RepartitionExec: partitioning=Hash([MobilePhoneModel@0, alias1@1], 4), input_partitions=4 -09)----------------AggregateExec: mode=Partial, gby=[MobilePhoneModel@1 as MobilePhoneModel, UserID@0 as alias1], aggr=[] -10)------------------FilterExec: MobilePhoneModel@1 != -11)--------------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 -12)----------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/clickbench_hits_10.parquet]]}, projection=[UserID, MobilePhoneModel], file_type=parquet, predicate=MobilePhoneModel@34 != , pruning_predicate=MobilePhoneModel_null_count@2 != row_count@3 AND (MobilePhoneModel_min@0 != OR != MobilePhoneModel_max@1), required_guarantees=[MobilePhoneModel not in ()] +06)----------AggregateExec: mode=Partial, gby=[MobilePhoneModel@1 as MobilePhoneModel], aggr=[count(DISTINCT hits.UserID)] +07)------------FilterExec: MobilePhoneModel@1 != +08)--------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +09)----------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/clickbench_hits_10.parquet]]}, projection=[UserID, MobilePhoneModel], file_type=parquet, predicate=MobilePhoneModel@34 != , pruning_predicate=MobilePhoneModel_null_count@2 != row_count@3 AND (MobilePhoneModel_min@0 != OR != MobilePhoneModel_max@1), required_guarantees=[MobilePhoneModel not in ()] query TI SELECT "MobilePhoneModel", COUNT(DISTINCT "UserID") AS u FROM hits WHERE "MobilePhoneModel" <> '' GROUP BY "MobilePhoneModel" ORDER BY u DESC LIMIT 10; @@ -312,25 +306,21 @@ EXPLAIN SELECT "MobilePhone", "MobilePhoneModel", COUNT(DISTINCT "UserID") AS u ---- logical_plan 01)Sort: u DESC NULLS FIRST, fetch=10 -02)--Projection: hits.MobilePhone, hits.MobilePhoneModel, count(alias1) AS u -03)----Aggregate: groupBy=[[hits.MobilePhone, hits.MobilePhoneModel]], aggr=[[count(alias1)]] -04)------Aggregate: groupBy=[[hits.MobilePhone, hits.MobilePhoneModel, hits.UserID AS alias1]], aggr=[[]] -05)--------SubqueryAlias: hits -06)----------Filter: hits_raw.MobilePhoneModel != Utf8View("") -07)------------TableScan: hits_raw projection=[UserID, MobilePhone, MobilePhoneModel], partial_filters=[hits_raw.MobilePhoneModel != Utf8View("")] +02)--Projection: hits.MobilePhone, hits.MobilePhoneModel, count(DISTINCT hits.UserID) AS u +03)----Aggregate: groupBy=[[hits.MobilePhone, hits.MobilePhoneModel]], aggr=[[count(DISTINCT hits.UserID)]] +04)------SubqueryAlias: hits +05)--------Filter: hits_raw.MobilePhoneModel != Utf8View("") +06)----------TableScan: hits_raw projection=[UserID, MobilePhone, MobilePhoneModel], partial_filters=[hits_raw.MobilePhoneModel != Utf8View("")] physical_plan 01)SortPreservingMergeExec: [u@2 DESC], fetch=10 -02)--ProjectionExec: expr=[MobilePhone@0 as MobilePhone, MobilePhoneModel@1 as MobilePhoneModel, count(alias1)@2 as u] -03)----SortExec: TopK(fetch=10), expr=[count(alias1)@2 DESC], preserve_partitioning=[true] -04)------AggregateExec: mode=FinalPartitioned, gby=[MobilePhone@0 as MobilePhone, MobilePhoneModel@1 as MobilePhoneModel], aggr=[count(alias1)] +02)--ProjectionExec: expr=[MobilePhone@0 as MobilePhone, MobilePhoneModel@1 as MobilePhoneModel, count(DISTINCT hits.UserID)@2 as u] +03)----SortExec: TopK(fetch=10), expr=[count(DISTINCT hits.UserID)@2 DESC], preserve_partitioning=[true] +04)------AggregateExec: mode=FinalPartitioned, gby=[MobilePhone@0 as MobilePhone, MobilePhoneModel@1 as MobilePhoneModel], aggr=[count(DISTINCT hits.UserID)] 05)--------RepartitionExec: partitioning=Hash([MobilePhone@0, MobilePhoneModel@1], 4), input_partitions=4 -06)----------AggregateExec: mode=Partial, gby=[MobilePhone@0 as MobilePhone, MobilePhoneModel@1 as MobilePhoneModel], aggr=[count(alias1)] -07)------------AggregateExec: mode=FinalPartitioned, gby=[MobilePhone@0 as MobilePhone, MobilePhoneModel@1 as MobilePhoneModel, alias1@2 as alias1], aggr=[] -08)--------------RepartitionExec: partitioning=Hash([MobilePhone@0, MobilePhoneModel@1, alias1@2], 4), input_partitions=4 -09)----------------AggregateExec: mode=Partial, gby=[MobilePhone@1 as MobilePhone, MobilePhoneModel@2 as MobilePhoneModel, UserID@0 as alias1], aggr=[] -10)------------------FilterExec: MobilePhoneModel@2 != -11)--------------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 -12)----------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/clickbench_hits_10.parquet]]}, projection=[UserID, MobilePhone, MobilePhoneModel], file_type=parquet, predicate=MobilePhoneModel@34 != , pruning_predicate=MobilePhoneModel_null_count@2 != row_count@3 AND (MobilePhoneModel_min@0 != OR != MobilePhoneModel_max@1), required_guarantees=[MobilePhoneModel not in ()] +06)----------AggregateExec: mode=Partial, gby=[MobilePhone@1 as MobilePhone, MobilePhoneModel@2 as MobilePhoneModel], aggr=[count(DISTINCT hits.UserID)] +07)------------FilterExec: MobilePhoneModel@2 != +08)--------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +09)----------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/clickbench_hits_10.parquet]]}, projection=[UserID, MobilePhone, MobilePhoneModel], file_type=parquet, predicate=MobilePhoneModel@34 != , pruning_predicate=MobilePhoneModel_null_count@2 != row_count@3 AND (MobilePhoneModel_min@0 != OR != MobilePhoneModel_max@1), required_guarantees=[MobilePhoneModel not in ()] query ITI SELECT "MobilePhone", "MobilePhoneModel", COUNT(DISTINCT "UserID") AS u FROM hits WHERE "MobilePhoneModel" <> '' GROUP BY "MobilePhone", "MobilePhoneModel" ORDER BY u DESC LIMIT 10; @@ -363,30 +353,28 @@ SELECT "SearchPhrase", COUNT(*) AS c FROM hits WHERE "SearchPhrase" <> '' GROUP ---- ## Q13 +# count(DISTINCT UserID) stays single-phase: UserID is UInt64 (integer), which +# has a fast native grouped accumulator, so the optimizer skips the two-phase rewrite. query TT EXPLAIN SELECT "SearchPhrase", COUNT(DISTINCT "UserID") AS u FROM hits WHERE "SearchPhrase" <> '' GROUP BY "SearchPhrase" ORDER BY u DESC LIMIT 10; ---- logical_plan 01)Sort: u DESC NULLS FIRST, fetch=10 -02)--Projection: hits.SearchPhrase, count(alias1) AS u -03)----Aggregate: groupBy=[[hits.SearchPhrase]], aggr=[[count(alias1)]] -04)------Aggregate: groupBy=[[hits.SearchPhrase, hits.UserID AS alias1]], aggr=[[]] -05)--------SubqueryAlias: hits -06)----------Filter: hits_raw.SearchPhrase != Utf8View("") -07)------------TableScan: hits_raw projection=[UserID, SearchPhrase], partial_filters=[hits_raw.SearchPhrase != Utf8View("")] +02)--Projection: hits.SearchPhrase, count(DISTINCT hits.UserID) AS u +03)----Aggregate: groupBy=[[hits.SearchPhrase]], aggr=[[count(DISTINCT hits.UserID)]] +04)------SubqueryAlias: hits +05)--------Filter: hits_raw.SearchPhrase != Utf8View("") +06)----------TableScan: hits_raw projection=[UserID, SearchPhrase], partial_filters=[hits_raw.SearchPhrase != Utf8View("")] physical_plan 01)SortPreservingMergeExec: [u@1 DESC], fetch=10 -02)--ProjectionExec: expr=[SearchPhrase@0 as SearchPhrase, count(alias1)@1 as u] -03)----SortExec: TopK(fetch=10), expr=[count(alias1)@1 DESC], preserve_partitioning=[true] -04)------AggregateExec: mode=FinalPartitioned, gby=[SearchPhrase@0 as SearchPhrase], aggr=[count(alias1)] +02)--ProjectionExec: expr=[SearchPhrase@0 as SearchPhrase, count(DISTINCT hits.UserID)@1 as u] +03)----SortExec: TopK(fetch=10), expr=[count(DISTINCT hits.UserID)@1 DESC], preserve_partitioning=[true] +04)------AggregateExec: mode=FinalPartitioned, gby=[SearchPhrase@0 as SearchPhrase], aggr=[count(DISTINCT hits.UserID)] 05)--------RepartitionExec: partitioning=Hash([SearchPhrase@0], 4), input_partitions=4 -06)----------AggregateExec: mode=Partial, gby=[SearchPhrase@0 as SearchPhrase], aggr=[count(alias1)] -07)------------AggregateExec: mode=FinalPartitioned, gby=[SearchPhrase@0 as SearchPhrase, alias1@1 as alias1], aggr=[] -08)--------------RepartitionExec: partitioning=Hash([SearchPhrase@0, alias1@1], 4), input_partitions=4 -09)----------------AggregateExec: mode=Partial, gby=[SearchPhrase@1 as SearchPhrase, UserID@0 as alias1], aggr=[] -10)------------------FilterExec: SearchPhrase@1 != -11)--------------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 -12)----------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/clickbench_hits_10.parquet]]}, projection=[UserID, SearchPhrase], file_type=parquet, predicate=SearchPhrase@39 != , pruning_predicate=SearchPhrase_null_count@2 != row_count@3 AND (SearchPhrase_min@0 != OR != SearchPhrase_max@1), required_guarantees=[SearchPhrase not in ()] +06)----------AggregateExec: mode=Partial, gby=[SearchPhrase@1 as SearchPhrase], aggr=[count(DISTINCT hits.UserID)] +07)------------FilterExec: SearchPhrase@1 != +08)--------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +09)----------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/clickbench_hits_10.parquet]]}, projection=[UserID, SearchPhrase], file_type=parquet, predicate=SearchPhrase@39 != , pruning_predicate=SearchPhrase_null_count@2 != row_count@3 AND (SearchPhrase_min@0 != OR != SearchPhrase_max@1), required_guarantees=[SearchPhrase not in ()] query TI SELECT "SearchPhrase", COUNT(DISTINCT "UserID") AS u FROM hits WHERE "SearchPhrase" <> '' GROUP BY "SearchPhrase" ORDER BY u DESC LIMIT 10; diff --git a/datafusion/sqllogictest/test_files/information_schema.slt b/datafusion/sqllogictest/test_files/information_schema.slt index bd0a8faedf12b..ae44aaf42700d 100644 --- a/datafusion/sqllogictest/test_files/information_schema.slt +++ b/datafusion/sqllogictest/test_files/information_schema.slt @@ -330,6 +330,7 @@ datafusion.optimizer.repartition_file_scans true datafusion.optimizer.repartition_joins true datafusion.optimizer.repartition_sorts true datafusion.optimizer.repartition_windows true +datafusion.optimizer.single_distinct_aggregation_prefer_memory false datafusion.optimizer.skip_failed_rules false datafusion.optimizer.subset_repartition_threshold 4 datafusion.optimizer.top_down_join_key_reordering true @@ -496,6 +497,7 @@ datafusion.optimizer.repartition_file_scans true When set to `true`, datasource datafusion.optimizer.repartition_joins true Should DataFusion repartition data using the join keys to execute joins in parallel using the provided `target_partitions` level datafusion.optimizer.repartition_sorts true Should DataFusion execute sorts in a per-partition fashion and merge afterwards instead of coalescing first and sorting globally. With this flag is enabled, plans in the form below ```text "SortExec: [a@0 ASC]", " CoalescePartitionsExec", " RepartitionExec: partitioning=RoundRobinBatch(8), input_partitions=1", ``` would turn into the plan below which performs better in multithreaded environments ```text "SortPreservingMergeExec: [a@0 ASC]", " SortExec: [a@0 ASC]", " RepartitionExec: partitioning=RoundRobinBatch(8), input_partitions=1", ``` datafusion.optimizer.repartition_windows true Should DataFusion repartition data using the partitions keys to execute window functions in parallel using the provided `target_partitions` level +datafusion.optimizer.single_distinct_aggregation_prefer_memory false When `true`, the `SingleDistinctToGroupBy` optimizer rule applies its cost guard universally — including to lone-distinct queries that have no companion non-distinct aggregates. The rewrite is skipped whenever the distinct function already has a native `GroupsAccumulator`, reducing peak memory by avoiding the larger intermediate `(group, distinct value)` result set that the two-phase plan produces. The trade-off is reduced parallelism: skipping the rewrite causes the engine to fall back to serial per-group accumulation, which is typically 2–3× slower for queries with high-cardinality distinct values. When `false` (the default), lone-distinct queries are always rewritten into the two-phase `GROUP BY` so they can fan out across partitions and achieve maximum throughput. See datafusion.optimizer.skip_failed_rules false When set to true, the logical plan optimizer will produce warning messages if any optimization rules produce errors and then proceed to the next rule. When set to false, any rules that produce errors will cause the query to fail datafusion.optimizer.subset_repartition_threshold 4 Partition count threshold for subset satisfaction optimization. When the current partition count is >= this threshold, DataFusion will skip repartitioning if the required partitioning expression is a subset of the current partition expression such as Hash(a) satisfies Hash(a, b). When the current partition count is < this threshold, DataFusion will repartition to increase parallelism even when subset satisfaction applies. Set to 0 to always repartition (disable subset satisfaction optimization). Set to a high value to always use subset satisfaction. Example (subset_repartition_threshold = 4): ```text Hash([a]) satisfies Hash([a, b]) because (Hash([a, b]) is subset of Hash([a]) If current partitions (3) < threshold (4), repartition: AggregateExec: mode=FinalPartitioned, gby=[a, b], aggr=[SUM(x)] RepartitionExec: partitioning=Hash([a, b], 8), input_partitions=3 AggregateExec: mode=Partial, gby=[a, b], aggr=[SUM(x)] DataSourceExec: file_groups={...}, output_partitioning=Hash([a], 3) If current partitions (8) >= threshold (4), use subset satisfaction: AggregateExec: mode=SinglePartitioned, gby=[a, b], aggr=[SUM(x)] DataSourceExec: file_groups={...}, output_partitioning=Hash([a], 8) ``` datafusion.optimizer.top_down_join_key_reordering true When set to true, the physical plan optimizer will run a top down process to reorder the join keys diff --git a/datafusion/sqllogictest/test_files/joins.slt b/datafusion/sqllogictest/test_files/joins.slt index b795b050f919f..ddb1b0089160c 100644 --- a/datafusion/sqllogictest/test_files/joins.slt +++ b/datafusion/sqllogictest/test_files/joins.slt @@ -1344,12 +1344,10 @@ inner join join_t2 on join_t1.t1_id = join_t2.t2_id group by join_t1.t1_id ---- logical_plan -01)Projection: join_t1.t1_id, count(alias1) AS count(DISTINCT join_t1.t1_int) -02)--Aggregate: groupBy=[[join_t1.t1_id]], aggr=[[count(alias1)]] -03)----Aggregate: groupBy=[[join_t1.t1_id, join_t1.t1_int AS alias1]], aggr=[[]] -04)------LeftSemi Join: join_t1.t1_id = join_t2.t2_id -05)--------TableScan: join_t1 projection=[t1_id, t1_int] -06)--------TableScan: join_t2 projection=[t2_id] +01)Aggregate: groupBy=[[join_t1.t1_id]], aggr=[[count(DISTINCT join_t1.t1_int)]] +02)--LeftSemi Join: join_t1.t1_id = join_t2.t2_id +03)----TableScan: join_t1 projection=[t1_id, t1_int] +04)----TableScan: join_t2 projection=[t2_id] # A similar query with two DISTINCT aggregates is also rewritten: each # `count(DISTINCT ...)` removes its own duplicates, so the join's duplicates diff --git a/datafusion/sqllogictest/test_files/single_distinct_to_groupby.slt b/datafusion/sqllogictest/test_files/single_distinct_to_groupby.slt index 78308f7897c6e..0b559d2ffd80a 100644 --- a/datafusion/sqllogictest/test_files/single_distinct_to_groupby.slt +++ b/datafusion/sqllogictest/test_files/single_distinct_to_groupby.slt @@ -98,11 +98,21 @@ logical_plan 02)--Aggregate: groupBy=[[t.g]], aggr=[[count(Int64(1)), count(DISTINCT t.xi)]] 03)----TableScan: t projection=[g, xi] -# `sum` reports nothing about `GroupsAccumulator` support from the argument -# types, and neither does any aggregate but `count`. Silence is not evidence -# that the rewrite pays, so the count keeps this plan unrewritten as well. -# Rewritten, this shape measured 3.15x the peak memory over 4,000,000 rows in -# 2,000 groups +# `sum(DISTINCT v)` lacks a fast grouped accumulator: by default it rewrites into +# the two-phase GROUP BY to avoid the 1.4–20x slower row-by-row GroupsAccumulatorAdapter. +query TT +EXPLAIN SELECT g, count(*) AS records, sum(DISTINCT v) AS distinct_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, sum(alias1) AS distinct_v +02)--Aggregate: groupBy=[[t.g]], aggr=[[sum(alias2), sum(alias1)]] +03)----Aggregate: groupBy=[[t.g, CAST(t.v AS Int64) AS alias1]], aggr=[[count(Int64(1)) AS alias2]] +04)------TableScan: t projection=[g, v] + +# When single_distinct_aggregation_prefer_memory = true, the rewrite is skipped to reduce peak memory +statement ok +set datafusion.optimizer.single_distinct_aggregation_prefer_memory = true; + query TT EXPLAIN SELECT g, count(*) AS records, sum(DISTINCT v) AS distinct_v FROM t GROUP BY g; ---- @@ -111,6 +121,9 @@ logical_plan 02)--Aggregate: groupBy=[[t.g]], aggr=[[count(Int64(1)), sum(DISTINCT CAST(t.v AS Int64))]] 03)----TableScan: t projection=[g, v] +statement ok +set datafusion.optimizer.single_distinct_aggregation_prefer_memory = false; + # `min(DISTINCT v)` is the same value as `min(v)`. EliminateAggregateDistinct # drops the flag before this rule runs, so there is no distinct aggregate left # to rewrite and the plan keeps one scalar per group instead of a row per @@ -129,14 +142,53 @@ 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] +01)Projection: t.g, sum(t.v) AS sum_v, count(DISTINCT t.xi) AS distinct_xi +02)--Aggregate: groupBy=[[t.g]], aggr=[[sum(CAST(t.v AS Int64)), count(DISTINCT t.xi)]] +03)----TableScan: t projection=[g, xi, v] + +######## +# Lone-distinct queries and single_distinct_aggregation_prefer_memory +######## + +# By default (prefer_memory = false), lone-distinct queries are always rewritten +# into the two-phase GROUP BY to allow parallel accumulation across partitions. +query TT +EXPLAIN SELECT g, sum(DISTINCT v) AS distinct_v FROM t GROUP BY g; +---- +logical_plan +01)Projection: t.g, sum(alias1) AS distinct_v +02)--Aggregate: groupBy=[[t.g]], aggr=[[sum(alias1)]] +03)----Aggregate: groupBy=[[t.g, CAST(t.v AS Int64) AS alias1]], aggr=[[]] +04)------TableScan: t projection=[g, v] + +# When single_distinct_aggregation_prefer_memory = true, lone-distinct queries +# whose aggregate has a native GroupsAccumulator skip the rewrite to reduce peak memory. +statement ok +set datafusion.optimizer.single_distinct_aggregation_prefer_memory = true; + +query TT +EXPLAIN SELECT g, sum(DISTINCT v) AS distinct_v FROM t GROUP BY g; +---- +logical_plan +01)Projection: t.g, sum(DISTINCT t.v) AS distinct_v +02)--Aggregate: groupBy=[[t.g]], aggr=[[sum(DISTINCT CAST(t.v AS Int64))]] +03)----TableScan: t projection=[g, v] + +# Both paths produce identical query results +query II +SELECT g, sum(DISTINCT v) AS distinct_v FROM t GROUP BY g ORDER BY g; +---- +1 112 +2 3 +3 1 + +statement ok +set datafusion.optimizer.single_distinct_aggregation_prefer_memory = false; statement ok set datafusion.explain.logical_plan_only = false; + ######## # count(*) vs count(col) vs count(1), grouped ######## diff --git a/datafusion/sqllogictest/test_files/tpch/plans/q16.slt.part b/datafusion/sqllogictest/test_files/tpch/plans/q16.slt.part index df7a26bf0a20d..5d85e3932e557 100644 --- a/datafusion/sqllogictest/test_files/tpch/plans/q16.slt.part +++ b/datafusion/sqllogictest/test_files/tpch/plans/q16.slt.part @@ -51,37 +51,33 @@ limit 10; ---- logical_plan 01)Sort: supplier_cnt DESC NULLS FIRST, part.p_brand ASC NULLS LAST, part.p_type ASC NULLS LAST, part.p_size ASC NULLS LAST, fetch=10 -02)--Projection: part.p_brand, part.p_type, part.p_size, count(alias1) AS supplier_cnt -03)----Aggregate: groupBy=[[part.p_brand, part.p_type, part.p_size]], aggr=[[count(alias1)]] -04)------Aggregate: groupBy=[[part.p_brand, part.p_type, part.p_size, partsupp.ps_suppkey AS alias1]], aggr=[[]] -05)--------LeftAnti Join: partsupp.ps_suppkey = __correlated_sq_1.s_suppkey null_aware -06)----------Projection: partsupp.ps_suppkey, part.p_brand, part.p_type, part.p_size -07)------------Inner Join: partsupp.ps_partkey = part.p_partkey -08)--------------TableScan: partsupp projection=[ps_partkey, ps_suppkey] -09)--------------Filter: part.p_size IN ([Int32(49), Int32(14), Int32(23), Int32(45), Int32(19), Int32(3), Int32(36), Int32(9)]) AND part.p_brand != Utf8View("Brand#45") AND part.p_type NOT LIKE Utf8View("MEDIUM POLISHED%") -10)----------------TableScan: part projection=[p_partkey, p_brand, p_type, p_size], partial_filters=[part.p_size IN ([Int32(49), Int32(14), Int32(23), Int32(45), Int32(19), Int32(3), Int32(36), Int32(9)]), part.p_brand != Utf8View("Brand#45"), part.p_type NOT LIKE Utf8View("MEDIUM POLISHED%")] -11)----------SubqueryAlias: __correlated_sq_1 -12)------------Projection: supplier.s_suppkey -13)--------------Filter: supplier.s_comment LIKE Utf8View("%Customer%Complaints%") -14)----------------TableScan: supplier projection=[s_suppkey, s_comment], partial_filters=[supplier.s_comment LIKE Utf8View("%Customer%Complaints%")] +02)--Projection: part.p_brand, part.p_type, part.p_size, count(DISTINCT partsupp.ps_suppkey) AS supplier_cnt +03)----Aggregate: groupBy=[[part.p_brand, part.p_type, part.p_size]], aggr=[[count(DISTINCT partsupp.ps_suppkey)]] +04)------LeftAnti Join: partsupp.ps_suppkey = __correlated_sq_1.s_suppkey null_aware +05)--------Projection: partsupp.ps_suppkey, part.p_brand, part.p_type, part.p_size +06)----------Inner Join: partsupp.ps_partkey = part.p_partkey +07)------------TableScan: partsupp projection=[ps_partkey, ps_suppkey] +08)------------Filter: part.p_size IN ([Int32(49), Int32(14), Int32(23), Int32(45), Int32(19), Int32(3), Int32(36), Int32(9)]) AND part.p_brand != Utf8View("Brand#45") AND part.p_type NOT LIKE Utf8View("MEDIUM POLISHED%") +09)--------------TableScan: part projection=[p_partkey, p_brand, p_type, p_size], partial_filters=[part.p_size IN ([Int32(49), Int32(14), Int32(23), Int32(45), Int32(19), Int32(3), Int32(36), Int32(9)]), part.p_brand != Utf8View("Brand#45"), part.p_type NOT LIKE Utf8View("MEDIUM POLISHED%")] +10)--------SubqueryAlias: __correlated_sq_1 +11)----------Projection: supplier.s_suppkey +12)------------Filter: supplier.s_comment LIKE Utf8View("%Customer%Complaints%") +13)--------------TableScan: supplier projection=[s_suppkey, s_comment], partial_filters=[supplier.s_comment LIKE Utf8View("%Customer%Complaints%")] physical_plan 01)SortPreservingMergeExec: [supplier_cnt@3 DESC, p_brand@0 ASC NULLS LAST, p_type@1 ASC NULLS LAST, p_size@2 ASC NULLS LAST], fetch=10 -02)--ProjectionExec: expr=[p_brand@0 as p_brand, p_type@1 as p_type, p_size@2 as p_size, count(alias1)@3 as supplier_cnt] -03)----SortExec: TopK(fetch=10), expr=[count(alias1)@3 DESC, p_brand@0 ASC NULLS LAST, p_type@1 ASC NULLS LAST, p_size@2 ASC NULLS LAST], preserve_partitioning=[true] -04)------AggregateExec: mode=FinalPartitioned, gby=[p_brand@0 as p_brand, p_type@1 as p_type, p_size@2 as p_size], aggr=[count(alias1)] +02)--ProjectionExec: expr=[p_brand@0 as p_brand, p_type@1 as p_type, p_size@2 as p_size, count(DISTINCT partsupp.ps_suppkey)@3 as supplier_cnt] +03)----SortExec: TopK(fetch=10), expr=[count(DISTINCT partsupp.ps_suppkey)@3 DESC, p_brand@0 ASC NULLS LAST, p_type@1 ASC NULLS LAST, p_size@2 ASC NULLS LAST], preserve_partitioning=[true] +04)------AggregateExec: mode=FinalPartitioned, gby=[p_brand@0 as p_brand, p_type@1 as p_type, p_size@2 as p_size], aggr=[count(DISTINCT partsupp.ps_suppkey)] 05)--------RepartitionExec: partitioning=Hash([p_brand@0, p_type@1, p_size@2], 4), input_partitions=4 -06)----------AggregateExec: mode=Partial, gby=[p_brand@0 as p_brand, p_type@1 as p_type, p_size@2 as p_size], aggr=[count(alias1)] -07)------------AggregateExec: mode=FinalPartitioned, gby=[p_brand@0 as p_brand, p_type@1 as p_type, p_size@2 as p_size, alias1@3 as alias1], aggr=[] -08)--------------RepartitionExec: partitioning=Hash([p_brand@0, p_type@1, p_size@2, alias1@3], 4), input_partitions=4 -09)----------------AggregateExec: mode=Partial, gby=[p_brand@1 as p_brand, p_type@2 as p_type, p_size@3 as p_size, ps_suppkey@0 as alias1], aggr=[] -10)------------------HashJoinExec: mode=CollectLeft, join_type=LeftAnti, on=[(ps_suppkey@0, s_suppkey@0)], null_aware -11)--------------------CoalescePartitionsExec -12)----------------------HashJoinExec: mode=Partitioned, join_type=Inner, on=[(ps_partkey@0, p_partkey@0)], projection=[ps_suppkey@1, p_brand@3, p_type@4, p_size@5] -13)------------------------RepartitionExec: partitioning=Hash([ps_partkey@0], 4), input_partitions=4 -14)--------------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:0..2932049], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:2932049..5864098], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:5864098..8796147], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:8796147..11728193]]}, projection=[ps_partkey, ps_suppkey], constraints=[PrimaryKey([0, 1])], file_type=csv, has_header=false -15)------------------------RepartitionExec: partitioning=Hash([p_partkey@0], 4), input_partitions=4 -16)--------------------------FilterExec: p_size@3 IN (SET) ([49, 14, 23, 45, 19, 3, 36, 9]) AND p_brand@1 != Brand#45 AND p_type@2 NOT LIKE MEDIUM POLISHED% -17)----------------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/part.tbl:0..597773], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/part.tbl:597773..1195546], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/part.tbl:1195546..1793319], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/part.tbl:1793319..2391090]]}, projection=[p_partkey, p_brand, p_type, p_size], constraints=[PrimaryKey([0])], file_type=csv, has_header=false -18)--------------------FilterExec: s_comment@1 LIKE %Customer%Complaints%, projection=[s_suppkey@0] -19)----------------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 -20)------------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/supplier.tbl]]}, projection=[s_suppkey, s_comment], constraints=[PrimaryKey([0])], file_type=csv, has_header=false +06)----------AggregateExec: mode=Partial, gby=[p_brand@1 as p_brand, p_type@2 as p_type, p_size@3 as p_size], aggr=[count(DISTINCT partsupp.ps_suppkey)] +07)------------HashJoinExec: mode=CollectLeft, join_type=LeftAnti, on=[(ps_suppkey@0, s_suppkey@0)], null_aware +08)--------------CoalescePartitionsExec +09)----------------HashJoinExec: mode=Partitioned, join_type=Inner, on=[(ps_partkey@0, p_partkey@0)], projection=[ps_suppkey@1, p_brand@3, p_type@4, p_size@5] +10)------------------RepartitionExec: partitioning=Hash([ps_partkey@0], 4), input_partitions=4 +11)--------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:0..2932049], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:2932049..5864098], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:5864098..8796147], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:8796147..11728193]]}, projection=[ps_partkey, ps_suppkey], constraints=[PrimaryKey([0, 1])], file_type=csv, has_header=false +12)------------------RepartitionExec: partitioning=Hash([p_partkey@0], 4), input_partitions=4 +13)--------------------FilterExec: p_size@3 IN (SET) ([49, 14, 23, 45, 19, 3, 36, 9]) AND p_brand@1 != Brand#45 AND p_type@2 NOT LIKE MEDIUM POLISHED% +14)----------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/part.tbl:0..597773], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/part.tbl:597773..1195546], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/part.tbl:1195546..1793319], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/part.tbl:1793319..2391090]]}, projection=[p_partkey, p_brand, p_type, p_size], constraints=[PrimaryKey([0])], file_type=csv, has_header=false +15)--------------FilterExec: s_comment@1 LIKE %Customer%Complaints%, projection=[s_suppkey@0] +16)----------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +17)------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/supplier.tbl]]}, projection=[s_suppkey, s_comment], constraints=[PrimaryKey([0])], file_type=csv, has_header=false diff --git a/docs/source/user-guide/configs.md b/docs/source/user-guide/configs.md index ca81398508b17..f4b4a38f74548 100644 --- a/docs/source/user-guide/configs.md +++ b/docs/source/user-guide/configs.md @@ -189,6 +189,7 @@ The following configuration settings are available: | datafusion.optimizer.enable_sort_pushdown | true | Enable sort pushdown optimization. When enabled, attempts to push sort requirements down to data sources that can natively handle them (e.g., by reversing file/row group read order). Returns **inexact ordering**: Sort operator is kept for correctness, but optimized input enables early termination for TopK queries (ORDER BY ... LIMIT N), providing significant speedup. Memory: No additional overhead (only changes read order). Future: Will add option to detect perfectly sorted data and eliminate Sort completely. Default: true | | datafusion.optimizer.enable_leaf_expression_pushdown | true | When set to true, the optimizer will extract leaf expressions (such as `get_field`) from filter/sort/join nodes into projections closer to the leaf table scans, and push those projections down towards the leaf nodes. | | datafusion.optimizer.enable_unions_to_filter | false | When set to true, the logical optimizer will rewrite `UNION DISTINCT` branches that read from the same source and differ only by filter predicates into a single branch with a combined filter. This optimization is conservative and only applies when the branches share the same source and compatible wrapper nodes such as identical projections or aliases. | +| datafusion.optimizer.single_distinct_aggregation_prefer_memory | false | When `true`, the `SingleDistinctToGroupBy` optimizer rule applies its cost guard universally — including to lone-distinct queries that have no companion non-distinct aggregates. The rewrite is skipped whenever the distinct function already has a native `GroupsAccumulator`, reducing peak memory by avoiding the larger intermediate `(group, distinct value)` result set that the two-phase plan produces. The trade-off is reduced parallelism: skipping the rewrite causes the engine to fall back to serial per-group accumulation, which is typically 2–3× slower for queries with high-cardinality distinct values. When `false` (the default), lone-distinct queries are always rewritten into the two-phase `GROUP BY` so they can fan out across partitions and achieve maximum throughput. See | | datafusion.explain.logical_plan_only | false | When set to true, the explain statement will only print logical plans | | datafusion.explain.physical_plan_only | false | When set to true, the explain statement will only print physical plans | | datafusion.explain.show_statistics | false | When set to true, the explain statement will print operator statistics for physical plans |