perf: add opt-in native grouped integer DISTINCT counts

benchmarks/bench.sh

@@ -163,6 +163,7 @@ hj:                     Benchmark for simple hash joins, testing various join sc
 smj:                    Benchmark for simple sort merge joins, testing various join scenarios
 dict:                   Benchmark for dictionary-encoded group-by scenarios
 array_agg_distinct:     1000K-group, two-row-per-group array_agg(DISTINCT) benchmark
+grouped_count_distinct: Grouped count(DISTINCT) by key type, group count and distinct cardinality
 compile_profile:        Compile and execute TPC-H across selected Cargo profiles, reporting timing and binary size
 
 
@@ -280,6 +281,9 @@ main() {
                     # Data is generated inline by the suite's load SQL.
                     echo "projection_subquery: no external data to generate"
                     ;;
+                grouped_count_distinct)
+                    echo "grouped_count_distinct: no external data to generate"
+                    ;;
                 asof_join)
                     data_asof_join
                     ;;
@@ -681,6 +685,9 @@ main() {
                 array_agg_distinct)
                     run_array_agg_distinct
                     ;;
+                grouped_count_distinct)
+                    run_grouped_count_distinct
+                    ;;
                 compile_profile)
                     run_compile_profile "${PROFILE_ARGS[@]}"
                     ;;
@@ -1813,6 +1820,15 @@ run_dict() {
     debug_run $CARGO_COMMAND --bin dfbench -- dict --iterations 5 -o "${RESULTS_FILE}" ${QUERY_ARG} ${LATENCY_ARG}
 }
 
+# Runs grouped DISTINCT counts over inline data with varying key types and cardinalities.
+run_grouped_count_distinct() {
+    debug_run env BENCH_NAME=grouped_count_distinct \
+      BENCH_RESULTS_FILE="$(sql_results_file grouped_count_distinct)" \
+      GROUPED_DISTINCT_ROWS="${GROUPED_DISTINCT_ROWS:-4000000}" \
+      ${QUERY:+BENCH_QUERY="${QUERY}"} \
+      bash -c "$SQL_CARGO_COMMAND"
+}
+
 # Runs the data-free high-cardinality array_agg(DISTINCT) SQL benchmark.
 run_array_agg_distinct() {
     echo "Running array_agg_distinct benchmark..."

benchmarks/sql_benchmarks/grouped_count_distinct/README.md

@@ -0,0 +1,30 @@
+<!--
+SPDX-License-Identifier: Apache-2.0
+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
+-->
+
+# Grouped COUNT(DISTINCT)
+
+Compare the default rewrite with the native accumulator opt-in at one target
+partition. The suite reads DataFusion settings from the environment; it does not
+enable the opt-in itself.
+
+```bash
+DATAFUSION_OPTIMIZER_SINGLE_DISTINCT_PREFER_NATIVE=false cargo run --release --bin benchmark_runner -- grouped_count_distinct --iterations 10 --partitions 1 --mem-pool-type greedy --memory-limit 1G
+DATAFUSION_OPTIMIZER_SINGLE_DISTINCT_PREFER_NATIVE=true cargo run --release --bin benchmark_runner -- grouped_count_distinct --iterations 10 --partitions 1 --mem-pool-type greedy --memory-limit 1G
+```
+
+Run from `benchmarks/`. Repeat at 4 and 12 target partitions to check the
+distributed rewrite, which remains selected with either flag value. The same
+environment variable works with `./bench.sh run grouped_count_distinct`.
+
+Q01-Q09 vary argument and grouping-key types, group count, distinct cardinality,
+and skew. Q10 covers the unfavorable native shape: at the default four million
+rows it has one million groups, each with four copies of the same value. Each
+native input row probes both the group table and the per-group distinct set;
+the rewrite's outer aggregate only sees the deduplicated pairs.
+
+Every query checks its result against an explicit distinct-pairs-then-count
+query before timing. Keep the memory pool and limit the same on both sides.
+The combined correctness query needs more memory than the measured aggregate.

benchmarks/sql_benchmarks/grouped_count_distinct/benchmarks/q01.benchmark

@@ -0,0 +1,25 @@
+name Q01
+group grouped_count_distinct
+
+load sql_benchmarks/grouped_count_distinct/init/load.sql
+
+assert I
+WITH actual AS (
+  SELECT g AS g, count(DISTINCT x) AS n FROM grouped_distinct GROUP BY g
+), expected AS (
+  SELECT g, count(x) AS n
+  FROM (SELECT DISTINCT g AS g, x AS x FROM grouped_distinct)
+  GROUP BY g
+)
+SELECT NOT EXISTS (SELECT * FROM actual EXCEPT SELECT * FROM expected)
+   AND NOT EXISTS (SELECT * FROM expected EXCEPT SELECT * FROM actual);
+----
+true
+
+expect_plan AggregateExec
+
+run
+-- Int64, 2,000 groups, high distinct cardinality (issue #25936).
+SELECT g, count(DISTINCT x) FROM grouped_distinct GROUP BY g;
+
+cleanup sql_benchmarks/grouped_count_distinct/init/cleanup.sql

benchmarks/sql_benchmarks/grouped_count_distinct/benchmarks/q02.benchmark

@@ -0,0 +1,25 @@
+name Q02
+group grouped_count_distinct
+
+load sql_benchmarks/grouped_count_distinct/init/load.sql
+
+assert I
+WITH actual AS (
+  SELECT g AS g, count(DISTINCT low_x) AS n FROM grouped_distinct GROUP BY g
+), expected AS (
+  SELECT g, count(x) AS n
+  FROM (SELECT DISTINCT g AS g, low_x AS x FROM grouped_distinct)
+  GROUP BY g
+)
+SELECT NOT EXISTS (SELECT * FROM actual EXCEPT SELECT * FROM expected)
+   AND NOT EXISTS (SELECT * FROM expected EXCEPT SELECT * FROM actual);
+----
+true
+
+expect_plan AggregateExec
+
+run
+-- Int64, 2,000 groups, at most eight distinct values per group.
+SELECT g, count(DISTINCT low_x) FROM grouped_distinct GROUP BY g;
+
+cleanup sql_benchmarks/grouped_count_distinct/init/cleanup.sql

benchmarks/sql_benchmarks/grouped_count_distinct/benchmarks/q03.benchmark

@@ -0,0 +1,25 @@
+name Q03
+group grouped_count_distinct
+
+load sql_benchmarks/grouped_count_distinct/init/load.sql
+
+assert I
+WITH actual AS (
+  SELECT few_g AS g, count(DISTINCT x) AS n FROM grouped_distinct GROUP BY few_g
+), expected AS (
+  SELECT g, count(x) AS n
+  FROM (SELECT DISTINCT few_g AS g, x AS x FROM grouped_distinct)
+  GROUP BY g
+)
+SELECT NOT EXISTS (SELECT * FROM actual EXCEPT SELECT * FROM expected)
+   AND NOT EXISTS (SELECT * FROM expected EXCEPT SELECT * FROM actual);
+----
+true
+
+expect_plan AggregateExec
+
+run
+-- Int64, eight groups, high distinct cardinality.
+SELECT few_g, count(DISTINCT x) FROM grouped_distinct GROUP BY few_g;
+
+cleanup sql_benchmarks/grouped_count_distinct/init/cleanup.sql

benchmarks/sql_benchmarks/grouped_count_distinct/benchmarks/q04.benchmark

@@ -0,0 +1,25 @@
+name Q04
+group grouped_count_distinct
+
+load sql_benchmarks/grouped_count_distinct/init/load.sql
+
+assert I
+WITH actual AS (
+  SELECT g AS g, count(DISTINCT unsigned_x) AS n FROM grouped_distinct GROUP BY g
+), expected AS (
+  SELECT g, count(x) AS n
+  FROM (SELECT DISTINCT g AS g, unsigned_x AS x FROM grouped_distinct)
+  GROUP BY g
+)
+SELECT NOT EXISTS (SELECT * FROM actual EXCEPT SELECT * FROM expected)
+   AND NOT EXISTS (SELECT * FROM expected EXCEPT SELECT * FROM actual);
+----
+true
+
+expect_plan AggregateExec
+
+run
+-- UInt64, 2,000 groups, high distinct cardinality.
+SELECT g, count(DISTINCT unsigned_x) FROM grouped_distinct GROUP BY g;
+
+cleanup sql_benchmarks/grouped_count_distinct/init/cleanup.sql

benchmarks/sql_benchmarks/grouped_count_distinct/benchmarks/q05.benchmark

@@ -0,0 +1,25 @@
+name Q05
+group grouped_count_distinct
+
+load sql_benchmarks/grouped_count_distinct/init/load.sql
+
+assert I
+WITH actual AS (
+  SELECT g AS g, count(DISTINCT string_x) AS n FROM grouped_distinct GROUP BY g
+), expected AS (
+  SELECT g, count(x) AS n
+  FROM (SELECT DISTINCT g AS g, string_x AS x FROM grouped_distinct)
+  GROUP BY g
+)
+SELECT NOT EXISTS (SELECT * FROM actual EXCEPT SELECT * FROM expected)
+   AND NOT EXISTS (SELECT * FROM expected EXCEPT SELECT * FROM actual);
+----
+true
+
+expect_plan AggregateExec
+
+run
+-- VARCHAR control: keep the existing two-phase rewrite.
+SELECT g, count(DISTINCT string_x) FROM grouped_distinct GROUP BY g;
+
+cleanup sql_benchmarks/grouped_count_distinct/init/cleanup.sql

benchmarks/sql_benchmarks/grouped_count_distinct/benchmarks/q06.benchmark

@@ -0,0 +1,25 @@
+name Q06
+group grouped_count_distinct
+
+load sql_benchmarks/grouped_count_distinct/init/load.sql
+
+assert I
+WITH actual AS (
+  SELECT g AS g, count(DISTINCT float_x) AS n FROM grouped_distinct GROUP BY g
+), expected AS (
+  SELECT g, count(x) AS n
+  FROM (SELECT DISTINCT g AS g, float_x AS x FROM grouped_distinct)
+  GROUP BY g
+)
+SELECT NOT EXISTS (SELECT * FROM actual EXCEPT SELECT * FROM expected)
+   AND NOT EXISTS (SELECT * FROM expected EXCEPT SELECT * FROM actual);
+----
+true
+
+expect_plan AggregateExec
+
+run
+-- Float64 control: no specialized DISTINCT GroupsAccumulator.
+SELECT g, count(DISTINCT float_x) FROM grouped_distinct GROUP BY g;
+
+cleanup sql_benchmarks/grouped_count_distinct/init/cleanup.sql

benchmarks/sql_benchmarks/grouped_count_distinct/benchmarks/q07.benchmark

@@ -0,0 +1,25 @@
+name Q07
+group grouped_count_distinct
+
+load sql_benchmarks/grouped_count_distinct/init/load.sql
+
+assert I
+WITH actual AS (
+  SELECT skew_g AS g, count(DISTINCT x) AS n FROM grouped_distinct GROUP BY skew_g
+), expected AS (
+  SELECT g, count(x) AS n
+  FROM (SELECT DISTINCT skew_g AS g, x AS x FROM grouped_distinct)
+  GROUP BY g
+)
+SELECT NOT EXISTS (SELECT * FROM actual EXCEPT SELECT * FROM expected)
+   AND NOT EXISTS (SELECT * FROM expected EXCEPT SELECT * FROM actual);
+----
+true
+
+expect_plan AggregateExec
+
+run
+-- Int64 argument, eight integer groups; at least 90% of rows in one hot group.
+SELECT skew_g, count(DISTINCT x) FROM grouped_distinct GROUP BY skew_g;
+
+cleanup sql_benchmarks/grouped_count_distinct/init/cleanup.sql

benchmarks/sql_benchmarks/grouped_count_distinct/benchmarks/q08.benchmark

@@ -0,0 +1,25 @@
+name Q08
+group grouped_count_distinct
+
+load sql_benchmarks/grouped_count_distinct/init/load.sql
+
+assert I
+WITH actual AS (
+  SELECT string_g AS g, count(DISTINCT x) AS n FROM grouped_distinct GROUP BY string_g
+), expected AS (
+  SELECT g, count(x) AS n
+  FROM (SELECT DISTINCT string_g AS g, x AS x FROM grouped_distinct)
+  GROUP BY g
+)
+SELECT NOT EXISTS (SELECT * FROM actual EXCEPT SELECT * FROM expected)
+   AND NOT EXISTS (SELECT * FROM expected EXCEPT SELECT * FROM actual);
+----
+true
+
+expect_plan AggregateExec
+
+run
+-- Int64 argument, eight string groups with a hot group (ClickBench Q10 shape).
+SELECT string_g, count(DISTINCT x) FROM grouped_distinct GROUP BY string_g;
+
+cleanup sql_benchmarks/grouped_count_distinct/init/cleanup.sql

benchmarks/sql_benchmarks/grouped_count_distinct/benchmarks/q09.benchmark

@@ -0,0 +1,27 @@
+name Q09
+group grouped_count_distinct
+
+load sql_benchmarks/grouped_count_distinct/init/load.sql
+
+assert I
+WITH actual AS (
+  SELECT skew_g AS g, string_g AS s, count(DISTINCT x) AS n
+  FROM grouped_distinct GROUP BY skew_g, string_g
+), expected AS (
+  SELECT g, s, count(x) AS n
+  FROM (SELECT DISTINCT skew_g AS g, string_g AS s, x AS x FROM grouped_distinct)
+  GROUP BY g, s
+)
+SELECT NOT EXISTS (SELECT * FROM actual EXCEPT SELECT * FROM expected)
+   AND NOT EXISTS (SELECT * FROM expected EXCEPT SELECT * FROM actual);
+----
+true
+
+expect_plan AggregateExec
+
+run
+-- Int64 argument, hot composite integer/string grouping (ClickBench Q11 shape).
+SELECT skew_g, string_g, count(DISTINCT x)
+FROM grouped_distinct GROUP BY skew_g, string_g;
+
+cleanup sql_benchmarks/grouped_count_distinct/init/cleanup.sql

benchmarks/sql_benchmarks/grouped_count_distinct/benchmarks/q10.benchmark

@@ -0,0 +1,26 @@
+name Q10
+group grouped_count_distinct
+
+load sql_benchmarks/grouped_count_distinct/init/load.sql
+
+assert I
+WITH actual AS (
+  SELECT many_g AS g, count(DISTINCT dup_x) AS n FROM grouped_distinct GROUP BY many_g
+), expected AS (
+  SELECT g, count(x) AS n
+  FROM (SELECT DISTINCT many_g AS g, dup_x AS x FROM grouped_distinct)
+  GROUP BY g
+)
+SELECT NOT EXISTS (SELECT * FROM actual EXCEPT SELECT * FROM expected)
+   AND NOT EXISTS (SELECT * FROM expected EXCEPT SELECT * FROM actual);
+----
+true
+
+expect_plan AggregateExec
+
+run
+-- Int64, 1M groups with four identical values per group at the default 4M rows.
+-- Native DISTINCT probes both the group table and its per-group set on each row.
+SELECT many_g, count(DISTINCT dup_x) FROM grouped_distinct GROUP BY many_g;
+
+cleanup sql_benchmarks/grouped_count_distinct/init/cleanup.sql

benchmarks/sql_benchmarks/grouped_count_distinct/grouped_count_distinct.suite

@@ -0,0 +1,15 @@
+description = "Grouped COUNT(DISTINCT) across types, cardinalities, skew and duplicate-heavy groups; compare --partitions 1, 4 and 12 with DATAFUSION_OPTIMIZER_SINGLE_DISTINCT_PREFER_NATIVE=false/true"
+
+[[options]]
+name = "rows"
+env = "GROUPED_DISTINCT_ROWS"
+default = "4000000"
+help = "Input rows generated outside the timed query."
+
+[[examples]]
+command = "cargo run --release --bin benchmark_runner -- grouped_count_distinct --iterations 10 --partitions 1"
+description = "Compare grouped DISTINCT counts at the issue's four-million-row scale."
+
+[[examples]]
+command = "cargo run --release --bin benchmark_runner -- grouped_count_distinct --iterations 10 --partitions 12"
+description = "Exercise parallel state merging, including hot integer, string and composite groups (Q07-Q09)."

benchmarks/sql_benchmarks/grouped_count_distinct/init/cleanup.sql

@@ -0,0 +1 @@
+DROP TABLE grouped_distinct;

benchmarks/sql_benchmarks/grouped_count_distinct/init/load.sql

@@ -0,0 +1,13 @@
+CREATE TABLE grouped_distinct AS
+SELECT value % 2000 AS g,
+       value % 8 AS few_g,
+       (value * 48271) % 999983 AS x,
+       (value / 2000) % 8 AS low_x,
+       arrow_cast((value * 48271) % 999983, 'UInt64') AS unsigned_x,
+       CAST((value * 48271) % 999983 AS VARCHAR) AS string_x,
+       CAST((value * 48271) % 999983 AS DOUBLE) AS float_x,
+       CASE WHEN value % 100 < 90 THEN 0 ELSE value % 8 END AS skew_g,
+       CAST(CASE WHEN value % 100 < 90 THEN 0 ELSE value % 8 END AS VARCHAR) AS string_g,
+       value % 1000000 AS many_g,
+       value % 2 AS dup_x
+FROM range(${GROUPED_DISTINCT_ROWS:-4000000});

datafusion/common/src/config.rs

@@ -1939,6 +1939,18 @@ 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
+
+        /// 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.
+        pub single_distinct_prefer_native: bool, default = false
     }
 }