You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/10/15 09:59:07 UTC
(datafusion) 01/01: perf: Partitioned topk emit row number (#26133)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-26133-59fb71d9bc06163a856a125a49949ef11ba334ec
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit b46cecfcf071a14c71f102559e14d22513ed0c9e
Author: Subham Singhal <Su...@users.noreply.github.com>
AuthorDate: Thu Oct 15 09:58:43 2026 +0000
perf: Partitioned topk emit row number (#26133)
## Which issue does this PR close?
- Part of [#24608](https://github.com/apache/datafusion/issues/24608)
## Rationale for this change
With `enable_window_topn`, `WindowTopN` rewrites `Filter(rn <= K) →
BoundedWindowAggExec → SortExec` into:
```text
BoundedWindowAggExec: row_number() PARTITION BY [pk] ORDER BY [val], mode=Sorted
PartitionedTopKExec: fn=row_number, fetch=2, partition=[pk], order=[val ASC]
```
The window node stays only to produce `rn`, but the operator already
knows every retained row's rank when it emits it: rows leave in
`(partition keys, order keys)` order, and the retained set is exactly
the rows with rank ≤ K. The node's cost scales with the operator's
output (K × partitions) and never amortizes, because every batch it sees
is dense in partition boundaries.
## What changes are included in this PR?
- `PartitionedTopKExec::try_new` takes `ranking_field:
Option<FieldRef>`. When set, the operator appends a `UInt64` column with
each row's rank, and `WindowTopN` drops the window node:
```text
PartitionedTopKExec: fn=row_number, fetch=2, partition=[pk], order=[val
ASC], emit=[rn]
```
- Each policy derives the rank from what it already holds. `ROW_NUMBER`:
position in the partition's drain. `RANK`: position of the first row
with the same key; boundary ties take the boundary key's. `DENSE_RANK`:
index of the distinct-key group. All three rely on the retained set
being a complete order-prefix of its partition.
- Applied only when the window node has exactly one expression, since
the operator appends one column. Otherwise the node is kept, as today.
The output schema is the node's (input fields plus that field), so
column indices above are unchanged and `schema_check()` holds.
- `compute_properties` declares the ordering the removed node published,
`[partition keys..., rank ASC NULLS LAST]`, so `ORDER BY pk, rn` stays
sort-free.
## Are these changes tested?
- `topk/mod.rs`: the rank column for each policy, including boundary
ties, an eviction that moves a row to the tie list, a `DENSE_RANK` group
gathered from several input batches, and output spanning several
`batch_size` chunks.
- `sorts/partitioned_topk.rs`: the widened schema, and the `[pk, rn]`
ordering (but not `[rn]` alone).
- `core/tests/physical_optimizer/window_topn.rs`: the node is removed
for one expression and kept for two; `ORDER BY pk, rn` plans no
`SortExec` for all three functions; end to end, the emitted column
equals the flag-off result (ties, 4 partitions, batch size 4).
- `window_topn.slt` and `range_partitioning.slt`: EXPLAIN output
updated, query results unchanged.
## Benchmark
h2o `window.sql` k=2 sweeps on `J1_1e7_1e7_NA.parquet` (10 M rows, 14
partitions), `enable_window_topn = true` in both arms.
| partitions | `ROW_NUMBER` | `RANK` | `DENSE_RANK` |
|---|---|---|---|
| 100 | 0.037 → 0.035 s (1.06×) | 0.036 → 0.037 s (0.97×) | 0.037 →
0.037 s (1.00×) |
| 1 K | 0.036 → 0.035 s (1.03×) | 0.038 → 0.039 s (0.97×) | 0.040 →
0.040 s (1.00×) |
| 10 K | 0.041 → 0.038 s (1.08×) | 0.034 → 0.032 s (1.06×) | 0.047 →
0.044 s (1.07×) |
| 100 K | **0.075 → 0.050 s (1.50×)** | **0.092 → 0.060 s (1.53×)** |
**0.147 → 0.118 s (1.25×)** |
## Are there any user-facing changes?
Only with enable_window_topn (default false): EXPLAIN no longer shows
the window node for single-expression ranking queries, and
PartitionedTopKExec shows emit=[...]. Results for the built-in functions
are unchanged.
A user-defined window function registered as row_number, rank or
dense_rank is no longer rewritten. WindowTopN used to match these
functions by name; it now compares the WindowUDF with the built-in
(row_number_udwf(), rank_udwf(), dense_rank_udwf()). Without this, the
operator's column would replace the custom evaluator's output, or the
query would fail at execution if the return type isn't UInt64. On main,
such a function already had its input pruned.
datafusion-functions-window becomes a regular dependency of
datafusion-physical-optimizer.
No breaking API change: PartitionedTopKExec::try_new keeps its
signature. with_ranking_field() and ranking_field() are new public
methods.
---
.../core/tests/physical_optimizer/window_topn.rs | 518 +++++++++++++--
datafusion/physical-optimizer/Cargo.toml | 2 +-
datafusion/physical-optimizer/src/optimizer.rs | 6 +-
datafusion/physical-optimizer/src/window_topn.rs | 98 ++-
.../physical-plan/src/sorts/partitioned_topk.rs | 305 ++++++++-
datafusion/physical-plan/src/topk/mod.rs | 709 +++++++++++++++++----
.../sqllogictest/test_files/range_partitioning.slt | 24 +-
datafusion/sqllogictest/test_files/window_topn.slt | 175 ++---
8 files changed, 1499 insertions(+), 338 deletions(-)
diff --git a/datafusion/core/tests/physical_optimizer/window_topn.rs b/datafusion/core/tests/physical_optimizer/window_topn.rs
index 85801e84c9..9ff80ab489 100644
--- a/datafusion/core/tests/physical_optimizer/window_topn.rs
+++ b/datafusion/core/tests/physical_optimizer/window_topn.rs
@@ -17,17 +17,24 @@
//! Tests for the WindowTopN physical optimizer rule.
+use std::ops::Range;
use std::sync::Arc;
-use arrow::datatypes::{DataType, Field, Schema};
+use arrow::array::{ArrayRef, RecordBatch};
+use arrow::datatypes::{DataType, Field, FieldRef, Schema};
use datafusion::prelude::{SessionConfig, SessionContext};
use datafusion_common::Result;
use datafusion_common::ScalarValue;
use datafusion_common::config::ConfigOptions;
use datafusion_expr::Operator;
-use datafusion_expr::{WindowFrame, WindowFrameBound, WindowFrameUnits};
+use datafusion_expr::{
+ PartitionEvaluator, Signature, Volatility, WindowFrame, WindowFrameBound,
+ WindowFrameUnits, WindowUDFImpl,
+};
use datafusion_functions_window::rank::{dense_rank_udwf, rank_udwf};
use datafusion_functions_window::row_number::row_number_udwf;
+use datafusion_functions_window_common::field::WindowUDFFieldArgs;
+use datafusion_functions_window_common::partition::PartitionEvaluatorArgs;
use datafusion_physical_expr::expressions::{BinaryExpr, Column, col, lit};
use datafusion_physical_expr::window::StandardWindowExpr;
use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr};
@@ -222,9 +229,8 @@ fn basic_row_number_rn_lteq_3() -> Result<()> {
let plan = build_window_topn_plan(3, Operator::LtEq)?;
let optimized = optimize(plan)?;
assert_snapshot!(plan_str(optimized.as_ref()), @r#"
- BoundedWindowAggExec: wdw=[row_number: Field { "row_number": UInt64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
- PartitionedTopKExec: fn=row_number, fetch=3, partition=[pk@0], order=[val@1 ASC]
- PlaceholderRowExec
+ PartitionedTopKExec: fn=row_number, fetch=3, partition=[pk@0], order=[val@1 ASC], emit=[row_number]
+ PlaceholderRowExec
"#);
Ok(())
}
@@ -234,9 +240,8 @@ fn rn_lt_3_becomes_fetch_2() -> Result<()> {
let plan = build_window_topn_plan(3, Operator::Lt)?;
let optimized = optimize(plan)?;
assert_snapshot!(plan_str(optimized.as_ref()), @r#"
- BoundedWindowAggExec: wdw=[row_number: Field { "row_number": UInt64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
- PartitionedTopKExec: fn=row_number, fetch=2, partition=[pk@0], order=[val@1 ASC]
- PlaceholderRowExec
+ PartitionedTopKExec: fn=row_number, fetch=2, partition=[pk@0], order=[val@1 ASC], emit=[row_number]
+ PlaceholderRowExec
"#);
Ok(())
}
@@ -286,9 +291,8 @@ fn flipped_3_gteq_rn() -> Result<()> {
let optimized = optimize(plan)?;
assert_snapshot!(plan_str(optimized.as_ref()), @r#"
- BoundedWindowAggExec: wdw=[row_number: Field { "row_number": UInt64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
- PartitionedTopKExec: fn=row_number, fetch=3, partition=[pk@0], order=[val@1 ASC]
- PlaceholderRowExec
+ PartitionedTopKExec: fn=row_number, fetch=3, partition=[pk@0], order=[val@1 ASC], emit=[row_number]
+ PlaceholderRowExec
"#);
Ok(())
}
@@ -395,9 +399,61 @@ fn with_projection_between() -> Result<()> {
let optimized = optimize(filter)?;
assert_snapshot!(plan_str(optimized.as_ref()), @r#"
ProjectionExec: expr=[pk@0 as pk, val@1 as val, row_number@2 as row_number]
- BoundedWindowAggExec: wdw=[row_number: Field { "row_number": UInt64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
- PartitionedTopKExec: fn=row_number, fetch=3, partition=[pk@0], order=[val@1 ASC]
- PlaceholderRowExec
+ PartitionedTopKExec: fn=row_number, fetch=3, partition=[pk@0], order=[val@1 ASC], emit=[row_number]
+ PlaceholderRowExec
+ "#);
+ Ok(())
+}
+
+/// Two `ROW_NUMBER()` expressions over the same keys: the operator can only
+/// append one column, so the window node has to stay and compute both.
+#[test]
+fn two_window_exprs_keep_the_window_node() -> Result<()> {
+ let s = schema();
+ let input: Arc<dyn ExecutionPlan> = Arc::new(PlaceholderRowExec::new(Arc::clone(&s)));
+
+ let partition_by = vec![col("pk", &s)?];
+ let order_by = vec![PhysicalSortExpr::new_default(col("val", &s)?).asc()];
+ let row_number_expr = |name: &str| -> Result<Arc<StandardWindowExpr>> {
+ Ok(Arc::new(StandardWindowExpr::new(
+ create_udwf_window_expr(
+ &row_number_udwf(),
+ &[],
+ &s,
+ name.to_string(),
+ false,
+ )?,
+ &partition_by,
+ &order_by,
+ Arc::new(WindowFrame::new_bounds(
+ WindowFrameUnits::Rows,
+ WindowFrameBound::Preceding(ScalarValue::UInt64(None)),
+ WindowFrameBound::CurrentRow,
+ )),
+ )))
+ };
+
+ let window: Arc<dyn ExecutionPlan> = Arc::new(BoundedWindowAggExec::try_new(
+ vec![row_number_expr("rn1")?, row_number_expr("rn2")?],
+ input,
+ InputOrderMode::Sorted,
+ true,
+ )?);
+
+ // Filter on the first of the two window columns (index 2).
+ let predicate = Arc::new(BinaryExpr::new(
+ Arc::new(Column::new("rn1", 2)),
+ Operator::LtEq,
+ lit(ScalarValue::UInt64(Some(3))),
+ ));
+ let filter: Arc<dyn ExecutionPlan> =
+ Arc::new(FilterExec::try_new(predicate, window)?);
+
+ let optimized = optimize(filter)?;
+ assert_snapshot!(plan_str(optimized.as_ref()), @r#"
+ BoundedWindowAggExec: wdw=[rn1: Field { "rn1": UInt64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW, rn2: Field { "rn2": UInt64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
+ PartitionedTopKExec: fn=row_number, fetch=3, partition=[pk@0], order=[val@1 ASC]
+ PlaceholderRowExec
"#);
Ok(())
}
@@ -405,7 +461,6 @@ fn with_projection_between() -> Result<()> {
// ----------------------------------------------------------------------
// RANK rule tests
// ----------------------------------------------------------------------
-
/// Build: FilterExec(rk op limit) → BoundedWindowAggExec(<udwf> PBY pk OBY val)
///
/// Matches the pre-`EnsureRequirements` plan shape (no `SortExec` under the window).
@@ -508,9 +563,8 @@ fn basic_rank_rk_lteq_3() -> Result<()> {
let plan = build_ranking_topn_plan(rank_udwf, "rank", 3, Operator::LtEq)?;
let optimized = optimize(plan)?;
assert_snapshot!(plan_str(optimized.as_ref()), @r#"
- BoundedWindowAggExec: wdw=[rank: Field { "rank": UInt64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
- PartitionedTopKExec: fn=rank, fetch=3, partition=[pk@0], order=[val@1 ASC]
- PlaceholderRowExec
+ PartitionedTopKExec: fn=rank, fetch=3, partition=[pk@0], order=[val@1 ASC], emit=[rank]
+ PlaceholderRowExec
"#);
Ok(())
}
@@ -520,9 +574,8 @@ fn rank_rk_lt_4_becomes_fetch_3() -> Result<()> {
let plan = build_ranking_topn_plan(rank_udwf, "rank", 4, Operator::Lt)?;
let optimized = optimize(plan)?;
assert_snapshot!(plan_str(optimized.as_ref()), @r#"
- BoundedWindowAggExec: wdw=[rank: Field { "rank": UInt64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
- PartitionedTopKExec: fn=rank, fetch=3, partition=[pk@0], order=[val@1 ASC]
- PlaceholderRowExec
+ PartitionedTopKExec: fn=rank, fetch=3, partition=[pk@0], order=[val@1 ASC], emit=[rank]
+ PlaceholderRowExec
"#);
Ok(())
}
@@ -532,9 +585,8 @@ fn rank_flipped_3_gteq_rk() -> Result<()> {
let plan = build_ranking_topn_plan(rank_udwf, "rank", 3, Operator::GtEq)?;
let optimized = optimize(plan)?;
assert_snapshot!(plan_str(optimized.as_ref()), @r#"
- BoundedWindowAggExec: wdw=[rank: Field { "rank": UInt64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
- PartitionedTopKExec: fn=rank, fetch=3, partition=[pk@0], order=[val@1 ASC]
- PlaceholderRowExec
+ PartitionedTopKExec: fn=rank, fetch=3, partition=[pk@0], order=[val@1 ASC], emit=[rank]
+ PlaceholderRowExec
"#);
Ok(())
}
@@ -544,9 +596,8 @@ fn rank_flipped_4_gt_rk_becomes_fetch_3() -> Result<()> {
let plan = build_ranking_topn_plan(rank_udwf, "rank", 4, Operator::Gt)?;
let optimized = optimize(plan)?;
assert_snapshot!(plan_str(optimized.as_ref()), @r#"
- BoundedWindowAggExec: wdw=[rank: Field { "rank": UInt64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
- PartitionedTopKExec: fn=rank, fetch=3, partition=[pk@0], order=[val@1 ASC]
- PlaceholderRowExec
+ PartitionedTopKExec: fn=rank, fetch=3, partition=[pk@0], order=[val@1 ASC], emit=[rank]
+ PlaceholderRowExec
"#);
Ok(())
}
@@ -576,9 +627,8 @@ fn basic_dense_rank_dr_lteq_3() -> Result<()> {
let plan = build_ranking_topn_plan(dense_rank_udwf, "dense_rank", 3, Operator::LtEq)?;
let optimized = optimize(plan)?;
assert_snapshot!(plan_str(optimized.as_ref()), @r#"
- BoundedWindowAggExec: wdw=[dense_rank: Field { "dense_rank": UInt64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
- PartitionedTopKExec: fn=dense_rank, fetch=3, partition=[pk@0], order=[val@1 ASC]
- PlaceholderRowExec
+ PartitionedTopKExec: fn=dense_rank, fetch=3, partition=[pk@0], order=[val@1 ASC], emit=[dense_rank]
+ PlaceholderRowExec
"#);
Ok(())
}
@@ -588,9 +638,8 @@ fn dense_rank_dr_lt_4_becomes_fetch_3() -> Result<()> {
let plan = build_ranking_topn_plan(dense_rank_udwf, "dense_rank", 4, Operator::Lt)?;
let optimized = optimize(plan)?;
assert_snapshot!(plan_str(optimized.as_ref()), @r#"
- BoundedWindowAggExec: wdw=[dense_rank: Field { "dense_rank": UInt64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
- PartitionedTopKExec: fn=dense_rank, fetch=3, partition=[pk@0], order=[val@1 ASC]
- PlaceholderRowExec
+ PartitionedTopKExec: fn=dense_rank, fetch=3, partition=[pk@0], order=[val@1 ASC], emit=[dense_rank]
+ PlaceholderRowExec
"#);
Ok(())
}
@@ -600,9 +649,8 @@ fn dense_rank_flipped_3_gteq_dr() -> Result<()> {
let plan = build_ranking_topn_plan(dense_rank_udwf, "dense_rank", 3, Operator::GtEq)?;
let optimized = optimize(plan)?;
assert_snapshot!(plan_str(optimized.as_ref()), @r#"
- BoundedWindowAggExec: wdw=[dense_rank: Field { "dense_rank": UInt64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
- PartitionedTopKExec: fn=dense_rank, fetch=3, partition=[pk@0], order=[val@1 ASC]
- PlaceholderRowExec
+ PartitionedTopKExec: fn=dense_rank, fetch=3, partition=[pk@0], order=[val@1 ASC], emit=[dense_rank]
+ PlaceholderRowExec
"#);
Ok(())
}
@@ -612,9 +660,8 @@ fn dense_rank_flipped_4_gt_dr_becomes_fetch_3() -> Result<()> {
let plan = build_ranking_topn_plan(dense_rank_udwf, "dense_rank", 4, Operator::Gt)?;
let optimized = optimize(plan)?;
assert_snapshot!(plan_str(optimized.as_ref()), @r#"
- BoundedWindowAggExec: wdw=[dense_rank: Field { "dense_rank": UInt64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
- PartitionedTopKExec: fn=dense_rank, fetch=3, partition=[pk@0], order=[val@1 ASC]
- PlaceholderRowExec
+ PartitionedTopKExec: fn=dense_rank, fetch=3, partition=[pk@0], order=[val@1 ASC], emit=[dense_rank]
+ PlaceholderRowExec
"#);
Ok(())
}
@@ -727,3 +774,396 @@ async fn partitioned_topk_exec_exposes_metrics() -> Result<()> {
Ok(())
}
+
+/// The operator emits the `ROW_NUMBER()` column itself, so the plan holds no
+/// `BoundedWindowAggExec` at all — and the column's values still restart at 1
+/// in every partition, across several execution partitions and with ties in
+/// the ORDER BY column.
+#[tokio::test]
+async fn row_number_column_is_emitted_without_a_window_node() -> Result<()> {
+ let mut config = SessionConfig::new()
+ .with_batch_size(4)
+ .with_target_partitions(4);
+ config.options_mut().optimizer.enable_window_topn = true;
+ let ctx = SessionContext::new_with_config(config);
+
+ // 5 partition keys; `val` ties within each key, so rn has to break the tie
+ // by position rather than by value.
+ ctx.sql(
+ "CREATE TABLE t AS
+ SELECT v % 5 AS pk, (v % 3) * 10 AS val FROM (SELECT unnest(range(0, 60)) AS v)",
+ )
+ .await?
+ .collect()
+ .await?;
+
+ let sql = "SELECT pk, val, rn FROM (
+ SELECT *, ROW_NUMBER() OVER (PARTITION BY pk ORDER BY val) AS rn FROM t
+ ) WHERE rn <= 5 ORDER BY pk, rn";
+ let plan = ctx.sql(sql).await?.create_physical_plan().await?;
+ assert!(
+ !plan_str(plan.as_ref()).contains("BoundedWindowAggExec"),
+ "the window node should be gone:\n{}",
+ plan_str(plan.as_ref())
+ );
+
+ let batches = collect(plan, ctx.task_ctx()).await?;
+ assert_snapshot!(datafusion_common::test_util::batches_to_string(&batches), @r"
+ +----+-----+----+
+ | pk | val | rn |
+ +----+-----+----+
+ | 0 | 0 | 1 |
+ | 0 | 0 | 2 |
+ | 0 | 0 | 3 |
+ | 0 | 0 | 4 |
+ | 0 | 10 | 5 |
+ | 1 | 0 | 1 |
+ | 1 | 0 | 2 |
+ | 1 | 0 | 3 |
+ | 1 | 0 | 4 |
+ | 1 | 10 | 5 |
+ | 2 | 0 | 1 |
+ | 2 | 0 | 2 |
+ | 2 | 0 | 3 |
+ | 2 | 0 | 4 |
+ | 2 | 10 | 5 |
+ | 3 | 0 | 1 |
+ | 3 | 0 | 2 |
+ | 3 | 0 | 3 |
+ | 3 | 0 | 4 |
+ | 3 | 10 | 5 |
+ | 4 | 0 | 1 |
+ | 4 | 0 | 2 |
+ | 4 | 0 | 3 |
+ | 4 | 0 | 4 |
+ | 4 | 10 | 5 |
+ +----+-----+----+
+ ");
+
+ // Same query with the flag off must produce the same rows, so the emitted
+ // column agrees with `BoundedWindowAggExec`'s own output.
+ let off = SessionContext::new_with_config(
+ SessionConfig::new()
+ .with_batch_size(4)
+ .with_target_partitions(4),
+ );
+ off.sql(
+ "CREATE TABLE t AS
+ SELECT v % 5 AS pk, (v % 3) * 10 AS val FROM (SELECT unnest(range(0, 60)) AS v)",
+ )
+ .await?
+ .collect()
+ .await?;
+ let expected = off.sql(sql).await?.collect().await?;
+ assert_eq!(
+ datafusion_common::test_util::batches_to_string(&batches),
+ datafusion_common::test_util::batches_to_string(&expected),
+ );
+
+ Ok(())
+}
+
+/// `RANK` and `DENSE_RANK` are emitted by the operator too, so their plans hold
+/// no `BoundedWindowAggExec` either. The values are checked against the same
+/// query with the flag off — i.e. against `BoundedWindowAggExec`'s own output
+/// over the unpruned input — on a tie-heavy table, which is where the two
+/// policies diverge from each other and from `ROW_NUMBER`.
+#[tokio::test]
+async fn rank_and_dense_rank_columns_are_emitted_without_a_window_node() -> Result<()> {
+ // `val` takes 3 distinct values with 4 rows each per partition, so within a
+ // partition RANK is 1, 5, 9 and DENSE_RANK is 1, 2, 3. Each limit keeps the
+ // first two values' 8 rows, so both columns hold a rank above 1 and, at
+ // batch size 4, span several output batches.
+ let create = "CREATE TABLE t AS
+ SELECT v % 5 AS pk, (v % 3) * 10 AS val FROM (SELECT unnest(range(0, 60)) AS v)";
+
+ for (func, limit) in [("RANK", 5), ("DENSE_RANK", 2)] {
+ let sql = format!(
+ "SELECT pk, val, rk FROM (
+ SELECT *, {func}() OVER (PARTITION BY pk ORDER BY val) AS rk FROM t
+ ) WHERE rk <= {limit} ORDER BY pk, rk, val"
+ );
+
+ let mut on_config = SessionConfig::new()
+ .with_batch_size(4)
+ .with_target_partitions(4);
+ on_config.options_mut().optimizer.enable_window_topn = true;
+ let on = SessionContext::new_with_config(on_config);
+ on.sql(create).await?.collect().await?;
+
+ let plan = on.sql(&sql).await?.create_physical_plan().await?;
+ assert!(
+ !plan_str(plan.as_ref()).contains("BoundedWindowAggExec"),
+ "{func}: the window node should be gone:\n{}",
+ plan_str(plan.as_ref())
+ );
+ let actual = collect(plan, on.task_ctx()).await?;
+
+ let off = SessionContext::new_with_config(
+ SessionConfig::new()
+ .with_batch_size(4)
+ .with_target_partitions(4),
+ );
+ off.sql(create).await?.collect().await?;
+ let expected = off.sql(&sql).await?.collect().await?;
+
+ assert_eq!(
+ datafusion_common::test_util::batches_to_string(&actual),
+ datafusion_common::test_util::batches_to_string(&expected),
+ "{func}: emitted column disagrees with BoundedWindowAggExec"
+ );
+
+ // Pin the values themselves, so a column that is wrong in the same way
+ // with the flag off and on cannot pass: per partition, 4 rows at rank 1
+ // and 4 at the second value's rank.
+ let second_rank = if func == "RANK" { 5 } else { 2 };
+ let mut ranks: Vec<u64> = Vec::new();
+ for batch in &actual {
+ let col = batch
+ .column(2)
+ .as_any()
+ .downcast_ref::<arrow::array::UInt64Array>()
+ .expect("ranking column is UInt64");
+ ranks.extend(col.values().iter().copied());
+ }
+ let per_partition = [[1; 4], [second_rank; 4]].concat();
+ assert_eq!(ranks, per_partition.repeat(5), "{func}: unexpected ranks");
+ }
+
+ Ok(())
+}
+
+/// The deleted `BoundedWindowAggExec` published the ordering
+/// `[partition keys..., ranking column ASC NULLS LAST]` (from each ranking
+/// function's `sort_options`). `PartitionedTopKExec` has to publish it too, or
+/// `ORDER BY <partition keys>, <ranking column>` regains a `SortExec` that the
+/// un-rewritten plan did not need.
+#[tokio::test]
+async fn ranking_column_ordering_keeps_order_by_sort_free() -> Result<()> {
+ for func in ["ROW_NUMBER", "RANK", "DENSE_RANK"] {
+ let mut config = SessionConfig::new().with_target_partitions(4);
+ config.options_mut().optimizer.enable_window_topn = true;
+ let ctx = SessionContext::new_with_config(config);
+ ctx.sql(
+ "CREATE TABLE t AS
+ SELECT arrow_cast(v % 5, 'Int64') AS pk, (v % 3) * 10 AS val
+ FROM (SELECT unnest(range(0, 60)) AS v)",
+ )
+ .await?
+ .collect()
+ .await?;
+
+ // NULLS FIRST on the partition key matches the nulls placement the
+ // operator emits, so the only thing that could force a sort here is a
+ // missing ordering on the ranking column itself.
+ let sql = format!(
+ "SELECT pk, val, rn FROM (
+ SELECT *, {func}() OVER (PARTITION BY pk ORDER BY val) AS rn FROM t
+ ) WHERE rn <= 5 ORDER BY pk NULLS FIRST, rn"
+ );
+ let plan = ctx.sql(&sql).await?.create_physical_plan().await?;
+ let displayed = plan_str(plan.as_ref());
+ assert!(
+ !displayed.contains("SortExec"),
+ "{func}: the ranking column's ordering should make this sort-free:\n{displayed}"
+ );
+ }
+ Ok(())
+}
+
+/// A user-defined window function that shares a built-in ranking function's
+/// name but computes something else: the constant 7, or (`frame_len`) the
+/// number of rows in its window frame.
+#[derive(Debug, PartialEq, Eq, Hash)]
+struct SameNameWindowUdf {
+ name: String,
+ return_type: DataType,
+ frame_len: bool,
+ signature: Signature,
+}
+
+impl SameNameWindowUdf {
+ fn new_udwf(
+ name: &str,
+ return_type: DataType,
+ frame_len: bool,
+ ) -> datafusion_expr::WindowUDF {
+ datafusion_expr::WindowUDF::from(Self {
+ name: name.to_string(),
+ return_type,
+ frame_len,
+ signature: Signature::nullary(Volatility::Immutable),
+ })
+ }
+}
+
+impl WindowUDFImpl for SameNameWindowUdf {
+ fn name(&self) -> &str {
+ &self.name
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn partition_evaluator(
+ &self,
+ _args: PartitionEvaluatorArgs,
+ ) -> Result<Box<dyn PartitionEvaluator>> {
+ Ok(Box::new(SameNameEvaluator {
+ return_type: self.return_type.clone(),
+ frame_len: self.frame_len,
+ }))
+ }
+
+ fn field(&self, field_args: WindowUDFFieldArgs) -> Result<FieldRef> {
+ Ok(Field::new(field_args.name(), self.return_type.clone(), false).into())
+ }
+}
+
+#[derive(Debug)]
+struct SameNameEvaluator {
+ return_type: DataType,
+ frame_len: bool,
+}
+
+impl PartitionEvaluator for SameNameEvaluator {
+ fn evaluate(
+ &mut self,
+ _values: &[ArrayRef],
+ range: &Range<usize>,
+ ) -> Result<ScalarValue> {
+ let value = if self.frame_len {
+ range.len() as i64
+ } else {
+ 7
+ };
+ ScalarValue::Int64(Some(value)).cast_to(&self.return_type)
+ }
+
+ fn supports_bounded_execution(&self) -> bool {
+ true
+ }
+
+ fn uses_window_frame(&self) -> bool {
+ self.frame_len
+ }
+
+ fn is_causal(&self) -> bool {
+ !self.frame_len
+ }
+}
+
+/// Runs `sql` over the tie-heavy test table with the flag on and off, with
+/// `udwfs` registered, asserts the flag-on plan has no `PartitionedTopKExec`
+/// and returns both results.
+async fn run_with_custom_udwfs(
+ udwfs: &[datafusion_expr::WindowUDF],
+ sql: &str,
+) -> Result<(Vec<RecordBatch>, Vec<RecordBatch>)> {
+ let create = "CREATE TABLE t AS
+ SELECT v % 5 AS pk, (v % 3) * 10 AS val FROM (SELECT unnest(range(0, 60)) AS v)";
+ let mut results = Vec::new();
+ for enabled in [true, false] {
+ let mut config = SessionConfig::new()
+ .with_batch_size(4)
+ .with_target_partitions(4);
+ config.options_mut().optimizer.enable_window_topn = enabled;
+ let ctx = SessionContext::new_with_config(config);
+ for udwf in udwfs {
+ ctx.register_udwf(udwf.clone());
+ }
+ ctx.sql(create).await?.collect().await?;
+
+ let plan = ctx.sql(sql).await?.create_physical_plan().await?;
+ assert!(
+ !plan_str(plan.as_ref()).contains("PartitionedTopKExec"),
+ "a custom window function must not be rewritten:\n{}",
+ plan_str(plan.as_ref())
+ );
+ results.push(collect(plan, ctx.task_ctx()).await?);
+ }
+ let off = results.pop().unwrap();
+ let on = results.pop().unwrap();
+ Ok((on, off))
+}
+
+/// A user-defined function registered under a built-in ranking function's name
+/// is not that function: the rule must leave it alone rather than prune its
+/// input and replace its evaluator with the operator's ranking column.
+#[tokio::test]
+async fn custom_udwf_with_ranking_name_is_not_rewritten() -> Result<()> {
+ for name in ["row_number", "rank", "dense_rank"] {
+ for return_type in [DataType::UInt64, DataType::Int64] {
+ let sql = format!(
+ "SELECT pk, val, rn FROM (
+ SELECT *, {name}() OVER (PARTITION BY pk ORDER BY val) AS rn FROM t
+ ) WHERE rn <= 7 ORDER BY pk, val"
+ );
+ let udwf = SameNameWindowUdf::new_udwf(name, return_type.clone(), false);
+ let (on, off) = run_with_custom_udwfs(&[udwf], &sql).await?;
+ assert_eq!(
+ datafusion_common::test_util::batches_to_string(&on),
+ datafusion_common::test_util::batches_to_string(&off),
+ "{name} ({return_type})"
+ );
+
+ // Every row passes `rn <= 7` and keeps the function's own value.
+ let rns = on
+ .iter()
+ .flat_map(|b| {
+ let col =
+ arrow::compute::cast(b.column(2), &DataType::Int64).unwrap();
+ col.as_any()
+ .downcast_ref::<arrow::array::Int64Array>()
+ .unwrap()
+ .values()
+ .to_vec()
+ })
+ .collect::<Vec<_>>();
+ assert_eq!(rns, vec![7; 60], "{name} ({return_type})");
+ }
+ }
+ Ok(())
+}
+
+/// The built-in `ROW_NUMBER` is filtered on, but its sibling is a user-defined
+/// function named `rank` that counts the rows in a `CURRENT ROW AND 1
+/// FOLLOWING` frame. Its value at the last retained row depends on the first
+/// pruned one, so the sibling has to block the rewrite.
+#[tokio::test]
+async fn custom_udwf_sibling_with_ranking_name_blocks_rewrite() -> Result<()> {
+ let sql = "SELECT pk, val, rn, rk FROM (
+ SELECT *,
+ ROW_NUMBER() OVER (PARTITION BY pk ORDER BY val) AS rn,
+ rank() OVER (
+ PARTITION BY pk ORDER BY val
+ ROWS BETWEEN CURRENT ROW AND 1 FOLLOWING
+ ) AS rk
+ FROM t
+ ) WHERE rn <= 2 ORDER BY pk, rn";
+ for return_type in [DataType::UInt64, DataType::Int64] {
+ let udwf = SameNameWindowUdf::new_udwf("rank", return_type.clone(), true);
+ let (on, off) = run_with_custom_udwfs(&[udwf], sql).await?;
+ assert_eq!(
+ datafusion_common::test_util::batches_to_string(&on),
+ datafusion_common::test_util::batches_to_string(&off),
+ "{return_type}"
+ );
+
+ // Both retained rows of each partition have a following row.
+ let rks = on
+ .iter()
+ .flat_map(|b| {
+ let col = arrow::compute::cast(b.column(3), &DataType::Int64).unwrap();
+ col.as_any()
+ .downcast_ref::<arrow::array::Int64Array>()
+ .unwrap()
+ .values()
+ .to_vec()
+ })
+ .collect::<Vec<_>>();
+ assert_eq!(rks, vec![2; 10], "{return_type}");
+ }
+ Ok(())
+}
diff --git a/datafusion/physical-optimizer/Cargo.toml b/datafusion/physical-optimizer/Cargo.toml
index cb03303ac3..7a9cb12b59 100644
--- a/datafusion/physical-optimizer/Cargo.toml
+++ b/datafusion/physical-optimizer/Cargo.toml
@@ -46,6 +46,7 @@ datafusion-common = { workspace = true }
datafusion-execution = { workspace = true }
datafusion-expr = { workspace = true }
datafusion-expr-common = { workspace = true, default-features = true }
+datafusion-functions-window = { workspace = true }
datafusion-physical-expr = { workspace = true }
datafusion-physical-expr-common = { workspace = true }
datafusion-physical-plan = { workspace = true }
@@ -57,6 +58,5 @@ recursive = { workspace = true, optional = true }
[dev-dependencies]
datafusion-expr = { workspace = true }
datafusion-functions = { workspace = true }
-datafusion-functions-window = { workspace = true }
insta = { workspace = true }
tokio = { workspace = true }
diff --git a/datafusion/physical-optimizer/src/optimizer.rs b/datafusion/physical-optimizer/src/optimizer.rs
index 81932594b8..5f3004bf9a 100644
--- a/datafusion/physical-optimizer/src/optimizer.rs
+++ b/datafusion/physical-optimizer/src/optimizer.rs
@@ -124,8 +124,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).
diff --git a/datafusion/physical-optimizer/src/window_topn.rs b/datafusion/physical-optimizer/src/window_topn.rs
index 758ad66de3..5834bdb45d 100644
--- a/datafusion/physical-optimizer/src/window_topn.rs
+++ b/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
@@ -56,6 +58,8 @@ use datafusion_common::config::ConfigOptions;
use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
use datafusion_common::{Result, ScalarValue};
use datafusion_expr::Operator;
+use datafusion_functions_window::rank::{dense_rank_udwf, rank_udwf};
+use datafusion_functions_window::row_number::row_number_udwf;
use datafusion_physical_expr::expressions::{BinaryExpr, Column, Literal};
use datafusion_physical_expr::window::StandardWindowExpr;
use datafusion_physical_expr::{LexOrdering, PhysicalSortExpr};
@@ -85,15 +89,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,7 +248,21 @@ 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,
@@ -241,11 +270,25 @@ impl WindowTopN {
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() {
@@ -362,23 +405,32 @@ fn scalar_to_usize(value: &ScalarValue) -> Option<usize> {
/// Identify which supported ranking window function `expr` is.
///
-/// Downcasts through `StandardWindowExpr` → `WindowUDFExpr` and checks
-/// the UDF name. Returns:
-/// - `Some(WindowFnKind::RowNumber)` for `"row_number"`
-/// - `Some(WindowFnKind::Rank)` for `"rank"`
-/// - `Some(WindowFnKind::DenseRank)` for `"dense_rank"`
+/// Downcasts through `StandardWindowExpr` → `WindowUDFExpr` and compares
+/// the UDF with the built-in one. Returns:
+/// - `Some(WindowFnKind::RowNumber)` for the built-in `row_number`
+/// - `Some(WindowFnKind::Rank)` for the built-in `rank`
+/// - `Some(WindowFnKind::DenseRank)` for the built-in `dense_rank`
/// - `None` for everything else
+///
+/// The comparison is `WindowUDF` equality rather than the name: a user-defined
+/// function registered as `row_number` (say) can compute anything, and the
+/// rewrite both prunes its input and, with a single window expression, drops
+/// its evaluator.
fn supported_window_fn(
expr: &Arc<dyn datafusion_physical_expr::window::WindowExpr>,
) -> Option<WindowFnKind> {
let swe = expr.as_any().downcast_ref::<StandardWindowExpr>()?;
let swfe = swe.get_standard_func_expr();
let udf = swfe.as_any().downcast_ref::<WindowUDFExpr>()?;
- match udf.fun().name() {
- "row_number" => Some(WindowFnKind::RowNumber),
- "rank" => Some(WindowFnKind::Rank),
- "dense_rank" => Some(WindowFnKind::DenseRank),
- _ => None,
+ let fun = udf.fun();
+ if fun == &row_number_udwf() {
+ Some(WindowFnKind::RowNumber)
+ } else if fun == &rank_udwf() {
+ Some(WindowFnKind::Rank)
+ } else if fun == &dense_rank_udwf() {
+ Some(WindowFnKind::DenseRank)
+ } else {
+ None
}
}
diff --git a/datafusion/physical-plan/src/sorts/partitioned_topk.rs b/datafusion/physical-plan/src/sorts/partitioned_topk.rs
index 04072fffaf..4eead3e1a4 100644
--- a/datafusion/physical-plan/src/sorts/partitioned_topk.rs
+++ b/datafusion/physical-plan/src/sorts/partitioned_topk.rs
@@ -34,14 +34,17 @@
use std::fmt::{self, Formatter};
use std::sync::Arc;
-use arrow::datatypes::SchemaRef;
+use arrow::compute::SortOptions;
+use arrow::datatypes::{FieldRef, SchemaRef};
use arrow::row::SortField;
use datafusion_common::Result;
use datafusion_common::tree_node::TreeNodeRecursion;
use datafusion_execution::TaskContext;
use datafusion_execution::runtime_env::RuntimeEnv;
use datafusion_physical_expr::PhysicalExpr;
-use datafusion_physical_expr_common::sort_expr::LexOrdering;
+use datafusion_physical_expr::equivalence::EquivalenceProperties;
+use datafusion_physical_expr::expressions::Column;
+use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr};
use futures::StreamExt;
use futures::TryStreamExt;
@@ -49,6 +52,7 @@ use crate::execution_plan::{Boundedness, EmissionType};
use crate::metrics::{ExecutionPlanMetricsSet, MetricsSet};
use crate::topk::{
PartitionedTopK, PartitionedTopKDenseRank, PartitionedTopKRank, build_sort_fields,
+ partitioned_topk_output_schema,
};
use crate::{ChildrenPropertiesMode, ReplaceChildrenOptions};
use crate::{
@@ -205,6 +209,18 @@ pub struct PartitionedTopKExec {
/// Which window function this operator is optimizing. Selects the
/// per-partition retention policy (see [`WindowFnKind`]).
fn_kind: WindowFnKind,
+ /// When set, the operator appends this column to its output, holding each
+ /// retained row's value for [`Self::fn_kind`]'s ranking function.
+ ///
+ /// Every policy already knows those values at emit time — the retained
+ /// rows leave in `(partition_keys, order_keys)` order, and the retained set
+ /// is a complete order-prefix of each partition, so `ROW_NUMBER` is the
+ /// emit position, `RANK` follows from comparing adjacent ORDER BY keys, and
+ /// `DENSE_RANK` is the index of the row's distinct-ORDER-BY-value group.
+ /// Filling them in here lets the rewrite drop the `BoundedWindowAggExec`
+ /// entirely instead of leaving it to re-derive the same numbers over the
+ /// operator's output.
+ ranking_field: Option<FieldRef>,
/// Execution metrics
metrics_set: ExecutionPlanMetricsSet,
/// Cached plan properties (output ordering, partitioning, etc.)
@@ -245,18 +261,40 @@ impl PartitionedTopKExec {
fetch: usize,
fn_kind: WindowFnKind,
) -> Result<Self> {
- let cache = Self::compute_properties(&input, expr.clone())?;
+ let cache =
+ Self::compute_properties(&input, expr.clone(), partition_prefix_len, None)?;
Ok(Self {
input,
expr,
partition_prefix_len,
fetch,
fn_kind,
+ ranking_field: None,
metrics_set: ExecutionPlanMetricsSet::new(),
cache: Arc::new(cache),
})
}
+ /// Append `field` to the output, holding each retained row's value for
+ /// [`Self::fn_kind`]'s ranking function, so the caller can drop the
+ /// `BoundedWindowAggExec` that would otherwise compute it.
+ ///
+ /// The output schema widens by that one column, so the plan properties
+ /// are recomputed under it.
+ pub fn with_ranking_field(self, field: FieldRef) -> Result<Self> {
+ let cache = Self::compute_properties(
+ &self.input,
+ self.expr.clone(),
+ self.partition_prefix_len,
+ Some(&field),
+ )?;
+ Ok(Self {
+ ranking_field: Some(field),
+ cache: Arc::new(cache),
+ ..self
+ })
+ }
+
/// Returns the child execution plan.
pub fn input(&self) -> &Arc<dyn ExecutionPlan> {
&self.input
@@ -283,18 +321,67 @@ impl PartitionedTopKExec {
self.fn_kind
}
+ /// Returns the ranking column this operator appends to its output, or
+ /// `None` when it emits only its input's columns.
+ pub fn ranking_field(&self) -> Option<&FieldRef> {
+ self.ranking_field.as_ref()
+ }
+
/// Compute [`PlanProperties`] for this operator.
///
/// The output is sorted by `sort_exprs` (partition keys then order keys),
/// uses the same partitioning as the input, emits all output at once
/// (`EmissionType::Final`), and is bounded.
+ ///
+ /// When a ranking column is appended, the output is *also* sorted by
+ /// `[partition keys..., ranking column ASC NULLS LAST]`: all three
+ /// retention policies assign ranks that are non-decreasing in emit order
+ /// within a partition. `BoundedWindowAggExec` declared the same ordering
+ /// (from each ranking function's `sort_options`), so declaring it here is
+ /// what lets `ORDER BY <partition keys>, <ranking column>` stay
+ /// sort-free after the window node is removed.
fn compute_properties(
input: &Arc<dyn ExecutionPlan>,
sort_exprs: LexOrdering,
+ partition_prefix_len: usize,
+ ranking_field: Option<&FieldRef>,
) -> Result<PlanProperties> {
- let mut eq_properties = input.equivalence_properties().clone();
+ // With a ranking column appended the output schema is wider than
+ // the input's, so the input's properties have to be re-hung under the
+ // new schema rather than cloned (the same reason
+ // `window_equivalence_properties` does this for `BoundedWindowAggExec`).
+ let mut eq_properties = match ranking_field {
+ None => input.equivalence_properties().clone(),
+ Some(_) => {
+ let output_schema =
+ partitioned_topk_output_schema(&input.schema(), ranking_field);
+ EquivalenceProperties::new(output_schema)
+ .extend(input.equivalence_properties().clone())?
+ }
+ };
+ let partition_exprs = sort_exprs[..partition_prefix_len].to_vec();
eq_properties.reorder(sort_exprs)?;
+ if let Some(field) = ranking_field {
+ // The ranking column is the last one, appended by
+ // `partitioned_topk_output_schema`.
+ let ranking_col =
+ Arc::new(Column::new(field.name(), input.schema().fields().len()))
+ as Arc<dyn PhysicalExpr>;
+ let ranking_sort = PhysicalSortExpr::new(
+ ranking_col,
+ SortOptions {
+ descending: false,
+ nulls_first: false,
+ },
+ );
+ // Mirrors `add_new_ordering_expr_with_partition_by`: the ranking
+ // value only ascends *within* a partition, so the ordering has to
+ // be prefixed by the partition keys. `reorder` above already
+ // established that prefix.
+ eq_properties.add_ordering(partition_exprs.into_iter().chain([ranking_sort]));
+ }
+
Ok(PlanProperties::new(
eq_properties,
input.output_partitioning().clone(),
@@ -311,16 +398,16 @@ impl DisplayAs for PartitionedTopKExec {
WindowFnKind::Rank => "rank",
WindowFnKind::DenseRank => "dense_rank",
};
+ let partition_exprs: Vec<String> = self.expr[..self.partition_prefix_len]
+ .iter()
+ .map(|e| format!("{}", e.expr))
+ .collect();
+ let order_exprs: Vec<String> = self.expr[self.partition_prefix_len..]
+ .iter()
+ .map(|e| format!("{e}"))
+ .collect();
match t {
DisplayFormatType::Default | DisplayFormatType::Verbose => {
- let partition_exprs: Vec<String> = self.expr[..self.partition_prefix_len]
- .iter()
- .map(|e| format!("{}", e.expr))
- .collect();
- let order_exprs: Vec<String> = self.expr[self.partition_prefix_len..]
- .iter()
- .map(|e| format!("{e}"))
- .collect();
write!(
f,
"PartitionedTopKExec: fn={}, fetch={}, partition=[{}], order=[{}]",
@@ -328,23 +415,22 @@ impl DisplayAs for PartitionedTopKExec {
self.fetch,
partition_exprs.join(", "),
order_exprs.join(", "),
- )
+ )?;
+ if let Some(field) = &self.ranking_field {
+ write!(f, ", emit=[{}]", field.name())?;
+ }
}
DisplayFormatType::TreeRender => {
- let partition_exprs: Vec<String> = self.expr[..self.partition_prefix_len]
- .iter()
- .map(|e| format!("{}", e.expr))
- .collect();
- let order_exprs: Vec<String> = self.expr[self.partition_prefix_len..]
- .iter()
- .map(|e| format!("{e}"))
- .collect();
writeln!(f, "fn={fn_label}")?;
writeln!(f, "fetch={}", self.fetch)?;
writeln!(f, "partition=[{}]", partition_exprs.join(", "))?;
- writeln!(f, "order=[{}]", order_exprs.join(", "))
+ writeln!(f, "order=[{}]", order_exprs.join(", "))?;
+ if let Some(field) = &self.ranking_field {
+ writeln!(f, "emit=[{}]", field.name())?;
+ }
}
}
+ Ok(())
}
}
@@ -386,13 +472,20 @@ impl ExecutionPlan for PartitionedTopKExec {
_: ReplaceChildrenOptions,
) -> Result<Arc<dyn ExecutionPlan>> {
assert_eq!(children.len(), 1);
- Ok(Arc::new(PartitionedTopKExec::try_new(
+ let mut exec = PartitionedTopKExec::try_new(
Arc::clone(&children[0]),
self.expr.clone(),
self.partition_prefix_len,
self.fetch,
self.fn_kind,
- )?))
+ )?;
+ // Not a `try_new` argument, so it has to be carried over by hand:
+ // dropping it would narrow the schema the nodes above were planned
+ // against.
+ if let Some(field) = &self.ranking_field {
+ exec = exec.with_ranking_field(Arc::clone(field))?;
+ }
+ Ok(Arc::new(exec))
}
fn apply_expressions(
@@ -439,6 +532,11 @@ impl ExecutionPlan for PartitionedTopKExec {
let batch_size = context.session_config().batch_size();
let runtime = Arc::clone(&context.runtime_env());
let metrics_set = self.metrics_set.clone();
+ let output_schema = self.schema();
+ let ranked_schema = self
+ .ranking_field
+ .is_some()
+ .then(|| Arc::clone(&output_schema));
let stream = futures::stream::once(async move {
do_partitioned_topk(
@@ -453,13 +551,14 @@ impl ExecutionPlan for PartitionedTopKExec {
batch_size,
runtime,
metrics_set,
+ ranked_schema,
)
.await
})
.try_flatten();
Ok(Box::pin(RecordBatchStreamAdapter::new(
- self.input.schema(),
+ output_schema,
stream,
)))
}
@@ -510,12 +609,13 @@ async fn do_partitioned_topk(
batch_size: usize,
runtime: Arc<RuntimeEnv>,
metrics_set: ExecutionPlanMetricsSet,
+ ranked_schema: Option<SchemaRef>,
) -> Result<SendableRecordBatchStream> {
match fn_kind {
WindowFnKind::RowNumber => {
let mut state = PartitionedTopK::try_new(
partition_id,
- schema,
+ &schema,
partition_exprs,
partition_sort_fields,
order_expr,
@@ -523,6 +623,7 @@ async fn do_partitioned_topk(
batch_size,
&runtime,
&metrics_set,
+ ranked_schema,
)?;
while let Some(batch) = input.next().await {
state.insert_batch(&batch?)?;
@@ -533,7 +634,7 @@ async fn do_partitioned_topk(
WindowFnKind::Rank => {
let mut state = PartitionedTopKRank::try_new(
partition_id,
- schema,
+ &schema,
partition_exprs,
partition_sort_fields,
order_expr,
@@ -541,6 +642,7 @@ async fn do_partitioned_topk(
batch_size,
&runtime,
&metrics_set,
+ ranked_schema,
)?;
while let Some(batch) = input.next().await {
state.insert_batch(&batch?)?;
@@ -551,7 +653,7 @@ async fn do_partitioned_topk(
WindowFnKind::DenseRank => {
let mut state = PartitionedTopKDenseRank::try_new(
partition_id,
- schema,
+ &schema,
partition_exprs,
partition_sort_fields,
order_expr,
@@ -559,6 +661,7 @@ async fn do_partitioned_topk(
batch_size,
&runtime,
&metrics_set,
+ ranked_schema,
)?;
while let Some(batch) = input.next().await {
state.insert_batch(&batch?)?;
@@ -568,3 +671,149 @@ async fn do_partitioned_topk(
}
}
}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use crate::placeholder_row::PlaceholderRowExec;
+ use arrow::datatypes::{DataType, Field, Schema};
+ use datafusion_physical_expr::expressions::col;
+ use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr;
+
+ fn pk_val_input() -> Result<(Arc<Schema>, Arc<dyn ExecutionPlan>)> {
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("pk", DataType::Int64, false),
+ Field::new("val", DataType::Int64, false),
+ ]));
+ let input: Arc<dyn ExecutionPlan> =
+ Arc::new(PlaceholderRowExec::new(Arc::clone(&schema)));
+ Ok((schema, input))
+ }
+
+ fn pk_then_val(schema: &Arc<Schema>) -> Result<LexOrdering> {
+ Ok(LexOrdering::new([
+ PhysicalSortExpr::new_default(col("pk", schema)?),
+ PhysicalSortExpr::new_default(col("val", schema)?),
+ ])
+ .expect("two sort expressions"))
+ }
+
+ /// Emitting the `ROW_NUMBER()` column widens the output schema by exactly
+ /// that field, and leaves it alone otherwise.
+ #[test]
+ fn ranking_field_widens_the_output_schema() -> Result<()> {
+ let (schema, input) = pk_val_input()?;
+ let rn_field = Arc::new(Field::new("rn", DataType::UInt64, false));
+
+ let without = PartitionedTopKExec::try_new(
+ input,
+ pk_then_val(&schema)?,
+ 1,
+ 3,
+ WindowFnKind::RowNumber,
+ )?;
+ assert_eq!(without.schema().fields().len(), 2);
+
+ let with = without.with_ranking_field(Arc::clone(&rn_field))?;
+ assert_eq!(with.schema().fields().len(), 3);
+ assert_eq!(with.schema().field(2), rn_field.as_ref());
+
+ // The ordering the operator guarantees survives the wider schema, and
+ // the appended column brings its own `[pk, rn]` ordering with it.
+ let eq = with.properties().equivalence_properties();
+ let pk = Arc::new(Column::new("pk", 0)) as Arc<dyn PhysicalExpr>;
+ let rn = Arc::new(Column::new("rn", 2)) as Arc<dyn PhysicalExpr>;
+ let asc_nulls_last = SortOptions {
+ descending: false,
+ nulls_first: false,
+ };
+ assert!(eq.ordering_satisfy([
+ PhysicalSortExpr::new(Arc::clone(&pk), asc_nulls_last),
+ PhysicalSortExpr::new(Arc::clone(&rn), asc_nulls_last),
+ ])?);
+ // ... but only within a partition: `rn` restarts at 1 per key, so it is
+ // not a global ordering on its own.
+ assert!(!eq.ordering_satisfy([PhysicalSortExpr::new(rn, asc_nulls_last)])?);
+ Ok(())
+ }
+
+ /// Every policy can emit its own ranking column, and each widens the
+ /// output schema the same way.
+ #[test]
+ fn every_policy_accepts_a_ranking_field() -> Result<()> {
+ let (schema, input) = pk_val_input()?;
+ let field = Arc::new(Field::new("rk", DataType::UInt64, false));
+
+ for fn_kind in [
+ WindowFnKind::RowNumber,
+ WindowFnKind::Rank,
+ WindowFnKind::DenseRank,
+ ] {
+ let exec = PartitionedTopKExec::try_new(
+ Arc::clone(&input),
+ pk_then_val(&schema)?,
+ 1,
+ 3,
+ fn_kind,
+ )?
+ .with_ranking_field(Arc::clone(&field))?;
+ assert_eq!(
+ exec.schema().field(2),
+ field.as_ref(),
+ "unexpected appended field for {fn_kind:?}"
+ );
+ assert_eq!(exec.ranking_field(), Some(&field));
+ }
+ Ok(())
+ }
+
+ /// The ranking field is set by `with_ranking_field` rather than passed to
+ /// `try_new`, so rebuilding the operator over a new child has to carry it
+ /// over. Dropping it would narrow the output schema under nodes that were
+ /// planned against the wider one.
+ #[test]
+ fn replace_children_keeps_the_ranking_field() -> Result<()> {
+ let (schema, input) = pk_val_input()?;
+ let rn_field = Arc::new(Field::new("rn", DataType::UInt64, false));
+ let exec: Arc<dyn ExecutionPlan> = Arc::new(
+ PartitionedTopKExec::try_new(
+ Arc::clone(&input),
+ pk_then_val(&schema)?,
+ 1,
+ 3,
+ WindowFnKind::RowNumber,
+ )?
+ .with_ranking_field(Arc::clone(&rn_field))?,
+ );
+
+ let rebuilt = Arc::clone(&exec).replace_children(
+ vec![input],
+ ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
+ )?;
+ assert_eq!(rebuilt.schema(), exec.schema());
+ assert_eq!(
+ rebuilt
+ .downcast_ref::<PartitionedTopKExec>()
+ .expect("still a PartitionedTopKExec")
+ .ranking_field(),
+ Some(&rn_field)
+ );
+ // The ordering on the appended column is recomputed with it.
+ let pk = Arc::new(Column::new("pk", 0)) as Arc<dyn PhysicalExpr>;
+ let rn = Arc::new(Column::new("rn", 2)) as Arc<dyn PhysicalExpr>;
+ let asc_nulls_last = SortOptions {
+ descending: false,
+ nulls_first: false,
+ };
+ assert!(
+ rebuilt
+ .properties()
+ .equivalence_properties()
+ .ordering_satisfy([
+ PhysicalSortExpr::new(pk, asc_nulls_last),
+ PhysicalSortExpr::new(rn, asc_nulls_last),
+ ])?
+ );
+ Ok(())
+ }
+}
diff --git a/datafusion/physical-plan/src/topk/mod.rs b/datafusion/physical-plan/src/topk/mod.rs
index 64b31f5e81..aee902f2b7 100644
--- a/datafusion/physical-plan/src/topk/mod.rs
+++ b/datafusion/physical-plan/src/topk/mod.rs
@@ -37,8 +37,8 @@ use super::metrics::{
use crate::spill::get_record_batch_memory_size;
use crate::{SendableRecordBatchStream, stream::RecordBatchStreamAdapter};
-use arrow::array::{ArrayRef, RecordBatch, UInt32Array};
-use arrow::datatypes::SchemaRef;
+use arrow::array::{ArrayRef, RecordBatch, UInt32Array, UInt64Array};
+use arrow::datatypes::{FieldRef, Schema, SchemaRef};
use datafusion_common::{
HashMap, Result, ScalarValue, internal_datafusion_err, internal_err,
};
@@ -1713,6 +1713,27 @@ impl StoreSlots for PartitionHeap {
}
}
+/// Output schema of `PartitionedTopKExec`: its input's columns, plus
+/// `ranking_field` when the operator emits the ranking column itself.
+///
+/// Built once, by the exec, and handed to each partition's state as is, so the
+/// batches a state emits and the schema the exec declares are the same
+/// `SchemaRef` rather than two schemas that happen to agree.
+pub(crate) fn partitioned_topk_output_schema(
+ input_schema: &SchemaRef,
+ ranking_field: Option<&FieldRef>,
+) -> SchemaRef {
+ let Some(field) = ranking_field else {
+ return Arc::clone(input_schema);
+ };
+ let mut fields = input_schema.fields().to_vec();
+ fields.push(Arc::clone(field));
+ Arc::new(Schema::new_with_metadata(
+ fields,
+ input_schema.metadata().clone(),
+ ))
+}
+
/// What [`PartitionedTopK`] and [`PartitionedTopKRank`] share: the two
/// [`RowConverter`]s — one for the partition key, one for the ORDER BY key —
/// and their scratch [`Rows`], the [`MemoryReservation`], [`TopKMetrics`], one
@@ -1725,7 +1746,14 @@ impl StoreSlots for PartitionHeap {
/// the retention rule — differs, so each operator writes its own row loop over
/// these fields between the two calls.
struct PartitionedTopKCore<S> {
- schema: SchemaRef,
+ /// The schema `emit` produces: the input's columns, plus the ranking
+ /// column when [`Self::emit_ranks`] is set.
+ output_schema: SchemaRef,
+ /// When set, `emit` appends a column holding each retained row's ranking
+ /// value, as computed by the operator's drain in [`Self::emit_with`]. That
+ /// lets the rewrite drop the `BoundedWindowAggExec` that would otherwise
+ /// re-derive the same values over this operator's output.
+ emit_ranks: bool,
metrics: TopKMetrics,
reservation: MemoryReservation,
/// ORDER BY expressions (excludes PARTITION BY).
@@ -1800,11 +1828,13 @@ struct PartitionedTopKCore<S> {
impl<S: StoreSlots> PartitionedTopKCore<S> {
/// `name` labels the memory consumer, as `name[partition_id]`.
+ /// `ranked_schema` is the exec's output schema when it emits the ranking
+ /// column, which makes `emit` append it, and `None` to emit `schema` as is.
#[expect(clippy::too_many_arguments)]
fn try_new(
name: &str,
partition_id: usize,
- schema: SchemaRef,
+ schema: &SchemaRef,
partition_exprs: Vec<Arc<dyn PhysicalExpr>>,
partition_sort_fields: Vec<SortField>,
order_expr: LexOrdering,
@@ -1812,6 +1842,7 @@ impl<S: StoreSlots> PartitionedTopKCore<S> {
batch_size: usize,
runtime: &Arc<RuntimeEnv>,
metrics: &ExecutionPlanMetricsSet,
+ ranked_schema: Option<SchemaRef>,
) -> Result<Self> {
assert!(k > 0, "{name} requires k > 0");
let reservation = MemoryConsumer::new(format!("{name}[{partition_id}]"))
@@ -1820,7 +1851,7 @@ impl<S: StoreSlots> PartitionedTopKCore<S> {
// Both encoders are shared by every partition, and each scratch buffer
// is sized to hold a whole batch so an `insert_batch` pass encodes once
// per column set with no regrowth.
- let order_sort_fields = build_sort_fields(&order_expr, &schema)?;
+ let order_sort_fields = build_sort_fields(&order_expr, schema)?;
let row_converter = RowConverter::new(order_sort_fields)?;
let scratch_rows =
row_converter.empty_rows(batch_size, ESTIMATED_BYTES_PER_ROW * batch_size);
@@ -1830,7 +1861,8 @@ impl<S: StoreSlots> PartitionedTopKCore<S> {
.empty_rows(batch_size, ESTIMATED_BYTES_PER_ROW * batch_size);
Ok(Self {
- schema,
+ emit_ranks: ranked_schema.is_some(),
+ output_schema: ranked_schema.unwrap_or_else(|| Arc::clone(schema)),
metrics: TopKMetrics::new(metrics, partition_id),
reservation,
expr: order_expr,
@@ -1907,15 +1939,20 @@ impl<S: StoreSlots> PartitionedTopKCore<S> {
/// its retained rows in output order, and return them as a stream of
/// `RecordBatch`es ordered by `(partition_keys, order_keys)`.
///
+ /// `drain` pairs each row with its ranking value, which is what the
+ /// operator emits when [`Self::emit_ranks`] is set and is ignored
+ /// otherwise.
+ ///
/// Only the order is resolved here; [`EmitState::stream`] interleaves the
/// rows out one `batch_size` chunk per poll, carrying the reservation until
/// the stream is dropped since the store's batches stay pinned until then.
- fn emit_with<I: IntoIterator<Item = StoreRef>>(
+ fn emit_with<I: IntoIterator<Item = (StoreRef, u64)>>(
self,
drain: impl FnMut(S) -> I,
) -> Result<SendableRecordBatchStream> {
let Self {
- schema,
+ output_schema,
+ emit_ranks,
metrics,
reservation,
partitions,
@@ -1937,7 +1974,8 @@ impl<S: StoreSlots> PartitionedTopKCore<S> {
.map(|(_key, state)| state)
.flat_map(drain);
EmitState::stream(
- schema,
+ output_schema,
+ emit_ranks,
metrics,
reservation,
batch_size,
@@ -1986,7 +2024,7 @@ impl PartitionedTopK {
#[expect(clippy::too_many_arguments)]
pub(crate) fn try_new(
partition_id: usize,
- schema: SchemaRef,
+ schema: &SchemaRef,
partition_exprs: Vec<Arc<dyn PhysicalExpr>>,
partition_sort_fields: Vec<SortField>,
order_expr: LexOrdering,
@@ -1994,6 +2032,7 @@ impl PartitionedTopK {
batch_size: usize,
runtime: &Arc<RuntimeEnv>,
metrics: &ExecutionPlanMetricsSet,
+ ranked_schema: Option<SchemaRef>,
) -> Result<Self> {
Ok(Self {
core: PartitionedTopKCore::try_new(
@@ -2007,6 +2046,7 @@ impl PartitionedTopK {
batch_size,
runtime,
metrics,
+ ranked_schema,
)?,
})
}
@@ -2120,8 +2160,15 @@ impl PartitionedTopK {
/// the reservation is carried into [`EmitState`] and released when the stream
/// is dropped rather than when this returns.
pub(crate) fn emit(self) -> Result<SendableRecordBatchStream> {
- self.core
- .emit_with(|mut heap| heap.drain_sorted().into_iter().map(|slot| slot.at))
+ // Each heap drains its partition's rows in ORDER BY order, low to
+ // high, so a row's `ROW_NUMBER()` is just its 1-based position in that
+ // drain.
+ self.core.emit_with(|mut heap| {
+ heap.drain_sorted()
+ .into_iter()
+ .zip(1u64..)
+ .map(|(slot, row_number)| (slot.at, row_number))
+ })
}
/// Total memory currently held by this operator, including all
@@ -2135,8 +2182,11 @@ impl PartitionedTopK {
/// Hands out `batch_size` rows at a time, interleaved out of the store batches
/// the heaps referenced.
struct EmitState {
+ /// The input's columns, plus the ranking column when `ranks` is set.
+ output_schema: SchemaRef,
metrics: TopKMetrics,
- /// Covers the pinned batches and `ordered` for as long as they are held.
+ /// Covers the pinned batches, `ordered` and `ranks` for as long as they
+ /// are held.
///
/// Carried here rather than dropped at the end of `emit` so the bytes stay
/// accounted for until the stream is, and released by this struct's drop.
@@ -2149,6 +2199,8 @@ struct EmitState {
batches: Vec<RecordBatch>,
/// `(array_pos, row)` pairs in `(partition_keys, order_keys)` order.
ordered: Vec<(usize, usize)>,
+ /// Each `ordered` row's ranking value, when the operator emits it.
+ ranks: Option<Vec<u64>>,
/// Scratch for [`interleave_referenced`], one entry per batch.
remap: Vec<usize>,
pos: usize,
@@ -2157,14 +2209,18 @@ struct EmitState {
impl EmitState {
/// Resolve `slots`, `len` of them in output order, against `store`, and
/// return the stream that emits them. Shared by [`PartitionedTopK::emit`]
- /// and [`PartitionedTopKRank::emit`], which differ only in that order.
+ /// and [`PartitionedTopKRank::emit`], which differ only in that order and
+ /// in how they derive each row's ranking value, which is kept only when
+ /// `emit_ranks` is set.
+ #[expect(clippy::too_many_arguments)]
fn stream(
- schema: SchemaRef,
+ output_schema: SchemaRef,
+ emit_ranks: bool,
metrics: TopKMetrics,
reservation: MemoryReservation,
batch_size: usize,
store: &RecordBatchStore,
- slots: impl Iterator<Item = StoreRef>,
+ slots: impl Iterator<Item = (StoreRef, u64)>,
len: usize,
) -> Result<SendableRecordBatchStream> {
let timer = metrics.baseline.elapsed_compute().timer();
@@ -2175,17 +2231,23 @@ impl EmitState {
// Flattened in output order, so the emit itself is a slice walk.
let mut ordered: Vec<(usize, usize)> = Vec::with_capacity(len);
- for slot in slots {
+ // Parallel to `ordered`.
+ let mut ranks: Option<Vec<u64>> = emit_ranks.then(|| Vec::with_capacity(len));
+ for (slot, rank) in slots {
let array_pos = *batch_id_array_pos
.get(&slot.batch_id)
.expect("a retained slot's batch_id is present in the store");
ordered.push((array_pos, slot.row as usize));
+ if let Some(ranks) = &mut ranks {
+ ranks.push(rank);
+ }
}
drop(timer);
// What survives this function is the store's batches — pinned until the
// returned stream is dropped — plus `ordered`, one 16-byte pair per
- // retained row, and `remap`, one entry per batch. Everything else the
+ // retained row, `ranks`, 8 more bytes per row when the ranking column
+ // is emitted, and `remap`, one entry per batch. Everything else the
// operator held (its scratch buffers, every partition and its interned
// key) is freed by now, so the resize below is normally a shrink. The
// reservation moves into the stream state rather than being dropped
@@ -2197,21 +2259,26 @@ impl EmitState {
+ store.batches_size
+ batches.capacity() * size_of::<RecordBatch>()
+ ordered.capacity() * size_of::<(usize, usize)>()
+ + ranks
+ .as_ref()
+ .map_or(0, |r| r.capacity() * size_of::<u64>())
+ batches.len() * size_of::<usize>(),
)?;
let state = EmitState {
+ output_schema: Arc::clone(&output_schema),
metrics,
_reservation: reservation,
batch_size,
remap: vec![usize::MAX; batches.len()],
batches,
ordered,
+ ranks,
pos: 0,
};
Ok(Box::pin(RecordBatchStreamAdapter::new(
- schema,
+ output_schema,
futures::stream::try_unfold(state, |mut state| async move {
Ok(state.next_batch()?.map(|batch| (batch, state)))
}),
@@ -2226,16 +2293,48 @@ impl EmitState {
}
let _timer = self.metrics.baseline.elapsed_compute().timer();
- let end = (self.pos + self.batch_size).min(self.ordered.len());
- let chunk = &self.ordered[self.pos..end];
+ let start = self.pos;
+ let end = (start + self.batch_size).min(self.ordered.len());
self.pos = end;
- let batch = interleave_referenced(&self.batches, chunk, &mut self.remap)?;
- (&batch).record_output(&self.metrics.baseline);
- Ok(Some(batch))
+ let batch = interleave_referenced(
+ &self.batches,
+ &self.ordered[start..end],
+ &mut self.remap,
+ )?;
+ let ranks = self.ranks.as_ref().map(|r| r[start..end].iter().copied());
+ finish_chunk(batch, ranks, &self.output_schema, &self.metrics).map(Some)
}
}
+/// Finish one output chunk: append `ranks` as its last column under
+/// `output_schema` when the operator emits the ranking column, then count the
+/// chunk in `output_rows`.
+///
+/// Each of the three retention policies derives the values differently —
+/// see [`PartitionedTopK::emit`], [`PartitionedTopKRank::emit`], and
+/// [`PartitionedTopKDenseRank::emit`] — but all three know them by the time
+/// they build an output batch, which is what makes the
+/// `BoundedWindowAggExec` above the operator removable. `RecordBatch::try_new`
+/// rejects a `ranks` whose length differs from the chunk's.
+fn finish_chunk(
+ batch: RecordBatch,
+ ranks: Option<impl IntoIterator<Item = u64>>,
+ output_schema: &SchemaRef,
+ metrics: &TopKMetrics,
+) -> Result<RecordBatch> {
+ let batch = match ranks {
+ None => batch,
+ Some(ranks) => {
+ let mut columns = batch.columns().to_vec();
+ columns.push(Arc::new(UInt64Array::from_iter_values(ranks)));
+ RecordBatch::try_new(Arc::clone(output_schema), columns)?
+ }
+ };
+ (&batch).record_output(&metrics.baseline);
+ Ok(batch)
+}
+
/// Per-partition state for `RANK()` semantics.
///
/// `RANK ≤ K` keeps the K best rows by ORDER BY plus every row tied at the
@@ -2318,7 +2417,7 @@ impl PartitionedTopKRank {
#[expect(clippy::too_many_arguments)]
pub(crate) fn try_new(
partition_id: usize,
- schema: SchemaRef,
+ schema: &SchemaRef,
partition_exprs: Vec<Arc<dyn PhysicalExpr>>,
partition_sort_fields: Vec<SortField>,
order_expr: LexOrdering,
@@ -2326,6 +2425,7 @@ impl PartitionedTopKRank {
batch_size: usize,
runtime: &Arc<RuntimeEnv>,
metrics: &ExecutionPlanMetricsSet,
+ ranked_schema: Option<SchemaRef>,
) -> Result<Self> {
Ok(Self {
core: PartitionedTopKCore::try_new(
@@ -2339,6 +2439,7 @@ impl PartitionedTopKRank {
batch_size,
runtime,
metrics,
+ ranked_schema,
)?,
evicted_key: Vec::new(),
ties_bytes: 0,
@@ -2475,8 +2576,31 @@ impl PartitionedTopKRank {
/// (all sharing the boundary ob).
pub(crate) fn emit(self) -> Result<SendableRecordBatchStream> {
self.core.emit_with(|mut state| {
- let heap = state.heap.drain_sorted().into_iter().map(|slot| slot.at);
- heap.chain(state.ties)
+ let heap = state.heap.drain_sorted();
+ // `RANK` is `1 + (rows with a strictly smaller ORDER BY value)`.
+ // Counting positions is only valid because the retained set is a
+ // *complete* order-prefix of the partition: the heap keeps the K
+ // smallest keys and the ties only add rows equal to the boundary, so
+ // no row with a smaller key than a retained row was ever dropped.
+ //
+ // The ties all sit at the boundary, the heap's largest key, so they
+ // share the rank of the first heap row holding it. Ties exist only
+ // while the heap is full, so a heap with no rows has none.
+ let boundary_rank = heap.last().map_or(0, |worst| {
+ heap.partition_point(|slot| slot.key < worst.key) as u64 + 1
+ });
+ // Over the sorted heap, a row either ties with its predecessor and
+ // repeats its rank or starts a new one at its own 1-based position.
+ let mut prev: Option<Vec<u8>> = None;
+ let mut rank = 0;
+ let heap = heap.into_iter().zip(1u64..).map(move |(slot, position)| {
+ if prev.as_ref() != Some(&slot.key) {
+ rank = position;
+ }
+ prev = Some(slot.key);
+ (slot.at, rank)
+ });
+ heap.chain(state.ties.into_iter().map(move |at| (at, boundary_rank)))
})
}
@@ -2632,7 +2756,6 @@ impl DenseRankPartitionState {
/// row count is added to the `row_replacements` metric.
/// - `ob_key >= max` → drop the whole run; no map mutation.
pub(crate) struct PartitionedTopKDenseRank {
- schema: SchemaRef,
metrics: TopKMetrics,
reservation: MemoryReservation,
/// ORDER BY expressions (excludes PARTITION BY).
@@ -2669,15 +2792,23 @@ pub(crate) struct PartitionedTopKDenseRank {
/// keeps the reservation proportional to the batches actually pinned
/// rather than to the number of entries pointing at them.
store: RecordBatchStore,
+ /// When set, `emit` appends a column holding each retained row's
+ /// `DENSE_RANK()`, letting the rewrite drop the `BoundedWindowAggExec`
+ /// that would otherwise re-derive it.
+ emit_ranks: bool,
+ /// The schema `emit` produces: the input's columns, plus the ranking
+ /// column when [`Self::emit_ranks`] is set.
+ output_schema: SchemaRef,
k: usize,
batch_size: usize,
}
impl PartitionedTopKDenseRank {
+ /// `ranked_schema` is as for [`PartitionedTopKCore::try_new`].
#[expect(clippy::too_many_arguments)]
pub(crate) fn try_new(
partition_id: usize,
- schema: SchemaRef,
+ schema: &SchemaRef,
partition_exprs: Vec<Arc<dyn PhysicalExpr>>,
partition_sort_fields: Vec<SortField>,
order_expr: LexOrdering,
@@ -2685,13 +2816,14 @@ impl PartitionedTopKDenseRank {
batch_size: usize,
runtime: &Arc<RuntimeEnv>,
metrics: &ExecutionPlanMetricsSet,
+ ranked_schema: Option<SchemaRef>,
) -> Result<Self> {
assert!(k > 0, "PartitionedTopKDenseRank requires k > 0");
let reservation =
MemoryConsumer::new(format!("PartitionedTopKDenseRank[{partition_id}]"))
.register(&runtime.memory_pool);
- let order_sort_fields = build_sort_fields(&order_expr, &schema)?;
+ let order_sort_fields = build_sort_fields(&order_expr, schema)?;
let row_converter = RowConverter::new(order_sort_fields)?;
let scratch_rows =
row_converter.empty_rows(batch_size, ESTIMATED_BYTES_PER_ROW * batch_size);
@@ -2701,7 +2833,8 @@ impl PartitionedTopKDenseRank {
.empty_rows(batch_size, ESTIMATED_BYTES_PER_ROW * batch_size);
Ok(Self {
- schema,
+ emit_ranks: ranked_schema.is_some(),
+ output_schema: ranked_schema.unwrap_or_else(|| Arc::clone(schema)),
metrics: TopKMetrics::new(metrics, partition_id),
reservation,
expr: order_expr,
@@ -2911,7 +3044,8 @@ impl PartitionedTopKDenseRank {
/// emitted rows are in ob-sorted order.
pub(crate) fn emit(self) -> Result<SendableRecordBatchStream> {
let Self {
- schema,
+ output_schema,
+ emit_ranks,
metrics,
reservation: _,
expr: _,
@@ -2959,6 +3093,9 @@ impl PartitionedTopKDenseRank {
let mut batch_id_pos: HashMap<u32, usize> = HashMap::new();
let mut indices: Vec<(usize, usize)> = Vec::with_capacity(batch_size);
let mut sorted_obs: Vec<(Vec<u8>, Vec<GroupEntry>)> = Vec::new();
+ // Parallel to `indices`: each row's `DENSE_RANK()`, when emitted.
+ let mut ranks: Option<Vec<u64>> =
+ emit_ranks.then(|| Vec::with_capacity(batch_size));
// Chunk at `batch_size` so the operator emits the same batch sizes
// as before and no single `interleave` output exceeds `batch_size`.
@@ -2972,7 +3109,11 @@ impl PartitionedTopKDenseRank {
// (byte-comparable encoding == sort order).
sorted_obs.extend(groups);
sorted_obs.sort_by(|a, b| a.0.cmp(&b.0));
- for (_ob, entries) in sorted_obs.drain(..) {
+ // `DENSE_RANK` is the 1-based index of a row's distinct ob value
+ // within its partition, and the state holds exactly the K
+ // distinct-smallest values — so the loop counter *is* the rank,
+ // shared by every row of the group.
+ for (dense_rank, (_ob, entries)) in (1u64..).zip(sorted_obs.drain(..)) {
for entry in entries {
let batch = &store
.get(entry.batch_id)
@@ -2989,10 +3130,18 @@ impl PartitionedTopKDenseRank {
})
});
indices.push((pos, row as usize));
+ if let Some(ranks) = &mut ranks {
+ ranks.push(dense_rank);
+ }
if indices.len() == batch_size {
let b = interleave_record_batch(&batch_refs, &indices)?;
- (&b).record_output(&metrics.baseline);
- out.push(Ok(b));
+ let ranks = ranks.as_mut().map(|r| r.drain(..));
+ out.push(Ok(finish_chunk(
+ b,
+ ranks,
+ &output_schema,
+ &metrics,
+ )?));
indices.clear();
batch_refs.clear();
batch_id_pos.clear();
@@ -3004,12 +3153,12 @@ impl PartitionedTopKDenseRank {
}
if !indices.is_empty() {
let b = interleave_record_batch(&batch_refs, &indices)?;
- (&b).record_output(&metrics.baseline);
- out.push(Ok(b));
+ let ranks = ranks.as_mut().map(|r| r.drain(..));
+ out.push(Ok(finish_chunk(b, ranks, &output_schema, &metrics)?));
}
Ok(Box::pin(RecordBatchStreamAdapter::new(
- schema,
+ output_schema,
futures::stream::iter(out),
)))
}
@@ -3053,7 +3202,7 @@ mod tests {
AsArray, BooleanArray, Float64Array, Int32Array, StringArray,
StringDictionaryBuilder,
};
- use arrow::datatypes::{DataType, Field, Int32Type, Schema};
+ use arrow::datatypes::{DataType, Field, Int32Type, Schema, UInt64Type};
use arrow_schema::SortOptions;
use datafusion_common::{assert_batches_eq, exec_datafusion_err};
use datafusion_execution::memory_pool::GreedyMemoryPool;
@@ -3943,11 +4092,13 @@ mod tests {
val_sort_options: SortOptions,
val_nullable: bool,
) -> Result<(Arc<Schema>, PartitionedTopK)> {
- build_partitioned_topk_inner(
+ build_pk_val_state(
+ PartitionedTopK::try_new,
k,
val_sort_options,
val_nullable,
&Arc::new(RuntimeEnv::default()),
+ None,
)
}
@@ -3957,15 +4108,43 @@ mod tests {
k: usize,
runtime: &Arc<RuntimeEnv>,
) -> Result<(Arc<Schema>, PartitionedTopK)> {
- build_partitioned_topk_inner(k, SortOptions::default(), false, runtime)
+ build_pk_val_state(
+ PartitionedTopK::try_new,
+ k,
+ SortOptions::default(),
+ false,
+ runtime,
+ None,
+ )
}
- fn build_partitioned_topk_inner(
+ /// The signature all three policies' `try_new` share, so one builder can
+ /// construct any of them.
+ type StateCtor<T> = fn(
+ usize,
+ &SchemaRef,
+ Vec<Arc<dyn PhysicalExpr>>,
+ Vec<SortField>,
+ LexOrdering,
+ usize,
+ usize,
+ &Arc<RuntimeEnv>,
+ &ExecutionPlanMetricsSet,
+ Option<SchemaRef>,
+ ) -> Result<T>;
+
+ /// Builds a `(pk Int32, val Int32)` schema and a state built by `ctor`,
+ /// partitioned by `pk ASC` and ordered by `val` under `val_sort_options`,
+ /// with batch size 8. With a `ranking_field` the state appends it, as the
+ /// exec would. Shared by every policy's builders.
+ fn build_pk_val_state<T>(
+ ctor: StateCtor<T>,
k: usize,
val_sort_options: SortOptions,
val_nullable: bool,
runtime: &Arc<RuntimeEnv>,
- ) -> Result<(Arc<Schema>, PartitionedTopK)> {
+ ranking_field: Option<FieldRef>,
+ ) -> Result<(Arc<Schema>, T)> {
let schema = pk_val_schema(val_nullable);
let pk_expr: Arc<dyn PhysicalExpr> = col("pk", schema.as_ref())?;
@@ -3978,23 +4157,40 @@ mod tests {
options: val_sort_options,
};
- let partition_ordering = vec![pk_sort_expr];
- let order_expr = LexOrdering::from([val_sort_expr]);
-
- let state = PartitionedTopK::try_new(
+ let ranked_schema = ranking_field
+ .map(|field| partitioned_topk_output_schema(&schema, Some(&field)));
+ let state = ctor(
0,
- Arc::clone(&schema),
+ &schema,
vec![pk_expr],
- build_sort_fields(&partition_ordering, &schema)?,
- order_expr,
+ build_sort_fields(&[pk_sort_expr], &schema)?,
+ LexOrdering::from([val_sort_expr]),
k,
8, // batch_size
runtime,
&ExecutionPlanMetricsSet::new(),
+ ranked_schema,
)?;
Ok((schema, state))
}
+ /// A state built by `ctor` with [`build_pk_val_state`]'s defaults that
+ /// emits `ranking_field`.
+ fn build_emitting<T>(
+ ctor: StateCtor<T>,
+ k: usize,
+ ranking_field: FieldRef,
+ ) -> Result<(Arc<Schema>, T)> {
+ build_pk_val_state(
+ ctor,
+ k,
+ SortOptions::default(),
+ false,
+ &Arc::new(RuntimeEnv::default()),
+ Some(ranking_field),
+ )
+ }
+
/// Partition-key nullability is a distinct concern from ORDER BY
/// nullability: the key layer decides whether NULL keys collapse
/// into one partition, and emit-time ordering decides where that
@@ -4020,7 +4216,7 @@ mod tests {
let state = PartitionedTopK::try_new(
0,
- Arc::clone(&schema),
+ &schema,
vec![pk_expr],
build_sort_fields(&partition_ordering, &schema)?,
order_expr,
@@ -4028,6 +4224,7 @@ mod tests {
8, // batch_size
&Arc::new(RuntimeEnv::default()),
&ExecutionPlanMetricsSet::new(),
+ None,
)?;
Ok((schema, state))
}
@@ -4201,7 +4398,7 @@ mod tests {
let mut state = PartitionedTopK::try_new(
0,
- Arc::clone(&schema),
+ &schema,
vec![pk_a, pk_b],
build_sort_fields(&partition_ordering, &schema)?,
order_expr,
@@ -4209,6 +4406,7 @@ mod tests {
8,
&Arc::new(RuntimeEnv::default()),
&ExecutionPlanMetricsSet::new(),
+ None,
)?;
let batch = RecordBatch::try_new(
@@ -5262,7 +5460,7 @@ mod tests {
}]);
let mut state = PartitionedTopK::try_new(
0,
- Arc::clone(&schema),
+ &schema,
vec![pk_expr],
build_sort_fields(&partition_ordering, &schema)?,
order_expr,
@@ -5270,6 +5468,7 @@ mod tests {
8,
&Arc::new(RuntimeEnv::default()),
&ExecutionPlanMetricsSet::new(),
+ None,
)?;
let empty = state.size();
@@ -5450,7 +5649,7 @@ mod tests {
}];
let mut state = PartitionedTopK::try_new(
0,
- Arc::clone(&schema),
+ &schema,
vec![pk_expr],
build_sort_fields(&partition_ordering, &schema)?,
LexOrdering::from(order_ordering.clone()),
@@ -5458,6 +5657,7 @@ mod tests {
3, // batch_size
&Arc::new(RuntimeEnv::default()),
&metrics,
+ None,
)?;
state.insert_batch(&pk_val_batch(
@@ -5474,6 +5674,175 @@ mod tests {
Ok(())
}
+ /// With a `ranking_field`, `emit` appends the retained rows'
+ /// `ROW_NUMBER()`: 1-based within each partition, restarting at every
+ /// partition boundary, and independent of how the partitions' rows were
+ /// interleaved on the way in.
+ #[tokio::test]
+ async fn test_partitioned_topk_emits_row_number_column() -> Result<()> {
+ let rn_field = Arc::new(Field::new("rn", DataType::UInt64, false));
+ let (schema, mut state) =
+ build_emitting(PartitionedTopK::try_new, 3, Arc::clone(&rn_field))?;
+
+ // pk=1 vals: 30, 10, 20, 40 → top-3 ASC = [10, 20, 30] (40 evicted)
+ // pk=2 vals: 5 → top-3 ASC = [5]
+ state.insert_batch(&pk_val_batch(
+ &schema,
+ vec![1, 2, 1, 1, 1],
+ vec![30, 5, 10, 20, 40],
+ )?)?;
+
+ let results: Vec<_> = state.emit()?.try_collect().await?;
+ assert_eq!(
+ results[0].schema().field(2),
+ rn_field.as_ref(),
+ "the appended field is the one the caller asked for"
+ );
+ assert_batches_eq!(
+ &[
+ "+----+-----+----+",
+ "| pk | val | rn |",
+ "+----+-----+----+",
+ "| 1 | 10 | 1 |",
+ "| 1 | 20 | 2 |",
+ "| 1 | 30 | 3 |",
+ "| 2 | 5 | 1 |",
+ "+----+-----+----+",
+ ],
+ &results
+ );
+ Ok(())
+ }
+
+ /// Input for the ranking-column chunk tests: partition 1 holds vals 1..=4
+ /// three times each, partition 2 holds two 5s, split over two input
+ /// batches. Every policy retains more rows than the builders' batch size
+ /// of 8, so the ranks have to be sliced correctly across output chunks.
+ fn ranking_chunk_batches(schema: &SchemaRef) -> Result<[RecordBatch; 2]> {
+ Ok([
+ pk_val_batch(schema, vec![1; 7], vec![1, 1, 1, 2, 2, 2, 3])?,
+ pk_val_batch(schema, vec![1, 1, 1, 1, 1, 2, 2], vec![3, 3, 4, 4, 4, 5, 5])?,
+ ])
+ }
+
+ /// `(val, rank)` for every emitted row, and how many batches carried them.
+ fn vals_and_ranks(batches: &[RecordBatch]) -> (Vec<(i32, u64)>, usize) {
+ let mut rows = Vec::new();
+ for batch in batches {
+ let vals = batch.column(1).as_primitive::<Int32Type>();
+ let ranks = batch.column(2).as_primitive::<UInt64Type>();
+ rows.extend(
+ vals.values()
+ .iter()
+ .copied()
+ .zip(ranks.values().iter().copied()),
+ );
+ }
+ (rows, batches.len())
+ }
+
+ /// `ROW_NUMBER()` keeps counting across an output chunk boundary, and
+ /// restarts at the next partition.
+ #[tokio::test]
+ async fn test_partitioned_topk_row_number_column_spans_output_chunks() -> Result<()> {
+ let rn = Arc::new(Field::new("rn", DataType::UInt64, false));
+ let (schema, mut state) = build_emitting(PartitionedTopK::try_new, 10, rn)?;
+ for batch in ranking_chunk_batches(&schema)? {
+ state.insert_batch(&batch)?;
+ }
+ let results: Vec<_> = state.emit()?.try_collect().await?;
+ let (rows, chunks) = vals_and_ranks(&results);
+ assert_eq!(chunks, 2);
+ assert_eq!(
+ rows,
+ vec![
+ (1, 1),
+ (1, 2),
+ (1, 3),
+ (2, 4),
+ (2, 5),
+ (2, 6),
+ (3, 7),
+ (3, 8),
+ // second chunk
+ (3, 9),
+ (4, 10),
+ (5, 1),
+ (5, 2),
+ ]
+ );
+ Ok(())
+ }
+
+ /// `RANK()` across an output chunk boundary: the second chunk opens inside
+ /// the 3s and ends with the boundary ties, so heap ranks and tie ranks are
+ /// both read past the first chunk.
+ #[tokio::test]
+ async fn test_partitioned_topk_rank_column_spans_output_chunks() -> Result<()> {
+ let rk = Arc::new(Field::new("rk", DataType::UInt64, false));
+ let (schema, mut state) = build_emitting(PartitionedTopKRank::try_new, 10, rk)?;
+ for batch in ranking_chunk_batches(&schema)? {
+ state.insert_batch(&batch)?;
+ }
+ let results: Vec<_> = state.emit()?.try_collect().await?;
+ let (rows, chunks) = vals_and_ranks(&results);
+ assert_eq!(chunks, 2);
+ assert_eq!(
+ rows,
+ vec![
+ (1, 1),
+ (1, 1),
+ (1, 1),
+ (2, 4),
+ (2, 4),
+ (2, 4),
+ (3, 7),
+ (3, 7),
+ // second chunk: the last 3, then the heap's 4 and its two ties
+ (3, 7),
+ (4, 10),
+ (4, 10),
+ (4, 10),
+ (5, 1),
+ (5, 1),
+ ]
+ );
+ Ok(())
+ }
+
+ /// `DENSE_RANK()` across an output chunk boundary, which its emit flushes
+ /// separately from the other two policies.
+ #[tokio::test]
+ async fn test_partitioned_topk_dense_rank_column_spans_output_chunks() -> Result<()> {
+ let dr = Arc::new(Field::new("dr", DataType::UInt64, false));
+ let (schema, mut state) =
+ build_emitting(PartitionedTopKDenseRank::try_new, 3, dr)?;
+ for batch in ranking_chunk_batches(&schema)? {
+ state.insert_batch(&batch)?;
+ }
+ let results: Vec<_> = state.emit()?.try_collect().await?;
+ let (rows, chunks) = vals_and_ranks(&results);
+ assert_eq!(chunks, 2);
+ assert_eq!(
+ rows,
+ vec![
+ (1, 1),
+ (1, 1),
+ (1, 1),
+ (2, 2),
+ (2, 2),
+ (2, 2),
+ (3, 3),
+ (3, 3),
+ // second chunk
+ (3, 3),
+ (5, 1),
+ (5, 1),
+ ]
+ );
+ Ok(())
+ }
+
/// Multiple distinct partition keys interleaved within a single
/// input batch — grouping rows by partition key, per-partition heap eviction,
/// and partition-key-ordered emit must all behave correctly.
@@ -5736,11 +6105,13 @@ mod tests {
val_sort_options: SortOptions,
val_nullable: bool,
) -> Result<(Arc<Schema>, PartitionedTopKRank)> {
- build_partitioned_topk_rank_inner(
+ build_pk_val_state(
+ PartitionedTopKRank::try_new,
k,
val_sort_options,
val_nullable,
&Arc::new(RuntimeEnv::default()),
+ None,
)
}
@@ -5751,42 +6122,14 @@ mod tests {
k: usize,
runtime: &Arc<RuntimeEnv>,
) -> Result<(Arc<Schema>, PartitionedTopKRank)> {
- build_partitioned_topk_rank_inner(k, SortOptions::default(), false, runtime)
- }
-
- fn build_partitioned_topk_rank_inner(
- k: usize,
- val_sort_options: SortOptions,
- val_nullable: bool,
- runtime: &Arc<RuntimeEnv>,
- ) -> Result<(Arc<Schema>, PartitionedTopKRank)> {
- let schema = pk_val_schema(val_nullable);
-
- let pk_expr: Arc<dyn PhysicalExpr> = col("pk", schema.as_ref())?;
- let pk_sort_expr = PhysicalSortExpr {
- expr: Arc::clone(&pk_expr),
- options: SortOptions::default(),
- };
- let val_sort_expr = PhysicalSortExpr {
- expr: col("val", schema.as_ref())?,
- options: val_sort_options,
- };
-
- let partition_sort_fields = build_sort_fields(&[pk_sort_expr], &schema)?;
- let order_expr = LexOrdering::from([val_sort_expr]);
-
- let state = PartitionedTopKRank::try_new(
- 0,
- Arc::clone(&schema),
- vec![pk_expr],
- partition_sort_fields,
- order_expr,
+ build_pk_val_state(
+ PartitionedTopKRank::try_new,
k,
- 8, // batch_size
+ SortOptions::default(),
+ false,
runtime,
- &ExecutionPlanMetricsSet::new(),
- )?;
- Ok((schema, state))
+ None,
+ )
}
/// Multiple distinct partition keys interleaved within a single
@@ -6058,6 +6401,82 @@ mod tests {
Ok(())
}
+ /// With a `ranking_field`, RANK's `emit` appends each retained row's
+ /// `RANK()`. The interesting rows are the boundary ties, which are emitted
+ /// *after* the heap's rows and must carry the rank of the first heap row
+ /// holding the boundary key rather than their own position.
+ #[tokio::test]
+ async fn test_partitioned_topk_rank_emits_rank_column() -> Result<()> {
+ let rk_field = Arc::new(Field::new("rk", DataType::UInt64, false));
+ let (schema, mut state) =
+ build_emitting(PartitionedTopKRank::try_new, 2, Arc::clone(&rk_field))?;
+
+ // pk=1 vals 10, 20, 20: heap fills with [10, 20], the second 20 ties at
+ // the boundary. True RANKs: 10→1, 20→2, 20→2.
+ // pk=2 vals 5, 5: both tie at rank 1 inside the heap.
+ // pk=3 vals 7, 7, 7: the boundary key fills the whole heap and the
+ // third 7 ties at it, so the tie's rank is the boundary's first
+ // position (1), not the heap's length (2).
+ state.insert_batch(&pk_val_batch(
+ &schema,
+ vec![1, 1, 1, 2, 2, 3, 3, 3],
+ vec![10, 20, 20, 5, 5, 7, 7, 7],
+ )?)?;
+
+ let results: Vec<_> = state.emit()?.try_collect().await?;
+ assert_eq!(results[0].schema().field(2), rk_field.as_ref());
+ assert_batches_eq!(
+ &[
+ "+----+-----+----+",
+ "| pk | val | rk |",
+ "+----+-----+----+",
+ "| 1 | 10 | 1 |",
+ "| 1 | 20 | 2 |",
+ "| 1 | 20 | 2 |",
+ "| 2 | 5 | 1 |",
+ "| 2 | 5 | 1 |",
+ "| 3 | 7 | 1 |",
+ "| 3 | 7 | 1 |",
+ "| 3 | 7 | 1 |",
+ "+----+-----+----+",
+ ],
+ &results
+ );
+ Ok(())
+ }
+
+ /// RANK's emitted column must agree with the window function it replaces
+ /// when an eviction leaves the boundary where it was: the evicted row
+ /// still ties with the new root, so it moves to the tie list, and its rank
+ /// then comes from the boundary rather than from a heap position.
+ #[tokio::test]
+ async fn test_partitioned_topk_rank_emits_rank_after_boundary_moves() -> Result<()> {
+ let rk_field = Arc::new(Field::new("rk", DataType::UInt64, false));
+ let (schema, mut state) =
+ build_emitting(PartitionedTopKRank::try_new, 2, Arc::clone(&rk_field))?;
+
+ // Arrive as 20, 20, 10: the heap fills with [20, 20], then 10 evicts one
+ // 20, whose bytes equal the new boundary (20) so it moves to the ties.
+ // True RANKs over {10, 20, 20}: 10→1, 20→2, 20→2.
+ state.insert_batch(&pk_val_batch(&schema, vec![1, 1], vec![20, 20])?)?;
+ state.insert_batch(&pk_val_batch(&schema, vec![1], vec![10])?)?;
+
+ let results: Vec<_> = state.emit()?.try_collect().await?;
+ assert_batches_eq!(
+ &[
+ "+----+-----+----+",
+ "| pk | val | rk |",
+ "+----+-----+----+",
+ "| 1 | 10 | 1 |",
+ "| 1 | 20 | 2 |",
+ "| 1 | 20 | 2 |",
+ "+----+-----+----+",
+ ],
+ &results
+ );
+ Ok(())
+ }
+
/// Tie rows are emitted in the same `EmitState` chunks as heap rows, so
/// they must be counted in `output_rows` once, not once as ties and again
/// as part of the chunk that holds them.
@@ -6076,7 +6495,7 @@ mod tests {
let metrics = ExecutionPlanMetricsSet::new();
let mut state = PartitionedTopKRank::try_new(
0,
- Arc::clone(&schema),
+ &schema,
vec![pk_expr],
build_sort_fields(&[pk_sort_expr], &schema)?,
LexOrdering::from([val_sort_expr]),
@@ -6084,6 +6503,7 @@ mod tests {
8, // batch_size
&Arc::new(RuntimeEnv::default()),
&metrics,
+ None,
)?;
// Two 5s fill the heap, the third 5 is retained as a tie.
@@ -6209,7 +6629,7 @@ mod tests {
}]);
PartitionedTopKRank::try_new(
0,
- Arc::clone(schema),
+ schema,
vec![pk_expr],
partition_sort_fields,
order_expr,
@@ -6217,6 +6637,7 @@ mod tests {
batch_size,
&Arc::new(RuntimeEnv::default()),
&ExecutionPlanMetricsSet::new(),
+ None,
)
}
@@ -6645,33 +7066,14 @@ mod tests {
val_sort_options: SortOptions,
val_nullable: bool,
) -> Result<(Arc<Schema>, PartitionedTopKDenseRank)> {
- let schema = pk_val_schema(val_nullable);
-
- let pk_expr: Arc<dyn PhysicalExpr> = col("pk", schema.as_ref())?;
- let pk_sort_expr = PhysicalSortExpr {
- expr: Arc::clone(&pk_expr),
- options: SortOptions::default(),
- };
- let val_sort_expr = PhysicalSortExpr {
- expr: col("val", schema.as_ref())?,
- options: val_sort_options,
- };
-
- let partition_sort_fields = build_sort_fields(&[pk_sort_expr], &schema)?;
- let order_expr = LexOrdering::from([val_sort_expr]);
-
- let state = PartitionedTopKDenseRank::try_new(
- 0,
- Arc::clone(&schema),
- vec![pk_expr],
- partition_sort_fields,
- order_expr,
+ build_pk_val_state(
+ PartitionedTopKDenseRank::try_new,
k,
- 8, // batch_size
+ val_sort_options,
+ val_nullable,
&Arc::new(RuntimeEnv::default()),
- &ExecutionPlanMetricsSet::new(),
- )?;
- Ok((schema, state))
+ None,
+ )
}
/// Single-batch DENSE_RANK top-2 across multiple partitions with
@@ -6707,6 +7109,75 @@ mod tests {
Ok(())
}
+ /// With a `ranking_field`, DENSE_RANK's `emit` appends each retained row's
+ /// `DENSE_RANK()`: the 1-based index of its distinct ORDER BY value, so
+ /// every row of a tied group shares one value and the counter does not skip
+ /// over group sizes the way `RANK` does.
+ #[tokio::test]
+ async fn test_partitioned_topk_dense_rank_emits_rank_column() -> Result<()> {
+ let dr_field = Arc::new(Field::new("dr", DataType::UInt64, false));
+ let (schema, mut state) =
+ build_emitting(PartitionedTopKDenseRank::try_new, 2, Arc::clone(&dr_field))?;
+
+ // pk=1 vals 10, 10, 20, 30: the 2 distinct-smallest are 10 and 20, so
+ // DENSE_RANKs are 10→1, 10→1, 20→2 and 30 is dropped. Note RANK would
+ // have given the 20 rank 3, not 2.
+ // pk=2 vals 7, 7: one distinct value, both dense rank 1.
+ state.insert_batch(&pk_val_batch(
+ &schema,
+ vec![1, 1, 1, 1, 2, 2],
+ vec![10, 10, 20, 30, 7, 7],
+ )?)?;
+
+ let results: Vec<_> = state.emit()?.try_collect().await?;
+ assert_eq!(results[0].schema().field(2), dr_field.as_ref());
+ assert_batches_eq!(
+ &[
+ "+----+-----+----+",
+ "| pk | val | dr |",
+ "+----+-----+----+",
+ "| 1 | 10 | 1 |",
+ "| 1 | 10 | 1 |",
+ "| 1 | 20 | 2 |",
+ "| 2 | 7 | 1 |",
+ "| 2 | 7 | 1 |",
+ "+----+-----+----+",
+ ],
+ &results
+ );
+ Ok(())
+ }
+
+ /// DENSE_RANK groups accumulate across batches, and a group's rows can come
+ /// from several source batches — each gathered separately at emit time. The
+ /// rank must be per *group*, not per gathered batch.
+ #[tokio::test]
+ async fn test_partitioned_topk_dense_rank_emits_rank_across_batches() -> Result<()> {
+ let dr_field = Arc::new(Field::new("dr", DataType::UInt64, false));
+ let (schema, mut state) =
+ build_emitting(PartitionedTopKDenseRank::try_new, 2, Arc::clone(&dr_field))?;
+
+ // Both groups (10 and 20) collect a row from each of the two batches.
+ state.insert_batch(&pk_val_batch(&schema, vec![1, 1], vec![20, 10])?)?;
+ state.insert_batch(&pk_val_batch(&schema, vec![1, 1], vec![10, 20])?)?;
+
+ let results: Vec<_> = state.emit()?.try_collect().await?;
+ assert_batches_eq!(
+ &[
+ "+----+-----+----+",
+ "| pk | val | dr |",
+ "+----+-----+----+",
+ "| 1 | 10 | 1 |",
+ "| 1 | 10 | 1 |",
+ "| 1 | 20 | 2 |",
+ "| 1 | 20 | 2 |",
+ "+----+-----+----+",
+ ],
+ &results
+ );
+ Ok(())
+ }
+
/// DENSE_RANK-specific: heavy ties within a batch. All rows at each
/// distinct ob value must be kept — within-call bucketing groups them
/// into one `GroupEntry` per distinct ob.
@@ -6862,7 +7333,7 @@ mod tests {
let mut state = PartitionedTopKDenseRank::try_new(
0,
- Arc::clone(&schema),
+ &schema,
vec![pk_expr],
partition_sort_fields,
order_expr,
@@ -6870,6 +7341,7 @@ mod tests {
8, // batch_size
&Arc::new(RuntimeEnv::default()),
&ExecutionPlanMetricsSet::new(),
+ None,
)?;
let dict_batch = |vals: &[&str]| -> Result<RecordBatch> {
@@ -7339,7 +7811,7 @@ mod tests {
}]);
let mut state = PartitionedTopKDenseRank::try_new(
0,
- Arc::clone(&schema),
+ &schema,
vec![pk_expr],
partition_sort_fields,
order_expr,
@@ -7347,6 +7819,7 @@ mod tests {
8, // batch_size
&Arc::new(RuntimeEnv::default()),
&ExecutionPlanMetricsSet::new(),
+ None,
)?;
// One row per partition, each with a wide key.
diff --git a/datafusion/sqllogictest/test_files/range_partitioning.slt b/datafusion/sqllogictest/test_files/range_partitioning.slt
index b3d2a1f442..56184845a6 100644
--- a/datafusion/sqllogictest/test_files/range_partitioning.slt
+++ b/datafusion/sqllogictest/test_files/range_partitioning.slt
@@ -1631,9 +1631,8 @@ EXPLAIN SELECT * FROM (
----
physical_plan
01)ProjectionExec: expr=[range_key@0 as range_key, value@1 as value, row_number() PARTITION BY [range_partitioned.range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@2 as rn]
-02)--BoundedWindowAggExec: wdw=[row_number() PARTITION BY [range_partitioned.range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [range_partitioned.range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
-03)----PartitionedTopKExec: fn=row_number, fetch=1, partition=[range_key@0], order=[value@1 DESC]
-04)------DataSourceExec: file_groups=<slt:ignore>, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet
+02)--PartitionedTopKExec: fn=row_number, fetch=1, partition=[range_key@0], order=[value@1 DESC], emit=[row_number() PARTITION BY [range_partitioned.range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]
+03)----DataSourceExec: file_groups=<slt:ignore>, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet
query III
SELECT * FROM (
@@ -1672,10 +1671,9 @@ EXPLAIN SELECT * FROM (
----
physical_plan
01)ProjectionExec: expr=[non_range_key@0 as non_range_key, value@1 as value, row_number() PARTITION BY [range_partitioned.non_range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@2 as rn]
-02)--BoundedWindowAggExec: wdw=[row_number() PARTITION BY [range_partitioned.non_range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [range_partitioned.non_range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
-03)----PartitionedTopKExec: fn=row_number, fetch=1, partition=[non_range_key@0], order=[value@1 DESC]
-04)------RepartitionExec: partitioning=Hash([non_range_key@0], 4), input_partitions=4
-05)--------DataSourceExec: file_groups=<slt:ignore>, projection=[non_range_key, value], output_partitioning=UnknownPartitioning(4), file_type=parquet
+02)--PartitionedTopKExec: fn=row_number, fetch=1, partition=[non_range_key@0], order=[value@1 DESC], emit=[row_number() PARTITION BY [range_partitioned.non_range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]
+03)----RepartitionExec: partitioning=Hash([non_range_key@0], 4), input_partitions=4
+04)------DataSourceExec: file_groups=<slt:ignore>, projection=[non_range_key, value], output_partitioning=UnknownPartitioning(4), file_type=parquet
query III
SELECT * FROM (
@@ -1708,9 +1706,8 @@ EXPLAIN SELECT * FROM (
----
physical_plan
01)ProjectionExec: expr=[range_key@0 as range_key, non_range_key@1 as non_range_key, value@2 as value, row_number() PARTITION BY [range_partitioned.range_key, range_partitioned.non_range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rn]
-02)--BoundedWindowAggExec: wdw=[row_number() PARTITION BY [range_partitioned.range_key, range_partitioned.non_range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [range_partitioned.range_key, range_partitioned.non_range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW [...]
-03)----PartitionedTopKExec: fn=row_number, fetch=1, partition=[range_key@0, non_range_key@1], order=[value@2 DESC]
-04)------DataSourceExec: file_groups=<slt:ignore>, projection=[range_key, non_range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet
+02)--PartitionedTopKExec: fn=row_number, fetch=1, partition=[range_key@0, non_range_key@1], order=[value@2 DESC], emit=[row_number() PARTITION BY [range_partitioned.range_key, range_partitioned.non_range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]
+03)----DataSourceExec: file_groups=<slt:ignore>, projection=[range_key, non_range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet
query IIII
SELECT * FROM (
@@ -1753,10 +1750,9 @@ EXPLAIN SELECT * FROM (
----
physical_plan
01)ProjectionExec: expr=[range_key@0 as range_key, non_range_key@1 as non_range_key, value@2 as value, row_number() PARTITION BY [range_partitioned.range_key, range_partitioned.non_range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rn]
-02)--BoundedWindowAggExec: wdw=[row_number() PARTITION BY [range_partitioned.range_key, range_partitioned.non_range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [range_partitioned.range_key, range_partitioned.non_range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW [...]
-03)----PartitionedTopKExec: fn=row_number, fetch=1, partition=[range_key@0, non_range_key@1], order=[value@2 DESC]
-04)------RepartitionExec: partitioning=Hash([range_key@0, non_range_key@1], 4), input_partitions=4
-05)--------DataSourceExec: file_groups=<slt:ignore>, projection=[range_key, non_range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet
+02)--PartitionedTopKExec: fn=row_number, fetch=1, partition=[range_key@0, non_range_key@1], order=[value@2 DESC], emit=[row_number() PARTITION BY [range_partitioned.range_key, range_partitioned.non_range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]
+03)----RepartitionExec: partitioning=Hash([range_key@0, non_range_key@1], 4), input_partitions=4
+04)------DataSourceExec: file_groups=<slt:ignore>, projection=[range_key, non_range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet
statement ok
reset datafusion.optimizer.subset_repartition_threshold;
diff --git a/datafusion/sqllogictest/test_files/window_topn.slt b/datafusion/sqllogictest/test_files/window_topn.slt
index 5c3757aa5f..637b6029c7 100644
--- a/datafusion/sqllogictest/test_files/window_topn.slt
+++ b/datafusion/sqllogictest/test_files/window_topn.slt
@@ -63,10 +63,8 @@ logical_plan
04)------TableScan: window_topn_t projection=[id, pk, val]
physical_plan
01)ProjectionExec: expr=[id@0 as id, pk@1 as pk, val@2 as val, row_number() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rn]
-02)--BoundedWindowAggExec: wdw=[row_number() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
-03)----RepartitionExec: partitioning=Hash([pk@1], 4), input_partitions=1, maintains_sort_order=true
-04)------PartitionedTopKExec: fn=row_number, fetch=3, partition=[pk@1], order=[val@2 ASC NULLS LAST]
-05)--------DataSourceExec: partitions=1, partition_sizes=[1]
+02)--PartitionedTopKExec: fn=row_number, fetch=3, partition=[pk@1], order=[val@2 ASC NULLS LAST], emit=[row_number() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]
+03)----DataSourceExec: partitions=1, partition_sizes=[1]
# Test 3: rn < 4 should give same results (fetch=3)
query III rowsort
@@ -154,10 +152,8 @@ logical_plan
04)------TableScan: window_topn_t projection=[id, pk, val]
physical_plan
01)ProjectionExec: expr=[id@0 as id, pk@1 as pk, val@2 as val, row_number() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rn]
-02)--BoundedWindowAggExec: wdw=[row_number() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
-03)----RepartitionExec: partitioning=Hash([pk@1], 4), input_partitions=1, maintains_sort_order=true
-04)------PartitionedTopKExec: fn=row_number, fetch=3, partition=[pk@1], order=[val@2 ASC NULLS LAST]
-05)--------DataSourceExec: partitions=1, partition_sizes=[1]
+02)--PartitionedTopKExec: fn=row_number, fetch=3, partition=[pk@1], order=[val@2 ASC NULLS LAST], emit=[row_number() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]
+03)----DataSourceExec: partitions=1, partition_sizes=[1]
# Test 7: Filter on data column (not window output) — should NOT optimize
query TT
@@ -188,10 +184,8 @@ EXPLAIN SELECT * FROM (
----
physical_plan
01)ProjectionExec: expr=[id@0 as id, pk@1 as pk, val@2 as val, row_number() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rn]
-02)--BoundedWindowAggExec: wdw=[row_number() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
-03)----RepartitionExec: partitioning=Hash([pk@1], 4), input_partitions=1, maintains_sort_order=true
-04)------PartitionedTopKExec: fn=row_number, fetch=3, partition=[pk@1], order=[val@2 ASC NULLS LAST]
-05)--------DataSourceExec: partitions=1, partition_sizes=[1]
+02)--PartitionedTopKExec: fn=row_number, fetch=3, partition=[pk@1], order=[val@2 ASC NULLS LAST], emit=[row_number() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]
+03)----DataSourceExec: partitions=1, partition_sizes=[1]
statement ok
SET datafusion.explain.physical_plan_only = false;
@@ -245,46 +239,32 @@ physical_plan
16)│ val: val │
17)└─────────────┬─────────────┘
18)┌─────────────┴─────────────┐
-19)│ BoundedWindowAggExec │
+19)│ PartitionedTopKExec │
20)│ -------------------- │
-21)│ mode: Sorted │
-22)│ │
-23)│ select_list: │
-24)│ row_number() PARTITION BY │
-25)│ [window_topn_t.pk] ORDER │
-26)│ BY [window_topn_t.val │
-27)│ ASC NULLS LAST] RANGE │
-28)│ BETWEEN UNBOUNDED │
-29)│ PRECEDING AND │
-30)│ CURRENT ROW │
-31)└─────────────┬─────────────┘
-32)┌─────────────┴─────────────┐
-33)│ RepartitionExec │
-34)│ -------------------- │
-35)│ partition_count(in->out): │
-36)│ 1 -> 4 │
-37)│ │
-38)│ partitioning_scheme: │
-39)│ Hash([pk@1], 4) │
-40)└─────────────┬─────────────┘
-41)┌─────────────┴─────────────┐
-42)│ PartitionedTopKExec │
-43)│ -------------------- │
-44)│ fetch: 3 │
-45)│ fn: row_number │
-46)│ │
-47)│ order: │
-48)│ [val@2 ASC NULLS LAST] │
-49)│ │
-50)│ partition: [pk@1] │
-51)└─────────────┬─────────────┘
-52)┌─────────────┴─────────────┐
-53)│ DataSourceExec │
-54)│ -------------------- │
-55)│ bytes: 480 │
-56)│ format: memory │
-57)│ rows: 10 │
-58)└───────────────────────────┘
+21)│ emit: │
+22)│ [row_number() PARTITION BY│
+23)│ [window_topn_t.pk] ORDER │
+24)│ BY [window_topn_t.val │
+25)│ ASC NULLS LAST] RANGE │
+26)│ BETWEEN UNBOUNDED │
+27)│ PRECEDING AND │
+28)│ CURRENT ROW] │
+29)│ │
+30)│ fetch: 3 │
+31)│ fn: row_number │
+32)│ │
+33)│ order: │
+34)│ [val@2 ASC NULLS LAST] │
+35)│ │
+36)│ partition: [pk@1] │
+37)└─────────────┬─────────────┘
+38)┌─────────────┴─────────────┐
+39)│ DataSourceExec │
+40)│ -------------------- │
+41)│ bytes: 480 │
+42)│ format: memory │
+43)│ rows: 10 │
+44)└───────────────────────────┘
statement ok
SET datafusion.explain.format = indent;
@@ -501,10 +481,8 @@ EXPLAIN SELECT * FROM (
----
physical_plan
01)ProjectionExec: expr=[id@0 as id, pk@1 as pk, val@2 as val, row_number() PARTITION BY [window_topn_t.pk, window_topn_t.id] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rn]
-02)--BoundedWindowAggExec: wdw=[row_number() PARTITION BY [window_topn_t.pk, window_topn_t.id] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [window_topn_t.pk, window_topn_t.id] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
-03)----RepartitionExec: partitioning=Hash([pk@1, id@0], 4), input_partitions=1, maintains_sort_order=true
-04)------PartitionedTopKExec: fn=row_number, fetch=3, partition=[pk@1, id@0], order=[val@2 ASC NULLS LAST]
-05)--------DataSourceExec: partitions=1, partition_sizes=[1]
+02)--PartitionedTopKExec: fn=row_number, fetch=3, partition=[pk@1, id@0], order=[val@2 ASC NULLS LAST], emit=[row_number() PARTITION BY [window_topn_t.pk, window_topn_t.id] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]
+03)----DataSourceExec: partitions=1, partition_sizes=[1]
statement ok
SET datafusion.explain.physical_plan_only = false;
@@ -533,10 +511,8 @@ EXPLAIN SELECT * FROM (
----
physical_plan
01)ProjectionExec: expr=[id@0 as id, pk@1 as pk, val@2 as val, row_number() PARTITION BY [window_topn_t.id] ORDER BY [window_topn_t.id ASC NULLS LAST, window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rn]
-02)--BoundedWindowAggExec: wdw=[row_number() PARTITION BY [window_topn_t.id] ORDER BY [window_topn_t.id ASC NULLS LAST, window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [window_topn_t.id] ORDER BY [window_topn_t.id ASC NULLS LAST, window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
-03)----RepartitionExec: partitioning=Hash([id@0], 4), input_partitions=1, maintains_sort_order=true
-04)------PartitionedTopKExec: fn=row_number, fetch=3, partition=[id@0], order=[val@2 ASC NULLS LAST]
-05)--------DataSourceExec: partitions=1, partition_sizes=[1]
+02)--PartitionedTopKExec: fn=row_number, fetch=3, partition=[id@0], order=[val@2 ASC NULLS LAST], emit=[row_number() PARTITION BY [window_topn_t.id] ORDER BY [window_topn_t.id ASC NULLS LAST, window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]
+03)----DataSourceExec: partitions=1, partition_sizes=[1]
# Test 19: Overlapping keys correctness (each id is unique, so rn=1 for all)
statement ok
@@ -569,10 +545,8 @@ EXPLAIN SELECT * FROM (
----
physical_plan
01)ProjectionExec: expr=[id@0 as id, pk@1 as pk, val@2 as val, row_number() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.pk ASC NULLS LAST, window_topn_t.val DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rn]
-02)--BoundedWindowAggExec: wdw=[row_number() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.pk ASC NULLS LAST, window_topn_t.val DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.pk ASC NULLS LAST, window_topn_t.val DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
-03)----RepartitionExec: partitioning=Hash([pk@1], 4), input_partitions=1, maintains_sort_order=true
-04)------PartitionedTopKExec: fn=row_number, fetch=3, partition=[pk@1], order=[val@2 DESC]
-05)--------DataSourceExec: partitions=1, partition_sizes=[1]
+02)--PartitionedTopKExec: fn=row_number, fetch=3, partition=[pk@1], order=[val@2 DESC], emit=[row_number() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.pk ASC NULLS LAST, window_topn_t.val DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]
+03)----DataSourceExec: partitions=1, partition_sizes=[1]
# Test 21: Correctness for PARTITION BY pk ORDER BY pk, val DESC
statement ok
@@ -604,10 +578,8 @@ EXPLAIN SELECT * FROM (
----
physical_plan
01)ProjectionExec: expr=[id@0 as id, pk@1 as pk, val@2 as val, row_number() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.pk DESC NULLS FIRST, window_topn_t.val DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rn]
-02)--BoundedWindowAggExec: wdw=[row_number() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.pk DESC NULLS FIRST, window_topn_t.val DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.pk DESC NULLS FIRST, window_topn_t.val DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
-03)----RepartitionExec: partitioning=Hash([pk@1], 4), input_partitions=1, maintains_sort_order=true
-04)------PartitionedTopKExec: fn=row_number, fetch=3, partition=[pk@1], order=[val@2 DESC]
-05)--------DataSourceExec: partitions=1, partition_sizes=[1]
+02)--PartitionedTopKExec: fn=row_number, fetch=3, partition=[pk@1], order=[val@2 DESC], emit=[row_number() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.pk DESC NULLS FIRST, window_topn_t.val DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]
+03)----DataSourceExec: partitions=1, partition_sizes=[1]
statement ok
SET datafusion.explain.physical_plan_only = false;
@@ -639,10 +611,8 @@ QUALIFY rn <= 3;
----
physical_plan
01)ProjectionExec: expr=[id@0 as id, pk@1 as pk, val@2 as val, row_number() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rn]
-02)--BoundedWindowAggExec: wdw=[row_number() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
-03)----RepartitionExec: partitioning=Hash([pk@1], 4), input_partitions=1, maintains_sort_order=true
-04)------PartitionedTopKExec: fn=row_number, fetch=3, partition=[pk@1], order=[val@2 ASC NULLS LAST]
-05)--------DataSourceExec: partitions=1, partition_sizes=[1]
+02)--PartitionedTopKExec: fn=row_number, fetch=3, partition=[pk@1], order=[val@2 ASC NULLS LAST], emit=[row_number() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]
+03)----DataSourceExec: partitions=1, partition_sizes=[1]
# Test 30: QUALIFY with < operator
statement ok
@@ -669,10 +639,8 @@ QUALIFY rnk <= 3;
----
physical_plan
01)ProjectionExec: expr=[id@0 as id, pk@1 as pk, val@2 as val, rank() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rnk]
-02)--BoundedWindowAggExec: wdw=[rank() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "rank() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
-03)----RepartitionExec: partitioning=Hash([pk@1], 4), input_partitions=1, maintains_sort_order=true
-04)------PartitionedTopKExec: fn=rank, fetch=3, partition=[pk@1], order=[val@2 ASC NULLS LAST]
-05)--------DataSourceExec: partitions=1, partition_sizes=[1]
+02)--PartitionedTopKExec: fn=rank, fetch=3, partition=[pk@1], order=[val@2 ASC NULLS LAST], emit=[rank() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]
+03)----DataSourceExec: partitions=1, partition_sizes=[1]
statement ok
SET datafusion.explain.physical_plan_only = false;
@@ -747,10 +715,8 @@ EXPLAIN SELECT * FROM (
----
physical_plan
01)ProjectionExec: expr=[id@0 as id, pk@1 as pk, val@2 as val, row_number() PARTITION BY [window_topn_nulls.pk] ORDER BY [window_topn_nulls.val ASC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rn]
-02)--BoundedWindowAggExec: wdw=[row_number() PARTITION BY [window_topn_nulls.pk] ORDER BY [window_topn_nulls.val ASC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [window_topn_nulls.pk] ORDER BY [window_topn_nulls.val ASC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
-03)----RepartitionExec: partitioning=Hash([pk@1], 4), input_partitions=1, maintains_sort_order=true
-04)------PartitionedTopKExec: fn=row_number, fetch=2, partition=[pk@1], order=[val@2 ASC]
-05)--------DataSourceExec: partitions=1, partition_sizes=[1]
+02)--PartitionedTopKExec: fn=row_number, fetch=2, partition=[pk@1], order=[val@2 ASC], emit=[row_number() PARTITION BY [window_topn_nulls.pk] ORDER BY [window_topn_nulls.val ASC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]
+03)----DataSourceExec: partitions=1, partition_sizes=[1]
query TT
EXPLAIN SELECT * FROM (
@@ -759,10 +725,8 @@ EXPLAIN SELECT * FROM (
----
physical_plan
01)ProjectionExec: expr=[id@0 as id, pk@1 as pk, val@2 as val, row_number() PARTITION BY [window_topn_nulls.pk] ORDER BY [window_topn_nulls.val DESC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rn]
-02)--BoundedWindowAggExec: wdw=[row_number() PARTITION BY [window_topn_nulls.pk] ORDER BY [window_topn_nulls.val DESC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [window_topn_nulls.pk] ORDER BY [window_topn_nulls.val DESC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
-03)----RepartitionExec: partitioning=Hash([pk@1], 4), input_partitions=1, maintains_sort_order=true
-04)------PartitionedTopKExec: fn=row_number, fetch=2, partition=[pk@1], order=[val@2 DESC NULLS LAST]
-05)--------DataSourceExec: partitions=1, partition_sizes=[1]
+02)--PartitionedTopKExec: fn=row_number, fetch=2, partition=[pk@1], order=[val@2 DESC NULLS LAST], emit=[row_number() PARTITION BY [window_topn_nulls.pk] ORDER BY [window_topn_nulls.val DESC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]
+03)----DataSourceExec: partitions=1, partition_sizes=[1]
statement ok
SET datafusion.explain.physical_plan_only = false;
@@ -842,10 +806,8 @@ logical_plan
04)------TableScan: window_topn_rank_t projection=[id, pk, val]
physical_plan
01)ProjectionExec: expr=[id@0 as id, pk@1 as pk, val@2 as val, rank() PARTITION BY [window_topn_rank_t.pk] ORDER BY [window_topn_rank_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rk]
-02)--BoundedWindowAggExec: wdw=[rank() PARTITION BY [window_topn_rank_t.pk] ORDER BY [window_topn_rank_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "rank() PARTITION BY [window_topn_rank_t.pk] ORDER BY [window_topn_rank_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
-03)----RepartitionExec: partitioning=Hash([pk@1], 4), input_partitions=1, maintains_sort_order=true
-04)------PartitionedTopKExec: fn=rank, fetch=3, partition=[pk@1], order=[val@2 ASC NULLS LAST]
-05)--------DataSourceExec: partitions=1, partition_sizes=[1]
+02)--PartitionedTopKExec: fn=rank, fetch=3, partition=[pk@1], order=[val@2 ASC NULLS LAST], emit=[rank() PARTITION BY [window_topn_rank_t.pk] ORDER BY [window_topn_rank_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]
+03)----DataSourceExec: partitions=1, partition_sizes=[1]
# Test R3: rk < 4 should give the same results (fetch = K-1 = 3)
query III rowsort
@@ -950,10 +912,8 @@ logical_plan
04)------TableScan: window_topn_rank_t projection=[id, pk, val]
physical_plan
01)ProjectionExec: expr=[id@0 as id, pk@1 as pk, val@2 as val, rank() PARTITION BY [window_topn_rank_t.pk, window_topn_rank_t.id] ORDER BY [window_topn_rank_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rk]
-02)--BoundedWindowAggExec: wdw=[rank() PARTITION BY [window_topn_rank_t.pk, window_topn_rank_t.id] ORDER BY [window_topn_rank_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "rank() PARTITION BY [window_topn_rank_t.pk, window_topn_rank_t.id] ORDER BY [window_topn_rank_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
-03)----RepartitionExec: partitioning=Hash([pk@1, id@0], 4), input_partitions=1, maintains_sort_order=true
-04)------PartitionedTopKExec: fn=rank, fetch=1, partition=[pk@1, id@0], order=[val@2 ASC NULLS LAST]
-05)--------DataSourceExec: partitions=1, partition_sizes=[1]
+02)--PartitionedTopKExec: fn=rank, fetch=1, partition=[pk@1, id@0], order=[val@2 ASC NULLS LAST], emit=[rank() PARTITION BY [window_topn_rank_t.pk, window_topn_rank_t.id] ORDER BY [window_topn_rank_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]
+03)----DataSourceExec: partitions=1, partition_sizes=[1]
# Test R9: RANK with DESC ordering
query III rowsort
@@ -1091,10 +1051,8 @@ logical_plan
04)------TableScan: window_topn_rank_dense_t projection=[id, pk, val]
physical_plan
01)ProjectionExec: expr=[id@0 as id, pk@1 as pk, val@2 as val, rank() PARTITION BY [window_topn_rank_dense_t.pk] ORDER BY [window_topn_rank_dense_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rk]
-02)--BoundedWindowAggExec: wdw=[rank() PARTITION BY [window_topn_rank_dense_t.pk] ORDER BY [window_topn_rank_dense_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "rank() PARTITION BY [window_topn_rank_dense_t.pk] ORDER BY [window_topn_rank_dense_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
-03)----RepartitionExec: partitioning=Hash([pk@1], 4), input_partitions=1, maintains_sort_order=true
-04)------PartitionedTopKExec: fn=rank, fetch=3, partition=[pk@1], order=[val@2 ASC NULLS LAST]
-05)--------DataSourceExec: partitions=1, partition_sizes=[1]
+02)--PartitionedTopKExec: fn=rank, fetch=3, partition=[pk@1], order=[val@2 ASC NULLS LAST], emit=[rank() PARTITION BY [window_topn_rank_dense_t.pk] ORDER BY [window_topn_rank_dense_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]
+03)----DataSourceExec: partitions=1, partition_sizes=[1]
statement ok
DROP TABLE window_topn_rank_dense_t;
@@ -1160,10 +1118,8 @@ logical_plan
04)------TableScan: window_topn_rank_null_t projection=[id, pk, val]
physical_plan
01)ProjectionExec: expr=[id@0 as id, pk@1 as pk, val@2 as val, rank() PARTITION BY [window_topn_rank_null_t.pk] ORDER BY [window_topn_rank_null_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rk]
-02)--BoundedWindowAggExec: wdw=[rank() PARTITION BY [window_topn_rank_null_t.pk] ORDER BY [window_topn_rank_null_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "rank() PARTITION BY [window_topn_rank_null_t.pk] ORDER BY [window_topn_rank_null_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
-03)----RepartitionExec: partitioning=Hash([pk@1], 4), input_partitions=1, maintains_sort_order=true
-04)------PartitionedTopKExec: fn=rank, fetch=2, partition=[pk@1], order=[val@2 ASC NULLS LAST]
-05)--------DataSourceExec: partitions=1, partition_sizes=[1]
+02)--PartitionedTopKExec: fn=rank, fetch=2, partition=[pk@1], order=[val@2 ASC NULLS LAST], emit=[rank() PARTITION BY [window_topn_rank_null_t.pk] ORDER BY [window_topn_rank_null_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]
+03)----DataSourceExec: partitions=1, partition_sizes=[1]
# Test R19: DESC NULLS LAST — pk=1: 3,2,1,NULL ranks 1,2,3,4; pk=2: 5,NULL,NULL ranks 1,2,2.
query III rowsort
@@ -1272,10 +1228,8 @@ logical_plan
04)------TableScan: window_topn_dense_rank_t projection=[id, pk, val]
physical_plan
01)ProjectionExec: expr=[id@0 as id, pk@1 as pk, val@2 as val, dense_rank() PARTITION BY [window_topn_dense_rank_t.pk] ORDER BY [window_topn_dense_rank_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as dr]
-02)--BoundedWindowAggExec: wdw=[dense_rank() PARTITION BY [window_topn_dense_rank_t.pk] ORDER BY [window_topn_dense_rank_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "dense_rank() PARTITION BY [window_topn_dense_rank_t.pk] ORDER BY [window_topn_dense_rank_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
-03)----RepartitionExec: partitioning=Hash([pk@1], 4), input_partitions=1, maintains_sort_order=true
-04)------PartitionedTopKExec: fn=dense_rank, fetch=2, partition=[pk@1], order=[val@2 ASC NULLS LAST]
-05)--------DataSourceExec: partitions=1, partition_sizes=[1]
+02)--PartitionedTopKExec: fn=dense_rank, fetch=2, partition=[pk@1], order=[val@2 ASC NULLS LAST], emit=[dense_rank() PARTITION BY [window_topn_dense_rank_t.pk] ORDER BY [window_topn_dense_rank_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]
+03)----DataSourceExec: partitions=1, partition_sizes=[1]
# Test DR3: dr < 3 should give the same results (fetch = K-1 = 2)
query III rowsort
@@ -1523,10 +1477,8 @@ logical_plan
04)------TableScan: window_topn_dr_null_t projection=[id, pk, val]
physical_plan
01)ProjectionExec: expr=[id@0 as id, pk@1 as pk, val@2 as val]
-02)--BoundedWindowAggExec: wdw=[dense_rank() PARTITION BY [window_topn_dr_null_t.pk] ORDER BY [window_topn_dr_null_t.val ASC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "dense_rank() PARTITION BY [window_topn_dr_null_t.pk] ORDER BY [window_topn_dr_null_t.val ASC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
-03)----RepartitionExec: partitioning=Hash([pk@1], 4), input_partitions=1, maintains_sort_order=true
-04)------PartitionedTopKExec: fn=dense_rank, fetch=2, partition=[pk@1], order=[val@2 ASC]
-05)--------DataSourceExec: partitions=1, partition_sizes=[1]
+02)--PartitionedTopKExec: fn=dense_rank, fetch=2, partition=[pk@1], order=[val@2 ASC], emit=[dense_rank() PARTITION BY [window_topn_dr_null_t.pk] ORDER BY [window_topn_dr_null_t.val ASC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]
+03)----DataSourceExec: partitions=1, partition_sizes=[1]
# Test DR14: NULLS LAST — NULL is the largest distinct ob (dr 4). 10 (dr
# 1) and 20 (dr 2) are the two smallest, so `dr <= 2` excludes the NULLs.
@@ -1574,12 +1526,9 @@ logical_plan
04)------TableScan: t projection=[c1, c2]
physical_plan
01)ProjectionExec: expr=[c1@0 as c1, c2@1 as c2, row_number() PARTITION BY [t.c1] ORDER BY [t.c2 DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@2 as rn]
-02)--BoundedWindowAggExec: wdw=[row_number() PARTITION BY [t.c1] ORDER BY [t.c2 DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [t.c1] ORDER BY [t.c2 DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted]
-03)----SortPreservingMergeExec: [c1@0 ASC NULLS LAST, c2@1 DESC]
-04)------SortExec: expr=[c1@0 ASC NULLS LAST, c2@1 DESC], preserve_partitioning=[true]
-05)--------PartitionedTopKExec: fn=row_number, fetch=1, partition=[c1@0], order=[c2@1 DESC]
-06)----------RepartitionExec: partitioning=Hash([c1@0], 5), input_partitions=1
-07)------------DataSourceExec: partitions=1, partition_sizes=[1]
+02)--PartitionedTopKExec: fn=row_number, fetch=1, partition=[c1@0], order=[c2@1 DESC], emit=[row_number() PARTITION BY [t.c1] ORDER BY [t.c2 DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]
+03)----RepartitionExec: partitioning=Hash([c1@0], 5), input_partitions=1
+04)------DataSourceExec: partitions=1, partition_sizes=[1]
statement ok
set datafusion.execution.target_partitions = 4;
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org