How DataFusion rewrites COUNT(DISTINCT x) into two passes

TL;DR: COUNT(DISTINCT b) is rewritten into a two-phase plan: an inner GROUP BY a, b that materializes the distinct values as rows, and an outer GROUP BY a that counts them. A plain count next to it is rolled up as sum of the inner partial counts, with a CASE to restore the 0 that sum reports as NULL over an empty input. DISTINCT is stripped entirely from aggregates whose merge is idempotent (min, max, bool_and, bit_or), because it cannot change the answer. The rewrite is deliberately conservative: it only fires when it is likely to pay for the extra group-by.

Why COUNT(DISTINCT x) is expensive

A hash aggregate is normally a beautiful thing: for each group it keeps one accumulator, updated in place, row by row. COUNT(x) is an i64 increment, SUM(x) is an f64 add. The aggregate is O(1) memory per group, streamable, and splittable into partial/final stages.

COUNT(DISTINCT x) breaks all of that. To count distinct values you must remember the values you have seen, so the accumulator becomes a hash set per group: a boxed, heap-allocated structure instead of a fixed-size value. The aggregate can no longer use the fast path GroupsAccumulator implementations that keep state in a flat Vec; it falls back to the GroupsAccumulatorAdapter that boxes a full Accumulator per group. The partial/final split becomes awkward too, because the partial stage would have to ship the whole hash set across the network.

It is also not decomposable: there is no way to compute the distinct count of a partition and combine the partition results, because the same value can appear in two partitions.

So DataFusion does the only thing that can work: turn the distinct set into rows.

The rewrite

    Before:
    Aggregate: groupBy=[[a]], aggr=[[count(DISTINCT b)]]
      TableScan: t

    After:
    Aggregate: groupBy=[[a]], aggr=[[count(alias1)]]
      Aggregate: groupBy=[[a, b AS alias1]], aggr=[[]]
        TableScan: t

The inner aggregate groups by (a, b), so it emits exactly one row per distinct (a, b) pair — the distinct values are now rows. The outer aggregate groups by a and counts those rows. The count is now a plain count with no DISTINCT flag, which means it can use the fast, flat, streamable accumulator path.

This is the rule SingleDistinctToGroupBy (datafusion/optimizer/src/single_distinct_to_groupby.rs), and the transformation adds the distinct argument to the inner group-by and removes the distinct flag from the outer aggregate:

// datafusion/optimizer/src/single_distinct_to_groupby.rs:299-330
if distinct {
    // `is_single_distinct_agg` checked that every distinct argument in the
    // plan is the same expression. A call may still take it more than once,
    // as `corr(DISTINCT x, x)` does, and keeps its arity here.
    assert_or_internal_err!(
        args.first().is_some_and(|f| args.iter().all(|a| a == f)),
        "DISTINCT aggregate arguments should be the same expression"
    );
    let arity = args.len();
    let arg = args.swap_remove(0);

    if group_fields_set.insert(arg.schema_name().to_string()) {
        inner_group_exprs.push(arg.alias(SINGLE_DISTINCT_ALIAS));
    }
    let outer = Expr::AggregateFunction(AggregateFunction::new_udf(
        func, vec![col(SINGLE_DISTINCT_ALIAS); arity],
        false, // intentional to remove distinct here
        filter, order_by, null_treatment,
    ));

Two details are worth pausing on. The arity handling exists because corr(DISTINCT x, x) takes the same argument twice — the tuple of distinct arguments is what gets deduplicated, so the outer aggregate still needs two arguments. And the inner group-by gets the argument aliased (SINGLE_DISTINCT_ALIAS = "alias1") rather than kept as a bare expression, because the outer aggregate has to be able to refer to it by name. The same aliasing trick is applied to non-trivial group-by expressions (a + 1 becomes a + 1 AS group_alias_0) so that the outer aggregate’s reference resolves against the inner aggregate’s output schema.

The preconditions

The rule is picky, and each precondition corresponds to a case where the rewrite would be wrong or merely useless (single_distinct_to_groupby.rs:126-170).

• One distinct field only. fields_set.len() != 1 disqualifies the node. count(DISTINCT a) + count(DISTINCT b) cannot be rewritten by this rule, because the inner group-by would need to be on (a, b) while the two counts need different grains.
• No FILTER and no ORDER BY on the aggregate. count(DISTINCT b) FILTER (WHERE ...) and count(DISTINCT b ORDER BY c) are not decomposable: the filter or the ordering is applied after deduplication, per group. The check is a one-liner: if filter.is_some() || !order_by.is_empty() { return Ok(false); }.
• No grouping sets. contains_grouping_set rejects ROLLUP/CUBE plans outright.

Rolling count up through sum

Now the count(*) case. If the query is SELECT a, count(*), count(DISTINCT b) FROM t GROUP BY a, the two-phase rewrite has to handle two aggregates with different grains. The inner group-by is on (a, b), so its count(*) counts the rows within each (a, b) pair — a partial count. The outer phase must add up those partial counts to get the total count per a. Adding is sum:

// datafusion/optimizer/src/single_distinct_to_groupby.rs:351-395
let outer_func = match rollup {
    Some(rollup) => Arc::clone(&rollup.sum),
    None => Arc::clone(&func),
};
inner_aggr_exprs.push(
    Expr::AggregateFunction(AggregateFunction::new_udf(
        Arc::clone(&func), args, false, filter, order_by, null_treatment,
    )).alias(&alias_str),
);
let outer = Expr::AggregateFunction(AggregateFunction::new_udf(
    outer_func, vec![col(&alias_str)], false, None, vec![], None,
));
// `sum` of no rows is NULL where `count` is 0.
// Restoring the 0 also keeps the column non-nullable, as `count` had it.
let proj = when(outer.clone().is_not_null(), outer.clone()).otherwise(lit(0_i64))?;

So count(*) in the presence of a count(DISTINCT …) becomes, in effect, sum of ones over the finer grain — with a CASE expression layered on top:

  After:
  Aggregate: groupBy=[[a]], aggr=[[count(alias1)]]
    Aggregate: groupBy=[[a, b AS alias1]], aggr=[[count(*) AS alias2]]
      TableScan: t
  Projection: CASE WHEN sum(alias2) IS NOT NULL THEN sum(alias2) ELSE 0 END, count(alias1)

The CASE is not decoration. SUM over zero rows is NULL; COUNT over zero rows is 0. The CASE restores the 0 and keeps the column non-nullable, which matters because the original plan’s output schema declared count as non-nullable and the optimizer may not change a plan’s output schema (the “logical schema is invariant under logical optimization” invariant, checked by assert_valid_optimization in datafusion/optimizer/src/optimizer.rs:775+).

Note also that the UDAFs are resolved from the function registry, not by name:

// datafusion/optimizer/src/single_distinct_to_groupby.rs:96-125
struct CountRollup { count: Arc<AggregateUDF>, sum: Arc<AggregateUDF> }

impl CountRollup {
    fn try_new(config: &dyn OptimizerConfig) -> Option<Self> {
        let registry = config.function_registry()?;
        Some(Self { count: registry.udaf("count").ok()?, sum: registry.udaf("sum").ok()? })
    }

    /// Whether `func` is the registry's `count`, and so decomposes into
    /// `sum` over per-partition counts.
    fn is_count(&self, func: &AggregateUDF) -> bool { self.count.as_ref() == func }
}

This is a deliberate design choice, spelled out in the OptimizerRule trait docs (datafusion/optimizer/src/optimizer.rs:112-125): the rule “should avoid function-specific transformations” because a user may have registered a count with different semantics (datafusion-spark does exactly that). Comparing Arc<AggregateUDF> identity against the registry’s count is the safe way to say “this really is the standard count”.

Does the rewrite pay for itself?

The rewrite is not free. The inner group-by has one row per (group, distinct value) pair instead of one row per group, and it keeps its state at that finer grain. Everything else in the plan now has to push more rows through the hash table. So the rule asks whether the benefit outweighs the cost (single_distinct_to_groupby.rs:180-213):

if func.groups_accumulator_supported_for_types(&arg_types, true) == Some(false) {
    return Ok(true);
}

The reasoning: the thing the rewrite buys is “the distinct aggregate no longer goes through GroupsAccumulatorAdapter with its one boxed accumulator per group”. If the aggregate already has a specialized GroupsAccumulator for its argument types, it never went near the adapter, and there is nothing to buy — only the inner group-by to pay for. groups_accumulator_supported_for_types(.., true) (the true is “is_distinct”) returning Some(false) means “this aggregate cannot use the fast path with DISTINCT”, which is exactly the situation where the rewrite helps.

The comment is refreshingly honest about the limits of the analysis:

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. … Measured in https://github.com/apache/datafusion/pull/24859.

And note the asymmetry: the pre-existing tolerance of a non-distinct sum, min or max is not gated on this heuristic, “since narrowing it would change plans that have always been rewritten”. Changing an optimization’s applicability is a behavior change; the project treats it as such.

Dropping DISTINCT where it cannot matter

The other half of the aggregate-rewrite story is removing DISTINCT where it is a no-op. min(DISTINCT x) and min(x) return the same value, because the merge operation is idempotent. DataFusion encodes this as a property of the UDAF, in datafusion/expr/src/udaf.rs:1745-1759:

pub enum DistinctHandling {
    /// The result is the same with or without `DISTINCT`, so the planner
    /// is free to drop it. `min`, `max`, `bool_and`, `bit_or`, ...
    Insensitive,
    /// The accumulator reads `AccumulatorArgs::is_distinct` and deduplicates
    /// its input, so the planner must leave the flag alone. `count`, `sum`,
    /// `avg`, `var_samp`, `array_agg`, ... This is the default.
    Sensitive,
    /// The accumulator does not implement `DISTINCT` at all; the planner has
    /// to deduplicate the input first (today `SingleDistinctToGroupBy` does
    /// that) or reject the query.
    Unsupported,
}

min and max declare it because “MAX is idempotent: duplicates cannot change the maximum” (datafusion/functions-aggregate/src/min_max.rs:326-330), and bool_and/bool_or because “Boolean AND/OR are idempotent: duplicates cannot change the result” (datafusion/functions-aggregate/src/bool_and_or.rs:187-191).

EliminateAggregateDistinct (datafusion/optimizer/src/eliminate_aggregate_distinct.rs) acts on it:

// datafusion/optimizer/src/eliminate_aggregate_distinct.rs:135-160
fn strip_insensitive_distinct(expr: Expr) -> Result<Transformed<Expr>> {
    Ok(match expr {
        Expr::AggregateFunction(mut agg)
            if agg.params.distinct
                && agg.func.distinct_handling() == DistinctHandling::Insensitive =>
        {
            agg.params.distinct = false;
            Transformed::yes(Expr::AggregateFunction(agg))
        }
        _ => Transformed::no(expr),
    })
}

The important subtlety is the all-or-nothing condition. The doc on can_strip_every_distinct (eliminate_aggregate_distinct.rs:101-133) explains it:

If only some of them were stripped, a Sensitive DISTINCT could stay beside a stripped aggregate. A stripped aggregate carries an alias, and SingleDistinctToGroupBy does not rewrite a node whose aggr_expr contains an alias. So min(DISTINCT x), count(DISTINCT x) would lose the rewrite that count needs. When every DISTINCT is Insensitive, none is left after stripping, so that rule has nothing to rewrite.

That is a lovely example of rules reasoning about each other. Stripping DISTINCT from min changes the expression’s schema_name (min(DISTINCT b) → min(b)), so the rule re-aliases it back to the original name to keep the schema stable. But that alias then makes the node ineligible for the SingleDistinctToGroupBy rewrite. Rather than let the two rules fight, EliminateAggregateDistinct declines to fire unless every DISTINCT in the node is Insensitive. The rule ordering in the optimizer list encodes the dependency (optimizer.rs:309-311): EliminateAggregateDistinct runs immediately before SingleDistinctToGroupBy.

SELECT DISTINCT is the same story

ReplaceDistinctWithAggregate (datafusion/optimizer/src/replace_distinct_aggregate.rs) applies the same idea one level up: SELECT DISTINCT a, b becomes GROUP BY a, b, because that is what a hash aggregate with no aggregate functions does — and it does it with the fast group-by path instead of a distinct-hash-set.

It even has a shortcut that removes the DISTINCT entirely when the input is already unique (replace_distinct_aggregate.rs:99-118): if the input schema carries a functional dependency with Dependency::Single mode covering exactly the projected columns, the DISTINCT is a no-op and is simply removed. Dependency::Single versus Dependency::Multi is exactly the distinction between “this is a primary key” and “this used to be a primary key but a join may have fanned it out”.

DISTINCT ON (a) b ... ORDER BY a DESC, c is rewritten into GROUP BY a with first_value(b ORDER BY a DESC, c) (replace_distinct_aggregate.rs:139-155) — again turning a “pick one row per group” operation into an aggregate the executor already knows how to run.

The physical side: stopping early

Once the plan is physical, there is a second trick for the SELECT DISTINCT a FROM t LIMIT 10 shape: once you have found ten distinct values, later rows cannot change the groups you have already found. So the aggregate can stop early — the rule turns Aggregate(Final, a) / Aggregate(Partial, a) under a Limit(10) into the same two aggregates carrying a soft limit of 10.

This is implemented as a new AggregateKind::DistinctLimit { group_by, limit } on both the partial and the final stage:

// datafusion/physical-plan/src/aggregates/mod.rs:1154-1173
pub fn try_optimize_distinct_soft_limit(
    mut self, limit: usize,
) -> Option<Transformed<Self>> {
    if !self.is_unordered_unfiltered_group_by_distinct() {
        return None;
    }
    self.kind = AggregateKind::DistinctLimit {
        group_by: Arc::clone(self.group_by()),
        limit,
    };
    Some(Transformed::yes(self))
}

The guard is_unordered_unfiltered_group_by_distinct is what makes the shortcut sound: with an ORDER BY, the “first ten groups” are not the “first ten rows’ groups”. The rule is also gated on optimizer.enable_distinct_aggregation_soft_limit (default true, datafusion/common/src/config.rs:1630).

Finally, CombinePartialFinalAggregate (datafusion/physical-optimizer/src/combine_partial_final_agg.rs) collapses the Partial/Final pair back into a single Single aggregate when the plan is not distributed — undoing the two-stage split when the two-stage split is not needed.

How the micro-rewrites compound

None of these rules is a big win alone; the win comes from the way they stack. EliminateAggregateDistinct (dropping a pointless DISTINCT), SingleDistinctToGroupBy (the two-phase distinct count), EliminateGroupByConstant, CommonSubexprEliminate (deduping repeated subexpressions) and OptimizeProjections (removing columns nobody reads) each do a small, safe, local rewrite. Together they turn a plan that materializes a per-group hash set into two flat group-bys and a projection.

That is the actual lesson of DataFusion’s optimizer design: the wins are not in the clever rules, they are in the composition of many boring, provably-safe local rewrites, plus the discipline to know exactly when a rule must not fire.

Trade-offs and traps

• Cardinality blowup. The inner group-by emits one row per (group, distinct value). If a group has 1000 distinct values, the plan now carries 1000 rows through a hash table where it used to carry 1.
• The heuristic is a proxy. rewrite_pays_for_count guesses based on whether the aggregate has a specialized group accumulator, and is deliberately biased towards not rewriting.
• Schema stability. Any rewrite that changes an aggregate’s name must re-alias it, or the plan’s output schema changes and the optimizer’s schema invariant check fails.
• Order and filter are excluded on purpose. count(DISTINCT b ORDER BY c) and count(DISTINCT b FILTER (WHERE …)) are not decomposable, and the rule returns Ok(false) rather than trying.
• The sum-of-counts trick changes nullability. sum is nullable, count is not; the CASE is what keeps the schema valid, and dropping it would be a silent correctness bug in the type system rather than in the data.