diff --git a/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q17.benchmark b/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q17.benchmark new file mode 100644 index 0000000000000..6708ca4ccacc0 --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q17.benchmark @@ -0,0 +1,16 @@ +# RANK top-2, ~100K partitions, with a sibling ROW_NUMBER expression. +# Project both ranks so the window stays above TopK and TopK's derived ranks +# are unused. Compare with q08's single-expression workload. + +subgroup window_sorted + +template sql_benchmarks/h2o/window_sorted.benchmark.template +QPAD=17 +FN=RANK +PK_MOD=100000 +OB_EXPR=v2 +OB_BOUND=count(DISTINCT ob) > 1000 +NAME=window_sorted_q17_rank_100k_partitions_two_windows +TWO_WINDOWS=true + +expect_plan BoundedWindowAggExec diff --git a/benchmarks/sql_benchmarks/h2o/window_sorted.benchmark.template b/benchmarks/sql_benchmarks/h2o/window_sorted.benchmark.template index 9170511492ea5..8091724148408 100644 --- a/benchmarks/sql_benchmarks/h2o/window_sorted.benchmark.template +++ b/benchmarks/sql_benchmarks/h2o/window_sorted.benchmark.template @@ -8,6 +8,7 @@ # OB_EXPR ORDER BY value expression (`v2` distinct, integer-cast for ties) # OB_BOUND SQL predicate pinning ob's cardinality, asserted below # NAME display name +# TWO_WINDOWS append and project ROW_NUMBER alongside FN (default false) # # Why this subgroup exists: the `window` subgroup registers `x` with no declared # ordering, so `output_ordering()` is None and the planner cannot know the input @@ -88,9 +89,9 @@ expect_plan output_ordering=[pk@0 ASC NULLS LAST, ob@1 DESC] run -- Top-2 per partition by ${FN} over the declared-sorted input. -SELECT pk, ob FROM ( +SELECT pk, ob ${TWO_WINDOWS:-false|, rnk, rn| } FROM ( SELECT pk, ob, - ${FN}() OVER (PARTITION BY pk ORDER BY ob DESC) AS rnk + ${FN}() OVER (PARTITION BY pk ORDER BY ob DESC) AS rnk ${TWO_WINDOWS:-false|, ROW_NUMBER() OVER (PARTITION BY pk ORDER BY ob DESC) AS rn| } FROM x_sorted ) sub_query WHERE rnk <= 2; diff --git a/datafusion/core/tests/physical_optimizer/window_topn.rs b/datafusion/core/tests/physical_optimizer/window_topn.rs index 85801e84c9189..950b3bfa2fc15 100644 --- a/datafusion/core/tests/physical_optimizer/window_topn.rs +++ b/datafusion/core/tests/physical_optimizer/window_topn.rs @@ -222,9 +222,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 +233,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 +284,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 +392,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 = 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> { + 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 = 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 = + 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 +454,6 @@ fn with_projection_between() -> Result<()> { // ---------------------------------------------------------------------- // RANK rule tests // ---------------------------------------------------------------------- - /// Build: FilterExec(rk op limit) → BoundedWindowAggExec( PBY pk OBY val) /// /// Matches the pre-`EnsureRequirements` plan shape (no `SortExec` under the window). @@ -508,9 +556,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 +567,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 +578,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 +589,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 +620,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 +631,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 +642,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 +653,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 +767,199 @@ 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 = Vec::new(); + for batch in &actual { + let col = batch + .column(2) + .as_any() + .downcast_ref::() + .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 , ` 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(()) +} diff --git a/datafusion/physical-optimizer/src/optimizer.rs b/datafusion/physical-optimizer/src/optimizer.rs index 80c33e1c08b6a..9fdaca6320b54 100644 --- a/datafusion/physical-optimizer/src/optimizer.rs +++ b/datafusion/physical-optimizer/src/optimizer.rs @@ -111,8 +111,10 @@ impl PhysicalOptimizer { // those are handled by the later `FilterPushdown` rule. // See `FilterPushdownPhase` for more details. Arc::new(FilterPushdown::new()), - // WindowTopN: replaces Filter(rn<=K) → Window(ROW_NUMBER) - // with Window(ROW_NUMBER) → PartitionedTopKExec(fetch=K). + // WindowTopN: replaces Filter(rn<=K) → Window(ROW_NUMBER) with + // PartitionedTopKExec(fetch=K), which emits `rn` itself; the + // window node is kept above it only when it has other window + // expressions to compute. // Must run before EnsureRequirements (so it can rewrite against the // window's declared ordering without pattern-matching a SortExec) // and before ProjectionPushdown (which embeds projections into FilterExec). diff --git a/datafusion/physical-optimizer/src/window_topn.rs b/datafusion/physical-optimizer/src/window_topn.rs index 758ad66de3017..28a6719a79070 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 @@ -85,15 +87,26 @@ use datafusion_physical_plan::windows::{BoundedWindowAggExec, WindowUDFExpr}; /// /// ```text /// [optional ProjectionExec] -/// BoundedWindowAggExec( PARTITION BY ... ORDER BY ...) +/// PartitionedTopKExec(fn=, partition_keys, order_keys, fetch=K, emit=[ 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( PARTITION BY ... ORDER BY ...) /// PartitionedTopKExec(fn=, 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 +246,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 +268,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 = 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() { diff --git a/datafusion/physical-plan/src/sorts/partitioned_topk.rs b/datafusion/physical-plan/src/sorts/partitioned_topk.rs index 04072fffaf8eb..4eead3e1a4ee6 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, /// Execution metrics metrics_set: ExecutionPlanMetricsSet, /// Cached plan properties (output ordering, partitioning, etc.) @@ -245,18 +261,40 @@ impl PartitionedTopKExec { fetch: usize, fn_kind: WindowFnKind, ) -> Result { - 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 { + 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 { &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 , ` stay + /// sort-free after the window node is removed. fn compute_properties( input: &Arc, sort_exprs: LexOrdering, + partition_prefix_len: usize, + ranking_field: Option<&FieldRef>, ) -> Result { - 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; + 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 = self.expr[..self.partition_prefix_len] + .iter() + .map(|e| format!("{}", e.expr)) + .collect(); + let order_exprs: Vec = self.expr[self.partition_prefix_len..] + .iter() + .map(|e| format!("{e}")) + .collect(); match t { DisplayFormatType::Default | DisplayFormatType::Verbose => { - let partition_exprs: Vec = self.expr[..self.partition_prefix_len] - .iter() - .map(|e| format!("{}", e.expr)) - .collect(); - let order_exprs: Vec = 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 = self.expr[..self.partition_prefix_len] - .iter() - .map(|e| format!("{}", e.expr)) - .collect(); - let order_exprs: Vec = 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> { 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, metrics_set: ExecutionPlanMetricsSet, + ranked_schema: Option, ) -> Result { 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, Arc)> { + let schema = Arc::new(Schema::new(vec![ + Field::new("pk", DataType::Int64, false), + Field::new("val", DataType::Int64, false), + ])); + let input: Arc = + Arc::new(PlaceholderRowExec::new(Arc::clone(&schema))); + Ok((schema, input)) + } + + fn pk_then_val(schema: &Arc) -> Result { + 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; + let rn = Arc::new(Column::new("rn", 2)) as Arc; + 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 = 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::() + .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; + let rn = Arc::new(Column::new("rn", 2)) as Arc; + 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 653875b8d356f..819bb713f7c1e 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, }; @@ -1714,6 +1714,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 @@ -1726,7 +1747,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 { - 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). @@ -1801,11 +1829,13 @@ struct PartitionedTopKCore { impl PartitionedTopKCore { /// `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>, partition_sort_fields: Vec, order_expr: LexOrdering, @@ -1813,6 +1843,7 @@ impl PartitionedTopKCore { batch_size: usize, runtime: &Arc, metrics: &ExecutionPlanMetricsSet, + ranked_schema: Option, ) -> Result { assert!(k > 0, "{name} requires k > 0"); let reservation = MemoryConsumer::new(format!("{name}[{partition_id}]")) @@ -1821,7 +1852,7 @@ impl PartitionedTopKCore { // 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); @@ -1831,7 +1862,8 @@ impl PartitionedTopKCore { .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, @@ -1908,15 +1940,20 @@ impl PartitionedTopKCore { /// 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>( + fn emit_with>( self, drain: impl FnMut(S) -> I, ) -> Result { let Self { - schema, + output_schema, + emit_ranks, metrics, reservation, partitions, @@ -1938,7 +1975,8 @@ impl PartitionedTopKCore { .map(|(_key, state)| state) .flat_map(drain); EmitState::stream( - schema, + output_schema, + emit_ranks, metrics, reservation, batch_size, @@ -1987,7 +2025,7 @@ impl PartitionedTopK { #[expect(clippy::too_many_arguments)] pub(crate) fn try_new( partition_id: usize, - schema: SchemaRef, + schema: &SchemaRef, partition_exprs: Vec>, partition_sort_fields: Vec, order_expr: LexOrdering, @@ -1995,6 +2033,7 @@ impl PartitionedTopK { batch_size: usize, runtime: &Arc, metrics: &ExecutionPlanMetricsSet, + ranked_schema: Option, ) -> Result { Ok(Self { core: PartitionedTopKCore::try_new( @@ -2008,6 +2047,7 @@ impl PartitionedTopK { batch_size, runtime, metrics, + ranked_schema, )?, }) } @@ -2121,8 +2161,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 { - 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 @@ -2136,8 +2183,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. @@ -2150,6 +2200,8 @@ struct EmitState { batches: Vec, /// `(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>, /// Scratch for [`interleave_referenced`], one entry per batch. remap: Vec, pos: usize, @@ -2158,14 +2210,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, + slots: impl Iterator, len: usize, ) -> Result { let timer = metrics.baseline.elapsed_compute().timer(); @@ -2176,17 +2232,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> = 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 @@ -2198,21 +2260,26 @@ impl EmitState { + store.batches_size + batches.capacity() * size_of::() + ordered.capacity() * size_of::<(usize, usize)>() + + ranks + .as_ref() + .map_or(0, |r| r.capacity() * size_of::()) + batches.len() * size_of::(), )?; 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))) }), @@ -2227,16 +2294,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>, + output_schema: &SchemaRef, + metrics: &TopKMetrics, +) -> Result { + 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 @@ -2319,7 +2418,7 @@ impl PartitionedTopKRank { #[expect(clippy::too_many_arguments)] pub(crate) fn try_new( partition_id: usize, - schema: SchemaRef, + schema: &SchemaRef, partition_exprs: Vec>, partition_sort_fields: Vec, order_expr: LexOrdering, @@ -2327,6 +2426,7 @@ impl PartitionedTopKRank { batch_size: usize, runtime: &Arc, metrics: &ExecutionPlanMetricsSet, + ranked_schema: Option, ) -> Result { Ok(Self { core: PartitionedTopKCore::try_new( @@ -2340,6 +2440,7 @@ impl PartitionedTopKRank { batch_size, runtime, metrics, + ranked_schema, )?, evicted_key: Vec::new(), ties_bytes: 0, @@ -2476,8 +2577,31 @@ impl PartitionedTopKRank { /// (all sharing the boundary ob). pub(crate) fn emit(self) -> Result { 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> = 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))) }) } @@ -2633,7 +2757,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). @@ -2670,15 +2793,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>, partition_sort_fields: Vec, order_expr: LexOrdering, @@ -2686,13 +2817,14 @@ impl PartitionedTopKDenseRank { batch_size: usize, runtime: &Arc, metrics: &ExecutionPlanMetricsSet, + ranked_schema: Option, ) -> Result { 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); @@ -2702,7 +2834,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, @@ -2912,7 +3045,8 @@ impl PartitionedTopKDenseRank { /// emitted rows are in ob-sorted order. pub(crate) fn emit(self) -> Result { let Self { - schema, + output_schema, + emit_ranks, metrics, reservation: _, expr: _, @@ -2959,6 +3093,9 @@ impl PartitionedTopKDenseRank { let mut batch_refs: Vec<&RecordBatch> = Vec::new(); let mut batch_id_pos: HashMap = HashMap::new(); let mut indices: Vec<(usize, usize)> = Vec::with_capacity(batch_size); + // Parallel to `indices`: each row's `DENSE_RANK()`, when emitted. + let mut ranks: Option> = + 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`. @@ -2973,7 +3110,11 @@ impl PartitionedTopKDenseRank { let mut sorted_obs: Vec<(Vec, Vec)> = groups.into_iter().collect(); sorted_obs.sort_by(|a, b| a.0.cmp(&b.0)); - for (_ob, entries) in sorted_obs { + // `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) { for entry in entries { let batch = &store .get(entry.batch_id) @@ -2990,10 +3131,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(); @@ -3005,12 +3154,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), ))) } @@ -3054,7 +3203,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; @@ -3944,11 +4093,13 @@ mod tests { val_sort_options: SortOptions, val_nullable: bool, ) -> Result<(Arc, PartitionedTopK)> { - build_partitioned_topk_inner( + build_pk_val_state( + PartitionedTopK::try_new, k, val_sort_options, val_nullable, &Arc::new(RuntimeEnv::default()), + None, ) } @@ -3958,15 +4109,43 @@ mod tests { k: usize, runtime: &Arc, ) -> Result<(Arc, 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 = fn( + usize, + &SchemaRef, + Vec>, + Vec, + LexOrdering, + usize, + usize, + &Arc, + &ExecutionPlanMetricsSet, + Option, + ) -> Result; + + /// 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( + ctor: StateCtor, k: usize, val_sort_options: SortOptions, val_nullable: bool, runtime: &Arc, - ) -> Result<(Arc, PartitionedTopK)> { + ranking_field: Option, + ) -> Result<(Arc, T)> { let schema = pk_val_schema(val_nullable); let pk_expr: Arc = col("pk", schema.as_ref())?; @@ -3979,23 +4158,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( + ctor: StateCtor, + k: usize, + ranking_field: FieldRef, + ) -> Result<(Arc, 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 @@ -4021,7 +4217,7 @@ mod tests { let state = PartitionedTopK::try_new( 0, - Arc::clone(&schema), + &schema, vec![pk_expr], build_sort_fields(&partition_ordering, &schema)?, order_expr, @@ -4029,6 +4225,7 @@ mod tests { 8, // batch_size &Arc::new(RuntimeEnv::default()), &ExecutionPlanMetricsSet::new(), + None, )?; Ok((schema, state)) } @@ -4202,7 +4399,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, @@ -4210,6 +4407,7 @@ mod tests { 8, &Arc::new(RuntimeEnv::default()), &ExecutionPlanMetricsSet::new(), + None, )?; let batch = RecordBatch::try_new( @@ -5263,7 +5461,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, @@ -5271,6 +5469,7 @@ mod tests { 8, &Arc::new(RuntimeEnv::default()), &ExecutionPlanMetricsSet::new(), + None, )?; let empty = state.size(); @@ -5451,7 +5650,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()), @@ -5459,6 +5658,7 @@ mod tests { 3, // batch_size &Arc::new(RuntimeEnv::default()), &metrics, + None, )?; state.insert_batch(&pk_val_batch( @@ -5475,6 +5675,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::(); + let ranks = batch.column(2).as_primitive::(); + 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. @@ -5737,11 +6106,13 @@ mod tests { val_sort_options: SortOptions, val_nullable: bool, ) -> Result<(Arc, PartitionedTopKRank)> { - build_partitioned_topk_rank_inner( + build_pk_val_state( + PartitionedTopKRank::try_new, k, val_sort_options, val_nullable, &Arc::new(RuntimeEnv::default()), + None, ) } @@ -5752,42 +6123,14 @@ mod tests { k: usize, runtime: &Arc, ) -> Result<(Arc, 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, - ) -> Result<(Arc, PartitionedTopKRank)> { - let schema = pk_val_schema(val_nullable); - - let pk_expr: Arc = 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 @@ -6059,6 +6402,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. @@ -6077,7 +6496,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]), @@ -6085,6 +6504,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. @@ -6210,7 +6630,7 @@ mod tests { }]); PartitionedTopKRank::try_new( 0, - Arc::clone(schema), + schema, vec![pk_expr], partition_sort_fields, order_expr, @@ -6218,6 +6638,7 @@ mod tests { batch_size, &Arc::new(RuntimeEnv::default()), &ExecutionPlanMetricsSet::new(), + None, ) } @@ -6646,33 +7067,14 @@ mod tests { val_sort_options: SortOptions, val_nullable: bool, ) -> Result<(Arc, PartitionedTopKDenseRank)> { - let schema = pk_val_schema(val_nullable); - - let pk_expr: Arc = 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 @@ -6708,6 +7110,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. @@ -6863,7 +7334,7 @@ mod tests { let mut state = PartitionedTopKDenseRank::try_new( 0, - Arc::clone(&schema), + &schema, vec![pk_expr], partition_sort_fields, order_expr, @@ -6871,6 +7342,7 @@ mod tests { 8, // batch_size &Arc::new(RuntimeEnv::default()), &ExecutionPlanMetricsSet::new(), + None, )?; let dict_batch = |vals: &[&str]| -> Result { @@ -7340,7 +7812,7 @@ mod tests { }]); let mut state = PartitionedTopKDenseRank::try_new( 0, - Arc::clone(&schema), + &schema, vec![pk_expr], partition_sort_fields, order_expr, @@ -7348,6 +7820,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 463cd1799e857..9ea1e184bee91 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=, 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=, 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=, 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=, 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], mode=[Sorted] -03)----PartitionedTopKExec: fn=row_number, fetch=1, partition=[range_key@0, non_range_key@1], order=[value@2 DESC] -04)------DataSourceExec: file_groups=, 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=, 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], mode=[Sorted] -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=, 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=, 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 5c3757aa5f770..637b6029c7e9e 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;