From d0c03145141a69799717f1fda0ce0f2e0048bd09 Mon Sep 17 00:00:00 2001 From: SubhamSinghal Date: Thu, 16 Jul 2026 19:17:14 +0530 Subject: [PATCH] perf: Extend WindowTopN to dense rank --- benchmarks/queries/h2o/window.sql | 50 ++ .../tests/physical_optimizer/window_topn.rs | 103 ++- .../physical-optimizer/src/window_topn.rs | 48 +- .../src/sorts/partitioned_topk.rs | 68 +- datafusion/physical-plan/src/topk/mod.rs | 849 +++++++++++++++++- .../sqllogictest/test_files/window_topn.slt | 316 +++++++ 6 files changed, 1387 insertions(+), 47 deletions(-) diff --git a/benchmarks/queries/h2o/window.sql b/benchmarks/queries/h2o/window.sql index 37df0a28ae614..b48a2aaabb67a 100644 --- a/benchmarks/queries/h2o/window.sql +++ b/benchmarks/queries/h2o/window.sql @@ -196,3 +196,53 @@ SELECT pk, largest_v2 FROM ( RANK() OVER (PARTITION BY (id3 % 100000) ORDER BY v2 DESC) AS rk_v2 FROM large WHERE v2 IS NOT NULL ) sub_query WHERE rk_v2 <= 2; + +-- Window Top-N (DENSE_RANK top-2 per partition, ~100 partitions) +-- The DENSE_RANK queries below mirror the RANK cardinality sweep above. +-- DENSE_RANK semantics keep every row whose ORDER BY value is among the +-- K distinct-smallest values in the partition, so total kept per partition +-- is unbounded in rows-per-distinct-value — exercises PartitionedTopKDenseRank's +-- HashMap-of-groups path. +SELECT pk, largest_v2 FROM ( + SELECT (id3 % 100) AS pk, v2 AS largest_v2, + DENSE_RANK() OVER (PARTITION BY (id3 % 100) ORDER BY v2 DESC) AS dr_v2 + FROM large WHERE v2 IS NOT NULL +) sub_query WHERE dr_v2 <= 2; + +-- Window Top-N (DENSE_RANK top-2 per partition, ~1K partitions) +SELECT pkey, largest_v2 FROM ( + SELECT (id3 % 1000) AS pkey, v2 AS largest_v2, + DENSE_RANK() OVER (PARTITION BY (id3 % 1000) ORDER BY v2 DESC) AS dr_v2 + FROM large WHERE v2 IS NOT NULL +) sub_query WHERE dr_v2 <= 2; + +-- Window Top-N (DENSE_RANK top-2 per partition, ~1K partitions, heavy ties) +-- v2 % 10 forces 10 distinct OBY values; most rows share the top-2 distinct +-- values so appends dominate — exercises the "Case A" append-to-existing-Vec +-- fast path in PartitionedTopKDenseRank. +SELECT pkey, largest_v2 FROM ( + SELECT (id3 % 1000) AS pkey, v2 AS largest_v2, + DENSE_RANK() OVER (PARTITION BY (id3 % 1000) ORDER BY (v2 % 10) DESC) AS dr_v2 + FROM large WHERE v2 IS NOT NULL +) sub_query WHERE dr_v2 <= 2; + +-- Window Top-N (DENSE_RANK top-2 per partition, ~10K partitions, low ties) +SELECT id2, largest_v2 FROM ( + SELECT id2, v2 AS largest_v2, + DENSE_RANK() OVER (PARTITION BY id2 ORDER BY v2 DESC) AS dr_v2 + FROM large WHERE v2 IS NOT NULL +) sub_query WHERE dr_v2 <= 2; + +-- Window Top-N (DENSE_RANK top-2 per partition, ~10K partitions, heavy ties) +SELECT id2, largest_v2 FROM ( + SELECT id2, v2 AS largest_v2, + DENSE_RANK() OVER (PARTITION BY id2 ORDER BY (v2 % 10) DESC) AS dr_v2 + FROM large WHERE v2 IS NOT NULL +) sub_query WHERE dr_v2 <= 2; + +-- Window Top-N (DENSE_RANK top-2 per partition, ~100K partitions) +SELECT pk, largest_v2 FROM ( + SELECT (id3 % 100000) AS pk, v2 AS largest_v2, + DENSE_RANK() OVER (PARTITION BY (id3 % 100000) ORDER BY v2 DESC) AS dr_v2 + FROM large WHERE v2 IS NOT NULL +) sub_query WHERE dr_v2 <= 2; diff --git a/datafusion/core/tests/physical_optimizer/window_topn.rs b/datafusion/core/tests/physical_optimizer/window_topn.rs index 07a1db127ec54..d1f052cbcf6de 100644 --- a/datafusion/core/tests/physical_optimizer/window_topn.rs +++ b/datafusion/core/tests/physical_optimizer/window_topn.rs @@ -488,8 +488,12 @@ fn build_ranking_topn_plan( Ok(filter) } -/// Build a RANK plan with NO ORDER BY: every row ties at rank 1 — degenerate. -fn build_rank_no_order_by_plan(limit_value: i64) -> Result> { +/// Build a RANK / DENSE_RANK plan with NO ORDER BY: every row ties at rank 1 — degenerate. +fn build_no_order_by_plan( + udwf_factory: fn() -> Arc, + udwf_name: &str, + limit_value: i64, +) -> Result> { let s = schema(); let input: Arc = Arc::new(PlaceholderRowExec::new(Arc::clone(&s))); @@ -503,7 +507,7 @@ fn build_rank_no_order_by_plan(limit_value: i64) -> Result Result = @@ -582,7 +586,7 @@ fn rank_no_order_by_no_change() -> Result<()> { // Without ORDER BY, every row ties at rank 1 — the optimization is // degenerate (entire input would be retained, ties storage unbounded). // The rule must skip. - let plan = build_rank_no_order_by_plan(3)?; + let plan = build_no_order_by_plan(rank_udwf, "rank", 3)?; let before = plan_str(plan.as_ref()); let optimized = optimize(plan)?; let after = plan_str(optimized.as_ref()); @@ -593,17 +597,98 @@ fn rank_no_order_by_no_change() -> Result<()> { Ok(()) } +// ---------------------------------------------------------------------- +// DENSE_RANK rule tests +// ---------------------------------------------------------------------- + #[test] -fn dense_rank_no_change() -> Result<()> { - // DENSE_RANK is not yet supported by the rule. The plan must pass - // through unchanged. +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 + "#); + Ok(()) +} + +#[test] +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 + "#); + Ok(()) +} + +#[test] +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 + "#); + Ok(()) +} + +#[test] +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 + "#); + Ok(()) +} + +#[test] +fn dense_rank_no_order_by_no_change() -> Result<()> { + // Without ORDER BY, every row ties at dense_rank 1 — the optimization + // is degenerate (entire input would be retained). The rule must skip. + let plan = build_no_order_by_plan(dense_rank_udwf, "dense_rank", 3)?; let before = plan_str(plan.as_ref()); let optimized = optimize(plan)?; let after = plan_str(optimized.as_ref()); assert_eq!( before, after, - "DENSE_RANK is unsupported and must not be rewritten" + "DENSE_RANK with empty ORDER BY must not be rewritten" ); Ok(()) } + +// ---------------------------------------------------------------------- +// Shared guard: `fn < 1` keeps nothing +// ---------------------------------------------------------------------- + +#[test] +fn predicate_lt_1_no_change() -> Result<()> { + // `fn < 1` (and the flipped `1 > fn`) yields limit_n = 0. Since + // ROW_NUMBER / RANK / DENSE_RANK are always >= 1, the predicate keeps + // nothing and the rule must skip — a fetch=0 PartitionedTopK* would + // otherwise panic on its `k > 0` assertion at execution time. + type UdwfFactory = fn() -> Arc; + let cases: [(UdwfFactory, &str); 3] = [ + (row_number_udwf, "row_number"), + (rank_udwf, "rank"), + (dense_rank_udwf, "dense_rank"), + ]; + for (factory, name) in cases { + let plan = build_ranking_topn_plan(factory, name, 1, Operator::Lt)?; + let before = plan_str(plan.as_ref()); + let optimized = optimize(plan)?; + let after = plan_str(optimized.as_ref()); + assert_eq!( + before, after, + "`{name} < 1` (limit 0) must not be rewritten" + ); + } + Ok(()) +} diff --git a/datafusion/physical-optimizer/src/window_topn.rs b/datafusion/physical-optimizer/src/window_topn.rs index c668608ca241b..7daa656849fce 100644 --- a/datafusion/physical-optimizer/src/window_topn.rs +++ b/datafusion/physical-optimizer/src/window_topn.rs @@ -26,7 +26,7 @@ //! ) WHERE rn <= K; //! ``` //! -//! or with `RANK()` in place of `ROW_NUMBER()`: +//! or with `RANK()` / `DENSE_RANK()` in place of `ROW_NUMBER()`: //! //! ```sql //! SELECT * FROM ( @@ -40,8 +40,8 @@ //! the `FilterExec` and `SortExec`. //! //! The appropriate [`WindowFnKind`] is forwarded to `PartitionedTopKExec`. -//! RANK requires a non-empty `ORDER BY` clause (otherwise all rows tie at -//! rank 1 and the optimization is degenerate). +//! `RANK` and `DENSE_RANK` require a non-empty `ORDER BY` clause (otherwise +//! all rows tie at rank 1 and the optimization is degenerate). //! //! See [`PartitionedTopKExec`] for details on the replacement operator. //! @@ -68,9 +68,9 @@ use datafusion_physical_plan::sorts::partitioned_topk::{ use datafusion_physical_plan::sorts::sort::SortExec; use datafusion_physical_plan::windows::{BoundedWindowAggExec, WindowUDFExpr}; -/// Physical optimizer rule that converts per-partition `ROW_NUMBER` and -/// `RANK` top-K queries into a more efficient plan using -/// [`PartitionedTopKExec`]. +/// Physical optimizer rule that converts per-partition `ROW_NUMBER`, +/// `RANK`, and `DENSE_RANK` top-K queries into a more efficient plan +/// using [`PartitionedTopKExec`]. /// /// # Pattern Detected /// @@ -86,12 +86,14 @@ use datafusion_physical_plan::windows::{BoundedWindowAggExec, WindowUDFExpr}; /// ```text /// [optional ProjectionExec] /// BoundedWindowAggExec( PARTITION BY ... ORDER BY ...) -/// PartitionedTopKExec(fn=, partition_keys, order_keys, fetch=K) +/// PartitionedTopKExec(fn=, partition_keys, order_keys, fetch=K) /// ``` /// /// The `FilterExec` is removed entirely. The `SortExec` is replaced by -/// `PartitionedTopKExec`, which maintains a per-partition top-K heap (and, -/// for `RANK`, a sibling ties `Vec`) instead of sorting the whole dataset. +/// `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. /// /// # Supported Predicates /// @@ -105,12 +107,12 @@ use datafusion_physical_plan::windows::{BoundedWindowAggExec, WindowUDFExpr}; /// All of the following must be true: /// - Config flag `enable_window_topn` is `true` /// - The plan matches `FilterExec → [ProjectionExec] → BoundedWindowAggExec → SortExec` -/// - The window function is `ROW_NUMBER` or `RANK` (not `DENSE_RANK`) +/// - The window function is `ROW_NUMBER`, `RANK`, or `DENSE_RANK` /// - The window function has a `PARTITION BY` clause (global top-K is /// already handled by `SortExec` with `fetch`) -/// - For `RANK`: a non-empty `ORDER BY` clause (otherwise all rows tie -/// at rank 1 — the optimization is useless and the boundary-tie storage -/// would be unbounded) +/// - For `RANK` / `DENSE_RANK`: a non-empty `ORDER BY` clause (otherwise +/// all rows tie at rank 1 — the optimization is useless and the +/// retained set would be unbounded) /// - The filter predicate compares the window output column to an integer /// literal using `<=`, `<`, `>=`, or `>` /// @@ -141,6 +143,14 @@ impl WindowTopN { // Step 2: Extract limit from predicate (rn <= K, rn < K, etc.) let (col_idx, limit_n) = extract_window_limit(filter.predicate())?; + // `fn < 1` (or the flipped `1 > fn`) yields limit_n = 0. Since + // ROW_NUMBER / RANK / DENSE_RANK are always >= 1, such a predicate + // keeps nothing — bail out and let the ordinary FilterExec return + // the empty result. (The PartitionedTopK* operators require k > 0.) + if limit_n == 0 { + return None; + } + // Step 3: Walk through optional ProjectionExec and RepartitionExec to find BoundedWindowAggExec let child = filter.input(); let (window_exec, intermediates) = find_window_below(child)?; @@ -172,10 +182,10 @@ impl WindowTopN { return None; } - // For RANK: an empty ORDER BY makes every row tie at rank 1 — - // the optimization is degenerate (we'd retain the entire input) - // and tie storage would be unbounded. - if matches!(fn_kind, WindowFnKind::Rank) + // For RANK / DENSE_RANK: an empty ORDER BY makes every row tie + // at rank 1 — the optimization is degenerate (we'd retain the + // entire input) and tie storage would be unbounded. + if matches!(fn_kind, WindowFnKind::Rank | WindowFnKind::DenseRank) && window_exprs[window_expr_idx].order_by().is_empty() { return None; @@ -315,7 +325,8 @@ fn scalar_to_usize(value: &ScalarValue) -> Option { /// the UDF name. Returns: /// - `Some(WindowFnKind::RowNumber)` for `"row_number"` /// - `Some(WindowFnKind::Rank)` for `"rank"` -/// - `None` for everything else (e.g. `dense_rank`) +/// - `Some(WindowFnKind::DenseRank)` for `"dense_rank"` +/// - `None` for everything else fn supported_window_fn( expr: &Arc, ) -> Option { @@ -325,6 +336,7 @@ fn supported_window_fn( match udf.fun().name() { "row_number" => Some(WindowFnKind::RowNumber), "rank" => Some(WindowFnKind::Rank), + "dense_rank" => Some(WindowFnKind::DenseRank), _ => None, } } diff --git a/datafusion/physical-plan/src/sorts/partitioned_topk.rs b/datafusion/physical-plan/src/sorts/partitioned_topk.rs index c250130341dc1..aea9a1f397130 100644 --- a/datafusion/physical-plan/src/sorts/partitioned_topk.rs +++ b/datafusion/physical-plan/src/sorts/partitioned_topk.rs @@ -24,9 +24,9 @@ //! ``` //! //! Instead of sorting the entire dataset, this operator delegates to a -//! per-partition heap-of-K implementation (one variant for `ROW_NUMBER` -//! and a sibling variant for `RANK`), both of which maintain one heap per -//! distinct partition key while sharing a single [`arrow::row::RowConverter`], +//! per-partition top-K implementation — one variant each for `ROW_NUMBER`, +//! `RANK`, and `DENSE_RANK` — all of which keep per-partition state while +//! sharing a single [`arrow::row::RowConverter`], //! [`MemoryReservation`](datafusion_execution::memory_pool::MemoryReservation), //! and metrics set across all partitions, and emit only the top-K rows //! per partition in sorted order `(partition_keys, order_keys)`. @@ -46,7 +46,9 @@ use futures::TryStreamExt; use crate::execution_plan::{Boundedness, EmissionType}; use crate::metrics::ExecutionPlanMetricsSet; -use crate::topk::{PartitionedTopK, PartitionedTopKRank, build_sort_fields}; +use crate::topk::{ + PartitionedTopK, PartitionedTopKDenseRank, PartitionedTopKRank, build_sort_fields, +}; use crate::{ DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, ExecutionPlanProperties, PlanProperties, SendableRecordBatchStream, stream::RecordBatchStreamAdapter, @@ -59,12 +61,19 @@ use crate::{ /// - [`Rank`](Self::Rank): K rows plus any rows tied at the boundary /// ORDER BY value (RANK semantics — `WHERE rk <= K` may keep more /// than K rows when ties straddle the boundary). +/// - [`DenseRank`](Self::DenseRank): every row whose ORDER BY value is +/// among the K distinct-smallest ORDER BY values in the partition +/// (DENSE_RANK semantics — total kept rows is unbounded in +/// rows-per-distinct-value). #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum WindowFnKind { /// `ROW_NUMBER()` — keep exactly K rows per partition. RowNumber, /// `RANK()` — keep K rows plus any rows tied at the boundary. Rank, + /// `DENSE_RANK()` — keep every row whose ob value is among the K + /// distinct-smallest ob values seen in the partition. + DenseRank, } /// Per-partition Top-K operator for window function queries. @@ -108,9 +117,10 @@ pub enum WindowFnKind { /// ``` /// /// Instead of sorting the entire dataset, this operator reads unsorted input -/// and delegates to a per-partition heap-of-K implementation (`PartitionedTopK` -/// for `ROW_NUMBER` and `PartitionedTopKRank` for `RANK`), each maintaining -/// one heap per distinct partition key while sharing a single +/// and delegates to a per-partition top-K implementation (`PartitionedTopK` +/// for `ROW_NUMBER`, `PartitionedTopKRank` for `RANK`, and +/// `PartitionedTopKDenseRank` for `DENSE_RANK`), each maintaining +/// per-partition state while sharing a single /// [`arrow::row::RowConverter`] / /// [`MemoryReservation`](datafusion_execution::memory_pool::MemoryReservation) /// across all partitions, and emits only the top-K rows per partition in @@ -162,11 +172,12 @@ pub enum WindowFnKind { /// /// # Limitations /// -/// - Only activated when the window function is `ROW_NUMBER` or `RANK` with -/// a `PARTITION BY` clause. `RANK` additionally requires a non-empty -/// `ORDER BY` (with an empty `ORDER BY`, every row ties at rank 1 and the -/// heap-of-K rewrite doesn't apply). Global top-K (no `PARTITION BY`) is -/// already handled efficiently by `SortExec` with `fetch`. +/// - Only activated when the window function is `ROW_NUMBER`, `RANK`, or +/// `DENSE_RANK` with a `PARTITION BY` clause. `RANK` and `DENSE_RANK` +/// additionally require a non-empty `ORDER BY` (with an empty `ORDER BY` +/// every row ties at rank 1 and the rewrite doesn't apply). Global top-K +/// (no `PARTITION BY`) is already handled efficiently by `SortExec` with +/// `fetch`. /// - For very high cardinality partition keys (millions of distinct values), /// both memory usage and runtime overhead can become significant. In such /// cases, the sort-based plan is more robust. Therefore, this optimization @@ -210,7 +221,8 @@ impl PartitionedTopKExec { /// that form the partition key. Must be >= 1. /// * `fetch` - Maximum rows to retain per partition (the K in "top-K"). /// * `fn_kind` - Which ranking window function this operator optimizes - /// ([`WindowFnKind::RowNumber`] or [`WindowFnKind::Rank`]). + /// ([`WindowFnKind::RowNumber`], [`WindowFnKind::Rank`], or + /// [`WindowFnKind::DenseRank`]). /// /// # Example /// @@ -295,6 +307,7 @@ impl DisplayAs for PartitionedTopKExec { let fn_label = match self.fn_kind { WindowFnKind::RowNumber => "row_number", WindowFnKind::Rank => "rank", + WindowFnKind::DenseRank => "dense_rank", }; match t { DisplayFormatType::Default | DisplayFormatType::Verbose => { @@ -431,9 +444,9 @@ impl ExecutionPlan for PartitionedTopKExec { } /// Read all input, feed each batch into a per-partition top-K state -/// (either [`PartitionedTopK`] for `ROW_NUMBER` or -/// [`PartitionedTopKRank`] for `RANK`), then emit results ordered by -/// `(partition_keys, order_keys)`. +/// ([`PartitionedTopK`] for `ROW_NUMBER`, [`PartitionedTopKRank`] for +/// `RANK`, or [`PartitionedTopKDenseRank`] for `DENSE_RANK`), then emit +/// results ordered by `(partition_keys, order_keys)`. /// /// # Phases /// @@ -443,10 +456,11 @@ impl ExecutionPlan for PartitionedTopKExec { /// `TopKMetrics` are shared across all distinct partition keys for /// this operator instance. /// -/// 2. **Emission** — `emit` drains all per-partition heaps in sorted +/// 2. **Emission** — `emit` drains all per-partition state in sorted /// partition-key order, returning a coalesced batch stream. For /// `RANK`, boundary-tied rows are materialized and emitted after -/// each partition's heap rows. +/// each partition's heap rows. For `DENSE_RANK`, rows are emitted +/// from a K-bounded map of distinct ob keys, sorted ascending. /// /// # Cost /// @@ -504,5 +518,23 @@ async fn do_partitioned_topk( drop(input); state.emit() } + WindowFnKind::DenseRank => { + let mut state = PartitionedTopKDenseRank::try_new( + partition_id, + schema, + partition_exprs, + partition_sort_fields, + order_expr, + fetch, + batch_size, + &runtime, + &metrics_set, + )?; + while let Some(batch) = input.next().await { + state.insert_batch(&batch?)?; + } + drop(input); + state.emit() + } } } diff --git a/datafusion/physical-plan/src/topk/mod.rs b/datafusion/physical-plan/src/topk/mod.rs index 1e3efff36b1d8..98531668d3695 100644 --- a/datafusion/physical-plan/src/topk/mod.rs +++ b/datafusion/physical-plan/src/topk/mod.rs @@ -1811,6 +1811,376 @@ impl PartitionedTopKRank { } } +/// Per-partition state for `DENSE_RANK()` semantics. +/// +/// A `HashMap, Vec>` keyed by the row-encoded ORDER BY +/// bytes, capped at `k` distinct keys. Each key's `Vec` holds +/// every row seen at that ob value, one entry per contributing source +/// `RecordBatch`. `TieEntry` is reused verbatim from [`RankPartitionState`]. +struct DenseRankPartitionState { + groups: HashMap, Vec>, +} + +impl DenseRankPartitionState { + fn size(&self) -> usize { + let table_overhead = + self.groups.capacity() * (size_of::>() + size_of::>()); + let contents: usize = self + .groups + .iter() + .map(|(key, entries)| { + key.capacity() + + entries.capacity() * size_of::() + + entries + .iter() + .map(|e| { + e.row_indices.capacity() * size_of::() + e.batch_bytes + }) + .sum::() + }) + .sum(); + table_overhead + contents + } +} + +/// Sibling to [`PartitionedTopK`] / [`PartitionedTopKRank`] implementing +/// `DENSE_RANK()` semantics. +/// +/// Per partition, retains every row whose ORDER BY value is among the K +/// distinct-smallest ob values seen for that partition. The total row +/// count kept per partition is unbounded in `rows_per_distinct_value` +/// (unlike `RANK`, which is bounded above by K + boundary ties). +/// +/// Like [`PartitionedTopK`], the [`RowConverter`], [`MemoryReservation`], +/// scratch [`Rows`] buffer, and [`TopKMetrics`] are shared across all +/// partitions for this operator instance. +/// +/// # Algorithm (per batch) +/// +/// Evaluate + encode partition-by and order-by columns once, then group +/// the batch's row indices by partition key. For each partition, bucket +/// that partition's rows by distinct ob value (a within-call +/// accumulation), then merge each bucket into the partition state. Every +/// bucket is built from the current batch's rows, so each `TieEntry` is +/// pinned to the batch its `row_indices` point into. +/// +/// For each partition, for each distinct `ob_key` run in this batch: +/// - `ob_key` already in `state.groups` → push this batch's run as a +/// new `TieEntry` (one entry per contributing batch). +/// - `ob_key` new, `state.groups.len() < k` → insert the run as a new +/// group. +/// - `ob_key` new, `state.groups.len() == k` → find the current max via +/// an O(K) scan of `state.groups.keys()`: +/// - `ob_key < max` → remove the max key (evict the entire max-key +/// group — up to many rows) and insert the run. The evicted group's +/// 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). + expr: LexOrdering, + /// Encoder for ORDER BY columns. Reused across partitions. + row_converter: RowConverter, + /// Scratch row buffer reused across `insert_batch` calls. + scratch_rows: Rows, + /// PARTITION BY expressions. + partition_exprs: Vec>, + /// Encoder for the partition key. + partition_converter: RowConverter, + /// Scratch row buffer for partition-key encoding. Reused across + /// `insert_batch` calls (cleared + appended each batch). + partition_scratch_rows: Rows, + /// One state per distinct partition key seen so far. Keyed by the + /// row-encoded PARTITION BY bytes (byte-comparable encoding, so the + /// `Vec` hashes, compares, and sorts identically to an + /// `OwnedRow`) which lets `insert_batch` look partitions up with + /// `entry_ref` — allocating a key only on first sight of a partition + /// rather than once per row. + states: HashMap, DenseRankPartitionState>, + /// Scratch map reused across `insert_batch` calls to group a batch's + /// row indices by partition key. Drained (not reallocated) each batch + /// so its backing table is allocated once, not per batch. + partition_groups: HashMap, Vec>, + /// Scratch map reused across partitions within a batch to bucket a + /// partition's rows by distinct ORDER BY value. Drained (not + /// reallocated) per partition so its backing table is allocated once, + /// not once per distinct partition key. + ob_runs: HashMap, Vec>, + k: usize, + batch_size: usize, +} + +impl PartitionedTopKDenseRank { + #[expect(clippy::too_many_arguments)] + pub(crate) fn try_new( + partition_id: usize, + schema: SchemaRef, + partition_exprs: Vec>, + partition_sort_fields: Vec, + order_expr: LexOrdering, + k: usize, + batch_size: usize, + runtime: &Arc, + metrics: &ExecutionPlanMetricsSet, + ) -> 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 row_converter = RowConverter::new(order_sort_fields)?; + let scratch_rows = + row_converter.empty_rows(batch_size, ESTIMATED_BYTES_PER_ROW * batch_size); + + let partition_converter = RowConverter::new(partition_sort_fields)?; + let partition_scratch_rows = partition_converter + .empty_rows(batch_size, ESTIMATED_BYTES_PER_ROW * batch_size); + + Ok(Self { + schema, + metrics: TopKMetrics::new(metrics, partition_id), + reservation, + expr: order_expr, + row_converter, + scratch_rows, + partition_exprs, + partition_converter, + partition_scratch_rows, + states: HashMap::new(), + partition_groups: HashMap::new(), + ob_runs: HashMap::new(), + k, + batch_size, + }) + } + + /// Encode PARTITION BY and ORDER BY columns once, demultiplex the + /// batch's rows by partition key, then per partition bucket the rows + /// by distinct ob value and merge each bucket into the partition + /// state as one [`TieEntry`]. + pub(crate) fn insert_batch(&mut self, batch: &RecordBatch) -> Result<()> { + let baseline = self.metrics.baseline.clone(); + let _timer = baseline.elapsed_compute().timer(); + + let num_rows = batch.num_rows(); + if num_rows == 0 { + return Ok(()); + } + + // Captured once so every `TieEntry` push from this batch can + // reuse it (avoids `get_record_batch_memory_size` per push). + let input_batch_bytes = get_record_batch_memory_size(batch); + + // 1. Encode partition columns. + let pk_arrays: Vec = self + .partition_exprs + .iter() + .map(|e| e.evaluate(batch).and_then(|v| v.into_array(num_rows))) + .collect::>()?; + self.partition_scratch_rows.clear(); + self.partition_converter + .append(&mut self.partition_scratch_rows, &pk_arrays)?; + + // 2. Group this batch's row indices by partition key. + // `partition_groups` is a reused scratch map: taken out here + // and drained below, so its backing table is allocated once + // for the operator, not once per batch. `entry_ref` owns the + // key only on Vacant, so it allocates one `Vec` per + // distinct partition rather than one per row. + let mut groups = std::mem::take(&mut self.partition_groups); + groups.clear(); + { + let pk_rows = &self.partition_scratch_rows; + for i in 0..num_rows { + groups + .entry_ref(pk_rows.row(i).as_ref()) + .or_default() + .push(i as u32); + } + } + + // 3. Evaluate ORDER BY columns and encode ONCE. + let ob_arrays: Vec = self + .expr + .iter() + .map(|e| e.expr.evaluate(batch).and_then(|v| v.into_array(num_rows))) + .collect::>()?; + self.scratch_rows.clear(); + self.row_converter + .append(&mut self.scratch_rows, &ob_arrays)?; + + let k = self.k; + let mut replacements: usize = 0; + + // 4. Per-partition: bucket this batch's rows by distinct ob value + // (within-call accumulation), then merge each bucket into the + // partition state as a single `TieEntry`. + for (pk, indices) in groups.drain() { + let state = + self.states + .entry(pk) + .or_insert_with(|| DenseRankPartitionState { + groups: HashMap::new(), + }); + + // Bucket by ob key. `ob_runs` is a reused scratch map (taken + // out and drained below) so its backing table is allocated + // once, not once per distinct partition key. `entry_ref` owns + // the key only on Vacant, so repeated rows of the same ob + // value don't re-allocate. + let mut runs = std::mem::take(&mut self.ob_runs); + runs.clear(); + for &orig_idx in &indices { + let ob_row = self.scratch_rows.row(orig_idx as usize); + runs.entry_ref(ob_row.as_ref()).or_default().push(orig_idx); + } + + for (ob_key, run_indices) in runs.drain() { + // Case A: ob already tracked — push this batch's run as a + // new `TieEntry` (one entry per contributing batch, exactly + // like RANK pushing one tie entry per batch). + if let Some(entries) = state.groups.get_mut(&ob_key) { + entries.push(TieEntry { + batch: batch.clone(), + row_indices: run_indices, + batch_bytes: input_batch_bytes, + }); + continue; + } + + // Case B: new ob, room available. + if state.groups.len() < k { + state.groups.insert( + ob_key, + vec![TieEntry { + batch: batch.clone(), + row_indices: run_indices, + batch_bytes: input_batch_bytes, + }], + ); + continue; + } + + // Case C: new ob, at K distinct keys. Find the current + // max (K-th distinct-best) via an O(K) scan — cold path. + // Scoped so the immutable borrow ends before mutation. + let evict_key: Option> = { + let max_key = state + .groups + .keys() + .map(|key| key.as_slice()) + .max() + .expect("state.groups has k >= 1 keys"); + (ob_key.as_slice() < max_key).then(|| max_key.to_vec()) + }; + if let Some(evicted_key) = evict_key { + let evicted = + state.groups.remove(&evicted_key).expect("max key present"); + replacements += + evicted.iter().map(|e| e.row_indices.len()).sum::(); + state.groups.insert( + ob_key, + vec![TieEntry { + batch: batch.clone(), + row_indices: run_indices, + batch_bytes: input_batch_bytes, + }], + ); + } + // else: ob >= max — drop the whole run. + } + + // Return the drained scratch map (capacity retained) for the + // next partition to reuse. + self.ob_runs = runs; + } + + // Return the drained scratch map (capacity retained) for the next + // batch to reuse. + self.partition_groups = groups; + + if replacements > 0 { + self.metrics.row_replacements.add(replacements); + } + self.reservation.try_resize(self.size())?; + Ok(()) + } + + /// Drain all per-partition maps in partition-key order and return + /// the rows as a stream of coalesced [`RecordBatch`]es ordered by + /// `(partition_keys, order_keys)`. Within a partition the distinct + /// ob keys are sorted (byte-comparable encoding == sort order) so + /// emitted rows are in ob-sorted order. + pub(crate) fn emit(self) -> Result { + let Self { + schema, + metrics, + reservation: _, + expr: _, + row_converter: _, + scratch_rows: _, + partition_exprs: _, + partition_converter: _, + partition_scratch_rows: _, + mut states, + partition_groups: _, + ob_runs: _, + k: _, + batch_size, + } = self; + let _timer = metrics.baseline.elapsed_compute().timer(); + + let mut sorted_pks: Vec> = states.keys().cloned().collect(); + sorted_pks.sort(); + + let mut coalescer = BatchCoalescer::new(Arc::clone(&schema), batch_size); + + for pk in sorted_pks { + let DenseRankPartitionState { groups } = + states.remove(&pk).expect("key from states.keys()"); + // HashMap is unordered — sort the <= K distinct ob keys so + // rows emit ascending (byte-comparable encoding == sort order). + 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 { + for entry in entries { + let indices = UInt32Array::from(entry.row_indices); + let sub = take_record_batch(&entry.batch, &indices)?; + (&sub).record_output(&metrics.baseline); + coalescer.push_batch(sub)?; + } + } + } + coalescer.finish_buffered_batch()?; + + let mut out: Vec> = Vec::new(); + while let Some(b) = coalescer.next_completed_batch() { + out.push(Ok(b)); + } + + Ok(Box::pin(RecordBatchStreamAdapter::new( + schema, + futures::stream::iter(out), + ))) + } + + /// Total memory currently held, including all per-partition states. + fn size(&self) -> usize { + size_of::() + + self.row_converter.size() + + self.partition_converter.size() + + self.scratch_rows.size() + + self.partition_scratch_rows.size() + + self.states.values().map(|s| s.size()).sum::() + + self.states.capacity() + * (size_of::>() + size_of::()) + } +} + #[cfg(test)] mod tests { use super::*; @@ -2537,7 +2907,7 @@ mod tests { } /// Multiple distinct partition keys interleaved within a single - /// input batch — the per-batch demux, per-partition heap eviction, + /// input batch — grouping rows by partition key, per-partition heap eviction, /// and partition-key-ordered emit must all behave correctly. #[tokio::test] async fn test_partitioned_topk_multi_partition_within_batch() -> Result<()> { @@ -2831,7 +3201,7 @@ mod tests { } /// Multiple distinct partition keys interleaved within a single - /// input batch — the per-batch demux, per-partition heap eviction, + /// input batch — grouping rows by partition key, per-partition heap eviction, /// and partition-key-ordered emit must all behave correctly. No /// ties: result should match a `ROW_NUMBER` top-K under the same K. #[tokio::test] @@ -3166,4 +3536,479 @@ mod tests { ); Ok(()) } + + // ==================================================================== + // PartitionedTopKDenseRank operator tests + // + // These mirror the RANK tests plus DENSE_RANK-specific cases: rows + // sharing an ob key coalesce into one `TieEntry`, unbounded + // rows-per-distinct-key, and eviction removes the entire max group + // when a strictly-smaller distinct ob arrives. + // ==================================================================== + + /// Builds a `(pk Int32, val Int32)` schema and a + /// `PartitionedTopKDenseRank` keyed on `pk ASC` (partition) and + /// `val ASC` (ORDER BY). + fn build_partitioned_topk_dense_rank( + k: usize, + ) -> Result<(Arc, PartitionedTopKDenseRank)> { + build_partitioned_topk_dense_rank_with_opts(k, SortOptions::default(), false) + } + + fn build_partitioned_topk_dense_rank_with_opts( + k: usize, + val_sort_options: SortOptions, + val_nullable: bool, + ) -> Result<(Arc, PartitionedTopKDenseRank)> { + let schema = Arc::new(Schema::new(vec![ + Field::new("pk", DataType::Int32, false), + Field::new("val", DataType::Int32, 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, + k, + 8, // batch_size + &Arc::new(RuntimeEnv::default()), + &ExecutionPlanMetricsSet::new(), + )?; + Ok((schema, state)) + } + + /// Single-batch DENSE_RANK top-2 across multiple partitions with + /// distinct ob values only — should behave identically to a + /// ROW_NUMBER top-2. Exercises per-partition grouping + emit order. + #[tokio::test] + async fn test_partitioned_topk_dense_rank_multi_partition_within_batch() -> Result<()> + { + let (schema, mut state) = build_partitioned_topk_dense_rank(2)?; + + // pk=1 vals: 10, 5, 8 → distinct-top-2 ASC = {5, 8} + // pk=2 vals: 20, 15 → distinct-top-2 ASC = {15, 20} + // pk=3 vals: 7 → distinct-top-2 ASC = {7} + let batch = + pk_val_batch(&schema, vec![1, 2, 1, 2, 1, 3], vec![10, 20, 5, 15, 8, 7])?; + state.insert_batch(&batch)?; + + let results: Vec<_> = state.emit()?.try_collect().await?; + assert_batches_eq!( + &[ + "+----+-----+", + "| pk | val |", + "+----+-----+", + "| 1 | 5 |", + "| 1 | 8 |", + "| 2 | 15 |", + "| 2 | 20 |", + "| 3 | 7 |", + "+----+-----+", + ], + &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 `TieEntry` per distinct ob. + /// + /// vals per partition (sorted logically): + /// pk=1: 1, 1, 1, 2, 2, 3, 3, 3, 4 + /// distinct-top-2 = {1, 2} → all 5 rows at those values retained. + #[tokio::test] + async fn test_partitioned_topk_dense_rank_heavy_ties_coalesced() -> Result<()> { + let (schema, mut state) = build_partitioned_topk_dense_rank(2)?; + + let batch = pk_val_batch( + &schema, + vec![1, 1, 1, 1, 1, 1, 1, 1, 1], + vec![1, 3, 1, 2, 3, 1, 2, 3, 4], + )?; + state.insert_batch(&batch)?; + + let results: Vec<_> = state.emit()?.try_collect().await?; + assert_batches_eq!( + &[ + "+----+-----+", + "| pk | val |", + "+----+-----+", + "| 1 | 1 |", + "| 1 | 1 |", + "| 1 | 1 |", + "| 1 | 2 |", + "| 1 | 2 |", + "+----+-----+", + ], + &results + ); + Ok(()) + } + + /// Rows tied at the same ob across two source batches must both + /// land under the same map key as separate `TieEntry`s — one per + /// source batch — but emit as a single contiguous run. + #[tokio::test] + async fn test_partitioned_topk_dense_rank_cross_batch_same_key() -> Result<()> { + let (schema, mut state) = build_partitioned_topk_dense_rank(2)?; + + // Batch 1: pk=1 with ob values {5, 5, 8}. groups after: {5→[..], 8→[..]}. + state.insert_batch(&pk_val_batch(&schema, vec![1, 1, 1], vec![5, 5, 8])?)?; + + // Batch 2: pk=1 with more 5s and an 8, plus a 20 that's dropped. + state.insert_batch(&pk_val_batch(&schema, vec![1, 1, 1], vec![5, 8, 20])?)?; + + let results: Vec<_> = state.emit()?.try_collect().await?; + assert_batches_eq!( + &[ + "+----+-----+", + "| pk | val |", + "+----+-----+", + "| 1 | 5 |", + "| 1 | 5 |", + "| 1 | 5 |", + "| 1 | 8 |", + "| 1 | 8 |", + "+----+-----+", + ], + &results + ); + Ok(()) + } + + /// Refactor guard: the full RANK-style path in one run — multi-partition + /// per-batch grouping, within-batch bucketing of scattered same-ob rows, + /// cross-batch append to an existing group, cross-batch new-key insert, + /// and cross-batch eviction of a whole max group. Every `TieEntry` is + /// built from its own source batch (no cross-batch coalescing), so the + /// retained rows must be exactly the K=2 smallest distinct ob values + /// per partition with all their rows, regardless of arrival order. + #[tokio::test] + async fn test_partitioned_topk_dense_rank_multi_batch_multi_partition() -> Result<()> + { + let (schema, mut state) = build_partitioned_topk_dense_rank(2)?; + + // Batch 1 interleaves pk=1 and pk=2, with same-ob rows scattered: + // pk=1 vals: 10, 20, 10, 20, 10 → {10:[×3], 20:[×2]} + // pk=2 vals: 100, 100 → {100:[×2]} + state.insert_batch(&pk_val_batch( + &schema, + vec![1, 2, 1, 1, 2, 1, 1], + vec![10, 100, 20, 20, 100, 10, 10], + )?)?; + + // Batch 2: + // pk=1 vals: 20, 5, 10 → append a 20, insert 5 (evicts the whole + // 20 group), append a 10 → retained distinct {5, 10}. + // pk=2 vals: 50 → insert 5th... new key, room → {50, 100}. + state.insert_batch(&pk_val_batch( + &schema, + vec![1, 2, 1, 1], + vec![20, 50, 5, 10], + )?)?; + + let results: Vec<_> = state.emit()?.try_collect().await?; + // pk=1: val=5 (×1 from batch 2), val=10 (×3 batch 1 + ×1 batch 2 = ×4). + // All 20s dropped (evicted). pk=2: val=50 (×1), val=100 (×2). + assert_batches_eq!( + &[ + "+----+-----+", + "| pk | val |", + "+----+-----+", + "| 1 | 5 |", + "| 1 | 10 |", + "| 1 | 10 |", + "| 1 | 10 |", + "| 1 | 10 |", + "| 2 | 50 |", + "| 2 | 100 |", + "| 2 | 100 |", + "+----+-----+", + ], + &results + ); + Ok(()) + } + + /// DENSE_RANK-specific: eviction removes the entire max group when + /// a strictly-smaller distinct ob arrives. Multiple rows at the + /// evicted key all disappear. + #[tokio::test] + async fn test_partitioned_topk_dense_rank_max_group_eviction() -> Result<()> { + let (schema, mut state) = build_partitioned_topk_dense_rank(2)?; + + // Batch 1: pk=1 with {10, 10, 20, 20}. groups={10→[..], 20→[..]}, at K. + state.insert_batch(&pk_val_batch( + &schema, + vec![1, 1, 1, 1], + vec![10, 10, 20, 20], + )?)?; + + // Batch 2: pk=1 with 5 — strictly smaller than max=20, evict entire + // 20 group; now groups={10, 5}. Then a 30 comes in and is dropped. + state.insert_batch(&pk_val_batch(&schema, vec![1, 1], vec![5, 30])?)?; + + let results: Vec<_> = state.emit()?.try_collect().await?; + assert_batches_eq!( + &[ + "+----+-----+", + "| pk | val |", + "+----+-----+", + "| 1 | 5 |", + "| 1 | 10 |", + "| 1 | 10 |", + "+----+-----+", + ], + &results + ); + Ok(()) + } + + /// Empty input must produce an empty output stream, not panic. + #[tokio::test] + async fn test_partitioned_topk_dense_rank_empty_input() -> Result<()> { + let (_schema, state) = build_partitioned_topk_dense_rank(3)?; + let results: Vec<_> = state.emit()?.try_collect().await?; + assert!(results.is_empty(), "empty input → empty output"); + Ok(()) + } + + /// `fetch = 1` retains only the smallest distinct ob per partition, + /// with all rows at that value kept. + #[tokio::test] + async fn test_partitioned_topk_dense_rank_fetch_one() -> Result<()> { + let (schema, mut state) = build_partitioned_topk_dense_rank(1)?; + + // pk=1 vals: 5, 3, 5, 3, 7 → distinct-top-1 = {3} → both 3s kept. + // pk=2 vals: 9, 4 → distinct-top-1 = {4} → single 4. + let batch = pk_val_batch( + &schema, + vec![1, 1, 1, 2, 1, 2, 1], + vec![5, 3, 5, 9, 3, 4, 7], + )?; + state.insert_batch(&batch)?; + + let results: Vec<_> = state.emit()?.try_collect().await?; + assert_batches_eq!( + &[ + "+----+-----+", + "| pk | val |", + "+----+-----+", + "| 1 | 3 |", + "| 1 | 3 |", + "| 2 | 4 |", + "+----+-----+", + ], + &results + ); + Ok(()) + } + + /// `K > distinct_ob_count` — nothing should be dropped. + #[tokio::test] + async fn test_partitioned_topk_dense_rank_k_exceeds_distinct() -> Result<()> { + let (schema, mut state) = build_partitioned_topk_dense_rank(10)?; + + // Only 3 distinct ob values under pk=1; all rows must be retained. + let batch = pk_val_batch(&schema, vec![1, 1, 1, 1], vec![5, 3, 3, 7])?; + state.insert_batch(&batch)?; + + let results: Vec<_> = state.emit()?.try_collect().await?; + assert_batches_eq!( + &[ + "+----+-----+", + "| pk | val |", + "+----+-----+", + "| 1 | 3 |", + "| 1 | 3 |", + "| 1 | 5 |", + "| 1 | 7 |", + "+----+-----+", + ], + &results + ); + Ok(()) + } + + /// `ORDER BY val DESC` — the row-encoded key ordering must reflect + /// the direction so the "distinct-K best" set is the K *largest* + /// distinct ob values. + #[tokio::test] + async fn test_partitioned_topk_dense_rank_desc_ordering() -> Result<()> { + let (schema, mut state) = build_partitioned_topk_dense_rank_with_opts( + 2, + SortOptions { + descending: true, + nulls_first: false, + }, + false, + )?; + + // pk=1 vals: 10, 5, 8, 12, 10 → distinct-top-2 DESC = {12, 10} + // → keep both 10s and 12. + let batch = pk_val_batch(&schema, vec![1, 1, 1, 1, 1], vec![10, 5, 8, 12, 10])?; + state.insert_batch(&batch)?; + + let results: Vec<_> = state.emit()?.try_collect().await?; + assert_batches_eq!( + &[ + "+----+-----+", + "| pk | val |", + "+----+-----+", + "| 1 | 12 |", + "| 1 | 10 |", + "| 1 | 10 |", + "+----+-----+", + ], + &results + ); + Ok(()) + } + + /// Cross-partition eviction independence — Case-C eviction in one + /// partition must not affect another partition's state. + #[tokio::test] + async fn test_partitioned_topk_dense_rank_partition_independence() -> Result<()> { + let (schema, mut state) = build_partitioned_topk_dense_rank(2)?; + + // Batch 1: pk=1 fills {10, 20}; pk=2 fills {30, 40}. + state.insert_batch(&pk_val_batch( + &schema, + vec![1, 1, 2, 2], + vec![10, 20, 30, 40], + )?)?; + + // Batch 2: pk=1 sees 5 (evicts 20). pk=2 sees 25 (evicts 40). + // Each partition's Case-C branch is independent. + state.insert_batch(&pk_val_batch(&schema, vec![1, 2], vec![5, 25])?)?; + + let results: Vec<_> = state.emit()?.try_collect().await?; + assert_batches_eq!( + &[ + "+----+-----+", + "| pk | val |", + "+----+-----+", + "| 1 | 5 |", + "| 1 | 10 |", + "| 2 | 25 |", + "| 2 | 30 |", + "+----+-----+", + ], + &results + ); + Ok(()) + } + + /// NULL sort values exercise the shared encoder's null-ordering + /// through the row-encoded key byte order. With `ASC NULLS + /// LAST`, a NULL is the *largest* distinct ob, so a partition with + /// >= K non-NULL distinct values evicts its NULLs, while a partition + /// whose only distinct value is NULL still emits it. + #[tokio::test] + async fn test_partitioned_topk_dense_rank_nulls_last_ordering() -> Result<()> { + let (schema, mut state) = build_partitioned_topk_dense_rank_with_opts( + 2, + SortOptions { + descending: false, + nulls_first: false, + }, + true, + )?; + + // pk=1 vals: NULL, 10, 20, NULL → distinct-top-2 NULLS LAST = {10, 20} + // pk=2 vals: NULL → distinct-top-2 = {NULL} + // pk=3 vals: 3, 3 → distinct-top-2 = {3} + let batch = nullable_pk_val_batch( + &schema, + vec![1, 1, 1, 1, 2, 3, 3], + vec![None, Some(10), Some(20), None, None, Some(3), Some(3)], + )?; + state.insert_batch(&batch)?; + + let results: Vec<_> = state.emit()?.try_collect().await?; + assert_batches_eq!( + &[ + "+----+-----+", + "| pk | val |", + "+----+-----+", + "| 1 | 10 |", + "| 1 | 20 |", + "| 2 | |", + "| 3 | 3 |", + "| 3 | 3 |", + "+----+-----+", + ], + &results + ); + Ok(()) + } + + /// `ASC NULLS FIRST` sorts NULLs *before* every non-NULL value, so a + /// NULL is the smallest distinct ob and is kept preferentially. Every + /// row at a retained distinct ob — including all tied NULLs — emits. + #[tokio::test] + async fn test_partitioned_topk_dense_rank_nulls_first_ordering() -> Result<()> { + let (schema, mut state) = build_partitioned_topk_dense_rank_with_opts( + 2, + SortOptions { + descending: false, + nulls_first: true, + }, + true, + )?; + + // pk=1 vals: NULL, 5, NULL, 8 → distinct-top-2 NULLS FIRST = {NULL, 5} + // pk=2 vals: 7, NULL → distinct-top-2 = {NULL, 7} + // pk=3 vals: 3, 1 → distinct-top-2 = {1, 3} + let batch = nullable_pk_val_batch( + &schema, + vec![1, 2, 1, 3, 1, 2, 1, 3], + vec![ + None, + Some(7), + Some(5), + Some(3), + None, + None, + Some(8), + Some(1), + ], + )?; + state.insert_batch(&batch)?; + + let results: Vec<_> = state.emit()?.try_collect().await?; + assert_batches_eq!( + &[ + "+----+-----+", + "| pk | val |", + "+----+-----+", + "| 1 | |", + "| 1 | |", + "| 1 | 5 |", + "| 2 | |", + "| 2 | 7 |", + "| 3 | 1 |", + "| 3 | 3 |", + "+----+-----+", + ], + &results + ); + Ok(()) + } } diff --git a/datafusion/sqllogictest/test_files/window_topn.slt b/datafusion/sqllogictest/test_files/window_topn.slt index 4dff4a779b385..67b74905501d6 100644 --- a/datafusion/sqllogictest/test_files/window_topn.slt +++ b/datafusion/sqllogictest/test_files/window_topn.slt @@ -1041,6 +1041,322 @@ SELECT id, pk, val FROM ( statement ok DROP TABLE window_topn_rank_null_t; +############################################################################### +# DENSE_RANK() tests +############################################################################### +# +# DENSE_RANK semantics: `WHERE dr <= K` keeps every row whose ORDER BY +# value is among the K distinct-smallest ORDER BY values in the +# partition. The total kept per partition is unbounded in +# rows-per-distinct-value (unlike RANK, which is bounded above by +# `K + ties at the boundary`). +# +# The tests below exercise: +# - heavy ties within a distinct ob key (all rows retained per key) +# - cross-batch appends under the same key +# - eviction of the entire max-key group when a strictly-smaller +# distinct ob arrives +# - the empty-ORDER-BY degenerate case (rule must NOT fire) + +statement ok +SET datafusion.optimizer.enable_window_topn = true; + +statement ok +CREATE TABLE window_topn_dense_rank_t (id INT, pk INT, val INT) AS VALUES + -- pk=1: distinct top-2 = {10, 20}. Every row at those keeps dr <= 2. + -- 10 appears once, 20 appears three times (heavy tie), 30/40 dropped. + (1, 1, 10), + (2, 1, 20), + (3, 1, 20), + (4, 1, 20), + (5, 1, 30), + (6, 1, 40), + -- pk=2: distinct top-2 = {5, 15}. Two rows at 5, one at 15. + (7, 2, 5), + (8, 2, 5), + (9, 2, 15), + (10, 2, 25), + -- pk=3: 100 then four 200s — 200 group is the "boundary-max" but with + -- dense_rank <= 2 both distinct values (100, 200) are retained. + (11, 3, 100), + (12, 3, 200), + (13, 3, 200), + (14, 3, 200), + (15, 3, 200), + (16, 3, 300); + +# Test DR1: Basic DENSE_RANK correctness — every row at the K distinct +# smallest values is retained. +query III rowsort +SELECT id, pk, val FROM ( + SELECT *, DENSE_RANK() OVER (PARTITION BY pk ORDER BY val) as dr FROM window_topn_dense_rank_t +) WHERE dr <= 2; +---- +1 1 10 +11 3 100 +12 3 200 +13 3 200 +14 3 200 +15 3 200 +2 1 20 +3 1 20 +4 1 20 +7 2 5 +8 2 5 +9 2 15 + +# Test DR2: EXPLAIN shows PartitionedTopKExec with fn=dense_rank +query TT +EXPLAIN SELECT * FROM ( + SELECT *, DENSE_RANK() OVER (PARTITION BY pk ORDER BY val) as dr FROM window_topn_dense_rank_t +) WHERE dr <= 2; +---- +logical_plan +01)Projection: window_topn_dense_rank_t.id, window_topn_dense_rank_t.pk, window_topn_dense_rank_t.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 AS dr +02)--Filter: 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(2) +03)----WindowAggr: windowExpr=[[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]] +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)----PartitionedTopKExec: fn=dense_rank, fetch=2, partition=[pk@1], order=[val@2 ASC NULLS LAST] +04)------DataSourceExec: partitions=1, partition_sizes=[1] + +# Test DR3: dr < 3 should give the same results (fetch = K-1 = 2) +query III rowsort +SELECT id, pk, val FROM ( + SELECT *, DENSE_RANK() OVER (PARTITION BY pk ORDER BY val) as dr FROM window_topn_dense_rank_t +) WHERE dr < 3; +---- +1 1 10 +11 3 100 +12 3 200 +13 3 200 +14 3 200 +15 3 200 +2 1 20 +3 1 20 +4 1 20 +7 2 5 +8 2 5 +9 2 15 + +# Test DR4: flipped predicate (2 >= dr) also fires the rule +query III rowsort +SELECT id, pk, val FROM ( + SELECT *, DENSE_RANK() OVER (PARTITION BY pk ORDER BY val) as dr FROM window_topn_dense_rank_t +) WHERE 2 >= dr; +---- +1 1 10 +11 3 100 +12 3 200 +13 3 200 +14 3 200 +15 3 200 +2 1 20 +3 1 20 +4 1 20 +7 2 5 +8 2 5 +9 2 15 + +# Test DR5: K exceeds every partition's distinct count — nothing dropped. +# pk=1: 4 distinct, pk=2: 3 distinct, pk=3: 3 distinct. dr <= 100 retains all. +query III rowsort +SELECT id, pk, val FROM ( + SELECT *, DENSE_RANK() OVER (PARTITION BY pk ORDER BY val) as dr FROM window_topn_dense_rank_t +) WHERE dr <= 100; +---- +1 1 10 +10 2 25 +11 3 100 +12 3 200 +13 3 200 +14 3 200 +15 3 200 +16 3 300 +2 1 20 +3 1 20 +4 1 20 +5 1 30 +6 1 40 +7 2 5 +8 2 5 +9 2 15 + +# Test DR6: DENSE_RANK without PARTITION BY — should NOT trigger the optimization +query TT +EXPLAIN SELECT * FROM ( + SELECT *, DENSE_RANK() OVER (ORDER BY val) as dr FROM window_topn_dense_rank_t +) WHERE dr <= 2; +---- +logical_plan +01)Projection: window_topn_dense_rank_t.id, window_topn_dense_rank_t.pk, window_topn_dense_rank_t.val, dense_rank() ORDER BY [window_topn_dense_rank_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW AS dr +02)--Filter: dense_rank() ORDER BY [window_topn_dense_rank_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW <= UInt64(2) +03)----WindowAggr: windowExpr=[[dense_rank() ORDER BY [window_topn_dense_rank_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]] +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() ORDER BY [window_topn_dense_rank_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as dr] +02)--FilterExec: dense_rank() ORDER BY [window_topn_dense_rank_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 <= 2 +03)----BoundedWindowAggExec: wdw=[dense_rank() ORDER BY [window_topn_dense_rank_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "dense_rank() 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] +04)------SortExec: expr=[val@2 ASC NULLS LAST], preserve_partitioning=[false] +05)--------DataSourceExec: partitions=1, partition_sizes=[1] + +# Test DR7: DENSE_RANK with empty ORDER BY — degenerate (every row at +# dense_rank 1), rule must NOT fire. +query TT +EXPLAIN SELECT * FROM ( + SELECT *, DENSE_RANK() OVER (PARTITION BY pk) as dr FROM window_topn_dense_rank_t +) WHERE dr <= 3; +---- +logical_plan +01)Projection: window_topn_dense_rank_t.id, window_topn_dense_rank_t.pk, window_topn_dense_rank_t.val, dense_rank() PARTITION BY [window_topn_dense_rank_t.pk] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING AS dr +02)--Filter: dense_rank() PARTITION BY [window_topn_dense_rank_t.pk] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING <= UInt64(3) +03)----WindowAggr: windowExpr=[[dense_rank() PARTITION BY [window_topn_dense_rank_t.pk] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING]] +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] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING@3 as dr] +02)--FilterExec: dense_rank() PARTITION BY [window_topn_dense_rank_t.pk] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING@3 <= 3 +03)----BoundedWindowAggExec: wdw=[dense_rank() PARTITION BY [window_topn_dense_rank_t.pk] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING: Field { "dense_rank() PARTITION BY [window_topn_dense_rank_t.pk] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING": UInt64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING], mode=[Sorted] +04)------SortExec: expr=[pk@1 ASC NULLS LAST], preserve_partitioning=[false] +05)--------DataSourceExec: partitions=1, partition_sizes=[1] + +# Test DR8: DESC ordering — distinct-top-2 DESC per partition. +# pk=1 DESC {40, 30, 20, 10}: top-2 = {40, 30} → 2 rows +# pk=2 DESC {25, 15, 5}: top-2 = {25, 15} → 2 rows +# pk=3 DESC {300, 200, 100}: top-2 = {300, 200} → 5 rows (200 appears 4x) +query III rowsort +SELECT id, pk, val FROM ( + SELECT *, DENSE_RANK() OVER (PARTITION BY pk ORDER BY val DESC) as dr FROM window_topn_dense_rank_t +) WHERE dr <= 2; +---- +10 2 25 +12 3 200 +13 3 200 +14 3 200 +15 3 200 +16 3 300 +5 1 30 +6 1 40 +9 2 15 + +# Test DR9: multi-column PARTITION BY +query III rowsort +SELECT id, pk, val FROM ( + SELECT *, DENSE_RANK() OVER (PARTITION BY pk, id ORDER BY val) as dr FROM window_topn_dense_rank_t +) WHERE dr <= 1; +---- +1 1 10 +10 2 25 +11 3 100 +12 3 200 +13 3 200 +14 3 200 +15 3 200 +16 3 300 +2 1 20 +3 1 20 +4 1 20 +5 1 30 +6 1 40 +7 2 5 +8 2 5 +9 2 15 + +# Test DR10: mixed ROW_NUMBER + DENSE_RANK — filter on DR should optimize +query TT +EXPLAIN SELECT * FROM ( + SELECT *, + ROW_NUMBER() OVER (PARTITION BY pk ORDER BY val) as rn, + DENSE_RANK() OVER (PARTITION BY pk ORDER BY val) as dr + FROM window_topn_dense_rank_t +) WHERE dr <= 2; +---- +logical_plan +01)Projection: window_topn_dense_rank_t.id, window_topn_dense_rank_t.pk, window_topn_dense_rank_t.val, row_number() 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 AS rn, 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 AS dr +02)--Filter: 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(2) +03)----WindowAggr: windowExpr=[[row_number() 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, 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]] +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, row_number() 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 rn, 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@4 as dr] +02)--BoundedWindowAggExec: wdw=[row_number() 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 { "row_number() 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, 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)----PartitionedTopKExec: fn=dense_rank, fetch=2, partition=[pk@1], order=[val@2 ASC NULLS LAST] +04)------DataSourceExec: partitions=1, partition_sizes=[1] + +# Test DR11: `dr < 1` keeps nothing (limit_n = 0). The rule must NOT +# fire (a fetch=0 PartitionedTopKExec would panic); the ordinary +# FilterExec returns the empty result. +query III rowsort +SELECT id, pk, val FROM ( + SELECT *, DENSE_RANK() OVER (PARTITION BY pk ORDER BY val) as dr FROM window_topn_dense_rank_t +) WHERE dr < 1; +---- + +statement ok +DROP TABLE window_topn_dense_rank_t; + +# --------------------------------------------------------------------------- +# DENSE_RANK NULL-in-ORDER-BY ordering +# +# Under DENSE_RANK a NULL is a distinct ob value occupying its own rank +# slot; whether it lands among the K distinct-smallest depends on +# NULLS FIRST / NULLS LAST. Correctness rests entirely on the +# byte-comparable row encoding driving the distinct-ob key ordering, so +# exercise both null placements explicitly. +# --------------------------------------------------------------------------- + +statement ok +CREATE TABLE window_topn_dr_null_t (id INT, pk INT, val INT) AS VALUES + (1, 1, NULL), + (2, 1, NULL), + (3, 1, 10), + (4, 1, 20), + (5, 1, 30); + +# Test DR12: NULLS FIRST — NULL is the smallest distinct ob (dr 1), then +# 10 (dr 2). `dr <= 2` keeps both NULL rows and val=10. +query III rowsort +SELECT id, pk, val FROM ( + SELECT *, DENSE_RANK() OVER (PARTITION BY pk ORDER BY val ASC NULLS FIRST) as dr + FROM window_topn_dr_null_t +) WHERE dr <= 2; +---- +1 1 NULL +2 1 NULL +3 1 10 + +# Test DR13: EXPLAIN confirms the NULLS FIRST case uses PartitionedTopKExec. +query TT +EXPLAIN SELECT id, pk, val FROM ( + SELECT *, DENSE_RANK() OVER (PARTITION BY pk ORDER BY val ASC NULLS FIRST) as dr + FROM window_topn_dr_null_t +) WHERE dr <= 2; +---- +logical_plan +01)Projection: window_topn_dr_null_t.id, window_topn_dr_null_t.pk, window_topn_dr_null_t.val +02)--Filter: 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(2) +03)----WindowAggr: windowExpr=[[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]] +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)----PartitionedTopKExec: fn=dense_rank, fetch=2, partition=[pk@1], order=[val@2 ASC] +04)------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. +query III rowsort +SELECT id, pk, val FROM ( + SELECT *, DENSE_RANK() OVER (PARTITION BY pk ORDER BY val ASC NULLS LAST) as dr + FROM window_topn_dr_null_t +) WHERE dr <= 2; +---- +3 1 10 +4 1 20 + +statement ok +DROP TABLE window_topn_dr_null_t; + # Reset config to default (false) statement ok SET datafusion.optimizer.enable_window_topn = false;