Skip to content

Commit ea33cce

Browse files
saadtajwarsaadtajwar
authored
feat(optimizer): logical rewrite Filter(row_number() = 1) over PARTITION BY → Aggregate(FIRST_VALUE(... ORDER BY)) (#23824)
## Which issue does this PR close? - Closes #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`
1 parent 181a87a commit ea33cce

9 files changed

Lines changed: 902 additions & 0 deletions

File tree

‎datafusion/common/src/config.rs‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1644,6 +1644,16 @@ config_namespace! {
16441644
/// cause regressions in both memory usage and runtime.
16451645
pub enable_window_topn: bool, default = false
16461646

1647+
/// When set to true, the optimizer will rewrite a "top-1 per group"
1648+
/// pattern of the form `Filter(row_number() = 1)` over a `PARTITION BY` window into an
1649+
/// `Aggregate(first_value(... ORDER BY ...) GROUP BY partition)`.
1650+
/// This avoids buffering / fully sorting the input, which is especially beneficial on wide payloads.
1651+
/// Disabled by default because on nested/wide value types the aggregate only becomes memory
1652+
/// efficient once the columnar nested-type `GroupsAccumulator` fast
1653+
/// path is available; enabling it without that support can route wide
1654+
/// payloads onto the slow per-group accumulator path.
1655+
pub enable_row_number_to_aggregate: bool, default = false
1656+
16471657
/// When set to true, the optimizer will push TopK (Sort with fetch)
16481658
/// below hash repartition when the partition key is a prefix of the
16491659
/// sort key, reducing data volume before the shuffle.

‎datafusion/core/src/optimizer_rule_reference.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,7 @@ Rule order matters. The default pipeline may change between releases.
5050
| 24 | `extract_leaf_expressions` | Pulls cheap leaf expressions closer to data sources so later pruning and filter rules can act earlier. |
5151
| 25 | `push_down_leaf_projections` | Pushes the helper projections created by leaf extraction toward leaf inputs. |
5252
| 26 | `optimize_projections` | Prunes unused columns and removes unnecessary logical projections. |
53+
| 27 | `replace_filter_top1` | Rewrites `Filter(row_number() = 1)` over a partitioned window into a `first_value` aggregate. |
5354

5455
### Physical Optimizer Rules
5556

‎datafusion/optimizer/src/lib.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,7 @@ pub mod propagate_empty_relation;
5454
pub mod push_down_filter;
5555
pub mod push_down_limit;
5656
pub mod replace_distinct_aggregate;
57+
pub mod replace_filter_top1;
5758
pub mod rewrite_set_comparison;
5859
pub mod scalar_subquery_to_join;
5960
pub mod simplify_expressions;

‎datafusion/optimizer/src/optimizer.rs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,7 @@ use crate::propagate_empty_relation::PropagateEmptyRelation;
4949
use crate::push_down_filter::PushDownFilter;
5050
use crate::push_down_limit::PushDownLimit;
5151
use crate::replace_distinct_aggregate::ReplaceDistinctWithAggregate;
52+
use crate::replace_filter_top1::ReplaceFilterTop1;
5253
use crate::rewrite_set_comparison::RewriteSetComparison;
5354
use crate::scalar_subquery_to_join::ScalarSubqueryToJoin;
5455
use crate::simplify_expressions::SimplifyExpressions;
@@ -302,6 +303,7 @@ impl Optimizer {
302303
Arc::new(ExtractLeafExpressions::new()),
303304
Arc::new(PushDownLeafProjections::new()),
304305
Arc::new(OptimizeProjections::new()),
306+
Arc::new(ReplaceFilterTop1::new()),
305307
];
306308

307309
Self::with_rules(rules)

0 commit comments

Comments
 (0)