diff --git a/datafusion/common/src/config.rs b/datafusion/common/src/config.rs index 63a7797e396ef..aefb0a5d30cfc 100644 --- a/datafusion/common/src/config.rs +++ b/datafusion/common/src/config.rs @@ -1941,15 +1941,22 @@ config_namespace! { pub enable_unions_to_filter: bool, default = false /// Prefer native grouped DISTINCT aggregation when target_partitions is 1 - /// and every aggregate explicitly supports a specialized GroupsAccumulator - /// for its DISTINCT argument types. Defaults to false, preserving the - /// two-phase GROUP BY rewrite. + /// and every DISTINCT argument is a fixed-width, non-float type (such as + /// integers, decimals, dates), instead of rewriting into a two-phase `GROUP BY`. + /// Defaults to false, preserving the two-phase GROUP BY rewrite. /// - /// The native plan can avoid materializing distinct (group, value) pairs, - /// but can be slower for many groups with duplicate values. It must fit - /// each group's distinct values in memory, so queries that succeed by - /// spilling with the rewrite may fail with this option enabled. One target - /// partition can still have multiple input partitions and merge native states. + /// The native plan avoids materializing intermediate `(group, distinct value)` + /// pairs, lowering peak memory when no memory limit is set. However, + /// native DISTINCT must fit each group's distinct values in memory: under + /// a memory limit, the rewrite can spill those pairs, but the native + /// accumulator cannot, so queries that succeed by spilling with the rewrite + /// may fail with this option enabled. + /// + /// Queries with multiple target partitions, non-distinct companion aggregates, + /// floating-point arguments, and aggregates with unsupported distinct handling + /// (such as `approx_median` or `stddev`) continue with the rewrite. + /// + /// See and pub single_distinct_prefer_native: bool, default = false } } diff --git a/datafusion/optimizer/src/single_distinct_to_groupby.rs b/datafusion/optimizer/src/single_distinct_to_groupby.rs index 73efe0de34475..c38b075c2930d 100644 --- a/datafusion/optimizer/src/single_distinct_to_groupby.rs +++ b/datafusion/optimizer/src/single_distinct_to_groupby.rs @@ -8,6 +8,7 @@ use std::sync::Arc; use crate::optimizer::ApplyOrder; use crate::{OptimizerConfig, OptimizerRule}; +use arrow::datatypes::DataType; use datafusion_common::{ DFSchema, DataFusionError, HashSet, Result, assert_or_internal_err, @@ -17,7 +18,7 @@ use datafusion_expr::builder::project; use datafusion_expr::expr::AggregateFunctionParams; use datafusion_expr::expr_schema::ExprSchemable; use datafusion_expr::{ - AggregateUDF, Expr, cast, col, + AggregateUDF, DistinctHandling, Expr, cast, col, expr::AggregateFunction, lit, logical_plan::{Aggregate, LogicalPlan}, @@ -69,8 +70,8 @@ use datafusion_expr::{ /// /// See `rewrite_pays_for_count` for when the `count` is allowed. /// With `single_distinct_prefer_native` enabled and one target partition, grouped -/// all-DISTINCT aggregates stay native if every function explicitly supports a -/// specialized `GroupsAccumulator` for its argument types. +/// all-DISTINCT aggregates stay native when every DISTINCT argument is a +/// fixed-width, non-float type (such as integers, decimals, dates). #[derive(Default, Debug)] pub struct SingleDistinctToGroupBy {} @@ -110,7 +111,56 @@ impl CountRollup { } } -/// Check whether all aggregate exprs are distinct on a single field. +/// Returns `true` for fixed-width primitive types where single-phase distinct +/// accumulation avoids heap allocation churn **and** where bit-identity equals +/// value-equality, so that skipping the GROUP BY rewrite cannot silently change +/// query results. +/// +/// Float types are deliberately excluded: IEEE 754 defines `-0.0 == 0.0`, so a +/// `GROUP BY` on floats merges those two bit-patterns into one group. The +/// native distinct accumulators hash raw bits and therefore count them +/// separately. Keeping floats in this list would let `single_distinct_prefer_native` +/// skip the rewrite for a float column, producing a different answer than the +/// non-rewritten path — a silent correctness regression. +fn is_fixed_width_type(dt: &DataType) -> bool { + matches!( + dt, + DataType::Boolean + | DataType::Int8 + | DataType::Int16 + | DataType::Int32 + | DataType::Int64 + | DataType::UInt8 + | DataType::UInt16 + | DataType::UInt32 + | DataType::UInt64 + | DataType::Date32 + | DataType::Date64 + | DataType::Time32(_) + | DataType::Time64(_) + | DataType::Timestamp(_, _) + | DataType::Duration(_) + | DataType::Interval(_) + | DataType::Decimal128(_, _) + | DataType::Decimal256(_, _) + ) +} + +/// Check whether all aggregate exprs are eligible for the single-distinct-to-group-by rewrite. +/// +/// With `use_native_grouped_distinct` enabled (which requires one target partition, +/// non-empty `GROUP BY`, and `single_distinct_prefer_native = true`), the rewrite +/// is skipped for all-DISTINCT queries on fixed-width primitive integer/decimal/date +/// types, where it avoids materializing intermediate `(group, distinct value)` pairs. +/// Float types (`Float16`, `Float32`, `Float64`) are **not** candidates for skipping: +/// their native accumulators hash raw bits, counting `-0.0` and `0.0` separately, +/// while `GROUP BY` treats them as equal under IEEE 754. Skipping the rewrite for +/// floats would silently produce wrong counts and averages. +/// Variable-length types (strings, binary), un-grouped aggregates, queries with +/// companion non-distinct aggregates, and multi-partition plans always continue +/// with the 2-phase rewrite. +/// Aggregates with [`DistinctHandling::Unsupported`] (e.g. `approx_median`) cannot +/// run natively with `DISTINCT` and are always rewritten. fn is_single_distinct_agg( aggr_expr: &[Expr], input_schema: &DFSchema, @@ -158,8 +208,17 @@ fn is_single_distinct_agg( if aggregate_count != aggr_expr.len() || fields_set.len() != 1 { return Ok(false); } - // With one target partition, a native GroupsAccumulator avoids building an - // extra aggregate with a row per (group, distinct value) pair. With multiple + // Aggregates with DistinctHandling::Unsupported (e.g. approx_median) cannot + // run natively with DISTINCT; they must always be rewritten. + if distinct_aggs + .iter() + .any(|(func, _)| func.distinct_handling() == DistinctHandling::Unsupported) + { + return Ok(true); + } + + // With one target partition, native accumulation avoids building an extra + // aggregate with a row per (group, distinct value) pair. With multiple // partitions, the rewrite also distributes deduplication by both keys; // keeping DISTINCT instead repartitions lists of values by group and merges // them in the final accumulator. That can concentrate work on a few groups. @@ -167,15 +226,23 @@ fn is_single_distinct_agg( // ALL or an existing partitioned table), requiring Partial/Final state merging. // This is opt-in: the rule cannot see memory limits or cost the input, and // native DISTINCT must fit each group's distinct values in memory. - if use_native_grouped_distinct - && distinct_aggs.len() == aggregate_count - // An unknown answer must preserve the existing rewrite. - && distinct_groups_support(&distinct_aggs, input_schema)? + // + // The rewrite is skipped only when every aggregate is a DISTINCT aggregate + // over fixed-width primitive types (integers, decimals, dates, etc.) where + // bit-identity equals value-equality. Floats are excluded to avoid -0.0 vs + // 0.0 IEEE 754 discrepancy, and variable-length types (strings, binary) are + // excluded to prevent heap allocation churn. + if use_native_grouped_distinct && distinct_aggs.len() == aggregate_count { + let arg_types = distinct_aggs .iter() - .all(|support| *support == Some(true)) - { - return Ok(false); + .flat_map(|(_, args)| args.iter()) + .map(|arg| arg.get_type(input_schema)) + .collect::>>()?; + if arg_types.iter().all(is_fixed_width_type) { + return Ok(false); + } } + if has_count_rollup && !rewrite_pays_for_count(&distinct_aggs, input_schema)? { return Ok(false); } @@ -591,6 +658,22 @@ mod tests { table_scan(Some("test"), &schema, None)?.build() } + /// `test` with a `Float64` `b`. + /// + /// Used to verify that `single_distinct_prefer_native` does **not** skip the rewrite for + /// float columns: `-0.0` and `0.0` are equal under IEEE 754 (GROUP BY + /// treats them as one group) but are bit-distinct (native accumulators + /// count them separately). Skipping the rewrite for floats would silently + /// change query results. + fn test_table_scan_float64_b() -> Result { + let schema = Schema::new(vec![ + Field::new("a", DataType::UInt32, false), + Field::new("b", DataType::Float64, 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`. @@ -602,17 +685,30 @@ mod tests { impl RegistryOptimizerContext { fn new() -> Self { + Self::with_prefer_native(false) + } + + fn with_prefer_native_and_partitions( + prefer_native: bool, + target_partitions: usize, + ) -> Self { let mut registry = MemoryFunctionRegistry::new(); registry.register_udaf(count_udaf()).unwrap(); registry.register_udaf(sum_udaf()).unwrap(); // Keep logical-plan snapshots independent of the available CPUs. let mut options = OptimizerContext::new().options(); - Arc::make_mut(&mut options).execution.target_partitions = 4; + let opts = Arc::make_mut(&mut options); + opts.execution.target_partitions = target_partitions; + opts.optimizer.single_distinct_prefer_native = prefer_native; Self { inner: OptimizerContext::new_with_config_options(options), registry, } } + + fn with_prefer_native(prefer_native: bool) -> Self { + Self::with_prefer_native_and_partitions(prefer_native, 1) + } } impl OptimizerConfig for RegistryOptimizerContext { @@ -650,6 +746,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()?; @@ -1366,4 +1480,220 @@ mod tests { assert_eq!(plan.schema().as_ref(), optimized_plan.schema().as_ref()); Ok(()) } + + #[derive(Debug, Clone, PartialEq, Eq, Hash)] + struct UnsupportedUdaf { + signature: Signature, + } + + impl UnsupportedUdaf { + fn new() -> Self { + Self { + signature: Signature::exact( + vec![DataType::UInt32], + Volatility::Immutable, + ), + } + } + } + + impl AggregateUDFImpl for UnsupportedUdaf { + fn name(&self) -> &str { + "unsupported" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Ok(DataType::UInt32) + } + + fn accumulator( + &self, + _acc_args: AccumulatorArgs, + ) -> Result> { + unimplemented!("not needed for optimizer tests") + } + + fn distinct_handling(&self) -> DistinctHandling { + DistinctHandling::Unsupported + } + } + + fn unsupported_distinct(expr: Expr) -> Expr { + Expr::AggregateFunction(AggregateFunction::new_udf( + Arc::new(AggregateUDF::from(UnsupportedUdaf::new())), + vec![expr], + true, + None, + vec![], + None, + )) + } + + #[test] + fn single_distinct_prefer_native_default_rewrites_lone_distinct() -> Result<()> { + let table_scan = test_table_scan()?; + + let plan = LogicalPlanBuilder::from(table_scan) + .aggregate(vec![col("a")], vec![count_distinct(col("b"))])? + .build()?; + + assert_optimized_plan_equal!( + plan, + @r" + 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] + " + ) + } + + #[test] + fn single_distinct_prefer_native_true_skips_rewrite_for_groups_accumulator() + -> Result<()> { + let table_scan = test_table_scan()?; + + let plan = LogicalPlanBuilder::from(table_scan) + .aggregate(vec![col("a")], vec![count_distinct(col("b"))])? + .build()?; + + let ctx = RegistryOptimizerContext::with_prefer_native(true); + 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] + // single_distinct_prefer_native=true must NOT skip the rewrite when there is no GROUP BY: + // without the 2-phase path DataFusion cannot parallelize un-grouped distinct + // work and the whole aggregation is forced onto a single thread. + fn single_distinct_prefer_native_true_rewrites_when_no_group_by() -> Result<()> { + let table_scan = test_table_scan()?; + + // No GROUP BY: only a bare count(DISTINCT b). + let plan = LogicalPlanBuilder::from(table_scan) + .aggregate(Vec::::new(), vec![count_distinct(col("b"))])? + .build()?; + + let ctx = RegistryOptimizerContext::with_prefer_native(true); + assert_optimized_plan_equal_with_context!( + ctx, + 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] + TableScan: test [a:UInt32, b:UInt32, c:UInt32] + " + ) + } + + #[test] + // single_distinct_prefer_native=true must NOT skip the rewrite when target_partitions > 1: + // parallel accumulation regresses throughput and raises memory under limits. + fn single_distinct_prefer_native_true_rewrites_when_target_partitions_greater_than_one() + -> Result<()> { + let table_scan = test_table_scan()?; + + let plan = LogicalPlanBuilder::from(table_scan) + .aggregate(vec![col("a")], vec![count_distinct(col("b"))])? + .build()?; + + let ctx = RegistryOptimizerContext::with_prefer_native_and_partitions(true, 4); + assert_optimized_plan_equal_with_context!( + ctx, + plan, + @r" + 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] + " + ) + } + + #[test] + // single_distinct_prefer_native=true must NOT skip the rewrite for Utf8 columns. + // Single-phase distinct on strings causes heavy per-value heap allocation churn + // and actually increases peak memory compared to the 2-phase rewrite. + fn single_distinct_prefer_native_true_rewrites_for_utf8() -> Result<()> { + // test_table_scan_utf8_b() has column `b` as Utf8. + let table_scan = test_table_scan_utf8_b()?; + + let plan = LogicalPlanBuilder::from(table_scan) + .aggregate(vec![col("a")], vec![count_distinct(col("b"))])? + .build()?; + + let ctx = RegistryOptimizerContext::with_prefer_native(true); + assert_optimized_plan_equal_with_context!( + ctx, + plan, + @r" + 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:Utf8] + TableScan: test [a:UInt32, b:Utf8, c:UInt32] + " + ) + } + + #[test] + fn single_distinct_prefer_native_true_still_rewrites_unsupported() -> Result<()> { + let table_scan = test_table_scan()?; + + let plan = LogicalPlanBuilder::from(table_scan) + .aggregate(vec![col("a")], vec![unsupported_distinct(col("b"))])? + .build()?; + + let ctx = RegistryOptimizerContext::with_prefer_native(true); + assert_optimized_plan_equal_with_context!( + ctx, + plan, + @r" + Projection: test.a, unsupported(alias1) AS unsupported(DISTINCT test.b) [a:UInt32, unsupported(DISTINCT test.b):UInt32;N] + Aggregate: groupBy=[[test.a]], aggr=[[unsupported(alias1)]] [a:UInt32, unsupported(alias1):UInt32;N] + Aggregate: groupBy=[[test.a, test.b AS alias1]], aggr=[[]] [a:UInt32, alias1:UInt32] + TableScan: test [a:UInt32, b:UInt32, c:UInt32] + " + ) + } + + #[test] + // single_distinct_prefer_native=true must NOT skip the rewrite for Float64 columns. + // + // The native distinct accumulators hash raw bits, so they count `-0.0` and + // `0.0` as two separate values. A GROUP BY on those same floats collapses + // them into one group under IEEE 754 equality. Skipping the rewrite for + // floats would silently change `count(DISTINCT f)` from 2 to 3 (and + // `avg(DISTINCT f)` from 0.75 to 0.5) on a table containing (0.0, -0.0, 1.5). + // The rewrite must still fire so the two-phase path's GROUP BY-based + // deduplication is what actually runs. + fn single_distinct_prefer_native_true_rewrites_for_float64() -> Result<()> { + let table_scan = test_table_scan_float64_b()?; + + let plan = LogicalPlanBuilder::from(table_scan) + .aggregate(vec![col("a")], vec![count_distinct(col("b"))])? + .build()?; + + let ctx = RegistryOptimizerContext::with_prefer_native(true); + assert_optimized_plan_equal_with_context!( + ctx, + plan, + @r" + 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:Float64] + TableScan: test [a:UInt32, b:Float64, c:UInt32] + " + ) + } } diff --git a/datafusion/sqllogictest/test_files/grouped_count_distinct.slt b/datafusion/sqllogictest/test_files/grouped_count_distinct.slt index 1490442ab383a..38c4dd42b11b1 100644 --- a/datafusion/sqllogictest/test_files/grouped_count_distinct.slt +++ b/datafusion/sqllogictest/test_files/grouped_count_distinct.slt @@ -2,8 +2,8 @@ # SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership # SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements -# With the opt-in enabled and one target partition, grouped integer DISTINCT -# counts keep their specialized accumulator. Multiple target partitions retain +# With the opt-in enabled and one target partition, grouped all-DISTINCT aggregates +# with fixed-width, non-float arguments stay native. Multiple target partitions retain # distributed deduplication. # String and floating-point DISTINCT arguments still use the two-phase rewrite. @@ -80,15 +80,20 @@ logical_plan 03)----Aggregate: groupBy=[[numeric_counts.g, numeric_counts.x AS alias1]], aggr=[[sum(CAST(numeric_counts.x AS Int64)) AS alias2]] 04)------TableScan: numeric_counts projection=[g, x] -# `sum` does not report specialized DISTINCT support, so it keeps the rewrite. +# `sum` with integer DISTINCT argument stays native. query TT EXPLAIN SELECT g, sum(DISTINCT x) FROM numeric_counts GROUP BY g; ---- logical_plan -01)Projection: numeric_counts.g, sum(alias1) AS sum(DISTINCT numeric_counts.x) -02)--Aggregate: groupBy=[[numeric_counts.g]], aggr=[[sum(alias1)]] -03)----Aggregate: groupBy=[[numeric_counts.g, CAST(numeric_counts.x AS Int64) AS alias1]], aggr=[[]] -04)------TableScan: numeric_counts projection=[g, x] +01)Aggregate: groupBy=[[numeric_counts.g]], aggr=[[sum(DISTINCT CAST(numeric_counts.x AS Int64))]] +02)--TableScan: numeric_counts projection=[g, x] + +query II rowsort +SELECT g, sum(DISTINCT x) FROM numeric_counts GROUP BY g; +---- +1 3 +2 NULL +3 3 # The global aggregate does not use a GroupsAccumulator. query TT diff --git a/datafusion/sqllogictest/test_files/information_schema.slt b/datafusion/sqllogictest/test_files/information_schema.slt index c1aabc2946ea6..0d3b315ebb640 100644 --- a/datafusion/sqllogictest/test_files/information_schema.slt +++ b/datafusion/sqllogictest/test_files/information_schema.slt @@ -505,7 +505,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_prefer_native false Prefer native grouped DISTINCT aggregation when target_partitions is 1 and every aggregate explicitly supports a specialized GroupsAccumulator for its DISTINCT argument types. Defaults to false, preserving the two-phase GROUP BY rewrite. The native plan can avoid materializing distinct (group, value) pairs, but can be slower for many groups with duplicate values. It must fit each group's distinct values in memory, so queries that succeed by spilling with the rewrite may fail with this option enabled. One target partition can still have multiple input partitions and merge native states. +datafusion.optimizer.single_distinct_prefer_native false Prefer native grouped DISTINCT aggregation when target_partitions is 1 and every DISTINCT argument is a fixed-width, non-float type (such as integers, decimals, dates), instead of rewriting into a two-phase `GROUP BY`. Defaults to false, preserving the two-phase GROUP BY rewrite. The native plan avoids materializing intermediate `(group, distinct value)` pairs, lowering peak memory when no memory limit is set. However, native DISTINCT must fit each group's distinct values in memory: under a memory limit, the rewrite can spill those pairs, but the native accumulator cannot, so queries that succeed by spilling with the rewrite may fail with this option enabled. Queries with multiple target partitions, non-distinct companion aggregates, floating-point arguments, and aggregates with unsupported distinct handling (such as `approx_median` or `stddev`) continue with the rewrite. See and 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/single_distinct_to_groupby.slt b/datafusion/sqllogictest/test_files/single_distinct_to_groupby.slt index 78308f7897c6e..a4335bedd8c25 100644 --- a/datafusion/sqllogictest/test_files/single_distinct_to_groupby.slt +++ b/datafusion/sqllogictest/test_files/single_distinct_to_groupby.slt @@ -470,3 +470,101 @@ physical_plan statement ok DROP TABLE t_decimal_distinct; + + +######## +# single_distinct_prefer_native config option +######## + +statement ok +CREATE TABLE t_prefer_mem AS SELECT * FROM (VALUES (1, 10), (1, 20), (2, 10)) AS v(g, x); + +# By default (single_distinct_prefer_native = false), lone-distinct queries rewrite to two-phase GROUP BY +query TT +EXPLAIN SELECT g, count(DISTINCT x) FROM t_prefer_mem GROUP BY g; +---- +logical_plan +01)Projection: t_prefer_mem.g, count(alias1) AS count(DISTINCT t_prefer_mem.x) +02)--Aggregate: groupBy=[[t_prefer_mem.g]], aggr=[[count(alias1)]] +03)----Aggregate: groupBy=[[t_prefer_mem.g, t_prefer_mem.x AS alias1]], aggr=[[]] +04)------TableScan: t_prefer_mem projection=[g, x] +physical_plan +01)ProjectionExec: expr=[g@0 as g, count(alias1)@1 as count(DISTINCT t_prefer_mem.x)] +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, x@1 as alias1], aggr=[] +08)--------------DataSourceExec: partitions=1, partition_sizes=[1] + +# stddev cannot run DISTINCT itself, so it is rewritten even next to count(*). +query IIR rowsort +SELECT g, count(*), stddev(DISTINCT x) FROM t_prefer_mem GROUP BY g; +---- +1 2 7.071067811865 +2 1 NULL + +# When single_distinct_prefer_native = true and target_partitions = 1, the rewrite is skipped to reduce peak memory +statement ok +set datafusion.execution.target_partitions = 1; + +statement ok +set datafusion.optimizer.single_distinct_prefer_native = true; + +query TT +EXPLAIN SELECT g, count(DISTINCT x) FROM t_prefer_mem GROUP BY g; +---- +logical_plan +01)Aggregate: groupBy=[[t_prefer_mem.g]], aggr=[[count(DISTINCT t_prefer_mem.x)]] +02)--TableScan: t_prefer_mem projection=[g, x] +physical_plan +01)AggregateExec: mode=Single, gby=[g@0 as g], aggr=[count(DISTINCT t_prefer_mem.x)] +02)--DataSourceExec: partitions=1, partition_sizes=[1] + +# Query results are identical under both settings +query II +SELECT g, count(DISTINCT x) FROM t_prefer_mem GROUP BY g ORDER BY g; +---- +1 2 +2 1 + +# Float DISTINCT keeps the rewrite: the native accumulator counts -0.0 and 0.0 separately. +statement ok +CREATE TABLE t_prefer_mem_float AS SELECT * FROM (VALUES (1, CAST('0.0' AS DOUBLE)), (1, CAST('-0.0' AS DOUBLE)), (1, CAST('1.5' AS DOUBLE))) AS v(g, f); + +query IIR +SELECT g, count(DISTINCT f), avg(DISTINCT f) FROM t_prefer_mem_float GROUP BY g; +---- +1 2 0.75 + +statement ok +DROP TABLE t_prefer_mem_float; + +# Multiple target partitions retain the rewrite even with single_distinct_prefer_native = true +statement ok +set datafusion.execution.target_partitions = 4; + +query TT +EXPLAIN SELECT g, count(DISTINCT x) FROM t_prefer_mem GROUP BY g; +---- +logical_plan +01)Projection: t_prefer_mem.g, count(alias1) AS count(DISTINCT t_prefer_mem.x) +02)--Aggregate: groupBy=[[t_prefer_mem.g]], aggr=[[count(alias1)]] +03)----Aggregate: groupBy=[[t_prefer_mem.g, t_prefer_mem.x AS alias1]], aggr=[[]] +04)------TableScan: t_prefer_mem projection=[g, x] +physical_plan +01)ProjectionExec: expr=[g@0 as g, count(alias1)@1 as count(DISTINCT t_prefer_mem.x)] +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, x@1 as alias1], aggr=[] +08)--------------DataSourceExec: partitions=1, partition_sizes=[1] + +statement ok +set datafusion.optimizer.single_distinct_prefer_native = false; + +statement ok +DROP TABLE t_prefer_mem; diff --git a/docs/source/user-guide/configs.md b/docs/source/user-guide/configs.md index 2c3ca887d4a38..19c6e5bffd78c 100644 --- a/docs/source/user-guide/configs.md +++ b/docs/source/user-guide/configs.md @@ -178,7 +178,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_prefer_native | false | Prefer native grouped DISTINCT aggregation when target_partitions is 1 and every aggregate explicitly supports a specialized GroupsAccumulator for its DISTINCT argument types. Defaults to false, preserving the two-phase GROUP BY rewrite. The native plan can avoid materializing distinct (group, value) pairs, but can be slower for many groups with duplicate values. It must fit each group's distinct values in memory, so queries that succeed by spilling with the rewrite may fail with this option enabled. One target partition can still have multiple input partitions and merge native states. | +| datafusion.optimizer.single_distinct_prefer_native | false | Prefer native grouped DISTINCT aggregation when target_partitions is 1 and every DISTINCT argument is a fixed-width, non-float type (such as integers, decimals, dates), instead of rewriting into a two-phase `GROUP BY`. Defaults to false, preserving the two-phase GROUP BY rewrite. The native plan avoids materializing intermediate `(group, distinct value)` pairs, lowering peak memory when no memory limit is set. However, native DISTINCT must fit each group's distinct values in memory: under a memory limit, the rewrite can spill those pairs, but the native accumulator cannot, so queries that succeed by spilling with the rewrite may fail with this option enabled. Queries with multiple target partitions, non-distinct companion aggregates, floating-point arguments, and aggregates with unsupported distinct handling (such as `approx_median` or `stddev`) continue with the rewrite. See and | | 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 |