DRAFT Bench 26133

benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q17.benchmark

@@ -0,0 +1,16 @@
+# RANK top-2, ~100K partitions, with a sibling ROW_NUMBER expression.
+# Project both ranks so the window stays above TopK and TopK's derived ranks
+# are unused. Compare with q08's single-expression workload.
+
+subgroup window_sorted
+
+template sql_benchmarks/h2o/window_sorted.benchmark.template
+QPAD=17
+FN=RANK
+PK_MOD=100000
+OB_EXPR=v2
+OB_BOUND=count(DISTINCT ob) > 1000
+NAME=window_sorted_q17_rank_100k_partitions_two_windows
+TWO_WINDOWS=true
+
+expect_plan BoundedWindowAggExec

benchmarks/sql_benchmarks/h2o/window_sorted.benchmark.template

@@ -8,6 +8,7 @@
 #   OB_EXPR  ORDER BY value expression  (`v2` distinct, integer-cast for ties)
 #   OB_BOUND SQL predicate pinning ob's cardinality, asserted below
 #   NAME     display name
+#   TWO_WINDOWS  append and project ROW_NUMBER alongside FN (default false)
 #
 # Why this subgroup exists: the `window` subgroup registers `x` with no declared
 # ordering, so `output_ordering()` is None and the planner cannot know the input
@@ -88,9 +89,9 @@ expect_plan output_ordering=[pk@0 ASC NULLS LAST, ob@1 DESC]
 
 run
 -- Top-2 per partition by ${FN} over the declared-sorted input.
-SELECT pk, ob FROM (
+SELECT pk, ob ${TWO_WINDOWS:-false|, rnk, rn| } FROM (
     SELECT pk, ob,
-           ${FN}() OVER (PARTITION BY pk ORDER BY ob DESC) AS rnk
+           ${FN}() OVER (PARTITION BY pk ORDER BY ob DESC) AS rnk ${TWO_WINDOWS:-false|, ROW_NUMBER() OVER (PARTITION BY pk ORDER BY ob DESC) AS rn| }
     FROM x_sorted
 ) sub_query WHERE rnk <= 2;
 

datafusion/physical-optimizer/src/optimizer.rs

@@ -111,8 +111,10 @@ impl PhysicalOptimizer {
             // those are handled by the later `FilterPushdown` rule.
             // See `FilterPushdownPhase` for more details.
             Arc::new(FilterPushdown::new()),
-            // WindowTopN: replaces Filter(rn<=K) → Window(ROW_NUMBER)
-            // with Window(ROW_NUMBER) → PartitionedTopKExec(fetch=K).
+            // WindowTopN: replaces Filter(rn<=K) → Window(ROW_NUMBER) with
+            // PartitionedTopKExec(fetch=K), which emits `rn` itself; the
+            // window node is kept above it only when it has other window
+            // expressions to compute.
             // Must run before EnsureRequirements (so it can rewrite against the
             // window's declared ordering without pattern-matching a SortExec)
             // and before ProjectionPushdown (which embeds projections into FilterExec).

datafusion/physical-optimizer/src/window_topn.rs

@@ -36,8 +36,10 @@
 //! ```
 //!
 //! And replaces the `FilterExec → BoundedWindowAggExec` pipeline with
-//! `BoundedWindowAggExec → PartitionedTopKExec(fetch=K)`, removing the
-//! `FilterExec` and inserting `PartitionedTopKExec` under the window.
+//! `PartitionedTopKExec(fetch=K)`, removing the `FilterExec`. When the window
+//! node holds a single ranking expression the operator emits that expression's
+//! column itself and the `BoundedWindowAggExec` is removed too; with two or
+//! more it is kept above the operator to compute its columns.
 //!
 //! The appropriate [`WindowFnKind`] is forwarded to `PartitionedTopKExec`.
 //! `RANK` and `DENSE_RANK` require a non-empty `ORDER BY` clause (otherwise
@@ -85,15 +87,26 @@ use datafusion_physical_plan::windows::{BoundedWindowAggExec, WindowUDFExpr};
 ///
 /// ```text
 /// [optional ProjectionExec]
-///   BoundedWindowAggExec(<ranking fn> PARTITION BY ... ORDER BY ...)
+///   PartitionedTopKExec(fn=<row_number|rank|dense_rank>, partition_keys, order_keys, fetch=K, emit=[<ranking fn> PARTITION BY ... ORDER BY ...])
+/// ```
+///
+/// or, when the window node holds two or more window expressions and so still
+/// has columns to compute:
+///
+/// ```text
+/// [optional ProjectionExec]
+///   BoundedWindowAggExec(<ranking fns> PARTITION BY ... ORDER BY ...)
 ///     PartitionedTopKExec(fn=<row_number|rank|dense_rank>, partition_keys, order_keys, fetch=K)
 /// ```
 ///
-/// The `FilterExec` is removed entirely. The child of `BoundedWindowAggExec` is now
-/// `PartitionedTopKExec`, which maintains per-partition top-K state (a
-/// heap for `ROW_NUMBER`, a heap plus boundary ties for `RANK`, a
-/// K-bounded distinct-ob map for `DENSE_RANK`) instead of sorting the
-/// whole dataset.
+/// The `FilterExec` is removed entirely, and `PartitionedTopKExec` maintains
+/// per-partition top-K state (a heap for `ROW_NUMBER`, a heap plus boundary
+/// ties for `RANK`, a K-bounded distinct-ob map for `DENSE_RANK`) instead of
+/// sorting the whole dataset. In the first shape it also appends the ranking
+/// column, which it can derive from the retained rows as it emits them, so
+/// nothing is left for the window node to do and it is removed as well. The
+/// column keeps the window expression's name, not the query's alias (`rn`):
+/// the alias is applied by the `ProjectionExec` above.
 ///
 /// # Supported Predicates
 ///
@@ -233,19 +246,47 @@ impl WindowTopN {
             return None;
         }
 
-        let partitioned_topk = PartitionedTopKExec::try_new(
+        // The operator can emit the ranking column itself, which makes the
+        // `BoundedWindowAggExec` above it redundant — see step 7. All three
+        // ranking functions are derivable from the retained set as it is
+        // emitted (`PartitionedTopKExec` documents how per policy), so the only
+        // condition is that there is exactly one window expression: the
+        // operator appends a single column, and it has to reproduce the whole
+        // of the window node's output schema. With two or more the window node
+        // stays and computes them as before.
+        let ranking_field = if window_exprs.len() == 1 {
+            Some(matched_expr.field().ok()?)
+        } else {
+            None
+        };
+
+        let mut partitioned_topk = PartitionedTopKExec::try_new(
             Arc::clone(window_exec_typed.input()),
             expr,
             partition_prefix_len,
             limit_n,
             fn_kind,
         )
         .ok()?;
-
-        // Step 7: Rebuild window with PartitionedTopKExec as its child
-        let mut result =
-            replace_children_if_necessary(window_exec, vec![Arc::new(partitioned_topk)])
+        if let Some(field) = &ranking_field {
+            partitioned_topk = partitioned_topk
+                .with_ranking_field(Arc::clone(field))
                 .ok()?;
+        }
+
+        // Step 7: Rebuild the pipeline above the operator.
+        //
+        // When the operator emits the ranking column itself, the
+        // `BoundedWindowAggExec` has nothing left to compute: it would
+        // re-derive, over the retained rows, exactly the values already in
+        // that column. Its output schema is the input's fields followed by
+        // its window expressions' fields, which is what the operator now
+        // produces, so the nodes above it keep referencing the same column
+        // indices. Drop it.
+        let mut result: Arc<dyn ExecutionPlan> = Arc::new(partitioned_topk);
+        if ranking_field.is_none() {
+            result = replace_children_if_necessary(window_exec, vec![result]).ok()?;
+        }
 
         // Step 8: Rebuild intermediate nodes (ProjectionExec/RepartitionExec)
         for node in intermediates.into_iter().rev() {