feat(optimizer): logical rewrite Filter(row_number() = 1) over PARTITION BY → Aggregate(FIRST_VALUE(... ORDER BY)) by saadtajwar · Pull Request #23824 · apache/datafusion

saadtajwar · 2026-07-22T23:49:13Z

Which issue does this PR close?

Closes Logical rewrite: Filter(row_number() = 1) over PARTITION BY → Aggregate(FIRST_VALUE(... ORDER BY)) #23603

Rationale for this change

DF plans the common "top-1 per group" pattern (ROW_NUMBER() OVER (PARTITION BY ... ORDER BY ...) WHERE rn = 1) as a full window sort followed by a filter - this is expensive on wide payloads because it materializes and ranks every row before discarding most.

This PR adds a logical optimizer rule that recognizes the pattern and rewrites it to GROUP BY partition keys with first_value(col ORDER BY ...), which can pick the winning row per group without a full sort. The rule is gated behind optimizer.enable_row_number_to_aggregate and defaults to off until nested-type GroupsAccumulator support lands in #23601/#23628, since the rewritten aggregate only becomes memory-efficient once that fast path exists.

What changes are included in this PR?

Adds ReplaceFilterTop1, a new logical optimizer rule that matches Filter(rn = 1 | rn <= 1 | rn < 2) over a single row_number() window with a non-empty PARTITION BY, and rewrites it to Projection -> Aggregate(first_value(...) ORDER BY ...) -> child, preserving the filter's output schema. It supports the direct Filter -> Window shape and Filter -> Projection -> Window when the projection passes the row_number column through unchanged.

A new config option datafusion.optimizer.enable_row_number_to_aggregate was added, and config docs and information_schema.slt are updated accordingly.

Are these changes tested?

Yes

Are there any user-facing changes?

Yes, users can opt in with SET datafusion.optimizer.enable_row_number_to_aggregate = true. The new setting is documented in docs/source/user-guide/configs.md and appears in information_schema

saadtajwar · 2026-07-22T23:50:42Z

+    let (name, op, val) = match (&**left, &**right) {
+        (
+            Expr::Column(Column { name, .. }),
+            Expr::Literal(ScalarValue::UInt64(Some(val)), _),


When I was testing it looked like the type was always UInt64 - is that a valid assumption? Or should I match on all int types?

saadtajwar · 2026-07-22T23:51:14Z

+
+/// Recognize either `Filter -> Window` or `Filter -> Projection -> Window`,
+/// where the window is a single `row_number()` with a non-empty `PARTITION BY`.
+fn validate_window_input(input: &Arc<LogicalPlan>) -> Option<WindowTop1<'_>> {


I was a bit unsure about this function name/the general structure and return values 😅 happy to receive any feedback here!

saadtajwar · 2026-07-22T23:53:11Z

+}
+
+/// Unalias expression reference
+fn unalias(mut expr: &Expr) -> &Expr {


BTW, looks like this unalias pattern is used in push_down_filter & eliminate_outer_join (although those two are recursive) - happy to consolidate this to a helper in a follow-up if we want!

saadtajwar · 2026-07-22T23:53:57Z

+/// - window has a `PARTITION BY` clause
+/// - gated behind the `optimizer.enable_row_number_to_aggregate` config option (off by default)
+#[derive(Default, Debug)]
+pub struct ReplaceFilterTop1 {}


Happy to rename this optimizer rule, I recognize that having a number in here might be a bit awkward (recurring theme throughout this PR of me not being good at naming things haha)

codecov-commenter · 2026-07-23T17:01:25Z

Codecov Report

❌ Patch coverage is 77.12082% with 89 lines in your changes missing coverage. Please review.
✅ Project coverage is 81.44%. Comparing base (1f5be0e) to head (302b9f6).
⚠️ Report is 36 commits behind head on main.

Files with missing lines	Patch %	Lines
datafusion/optimizer/src/replace_filter_top1.rs	77.06%	47 Missing and 42 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main   #23824      +/-   ##
==========================================
- Coverage   81.45%   81.44%   -0.01%     
==========================================
  Files        1118     1119       +1     
  Lines      399555   399944     +389     
  Branches   399555   399944     +389     
==========================================
+ Hits       325462   325754     +292     
- Misses      55097    55147      +50     
- Partials    18996    19043      +47

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

saadtajwar · 2026-07-23T17:18:05Z

@zhuqi-lucas & @tohuya6 - ready for review for #23603! Looking forward to the feedback! 😁 🎉

saadtajwar · 2026-07-23T17:19:45Z

Another note: I'll make a follow-up issue for this comment from @2010YOUY01 and address in a subsequent PR!

saadtajwar · 2026-08-03T22:31:16Z

@zhuqi-lucas - thanks again for taking the time to review! Have pushed a couple of commits that should address your comments - please let me know your thoughts! Thanks!

saadtajwar · 2026-08-18T14:47:12Z

@zhuqi-lucas - just wanted to follow up here! Thanks!

zhuqi-lucas

LGTM

alamb · 2026-08-27T17:39:07Z

@tohuya6 does this PR look good to you?

alamb · 2026-10-24T10:44:37Z

This PR appears to be waiting on the author to respond to reviewer feedback so marking as draft. Please mark ready to review when it is ready for another look

saadtajwar · 2026-10-28T14:30:11Z

This PR appears to be waiting on the author to respond to reviewer feedback so marking as draft. Please mark ready to review when it is ready for another look

Hi @alamb - thanks for checking in here! Just fixed the merge conflict & marked the outstanding reviews as resolved - this should be good to go! Thanks in advance for your time! :)

alamb · 2026-11-02T20:32:06Z

@zhuqi-lucas do you think this one is ready to go?

zhuqi-lucas · 2026-11-03T02:43:22Z

Thanks @alamb @tohuya6 , I think it's ready after resolve conflicts, and turning it on by default is tracked in #24761.

…tajwar/logical-rewrite-filter-top1 Resolve configs.md conflict by regenerating from merged config.

saadtajwar · 2026-11-03T13:36:16Z

Thanks @alamb @tohuya6 , I think it's ready after resolve conflicts, and turning it on by default is tracked in #24761.

Thanks @zhuqi-lucas ! Just resolved those merge conflicts, and I'm happy to fast follow with a PR for #24761 when this is in! Appreciate you all :)

saadtajwar added 9 commits July 21, 2026 19:42

feat: logical rewrite for filter on rn=1 to aggregate first_value

9293195

helper functions for predicate and window input validation

7549d87

finished rewrite

7e741ec

starting on tests

eb4bead

added more tests

c5994d5

added docs

ee28107

adding more comments

271d291

more comments

02a4c34

removing debug stuff

daa3858

github-actions Bot added documentation Improvements or additions to documentation optimizer Optimizer rules sqllogictest SQL Logic Tests (.slt) common Related to common crate labels Jul 22, 2026

saadtajwar commented Jul 22, 2026

View reviewed changes

github-actions Bot added the auto detected api change Auto detected API change label Jul 22, 2026

fixing CI bugs;

9ae64e2

github-actions Bot added the core Core DataFusion crate label Jul 23, 2026

config change

3ac1253

Merge branch 'main' into saadtajwar/logical-rewrite-filter-top1

e4b7e19

saadtajwar marked this pull request as ready for review July 23, 2026 17:17

saadtajwar mentioned this pull request Jul 23, 2026

Logical rewrite: Filter(row_number() = 1) over PARTITION BY → Aggregate(FIRST_VALUE(... ORDER BY)) #23603

Open

saadtajwar changed the title ~~feat(optimizer): logical rewrite Filter(row_number() = 1) over PARTITION BY → Aggregate(FIRST_VALUE(... ORDER BY)) (DRAFT)~~ feat(optimizer): logical rewrite Filter(row_number() = 1) over PARTITION BY → Aggregate(FIRST_VALUE(... ORDER BY)) Jul 23, 2026

tohuya6 suggested changes Jul 23, 2026

View reviewed changes

Comment thread datafusion/optimizer/src/replace_filter_top1.rs Resolved

saadtajwar and others added 4 commits July 23, 2026 15:02

making access of config function registry safer

6d4beb5

Merge branch 'main' into saadtajwar/logical-rewrite-filter-top1

45c0461

Merge branch 'main' into saadtajwar/logical-rewrite-filter-top1

e92b951

Merge branch 'main' into saadtajwar/logical-rewrite-filter-top1

f1275a8

adding comment for function name being safe

cc83011

spelling fix

52d0ebe

saadtajwar requested review from tohuya6 and zhuqi-lucas August 7, 2026 00:35

Merge branch 'main' into saadtajwar/logical-rewrite-filter-top1

c525f66

zhuqi-lucas reviewed Aug 7, 2026

View reviewed changes

Comment thread datafusion/sqllogictest/test_files/row_number_to_aggregate.slt Resolved

saadtajwar and others added 2 commits August 7, 2026 20:09

more tests cases added

7823342

Merge branch 'main' into saadtajwar/logical-rewrite-filter-top1

378777a

saadtajwar requested a review from zhuqi-lucas August 9, 2026 16:57

Merge branch 'main' into saadtajwar/logical-rewrite-filter-top1

8b62766

Merge branch 'main' into saadtajwar/logical-rewrite-filter-top1

302b9f6

zhuqi-lucas approved these changes Aug 26, 2026

View reviewed changes

saadtajwar mentioned this pull request Aug 28, 2026

Enable enable_row_number_to_aggregate by default #24761

Open

alamb marked this pull request as draft October 24, 2026 10:44

Merge upstream/main into saadtajwar/logical-rewrite-filter-top1

aa71da4

saadtajwar marked this pull request as ready for review October 28, 2026 14:29

saadtajwar requested a review from alamb October 28, 2026 14:29

saadtajwar added 2 commits October 28, 2026 10:47

chore: update license headers

c254276

rerun CI

a90b0f5

akupchinskiy mentioned this pull request Nov 3, 2026

first_value/last_value ORDER BY a time zone aware timestamp gets slower with more partitions (Arc<str> refcount contention in FirstLastGroupsAccumulator) #27046

Open

Merge branch 'main' of https://github.com/apache/datafusion into saad…

5c1fd73

…tajwar/logical-rewrite-filter-top1 Resolve configs.md conflict by regenerating from merged config.

github-actions Bot removed the auto detected api change Auto detected API change label Nov 3, 2026

Conversation

saadtajwar commented Jul 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

Uh oh!

Uh oh!

saadtajwar Jul 22, 2026

Choose a reason for hiding this comment

Uh oh!

saadtajwar Jul 22, 2026

Choose a reason for hiding this comment

Uh oh!

saadtajwar Jul 22, 2026

Choose a reason for hiding this comment

Uh oh!

saadtajwar Jul 22, 2026

Choose a reason for hiding this comment

Uh oh!

codecov-commenter commented Jul 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

saadtajwar commented Jul 23, 2026

Uh oh!

saadtajwar commented Jul 23, 2026

Uh oh!

Uh oh!

saadtajwar commented Aug 3, 2026

Uh oh!

Uh oh!

saadtajwar commented Aug 18, 2026

Uh oh!

zhuqi-lucas left a comment

Choose a reason for hiding this comment

Uh oh!

alamb commented Aug 27, 2026

Uh oh!

alamb commented Oct 24, 2026

Uh oh!

saadtajwar commented Oct 28, 2026

Uh oh!

alamb commented Nov 2, 2026

Uh oh!

zhuqi-lucas commented Nov 3, 2026

Uh oh!

saadtajwar commented Nov 3, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

saadtajwar commented Jul 22, 2026 •

edited

Loading

codecov-commenter commented Jul 23, 2026 •

edited

Loading