Skip to content

feat: Optimize ORDER BY by Pruning Functionally Redundant Sort Keys - #21362

Merged
alamb merged 10 commits into
apache:mainfrom
xiedeyantu:sortkey
Apr 19, 2026
Merged

alamb merged 10 commits into
apache:mainfrom
xiedeyantu:sortkey

Conversation

@xiedeyantu

Copy link
Copy Markdown
Member

Which issue does this PR close?

Rationale for this change

This PR adds functional-dependency-based simplification for ORDER BY clauses. When an earlier sort key already functionally determines a later key, the later key is redundant and can be removed without changing query semantics. This reduces unnecessary sorting work and avoids carrying extra sort keys through planning and execution.

What changes are included in this PR?

This PR extends the existing functional dependency utilities with a helper for pruning redundant sort keys, and wires that helper into eliminate_duplicated_expr so Sort nodes can be simplified during optimization. It also adds regression coverage for both the positive case, where a trailing sort key is removed, and the negative case, where sort order prevents pruning.

Are these changes tested?

Yes. I added unit tests covering:

  • removal of a functionally redundant trailing ORDER BY key
  • preservation of ordering when the dependent column appears before its determinant

I also ran cargo test -p datafusion-optimizer eliminate_duplicated_expr -- --nocapture successfully, and cargo fmt --all passes.

Are there any user-facing changes?

Yes, but only in query planning behavior. Some queries with redundant ORDER BY keys may produce simpler plans and run more efficiently. There are no public API changes.

@github-actions github-actions Bot added optimizer Optimizer rules common Related to common crate labels Apr 4, 2026
@github-actions github-actions Bot added the sqllogictest SQL Logic Tests (.slt) label Apr 4, 2026

@neilconway neilconway left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks good overall! A few minor suggestions.

@xiedeyantu

Copy link
Copy Markdown
Member Author

Looks good overall! A few minor suggestions.

@neilconway Thank you for your comments and suggestions; I think the scenario you mentioned—changing order by deptno, total_sal, abs(deptno) to order by deptno, abs(deptno)—is an excellent point. We should definitely support this; our previous restrictions were indeed a bit too strict. We could even go a step further in the future by leveraging injective functions for additional optimizations (though it might be better to implement this in a separate PR). Regarding the rest of your comments, I have attempted to address each one by submitting a corresponding commit. Could you please take another look and review them? Thank you once again for your thorough reviews of every single one of my PRs!

@neilconway neilconway left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for iterating on this!

Can you "resolve" comment threads for review comments you believe have been addressed, please?

@xiedeyantu

Copy link
Copy Markdown
Member Author

Thanks for iterating on this!

Can you "resolve" comment threads for review comments you believe have been addressed, please?

@neilconway My previous understanding was that the issue would only be marked as "resolved" after I had fixed it and the reviewer had confirmed that everything was in order; that is why I didn't click the button. Going forward, I will follow your suggestion. Thank you!

@neilconway

Copy link
Copy Markdown
Contributor

My previous understanding was that the issue would only be marked as "resolved" after I had fixed it and the reviewer had confirmed that everything was in order; that is why I didn't click the button.

I think it's simpler if the PR submitter just proactively "resolves" comments they believe have been resolved; if the reviewer disagrees, they can always reopen the comment thread.

@neilconway neilconway left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks good to me! Nice work.

@alamb PR looks reasonable to me.

@xiedeyantu

Copy link
Copy Markdown
Member Author

My previous understanding was that the issue would only be marked as "resolved" after I had fixed it and the reviewer had confirmed that everything was in order; that is why I didn't click the button.

I think it's simpler if the PR submitter just proactively "resolves" comments they believe have been resolved; if the reviewer disagrees, they can always reopen the comment thread.

I completely agree; this will make it much easier for the reviewer to conduct the next round of reviews. Thank you!

@alamb alamb left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I am not sure about the logic in this one. Thanks @neilconway and @xiedeyantu

----
logical_plan
01)Sort: table_with_ordered_pk.c1 ASC NULLS LAST, table_with_ordered_pk.c2 ASC NULLS LAST
01)Sort: table_with_ordered_pk.c1 ASC NULLS LAST

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't understand this change -- the query requires ORDR BY c1, c2 but now the query only sorts on c1. The primary key on c1 means there are no duplicates, but how does that ensure it is also ordered by c2

For example what about

INSERT INTO table VALUES (1,2)
INSERT INTO table VALUES (2,1)

That is still ordered by c1, but if you don't also sort of c2, you'll end up with the wrong sort

@neilconway neilconway Apr 10, 2026 •

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hmm, the optimization seems sound to me. If c1 functionally determines c2, we know that each distinct c1 value is associated with exactly one c2 value. So sorting by c1 is sufficient; adding in c2 as a tiebreaker / secondary sort key is never useful.

In the example, if c1 is a PK of table, sorting by c1 is sufficient to get the right sort order -- there will never be two rows with the same c1 value.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@alamb I fully agree with Neil's explanation: functional dependency serves precisely to guarantee this property—for instance, an injective function inherently exhibits the characteristics of functional dependency. In the context of ORDER BY, an additional aspect of ordering comes into play, requiring evaluation based on the sequence of the original fields (as explained in the accompanying code); this ensures that the semantics remain correct after the ORDER BY fields have been eliminated.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

oh, I see -- that makes sense -- the fact that there is a PK means there will never be a tie on c1.

Thank you for the clarification

@alamb alamb left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you @xiedeyantu and @neilconway

@alamb alamb added the performance Make DataFusion faster label Apr 16, 2026
@xiedeyantu

Copy link
Copy Markdown
Member Author

@alamb Thanks for your review, do I need to do something else?

@xiedeyantu

Copy link
Copy Markdown
Member Author

@alamb Hi, is this PR ready to be merged?

@alamb
alamb added this pull request to the merge queue Apr 19, 2026
@alamb

alamb commented Apr 19, 2026

Copy link
Copy Markdown
Contributor

@alamb Thanks for your review, do I need to do something else?

Nope -- sorry @xiedeyantu -- we normally approve PRs and then wait a day or two in case other committers want to review the PRs as well.

I am sorry for the delay

Merged via the queue into apache:main with commit fd882fb Apr 19, 2026
35 checks passed
@xiedeyantu

Copy link
Copy Markdown
Member Author

@alamb Thanks for your review, do I need to do something else?

Nope -- sorry @xiedeyantu -- we normally approve PRs and then wait a day or two in case other committers want to review the PRs as well.

I am sorry for the delay

Understood. Thank you for your review.

Rich-T-kid pushed a commit to Rich-T-kid/datafusion that referenced this pull request Apr 21, 2026
…pache#21362)

## Which issue does this PR close?

- Closes apache#21361 .

## Rationale for this change

This PR adds functional-dependency-based simplification for `ORDER BY`
clauses. When an earlier sort key already functionally determines a
later key, the later key is redundant and can be removed without
changing query semantics. This reduces unnecessary sorting work and
avoids carrying extra sort keys through planning and execution.

## What changes are included in this PR?

This PR extends the existing functional dependency utilities with a
helper for pruning redundant sort keys, and wires that helper into
`eliminate_duplicated_expr` so `Sort` nodes can be simplified during
optimization. It also adds regression coverage for both the positive
case, where a trailing sort key is removed, and the negative case, where
sort order prevents pruning.

## Are these changes tested?

Yes. I added unit tests covering:
- removal of a functionally redundant trailing `ORDER BY` key
- preservation of ordering when the dependent column appears before its
determinant

I also ran `cargo test -p datafusion-optimizer eliminate_duplicated_expr
-- --nocapture` successfully, and `cargo fmt --all` passes.

## Are there any user-facing changes?

Yes, but only in query planning behavior. Some queries with redundant
`ORDER BY` keys may produce simpler plans and run more efficiently.
There are no public API changes.

---------

Co-authored-by: Andrew Lamb <[email protected]>
comphead pushed a commit to comphead/arrow-datafusion that referenced this pull request Oct 14, 2026
…pache#26150)

## Which issue does this PR close?

- None yet. The bug and its repro are below.

## Rationale for this change

If `x` is a primary key, the optimizer treats `CAST(x AS ...)` as `x`.
It then drops the other ORDER BY or GROUP BY keys that `x` determines,
so the query returns wrong results:

```sql
CREATE TABLE t (x DOUBLE, y VARCHAR, PRIMARY KEY (x))
  AS VALUES (1.1, 'b'), (1.2, 'a'), (2.1, 'c');

SELECT x, y FROM t ORDER BY CAST(x AS INT), y;
-- main:     1.1 b / 1.2 a / 2.1 c   (`y` is dropped from the sort)
-- expected: 1.2 a / 1.1 b / 2.1 c

SELECT CAST(x AS INT) k, count(*) n FROM t GROUP BY CAST(x AS INT), y;
-- main:     (1, 2) / (2, 1)         (`y` is dropped from the grouping)
-- expected: (1, 1) / (1, 1) / (2, 1)

SELECT y, count(*) FROM t GROUP BY CAST(x AS INT);
-- main:     accepted
-- expected: Error: Column in SELECT must be in GROUP BY or an aggregate function
```

`CAST(x AS INT)` is named `t.x`, because casts are left out of
expression names. The functional dependency helpers matched GROUP BY and
ORDER BY expressions to key columns by comparing names, so the cast
matched the key `t.x`.

The same happens with `TRY_CAST`, with a `UNIQUE NOT NULL` key, and with
a key that comes from an inner GROUP BY. The ORDER BY case comes from
the sort key pruning added in apache#21362 (54.0.0).

## What changes are included in this PR?

- **Index-based helpers.** The four helpers in
`functional_dependencies.rs` now take, for each GROUP BY or ORDER BY
expression, the index of the input field it references (`Option<usize>`)
instead of its name. A computed expression has no index, so it never
matches a key. The new `datafusion_expr::utils::passthrough_field_index`
returns the index for a column reference, aliased or not.
- **One GROUP BY list.** The dependencies of an `Aggregate`'s output use
the GROUP BY list that its schema is built from
(`grouping_set_to_exprlist`). Before, they used a list de-duplicated by
name, which merged `CAST(x AS INT)` and `x`.
- **Projections resolve columns.** A projection now finds each input
column with `DFSchema::index_of_column_by_name` instead of comparing
`"qualifier.name"` strings. The unit test
`projection_duplicate_flattened_name_uses_first_input_index` asserted
the old string behaviour: a column named `"orders.id"` got the
dependency of the different field `orders.id`. It is renamed and now
asserts that each column keeps its own dependency.

## What is the testing strategy for this PR?

- **New tests.** Section 6 of `functional_dependencies.slt` has one
query for each user of the helpers:
  - ORDER BY pruning;
  - GROUP BY pruning;
  - `Aggregate` output dependencies;
  - GROUP BY expansion.

On `main`, each of these returns a wrong result or accepts an invalid
query. The same queries also run over `TRY_CAST` on a `UNIQUE NOT NULL`
key and over a key from an inner GROUP BY. One more query pins that a
widening cast (`INT` to `BIGINT`) of a key does not carry its
dependencies.
- **New unit tests.** `passthrough_field_index` resolves an ambiguous
unqualified column to nothing. A key that is in the GROUP BY list two
times (`id, id AS k`, builder API only) still determines the row.
- **Existing tests** are unchanged, apart from the rewritten unit test.
- **Planning time.** Planning is not slower. The `sql_planner` TPC-H and
TPC-DS benchmarks, whose tables have primary keys, are about 5% faster,
because the new code no longer renders expression names.

<details><summary>Benchmark numbers (3 interleaved runs per
side)</summary>

| Benchmark | `main` | This PR |
|---|---|---|
| `physical_plan_tpch_all` | 17.3 / 17.8 / 18.0 ms | 16.1 / 17.0 / 17.2
ms |
| `physical_plan_tpcds_all` | 291 / 300 / 301 ms | 276 / 280 / 282 ms |

</details>

## Are there any user-facing changes?

- **Results.** The queries above return correct results.
- **Rejected queries.** Only a column reference to a key carries its
dependencies. A query that selects other columns while grouping by an
expression over the key is now rejected, even when the expression keeps
distinct values distinct, like `GROUP BY CAST(x AS BIGINT)` over an
`INT` key. `main` accepts it and its result is correct. PostgreSQL and
DuckDB reject it too. Group by the key column itself instead.
- **API change.** Four `pub` functions in `datafusion_common` take
`&[Option<usize>]` instead of `&[String]`:
  - `aggregate_functional_dependencies`
  - `get_target_functional_dependencies`
  - `get_required_group_by_exprs_indices`
  - `get_required_sort_exprs_indices`

  The upgrade guide has a migration note.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

common Related to common crate optimizer Optimizer rules performance Make DataFusion faster sqllogictest SQL Logic Tests (.slt)

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Optimize ORDER BY by Pruning Functionally Redundant Sort Keys

3 participants