RankExec to optimize ROW_NUMBER <= K

Author: Vedin

Cargo.lock

@@ -737,6 +737,11 @@ config_namespace! {
         /// past window functions, if possible
         pub enable_window_limits: bool, default = true
 
+        /// When set to true, the physical optimizer will rewrite `ROW_NUMBER() OVER
+        /// (PARTITION BY .. ORDER BY ..) <= K` into a dedicated rank/TopN-per-partition
+        /// operator that can spill instead of using a full window aggregate.
+        pub enable_rank_rewrite: bool, default = true
+
         /// When set to true attempts to push down dynamic filters generated by operators into the file scan phase.
         /// For example, for a query such as `SELECT * FROM t ORDER BY timestamp DESC LIMIT 10`, the optimizer
         /// will attempt to push down the current top 10 timestamps that the TopK operator references into the file scans.

datafusion/common/src/config.rs

@@ -53,6 +53,7 @@ recursive = { workspace = true, optional = true }
 
 [dev-dependencies]
 datafusion-expr = { workspace = true }
+datafusion-functions-window = { workspace = true }
 datafusion-functions-nested = { workspace = true }
 insta = { workspace = true }
 tokio = { workspace = true }

datafusion/physical-optimizer/Cargo.toml

@@ -38,6 +38,7 @@ pub mod limited_distinct_aggregation;
 pub mod optimizer;
 pub mod output_requirements;
 pub mod projection_pushdown;
+pub mod rank_topn;
 pub use datafusion_pruning as pruning;
 mod limit_pushdown_past_window;
 pub mod sanity_checker;

datafusion/physical-optimizer/src/lib.rs

@@ -32,6 +32,7 @@ use crate::limit_pushdown::LimitPushdown;
 use crate::limited_distinct_aggregation::LimitedDistinctAggregation;
 use crate::output_requirements::OutputRequirements;
 use crate::projection_pushdown::ProjectionPushdown;
+use crate::rank_topn::RankTopNPerPartition;
 use crate::sanity_checker::SanityCheckPlan;
 use crate::topk_aggregation::TopKAggregation;
 use crate::update_aggr_exprs::OptimizeAggregateOrder;
@@ -104,6 +105,7 @@ impl PhysicalOptimizer {
             // those are handled by the later `FilterPushdown` rule.
             // See `FilterPushdownPhase` for more details.
             Arc::new(FilterPushdown::new()),
+            Arc::new(RankTopNPerPartition::new()),
             // The EnforceDistribution rule is for adding essential repartitioning to satisfy distribution
             // requirements. Please make sure that the whole plan tree is determined before this rule.
             // This rule increases parallelism if doing so is beneficial to the physical plan; i.e. at