(datafusion) 01/01: fix: do not use the top-k aggregate stream for a DISTINCT limit with no ordering direction (#25349)

[gi...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-25349-be5a96e76df9e0111cd4f21998b60a8db7212ba3
in repository https://gitbox.apache.org/repos/asf/datafusion.git

commit 18cbb877ba9a5cae7fc82313e4569a982d202b83
Author: Adrian Garcia Badaracco <17...@users.noreply.github.com>
AuthorDate: Wed Sep 16 02:14:44 2026 +0000

    fix: do not use the top-k aggregate stream for a DISTINCT limit with no ordering direction (#25349)
    
    ## Which issue does this PR close?
    
    - Closes https://github.com/apache/datafusion/issues/25350.
    
    ## Rationale for this change
    
    `SELECT DISTINCT ... LIMIT n` with no `ORDER BY` fails at execution with
    `Internal error: Ordering direction required for DISTINCT with limit`
    when one of the projected columns is a window function:
    
    ```sql
    CREATE TABLE t (env VARCHAR, region VARCHAR, deployment VARCHAR);
    INSERT INTO t VALUES ('prod','a','x'), ('prod','a','y'), ('staging','b','x');
    SELECT DISTINCT env, region, count(*) OVER (PARTITION BY env, deployment) AS n FROM t LIMIT 30;
    ```
    
    The mechanism:
    
    1. `LimitedDistinctAggregation` pushes `LimitOptions::new(limit)` (no
    direction, `descending: None`) into the `AggregateExec` while
    `is_unordered_unfiltered_group_by_distinct()` is true.
    2. `EnsureRequirements` then inserts the sort that the window function
    needs below the aggregation. `replace_children_if_necessary` rebuilds
    the aggregation through `replace_children` with
    `ChildrenPropertiesMode::Recompute`. This derives the properties again
    from the sorted child but keeps `limit_options`. The rebuilt node has
    `ordering_mode=PartiallySorted(..)` and an output ordering, so the
    predicate is now false.
    3. `execute_typed` selects `GroupedTopKAggregateStream`. Its constructor
    needs an ordering direction, and no rule supplied one, so it returns the
    error.
    
    The top-k stream also reads only the first group key
    (`aggr.group_expr().expr()[0]`). Thus it would give wrong results for a
    multi-column `DISTINCT` even with a direction. The correct behavior is
    to not select that stream in this case.
    
    ## What changes are included in this PR?
    
    `AggregateExec::execute_typed` selects the top-k path only when an
    ordering direction is available, which is what
    `GroupedTopKAggregateStream::new` requires:
    
    - `limit_options().descending` is set. `TopKAggregation` sets it, and
    that rule accepts only a single group key.
    - Or `get_minmax_desc()` supplies a direction.
    
    Otherwise, execution uses the regular grouped streams. This is correct
    because they treat the pushed limit as a soft limit, and the `LIMIT`
    above the aggregation still truncates the result.
    
    The limit pushdown and the plans do not change. Only the stream
    selection at execution changes.
    
    ## What is the testing strategy for this PR?
    
    A new case in `datafusion/sqllogictest/test_files/aggregate.slt` runs
    the multi-column and single-column forms of the query and checks the
    rows (the same rows as PostgreSQL 16.15 and DuckDB 1.5.2). Without the
    `mod.rs` change, both queries fail with the internal error. With the
    change, they pass.
    
    Also run locally, on this branch rebased onto current `main`:
    
    - `cargo test -p datafusion-physical-plan -p
    datafusion-physical-optimizer`
    - `cargo test --profile ci -p datafusion-sqllogictest --test
    sqllogictests` (all files)
    
    ## Are there any user-facing changes?
    
    Queries that failed with `Internal error: Ordering direction required
    for DISTINCT with limit` now return the correct rows. There are no API
    changes and no plan changes.
    
    🤖 Generated with [Claude Code](https://claude.com/claude-code)
    
    ---------
    
    Co-authored-by: Claude <no...@anthropic.com>
---
 datafusion/physical-plan/src/aggregates/mod.rs   | 19 +++++++++++++++-
 datafusion/sqllogictest/test_files/aggregate.slt | 29 ++++++++++++++++++++++++
 2 files changed, 47 insertions(+), 1 deletion(-)

diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs
index 5b0d2c2d2e..49a726f9be 100644
--- a/datafusion/physical-plan/src/aggregates/mod.rs
+++ b/datafusion/physical-plan/src/aggregates/mod.rs
@@ -1219,9 +1219,26 @@ impl AggregateExec {
             )?));
         }
 
-        // grouping by an expression that has a sort/limit upstream
+        // Grouping by an expression that has a sort/limit upstream.
+        //
+        // `GroupedTopKAggregateStream` keeps a priority queue, so it only works
+        // when an ordering direction is available: either this is a MIN/MAX
+        // aggregate, whose direction is implied by the accumulator, or the
+        // optimizer pushed a direction down next to the limit.
+        //
+        // A direction-less limit is a soft hint pushed by
+        // `LimitedDistinctAggregation`, which only pushes it while
+        // `is_unordered_unfiltered_group_by_distinct()` holds. That predicate is
+        // not stable: a later rule such as `EnsureRequirements` can insert the
+        // sort a window function requires below this aggregate, which gives the
+        // rebuilt aggregate an output ordering and makes the predicate false
+        // without ever supplying a direction. Falling back to the regular
+        // grouped streams is correct in that case, because they treat the limit
+        // as a soft limit and the `LIMIT` above this aggregate still truncates
+        // the result.
         if let Some(config) = self.limit_options
             && !self.is_unordered_unfiltered_group_by_distinct()
+            && (config.descending.is_some() || self.get_minmax_desc().is_some())
         {
             return Ok(StreamType::GroupedPriorityQueue(
                 GroupedTopKAggregateStream::new(self, context, partition, config.limit)?,
diff --git a/datafusion/sqllogictest/test_files/aggregate.slt b/datafusion/sqllogictest/test_files/aggregate.slt
index 3eabd5c4ff..a4a54cdee2 100644
--- a/datafusion/sqllogictest/test_files/aggregate.slt
+++ b/datafusion/sqllogictest/test_files/aggregate.slt
@@ -7522,6 +7522,35 @@ set datafusion.optimizer.enable_distinct_aggregation_soft_limit = true;
 statement ok
 set datafusion.optimizer.repartition_aggregations = true;
 
+# A window function in the projection makes `EnsureRequirements` add a sort below
+# the distinct aggregation, so the aggregation ends up with an output ordering after
+# `LimitedDistinctAggregation` pushed a direction-less limit into it. The query
+# must still execute and return the distinct rows.
+statement ok
+CREATE TABLE distinct_limit_window_table(env VARCHAR, region VARCHAR, deployment VARCHAR) AS VALUES
+  ('prod', 'a', 'x'),
+  ('prod', 'a', 'y'),
+  ('staging', 'b', 'x');
+
+query TTI rowsort
+SELECT DISTINCT env, region, count(*) OVER (PARTITION BY env, deployment) AS n
+FROM distinct_limit_window_table
+LIMIT 30;
+----
+prod a 1
+staging b 1
+
+query TI rowsort
+SELECT DISTINCT env, count(*) OVER (PARTITION BY env) AS n
+FROM distinct_limit_window_table
+LIMIT 30;
+----
+prod 2
+staging 1
+
+statement ok
+DROP TABLE distinct_limit_window_table;
+
 #
 # regr_*() tests
 #


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org