From d0c03145141a69799717f1fda0ce0f2e0048bd09 Mon Sep 17 00:00:00 2001 From: SubhamSinghal Date: Thu, 16 Jul 2026 19:17:14 +0530 Subject: [PATCH 1/6] 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; From 4530bba18ccbf135cf2b1ef85e37fd321cb946d7 Mon Sep 17 00:00:00 2001 From: SubhamSinghal Date: Tue, 28 Jul 2026 21:40:35 +0530 Subject: [PATCH 2/6] address review comments --- .../physical-optimizer/src/window_topn.rs | 56 ++++++-- datafusion/physical-plan/src/topk/mod.rs | 52 ++++--- .../sqllogictest/test_files/window_topn.slt | 134 ++++++++++++++++++ 3 files changed, 216 insertions(+), 26 deletions(-) diff --git a/datafusion/physical-optimizer/src/window_topn.rs b/datafusion/physical-optimizer/src/window_topn.rs index 7daa656849fce..9e91739ae15ab 100644 --- a/datafusion/physical-optimizer/src/window_topn.rs +++ b/datafusion/physical-optimizer/src/window_topn.rs @@ -108,11 +108,18 @@ use datafusion_physical_plan::windows::{BoundedWindowAggExec, WindowUDFExpr}; /// - Config flag `enable_window_topn` is `true` /// - The plan matches `FilterExec → [ProjectionExec] → BoundedWindowAggExec → SortExec` /// - The window function is `ROW_NUMBER`, `RANK`, or `DENSE_RANK` +/// - Every window expression in the `BoundedWindowAggExec` is `ROW_NUMBER`, +/// `RANK`, or `DENSE_RANK` over the same `PARTITION BY` / `ORDER BY`. A +/// sibling that reads pruned rows (e.g. `LEAD`, or an aggregate whose +/// frame is not strictly backward-looking) would be computed over the +/// pruned input and give wrong results. /// - The window function has a `PARTITION BY` clause (global top-K is /// already handled by `SortExec` with `fetch`) -/// - 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) +/// - At least one `ORDER BY` key survives past the `PARTITION BY` prefix +/// (so the operator has a non-empty ORDER BY). This rejects both a +/// missing `ORDER BY` and one fully covered by the partition prefix +/// such as `PARTITION BY pk ORDER BY pk`; for `RANK` / `DENSE_RANK` +/// such orderings also make every row tie at rank 1 (degenerate). /// - The filter predicate compares the window output column to an integer /// literal using `<=`, `<`, `>=`, or `>` /// @@ -169,6 +176,28 @@ impl WindowTopN { } let fn_kind = supported_window_fn(&window_exprs[window_expr_idx])?; + // Tail-pruning drops the rows that rank after the retained top-K, + // and every window expression in this `BoundedWindowAggExec` is + // then evaluated over the *pruned* input. The rewrite is only valid + // if each expression's value for a *retained* row is unaffected by + // the dropped rows. ROW_NUMBER / RANK / DENSE_RANK over the same + // PARTITION BY / ORDER BY satisfy this — each depends only on rows + // at or before the current row in that order, all of which are + // retained. A sibling like `LEAD(x)` reads following (pruned) rows, + // so at the retained boundary it would resolve to a dropped row and + // give a wrong result. Bail out unless every window expression is a + // supported ranking function sharing the matched expression's + // partition/order keys. + let matched_expr = &window_exprs[window_expr_idx]; + let all_prune_safe = window_exprs.iter().all(|e| { + supported_window_fn(e).is_some() + && e.partition_by() == matched_expr.partition_by() + && e.order_by() == matched_expr.order_by() + }); + if !all_prune_safe { + return None; + } + // Step 5: child of window is SortExec (verified above) let sort_child = sort_exec.input(); @@ -182,12 +211,21 @@ impl WindowTopN { return None; } - // 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() - { + // `PartitionedTopKExec` derives its ORDER BY keys from the + // SortExec ordering *beyond* the partition prefix + // (`expr[partition_prefix_len..]`). That slice is empty in two + // cases: + // * no ORDER BY at all (e.g. `ROW_NUMBER() OVER (PARTITION BY pk)`); + // * ORDER BY keys fully covered by the partition prefix (e.g. + // `DENSE_RANK() OVER (PARTITION BY pk ORDER BY pk)`, whose + // deduplicated sort ordering is just `[pk]`). + // With zero order keys the operator panics on execution (it + // requires at least one), and for RANK / DENSE_RANK every row + // would tie at rank 1 (a degenerate, unbounded retained set). + // `order_by()` alone does not catch the second case — it reports + // `[pk]` even though no order key survives past the partition + // prefix — so guard on the effective order-key count instead. + if sort_exec.expr().len() <= partition_prefix_len { return None; } diff --git a/datafusion/physical-plan/src/topk/mod.rs b/datafusion/physical-plan/src/topk/mod.rs index 98531668d3695..80d1830347b51 100644 --- a/datafusion/physical-plan/src/topk/mod.rs +++ b/datafusion/physical-plan/src/topk/mod.rs @@ -1819,6 +1819,14 @@ impl PartitionedTopKRank { /// `RecordBatch`. `TieEntry` is reused verbatim from [`RankPartitionState`]. struct DenseRankPartitionState { groups: HashMap, Vec>, + /// Cached admission boundary: the largest tracked ob value once + /// `groups` is full. A `HashMap` is unordered, so finding it otherwise + /// costs an O(K) scan of the keys on *every* new distinct ob value — + /// O(N·K) overall on mostly-distinct input. Caching makes the common + /// "new ob is worse than the boundary → drop" path O(1); the O(K) scan + /// is paid only when the cache is stale (`None`) — after a fill or an + /// eviction changes the key set. + max_key: Option>, } impl DenseRankPartitionState { @@ -1839,7 +1847,7 @@ impl DenseRankPartitionState { .sum::() }) .sum(); - table_overhead + contents + table_overhead + contents + self.max_key.as_ref().map_or(0, |k| k.capacity()) } } @@ -1869,8 +1877,9 @@ impl DenseRankPartitionState { /// 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` new, `state.groups.len() == k` → the largest tracked ob +/// value is the admission boundary, cached in `max_key` (recomputed via +/// an O(K) scan of `state.groups.keys()` only when the cache is stale): /// - `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. @@ -2024,6 +2033,7 @@ impl PartitionedTopKDenseRank { .entry(pk) .or_insert_with(|| DenseRankPartitionState { groups: HashMap::new(), + max_key: None, }); // Bucket by ob key. `ob_runs` is a reused scratch map (taken @@ -2061,22 +2071,30 @@ impl PartitionedTopKDenseRank { batch_bytes: input_batch_bytes, }], ); + // A new distinct key may raise the boundary; drop the + // cached max so it is recomputed on the next Case C. + state.max_key = None; 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 { + // Case C: new ob, at K distinct keys. The largest tracked + // ob value is the admission boundary; it is cached in + // `max_key` and recomputed via an O(K) scan only when the + // cache is stale, so the common "ob >= max → drop" path is + // O(1). Scoped so the immutable borrow ends before mutation. + if state.max_key.is_none() { + state.max_key = state.groups.keys().max().cloned(); + } + let is_smaller = state + .max_key + .as_deref() + .map(|max_key| ob_key.as_slice() < max_key) + .expect("state.groups has k >= 1 keys"); + + if is_smaller { + // Evict the entire max-key group and drop the cached + // max (recomputed on the next Case C). + let evicted_key = state.max_key.take().expect("max key present"); let evicted = state.groups.remove(&evicted_key).expect("max key present"); replacements += @@ -2139,7 +2157,7 @@ impl PartitionedTopKDenseRank { let mut coalescer = BatchCoalescer::new(Arc::clone(&schema), batch_size); for pk in sorted_pks { - let DenseRankPartitionState { groups } = + let DenseRankPartitionState { groups, max_key: _ } = 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). diff --git a/datafusion/sqllogictest/test_files/window_topn.slt b/datafusion/sqllogictest/test_files/window_topn.slt index 67b74905501d6..fed642f1b0b51 100644 --- a/datafusion/sqllogictest/test_files/window_topn.slt +++ b/datafusion/sqllogictest/test_files/window_topn.slt @@ -329,9 +329,115 @@ physical_plan 04)------SortExec: expr=[pk@1 ASC NULLS LAST, val@2 ASC NULLS LAST], preserve_partitioning=[false] 05)--------DataSourceExec: partitions=1, partition_sizes=[1] +# Sibling safety: an aggregate sibling (SUM) is computed over the pruned +# input, so the rule must NOT fire even though the filter is on ROW_NUMBER. +# SUM over the default RANGE frame includes tie-peers that pruning would +# drop, giving a wrong sum. The plan keeps SortExec + FilterExec — no +# PartitionedTopKExec. +query TT +EXPLAIN SELECT * FROM ( + SELECT *, + ROW_NUMBER() OVER (PARTITION BY pk ORDER BY val) as rn, + SUM(val) OVER (PARTITION BY pk ORDER BY val) as running_sum + FROM window_topn_t +) WHERE 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, sum(window_topn_t.val) PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@4 as running_sum] +02)--FilterExec: 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 <= 3 +03)----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, sum(window_topn_t.val) PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "sum(window_topn_t.val) PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": nullable Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] +04)------SortExec: expr=[pk@1 ASC NULLS LAST, val@2 ASC NULLS LAST], preserve_partitioning=[false] +05)--------DataSourceExec: partitions=1, partition_sizes=[1] + +# Sibling safety: a LEAD sibling reads following (pruned) rows, so the +# rule must NOT fire — the boundary row's LEAD would resolve to a pruned +# row. The plan keeps SortExec + FilterExec. +query TT +EXPLAIN SELECT * FROM ( + SELECT *, + ROW_NUMBER() OVER (PARTITION BY pk ORDER BY val) as rn, + LEAD(val) OVER (PARTITION BY pk ORDER BY val) as next_val + FROM window_topn_t +) WHERE 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, lead(window_topn_t.val) PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@4 as next_val] +02)--FilterExec: 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 <= 3 +03)----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, lead(window_topn_t.val) PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "lead(window_topn_t.val) PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": nullable Int32 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] +04)------SortExec: expr=[pk@1 ASC NULLS LAST, val@2 ASC NULLS LAST], preserve_partitioning=[false] +05)--------DataSourceExec: partitions=1, partition_sizes=[1] + statement ok SET datafusion.explain.physical_plan_only = false; +# Sibling safety (correctness): LEAD reads following rows, so if the rule +# fired, the pruned input would give a wrong LEAD at the retained boundary +# — each rn=2 row's next_val would become NULL instead of the (pruned) +# rn=3 row's value. The guard keeps the normal plan, so LEAD sees the full +# partition and next_val is correct (30 / 25 / 100 for the rn=2 rows). +query IIIII rowsort +SELECT id, pk, val, rn, next_val FROM ( + SELECT *, + ROW_NUMBER() OVER (PARTITION BY pk ORDER BY val) as rn, + LEAD(val) OVER (PARTITION BY pk ORDER BY val) as next_val + FROM window_topn_t +) WHERE rn <= 2; +---- +1 1 10 1 20 +10 3 75 2 100 +2 1 20 2 30 +5 2 5 1 15 +6 2 15 2 25 +9 3 50 1 75 + +# A tie table for sibling-safety correctness: pk=1 has three peers at +# val=5, one 10, one 20 (row_number 1..5; rank 1,1,1,4,5; dense_rank +# 1,1,1,2,3). +statement ok +CREATE TABLE window_topn_sib_t (id INT, pk INT, val INT) AS VALUES + (1, 1, 5), + (2, 1, 5), + (3, 1, 5), + (4, 1, 10), + (5, 1, 20); + +# SUM sibling (correctness): fetch=2 but there are three peers at val=5, +# so one peer would be pruned. The guard keeps the normal plan, so the +# RANGE frame sees all three 5s and running_sum = 15. If the rule fired, +# the pruned 2-row input would sum to 10. (id excluded — which two of the +# three tied rows get row_number 1/2 is not deterministic.) +query III rowsort +SELECT pk, val, running_sum FROM ( + SELECT *, + ROW_NUMBER() OVER (PARTITION BY pk ORDER BY val) as rn, + SUM(val) OVER (PARTITION BY pk ORDER BY val) as running_sum + FROM window_topn_sib_t +) WHERE rn <= 2; +---- +1 5 15 +1 5 15 + +# All three ranking siblings (ROW_NUMBER + RANK + DENSE_RANK) over the same +# window are prune-safe, so the rule DOES fire (filter on rn). Verify the +# rank/dense_rank values stay correct under the optimized plan — including +# the val=10 row where row_number=4, rank=4, dense_rank=2 all differ. +query IIII rowsort +SELECT pk, val, rnk, dr FROM ( + SELECT *, + ROW_NUMBER() OVER (PARTITION BY pk ORDER BY val) as rn, + RANK() OVER (PARTITION BY pk ORDER BY val) as rnk, + DENSE_RANK() OVER (PARTITION BY pk ORDER BY val) as dr + FROM window_topn_sib_t +) WHERE rn <= 4; +---- +1 10 4 2 +1 5 1 1 +1 5 1 1 +1 5 1 1 + +statement ok +DROP TABLE window_topn_sib_t; + # Test 15: ROW_NUMBER with DESC ordering — correctness query III rowsort SELECT id, pk, val FROM ( @@ -1292,6 +1398,34 @@ SELECT id, pk, val FROM ( ) WHERE dr < 1; ---- +# Test DR11b: ORDER BY keys fully covered by the PARTITION BY prefix +# (`PARTITION BY pk ORDER BY pk`). The deduplicated sort ordering is just +# `[pk]`, so no order key survives past the partition prefix. The rule +# must NOT fire — a PartitionedTopKExec built with zero order expressions +# panics on execution. Every row ties at dense_rank 1, so `dr <= 2` keeps +# the whole table (returned via the ordinary window + filter plan). +query III rowsort +SELECT id, pk, val FROM ( + SELECT *, DENSE_RANK() OVER (PARTITION BY pk ORDER BY pk) as dr FROM window_topn_dense_rank_t +) WHERE dr <= 2; +---- +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 + statement ok DROP TABLE window_topn_dense_rank_t; From a74ed66cf4381094ffea4d44ee5ee25a9b02fbf8 Mon Sep 17 00:00:00 2001 From: SubhamSinghal Date: Thu, 30 Jul 2026 15:20:38 +0530 Subject: [PATCH 3/6] Replace max_keys with binaryHeap --- datafusion/physical-plan/src/topk/mod.rs | 67 +++++++++++------------- 1 file changed, 30 insertions(+), 37 deletions(-) diff --git a/datafusion/physical-plan/src/topk/mod.rs b/datafusion/physical-plan/src/topk/mod.rs index 80d1830347b51..04f6aadd0f643 100644 --- a/datafusion/physical-plan/src/topk/mod.rs +++ b/datafusion/physical-plan/src/topk/mod.rs @@ -1819,14 +1819,13 @@ impl PartitionedTopKRank { /// `RecordBatch`. `TieEntry` is reused verbatim from [`RankPartitionState`]. struct DenseRankPartitionState { groups: HashMap, Vec>, - /// Cached admission boundary: the largest tracked ob value once - /// `groups` is full. A `HashMap` is unordered, so finding it otherwise - /// costs an O(K) scan of the keys on *every* new distinct ob value — - /// O(N·K) overall on mostly-distinct input. Caching makes the common - /// "new ob is worse than the boundary → drop" path O(1); the O(K) scan - /// is paid only when the cache is stale (`None`) — after a fill or an - /// eviction changes the key set. - max_key: Option>, + /// The same keys as `groups`, in a max-heap: the admission boundary + /// (the largest tracked ob value) is an O(1) `peek()` check, and + /// admission / removal are O(log K). + /// + /// INVARIANT: `keys` and `groups.keys()` hold the same set. Every + /// insertion into / removal from `groups` must mirror into `keys`. + keys: BinaryHeap>, } impl DenseRankPartitionState { @@ -1847,7 +1846,11 @@ impl DenseRankPartitionState { .sum::() }) .sum(); - table_overhead + contents + self.max_key.as_ref().map_or(0, |k| k.capacity()) + // `keys` duplicates every key's bytes; charge for them plus the + // heap's backing Vec (one `Vec` slot per reserved element). + let keys_overhead: usize = self.keys.capacity() * size_of::>() + + self.keys.iter().map(|k| k.capacity()).sum::(); + table_overhead + contents + keys_overhead } } @@ -1878,8 +1881,8 @@ impl DenseRankPartitionState { /// - `ob_key` new, `state.groups.len() < k` → insert the run as a new /// group. /// - `ob_key` new, `state.groups.len() == k` → the largest tracked ob -/// value is the admission boundary, cached in `max_key` (recomputed via -/// an O(K) scan of `state.groups.keys()` only when the cache is stale): +/// value is the admission boundary, read from the `state.keys` max-heap +/// in O(1): /// - `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. @@ -2033,7 +2036,7 @@ impl PartitionedTopKDenseRank { .entry(pk) .or_insert_with(|| DenseRankPartitionState { groups: HashMap::new(), - max_key: None, + keys: BinaryHeap::with_capacity(k), }); // Bucket by ob key. `ob_runs` is a reused scratch map (taken @@ -2063,6 +2066,7 @@ impl PartitionedTopKDenseRank { // Case B: new ob, room available. if state.groups.len() < k { + state.keys.push(ob_key.clone()); state.groups.insert( ob_key, vec![TieEntry { @@ -2071,34 +2075,23 @@ impl PartitionedTopKDenseRank { batch_bytes: input_batch_bytes, }], ); - // A new distinct key may raise the boundary; drop the - // cached max so it is recomputed on the next Case C. - state.max_key = None; continue; } // Case C: new ob, at K distinct keys. The largest tracked - // ob value is the admission boundary; it is cached in - // `max_key` and recomputed via an O(K) scan only when the - // cache is stale, so the common "ob >= max → drop" path is - // O(1). Scoped so the immutable borrow ends before mutation. - if state.max_key.is_none() { - state.max_key = state.groups.keys().max().cloned(); - } - let is_smaller = state - .max_key - .as_deref() - .map(|max_key| ob_key.as_slice() < max_key) - .expect("state.groups has k >= 1 keys"); - - if is_smaller { - // Evict the entire max-key group and drop the cached - // max (recomputed on the next Case C). - let evicted_key = state.max_key.take().expect("max key present"); - let evicted = - state.groups.remove(&evicted_key).expect("max key present"); + // ob value is the admission boundary. + let max_key = state.keys.peek().expect("state.groups has k >= 1 keys"); + if ob_key.as_slice() < max_key.as_slice() { + // Evict the entire max-key group, from both the map + // and its ordered mirror. + let evicted_key = state.keys.pop().expect("max key present"); + let evicted = state + .groups + .remove(&evicted_key) + .expect("keys mirrors groups"); replacements += evicted.iter().map(|e| e.row_indices.len()).sum::(); + state.keys.push(ob_key.clone()); state.groups.insert( ob_key, vec![TieEntry { @@ -2157,10 +2150,10 @@ impl PartitionedTopKDenseRank { let mut coalescer = BatchCoalescer::new(Arc::clone(&schema), batch_size); for pk in sorted_pks { - let DenseRankPartitionState { groups, max_key: _ } = + let DenseRankPartitionState { groups, keys: _ } = 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). + // 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)); From 8e97623c98969a34d9991c4920e479985b0b89e8 Mon Sep 17 00:00:00 2001 From: SubhamSinghal Date: Thu, 30 Jul 2026 19:20:34 +0530 Subject: [PATCH 4/6] Prevents memory over accounting --- datafusion/physical-plan/src/topk/mod.rs | 257 +++++++++++++++++++---- 1 file changed, 219 insertions(+), 38 deletions(-) diff --git a/datafusion/physical-plan/src/topk/mod.rs b/datafusion/physical-plan/src/topk/mod.rs index 04f6aadd0f643..bab5d7d27f502 100644 --- a/datafusion/physical-plan/src/topk/mod.rs +++ b/datafusion/physical-plan/src/topk/mod.rs @@ -1811,14 +1811,33 @@ impl PartitionedTopKRank { } } +/// A run of rows from a single source [`RecordBatch`] sharing one +/// distinct ORDER BY value. Materialized at emit time via +/// [`take_record_batch`]. +/// +/// The batch is referenced by `batch_id` rather than held directly, so the +/// operator-scoped [`RecordBatchStore`] can charge each distinct source +/// batch once however many entries reference it. Holding a batch per entry +/// and charging its bytes per entry would inflate the reservation by a +/// factor of (partitions × K), since a single batch contributes an entry to +/// every partition and ob group it touches. +#[derive(Debug)] +struct GroupEntry { + /// Indices into the batch identified by `batch_id`. Always non-empty + /// by construction. + row_indices: Vec, + /// Key into `PartitionedTopKDenseRank::store`. + batch_id: u32, +} + /// 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`]. +/// 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`. struct DenseRankPartitionState { - groups: HashMap, Vec>, + groups: HashMap, Vec>, /// The same keys as `groups`, in a max-heap: the admission boundary /// (the largest tracked ob value) is an O(1) `peek()` check, and /// admission / removal are O(log K). @@ -1830,19 +1849,17 @@ struct DenseRankPartitionState { impl DenseRankPartitionState { fn size(&self) -> usize { - let table_overhead = - self.groups.capacity() * (size_of::>() + size_of::>()); + 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.capacity() * size_of::() + entries .iter() - .map(|e| { - e.row_indices.capacity() * size_of::() + e.batch_bytes - }) + .map(|e| e.row_indices.capacity() * size_of::()) .sum::() }) .sum(); @@ -1864,7 +1881,10 @@ impl DenseRankPartitionState { /// /// Like [`PartitionedTopK`], the [`RowConverter`], [`MemoryReservation`], /// scratch [`Rows`] buffer, and [`TopKMetrics`] are shared across all -/// partitions for this operator instance. +/// partitions for this operator instance. So is the +/// [`RecordBatchStore`]: retained rows are held as `(batch_id, indices)` +/// so each source batch is charged once for the whole operator, however +/// many partitions and ob groups reference it. /// /// # Algorithm (per batch) /// @@ -1872,12 +1892,12 @@ impl DenseRankPartitionState { /// 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 +/// bucket is built from the current batch's rows, so each `GroupEntry` 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). +/// new `GroupEntry` (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` → the largest tracked ob @@ -1920,6 +1940,11 @@ pub(crate) struct PartitionedTopKDenseRank { /// reallocated) per partition so its backing table is allocated once, /// not once per distinct partition key. ob_runs: HashMap, Vec>, + /// Source batches referenced by the retained `GroupEntry`s, held once + /// for the whole operator with a use count per batch. This is what + /// keeps the reservation proportional to the batches actually pinned + /// rather than to the number of entries pointing at them. + store: RecordBatchStore, k: usize, batch_size: usize, } @@ -1964,6 +1989,7 @@ impl PartitionedTopKDenseRank { states: HashMap::new(), partition_groups: HashMap::new(), ob_runs: HashMap::new(), + store: RecordBatchStore::new(), k, batch_size, }) @@ -1972,7 +1998,7 @@ impl PartitionedTopKDenseRank { /// 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`]. + /// state as one [`GroupEntry`]. pub(crate) fn insert_batch(&mut self, batch: &RecordBatch) -> Result<()> { let baseline = self.metrics.baseline.clone(); let _timer = baseline.elapsed_compute().timer(); @@ -1982,9 +2008,13 @@ impl PartitionedTopKDenseRank { 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); + // Register this batch with the store up front. `uses` is bumped + // once per `GroupEntry` created below but lives on this local + // entry, so the store sees a single insert per batch rather than + // one per group — and `insert` drops batches that retained + // nothing without ever charging for them. + let mut batch_entry = self.store.register(batch.clone()); + let batch_id = batch_entry.id; // 1. Encode partition columns. let pk_arrays: Vec = self @@ -2029,7 +2059,7 @@ impl PartitionedTopKDenseRank { // 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`. + // partition state as a single `GroupEntry`. for (pk, indices) in groups.drain() { let state = self.states @@ -2053,26 +2083,25 @@ impl PartitionedTopKDenseRank { 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). + // new `GroupEntry` (one entry per contributing batch). if let Some(entries) = state.groups.get_mut(&ob_key) { - entries.push(TieEntry { - batch: batch.clone(), + batch_entry.uses += 1; + entries.push(GroupEntry { row_indices: run_indices, - batch_bytes: input_batch_bytes, + batch_id, }); continue; } // Case B: new ob, room available. if state.groups.len() < k { + batch_entry.uses += 1; state.keys.push(ob_key.clone()); state.groups.insert( ob_key, - vec![TieEntry { - batch: batch.clone(), + vec![GroupEntry { row_indices: run_indices, - batch_bytes: input_batch_bytes, + batch_id, }], ); continue; @@ -2089,15 +2118,25 @@ impl PartitionedTopKDenseRank { .groups .remove(&evicted_key) .expect("keys mirrors groups"); - replacements += - evicted.iter().map(|e| e.row_indices.len()).sum::(); + for e in &evicted { + replacements += e.row_indices.len(); + if e.batch_id == batch_id { + // Admitted earlier in this same call, so the + // store has not seen `batch_entry` yet — drop + // the pending use rather than calling `unuse`, + // which panics on an unregistered id. + batch_entry.uses -= 1; + } else { + self.store.unuse(e.batch_id); + } + } + batch_entry.uses += 1; state.keys.push(ob_key.clone()); state.groups.insert( ob_key, - vec![TieEntry { - batch: batch.clone(), + vec![GroupEntry { row_indices: run_indices, - batch_bytes: input_batch_bytes, + batch_id, }], ); } @@ -2113,6 +2152,9 @@ impl PartitionedTopKDenseRank { // batch to reuse. self.partition_groups = groups; + // Charges `batch` once if any group retained rows from it. + self.store.insert(batch_entry); + if replacements > 0 { self.metrics.row_replacements.add(replacements); } @@ -2139,6 +2181,7 @@ impl PartitionedTopKDenseRank { mut states, partition_groups: _, ob_runs: _, + store, k: _, batch_size, } = self; @@ -2154,13 +2197,17 @@ impl PartitionedTopKDenseRank { states.remove(&pk).expect("key from states.keys()"); // Sort the <= K distinct ob keys so rows emit ascending // (byte-comparable encoding == sort order). - let mut sorted_obs: Vec<(Vec, Vec)> = + 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 batch = &store + .get(entry.batch_id) + .expect("retained batch_id present in store") + .batch; let indices = UInt32Array::from(entry.row_indices); - let sub = take_record_batch(&entry.batch, &indices)?; + let sub = take_record_batch(batch, &indices)?; (&sub).record_output(&metrics.baseline); coalescer.push_batch(sub)?; } @@ -2189,6 +2236,7 @@ impl PartitionedTopKDenseRank { + self.states.values().map(|s| s.size()).sum::() + self.states.capacity() * (size_of::>() + size_of::()) + + self.store.size() } } @@ -3552,7 +3600,7 @@ mod tests { // PartitionedTopKDenseRank operator tests // // These mirror the RANK tests plus DENSE_RANK-specific cases: rows - // sharing an ob key coalesce into one `TieEntry`, unbounded + // sharing an ob key coalesce into one `GroupEntry`, unbounded // rows-per-distinct-key, and eviction removes the entire max group // when a strictly-smaller distinct ob arrives. // ==================================================================== @@ -3638,7 +3686,7 @@ mod tests { /// 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. + /// into one `GroupEntry` per distinct ob. /// /// vals per partition (sorted logically): /// pk=1: 1, 1, 1, 2, 2, 3, 3, 3, 4 @@ -3673,7 +3721,7 @@ mod tests { } /// Rows tied at the same ob across two source batches must both - /// land under the same map key as separate `TieEntry`s — one per + /// land under the same map key as separate `GroupEntry`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<()> { @@ -3706,7 +3754,7 @@ mod tests { /// 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 + /// and cross-batch eviction of a whole max group. Every `GroupEntry` 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. @@ -4022,4 +4070,137 @@ mod tests { ); Ok(()) } + + /// Total `GroupEntry` count across all partitions. + fn dense_rank_entry_count(state: &PartitionedTopKDenseRank) -> usize { + state + .states + .values() + .flat_map(|s| s.groups.values()) + .map(|entries| entries.len()) + .sum() + } + + /// One source batch feeding many retained groups must be charged + /// once, not once per group. + /// + /// Dense-rank retains up to K distinct-ob groups per partition and + /// each can draw rows from the same batch, so charging per entry + /// inflates the reservation by (partitions × K) — here 6× — and can + /// trip a spurious `ResourcesExhausted`. + #[tokio::test] + async fn test_partitioned_topk_dense_rank_charges_batch_once() -> Result<()> { + let (schema, mut state) = build_partitioned_topk_dense_rank(3)?; + + // pk=1 retains {1,2,3}, pk=2 retains {10,20,30}: 6 groups, all + // from this one batch. + let batch = pk_val_batch( + &schema, + vec![1, 1, 1, 1, 2, 2, 2, 2], + vec![1, 2, 3, 4, 10, 20, 30, 40], + )?; + let batch_bytes = get_record_batch_memory_size(&batch); + state.insert_batch(&batch)?; + + assert_eq!(dense_rank_entry_count(&state), 6); + assert_eq!(state.store.len(), 1); + assert_eq!(state.store.batches_size, batch_bytes); + Ok(()) + } + + /// Evicting the last group referencing a batch must release the + /// batch's bytes, or the reservation only ever grows. + #[tokio::test] + async fn test_partitioned_topk_dense_rank_releases_evicted_batch() -> Result<()> { + let (schema, mut state) = build_partitioned_topk_dense_rank(2)?; + + let first = pk_val_batch(&schema, vec![1, 1], vec![50, 60])?; + state.insert_batch(&first)?; + assert_eq!(state.store.len(), 1); + + // Both values beat {50, 60}, so every group from `first` is + // evicted and only `second` remains charged. + let second = pk_val_batch(&schema, vec![1, 1], vec![5, 6])?; + let second_bytes = get_record_batch_memory_size(&second); + state.insert_batch(&second)?; + + assert_eq!(state.store.len(), 1); + assert_eq!(state.store.batches_size, second_bytes); + + let results: Vec<_> = state.emit()?.try_collect().await?; + assert_batches_eq!( + &[ + "+----+-----+", + "| pk | val |", + "+----+-----+", + "| 1 | 5 |", + "| 1 | 6 |", + "+----+-----+", + ], + &results + ); + Ok(()) + } + + /// The mirror of the above: a batch whose every run is rejected must + /// not be charged at all. + /// + /// Nothing references it, so nothing would ever release it — charging + /// it would pin both the bytes and the batch for the operator's + /// lifetime. + #[tokio::test] + async fn test_partitioned_topk_dense_rank_ignores_fully_rejected_batch() -> Result<()> + { + let (schema, mut state) = build_partitioned_topk_dense_rank(2)?; + + let first = pk_val_batch(&schema, vec![1, 1], vec![5, 6])?; + let first_bytes = get_record_batch_memory_size(&first); + state.insert_batch(&first)?; + assert_eq!(state.store.len(), 1); + + // At K=2 with {5, 6} tracked, both values lose to the boundary, so + // no `GroupEntry` points at `second`. + state.insert_batch(&pk_val_batch(&schema, vec![1, 1], vec![50, 60])?)?; + + assert_eq!(dense_rank_entry_count(&state), 2); + assert_eq!(state.store.len(), 1); + assert_eq!(state.store.batches_size, first_bytes); + Ok(()) + } + + /// A group admitted and then evicted within the *same* + /// `insert_batch` call: the batch is still pending (not yet handed to + /// the store), so releasing it must decrement the in-flight use count + /// rather than call `unuse` on an unregistered id. + /// + /// `ob_runs` drains in hash order, so with K=1 and many distinct + /// values the minimum is almost never seen first and the run + /// admit-then-evict path is taken repeatedly. + #[tokio::test] + async fn test_partitioned_topk_dense_rank_evicts_same_call_group() -> Result<()> { + let (schema, mut state) = build_partitioned_topk_dense_rank(1)?; + + let pks = vec![1; 32]; + let vals: Vec = (0..32).rev().collect(); + let batch = pk_val_batch(&schema, pks, vals)?; + let batch_bytes = get_record_batch_memory_size(&batch); + state.insert_batch(&batch)?; + + assert_eq!(dense_rank_entry_count(&state), 1); + assert_eq!(state.store.len(), 1); + assert_eq!(state.store.batches_size, batch_bytes); + + let results: Vec<_> = state.emit()?.try_collect().await?; + assert_batches_eq!( + &[ + "+----+-----+", + "| pk | val |", + "+----+-----+", + "| 1 | 0 |", + "+----+-----+", + ], + &results + ); + Ok(()) + } } From afc931c10aa0d15445bb8b49f4779e03a7d79ac0 Mon Sep 17 00:00:00 2001 From: SubhamSinghal Date: Thu, 30 Jul 2026 19:42:49 +0530 Subject: [PATCH 5/6] Remove eager heap allocation --- datafusion/physical-plan/src/topk/mod.rs | 39 +++++++++++++++++++----- 1 file changed, 32 insertions(+), 7 deletions(-) diff --git a/datafusion/physical-plan/src/topk/mod.rs b/datafusion/physical-plan/src/topk/mod.rs index bab5d7d27f502..1fb64c15b280f 100644 --- a/datafusion/physical-plan/src/topk/mod.rs +++ b/datafusion/physical-plan/src/topk/mod.rs @@ -1836,6 +1836,7 @@ struct GroupEntry { /// 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`. +#[derive(Default)] struct DenseRankPartitionState { groups: HashMap, Vec>, /// The same keys as `groups`, in a max-heap: the admission boundary @@ -2061,13 +2062,10 @@ impl PartitionedTopKDenseRank { // (within-call accumulation), then merge each bucket into the // partition state as a single `GroupEntry`. for (pk, indices) in groups.drain() { - let state = - self.states - .entry(pk) - .or_insert_with(|| DenseRankPartitionState { - groups: HashMap::new(), - keys: BinaryHeap::with_capacity(k), - }); + let state = self + .states + .entry(pk) + .or_insert_with(DenseRankPartitionState::default); // Bucket by ob key. `ob_runs` is a reused scratch map (taken // out and drained below) so its backing table is allocated @@ -4168,6 +4166,33 @@ mod tests { Ok(()) } + /// `keys` must grow on demand rather than reserve K slots when a + /// partition is first seen. `size()` charges `keys.capacity()`, so + /// eager sizing reserves O(partitions * K) for slots that never hold + /// a key — enough to fail a memory limit on a high-cardinality input + /// whose partitions each keep a handful of distinct values. + #[tokio::test] + async fn test_partitioned_topk_dense_rank_heap_grows_on_demand() -> Result<()> { + const K: usize = 512; + const PARTITIONS: usize = 64; + let (schema, mut state) = build_partitioned_topk_dense_rank(K)?; + + // One row per partition: every partition holds exactly one + // distinct ob value, K - 1 slots short of capacity. + let pks: Vec = (0..PARTITIONS as i32).collect(); + let vals = pks.clone(); + state.insert_batch(&pk_val_batch(&schema, pks, vals)?)?; + + assert_eq!(state.states.len(), PARTITIONS); + let heap_slots: usize = state.states.values().map(|s| s.keys.capacity()).sum(); + // Eager `with_capacity(K)` would reserve PARTITIONS * K = 32768. + assert!( + heap_slots <= PARTITIONS * 8, + "reserved {heap_slots} heap slots to hold {PARTITIONS} keys" + ); + Ok(()) + } + /// A group admitted and then evicted within the *same* /// `insert_batch` call: the batch is still pending (not yet handed to /// the store), so releasing it must decrement the in-flight use count From c7e67c9f42330563ad21f8ee8e9d82de360b6b6c Mon Sep 17 00:00:00 2001 From: SubhamSinghal Date: Fri, 31 Jul 2026 15:30:25 +0530 Subject: [PATCH 6/6] Update size allocation --- datafusion/physical-plan/src/topk/mod.rs | 101 ++++++++++++++++++++++- 1 file changed, 99 insertions(+), 2 deletions(-) diff --git a/datafusion/physical-plan/src/topk/mod.rs b/datafusion/physical-plan/src/topk/mod.rs index 1fb64c15b280f..1d02397b92229 100644 --- a/datafusion/physical-plan/src/topk/mod.rs +++ b/datafusion/physical-plan/src/topk/mod.rs @@ -2226,14 +2226,31 @@ impl PartitionedTopKDenseRank { /// Total memory currently held, including all per-partition states. fn size(&self) -> usize { + // Per partition: the state itself plus the encoded partition key + // owned by the map. The key bytes are a heap allocation the table + // slot doesn't cover, and with wide or numerous partition keys + // they dominate the fixed-size slots. + let states_contents: usize = self + .states + .iter() + .map(|(pk, state)| pk.capacity() + state.size()) + .sum(); + // `partition_groups` and `ob_runs` are drained, not dropped, so + // their backing tables outlive every `insert_batch` call. Both are + // empty by the time `size()` runs (drained above), so only the + // retained capacity is charged. + let scratch_tables = self.partition_groups.capacity() + * (size_of::>() + size_of::>()) + + self.ob_runs.capacity() * (size_of::>() + size_of::>()); 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::() + + states_contents + self.states.capacity() * (size_of::>() + size_of::()) + + scratch_tables + self.store.size() } } @@ -2241,7 +2258,7 @@ impl PartitionedTopKDenseRank { #[cfg(test)] mod tests { use super::*; - use arrow::array::{BooleanArray, Float64Array, Int32Array}; + use arrow::array::{BooleanArray, Float64Array, Int32Array, StringArray}; use arrow::datatypes::{DataType, Field, Schema}; use arrow_schema::SortOptions; use datafusion_common::assert_batches_eq; @@ -4193,6 +4210,86 @@ mod tests { Ok(()) } + /// The encoded partition keys owned by `states`, and the backing + /// tables the drained scratch maps keep, are long-lived heap + /// allocations `size()` must charge: they persist for the operator's + /// life yet belong to no `GroupEntry`, so per-entry accounting can't + /// see them. With numerous or wide partition keys the key bytes are + /// the larger term. + #[tokio::test] + async fn test_partitioned_topk_dense_rank_size_covers_keys_and_scratch() -> Result<()> + { + const PARTITIONS: usize = 64; + const KEY_WIDTH: usize = 1024; + + let schema = Arc::new(Schema::new(vec![ + Field::new("pk", DataType::Utf8, false), + Field::new("val", DataType::Int32, false), + ])); + let pk_expr: Arc = col("pk", schema.as_ref())?; + let partition_sort_fields = build_sort_fields( + &[PhysicalSortExpr { + expr: Arc::clone(&pk_expr), + options: SortOptions::default(), + }], + &schema, + )?; + let order_expr = LexOrdering::from([PhysicalSortExpr { + expr: col("val", schema.as_ref())?, + options: SortOptions::default(), + }]); + let mut state = PartitionedTopKDenseRank::try_new( + 0, + Arc::clone(&schema), + vec![pk_expr], + partition_sort_fields, + order_expr, + 4, + 8, // batch_size + &Arc::new(RuntimeEnv::default()), + &ExecutionPlanMetricsSet::new(), + )?; + + // One row per partition, each with a wide key. + let pks: Vec = (0..PARTITIONS) + .map(|i| format!("{}{i:04}", "p".repeat(KEY_WIDTH - 4))) + .collect(); + let vals: Vec = (0..PARTITIONS as i32).collect(); + state.insert_batch(&RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(StringArray::from(pks)), + Arc::new(Int32Array::from(vals)), + ], + )?)?; + assert_eq!(state.states.len(), PARTITIONS); + + let key_bytes: usize = state.states.keys().map(|pk| pk.capacity()).sum(); + let scratch_bytes = state.partition_groups.capacity() + * (size_of::>() + size_of::>()) + + state.ob_runs.capacity() * (size_of::>() + size_of::>()); + assert!(key_bytes >= PARTITIONS * KEY_WIDTH, "key bytes {key_bytes}"); + assert!(scratch_bytes > 0, "scratch tables never allocated"); + + // Reconstruct the total from its parts. Both terms above have to + // appear for this to balance, so dropping either from `size()` + // fails here rather than being absorbed by the slack in some + // other term. + let expected = size_of::() + + state.row_converter.size() + + state.partition_converter.size() + + state.scratch_rows.size() + + state.partition_scratch_rows.size() + + key_bytes + + state.states.values().map(|s| s.size()).sum::() + + state.states.capacity() + * (size_of::>() + size_of::()) + + scratch_bytes + + state.store.size(); + assert_eq!(state.size(), expected); + Ok(()) + } + /// A group admitted and then evicted within the *same* /// `insert_batch` call: the batch is still pending (not yet handed to /// the store), so releasing it must decrement the in-flight use count