diff --git a/datafusion/physical-plan/src/topk/mod.rs b/datafusion/physical-plan/src/topk/mod.rs index 030614eed02e5..79fd2eb7a8346 100644 --- a/datafusion/physical-plan/src/topk/mod.rs +++ b/datafusion/physical-plan/src/topk/mod.rs @@ -23,7 +23,7 @@ use arrow::{ BatchCoalescer, FilterBuilder, interleave_record_batch, prep_null_mask_filter, take_record_batch, }, - row::{OwnedRow, RowConverter, Rows, SortField}, + row::{RowConverter, Rows, SortField}, }; use datafusion_expr::{ColumnarValue, Operator}; use std::mem::size_of; @@ -1272,8 +1272,20 @@ pub(crate) struct PartitionedTopK { partition_exprs: Vec>, /// Encoder for the partition key. partition_converter: RowConverter, - /// One heap per distinct partition key seen so far. - heaps: HashMap, + /// Scratch row buffer for partition-key encoding. Reused across + /// `insert_batch` calls (cleared + appended each batch). + partition_scratch_rows: Rows, + /// One heap per distinct partition key seen so far. Keyed by the + /// row-encoded PARTITION BY bytes (a 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. + heaps: HashMap, TopKHeap>, + /// 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>, k: usize, batch_size: usize, } @@ -1301,6 +1313,8 @@ impl PartitionedTopK { 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, @@ -1311,7 +1325,9 @@ impl PartitionedTopK { scratch_rows, partition_exprs, partition_converter, + partition_scratch_rows, heaps: HashMap::new(), + partition_groups: HashMap::new(), k, batch_size, }) @@ -1335,15 +1351,26 @@ impl PartitionedTopK { .iter() .map(|e| e.evaluate(batch).and_then(|v| v.into_array(num_rows))) .collect::>()?; - let pk_rows = self.partition_converter.convert_columns(&pk_arrays)?; + self.partition_scratch_rows.clear(); + self.partition_converter + .append(&mut self.partition_scratch_rows, &pk_arrays)?; // 2. Demultiplex row indices by partition key (per-batch). - let mut groups: HashMap> = HashMap::new(); - for i in 0..num_rows { - groups - .entry(pk_rows.row(i).owned()) - .or_default() - .push(i as u32); + // `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 on the full batch and encode ONCE. @@ -1360,7 +1387,7 @@ impl PartitionedTopK { // qualifying rows into the partition's heap. let k = self.k; let mut replacements: usize = 0; - for (pk, indices) in groups { + for (pk, indices) in groups.drain() { let heap = self.heaps.entry(pk).or_insert_with(|| TopKHeap::new(k)); // Once a heap is full, most rows at high partition cardinality @@ -1396,6 +1423,10 @@ impl PartitionedTopK { heap.maybe_compact()?; } + // 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); } @@ -1416,13 +1447,15 @@ impl PartitionedTopK { scratch_rows: _, partition_exprs: _, partition_converter: _, + partition_scratch_rows: _, mut heaps, + partition_groups: _, k: _, batch_size, } = self; let _timer = metrics.baseline.elapsed_compute().timer(); - let mut sorted_pks: Vec = heaps.keys().cloned().collect(); + let mut sorted_pks: Vec> = heaps.keys().cloned().collect(); sorted_pks.sort(); let mut coalescer = BatchCoalescer::new(Arc::clone(&schema), batch_size); @@ -1450,24 +1483,39 @@ impl PartitionedTopK { /// Total memory currently held by this operator, including all /// per-partition heaps. fn size(&self) -> usize { + // Per partition: the heap plus the encoded partition key owned by + // the map. The key bytes are a heap allocation the table slot + // doesn't cover. + let heaps_contents: usize = self + .heaps + .iter() + .map(|(pk, heap)| pk.capacity() + heap.size()) + .sum(); size_of::() + self.row_converter.size() + self.partition_converter.size() + self.scratch_rows.size() - + self.heaps.values().map(|h| h.size()).sum::() - + self.heaps.capacity() * (size_of::() + size_of::()) + + self.partition_scratch_rows.size() + + heaps_contents + + self.heaps.capacity() * (size_of::>() + size_of::()) + + self.partition_groups.capacity() + * (size_of::>() + size_of::>()) } } -/// A run of rows from a single source [`RecordBatch`] that tied at the -/// boundary when inserted. Stored as `(batch, indices)` and materialized -/// at emit time via [`take_record_batch`]. +/// Rows that tied at the boundary when inserted, materialized into a +/// batch holding *only* those rows. +/// +/// The rows are gathered eagerly rather than kept as `(source_batch, +/// indices)`: a tie entry lives until the boundary moves, so holding the +/// source batch would pin an entire input batch — and charge for it — +/// for as long as a single row of it stays tied. With ties spread across +/// many input batches that makes retained memory grow with the *input* +/// size instead of with `K + ties`. #[derive(Debug)] struct TieEntry { + /// The tied rows, and nothing else. Always non-empty by construction. batch: RecordBatch, - /// Indices into `batch` of the rows tied at the (then-current) - /// boundary. Always non-empty by construction. - row_indices: Vec, /// `get_record_batch_memory_size(&batch)` captured at push time so /// `RankPartitionState::size()` doesn't recurse through `batch`'s /// columns on every `try_resize` call. @@ -1489,13 +1537,15 @@ struct RankPartitionState { impl RankPartitionState { fn size(&self) -> usize { let ties_buffer = self.ties.capacity() * size_of::(); - let ties_contents: usize = self - .ties - .iter() - .map(|t| t.row_indices.capacity() * size_of::() + t.batch_bytes) - .sum(); + let ties_contents: usize = self.ties.iter().map(|t| t.batch_bytes).sum(); self.heap.size() + ties_buffer + ties_contents } + + /// Push `batch`'s rows onto the tie list, charging exactly their bytes. + fn push_ties(&mut self, batch: RecordBatch) { + let batch_bytes = get_record_batch_memory_size(&batch); + self.ties.push(TieEntry { batch, batch_bytes }); + } } /// Sibling to [`PartitionedTopK`] implementing `RANK()` semantics. @@ -1540,8 +1590,17 @@ pub(crate) struct PartitionedTopKRank { /// `insert_batch` calls (cleared + appended each batch) so we /// avoid allocating a fresh `Rows` buffer every batch. partition_scratch_rows: Rows, - /// One rank state per distinct partition key seen so far. - states: HashMap, + /// One rank state per distinct partition key seen so far. Keyed by + /// the row-encoded PARTITION BY bytes (a 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, RankPartitionState>, + /// 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>, k: usize, batch_size: usize, } @@ -1584,6 +1643,7 @@ impl PartitionedTopKRank { partition_converter, partition_scratch_rows, states: HashMap::new(), + partition_groups: HashMap::new(), k, batch_size, }) @@ -1602,12 +1662,6 @@ impl PartitionedTopKRank { return Ok(()); } - // Captured once so the per-tie push from this batch can reuse - // it (computing `get_record_batch_memory_size` is O(cols × - // buffer walk) and we'd otherwise pay it per push and again - // per `try_resize` call). - let input_batch_bytes = get_record_batch_memory_size(batch); - // 1. Evaluate + encode partition columns into the reusable // scratch (cleared then appended). let pk_arrays: Vec = self @@ -1618,15 +1672,23 @@ impl PartitionedTopKRank { self.partition_scratch_rows.clear(); self.partition_converter .append(&mut self.partition_scratch_rows, &pk_arrays)?; - let pk_rows = &self.partition_scratch_rows; // 2. Demultiplex row indices by partition key (per-batch). - let mut groups: HashMap> = HashMap::new(); - for i in 0..num_rows { - groups - .entry(pk_rows.row(i).owned()) - .or_default() - .push(i as u32); + // `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 on the full batch and encode ONCE. @@ -1643,22 +1705,45 @@ impl PartitionedTopKRank { let k = self.k; let mut replacements: usize = 0; - for (pk, indices) in groups { + for (pk, indices) in groups.drain() { let state = self.states.entry(pk).or_insert_with(|| RankPartitionState { heap: TopKHeap::new(k), ties: Vec::new(), }); - // Equal indices for THIS batch only. Coalesced into a single - // tie entry at the end of the partition's loop. Discarded if - // the boundary moves up mid-loop (those rows were tied to the - // old boundary, which is now strictly worse than the new K-th). + // Once the heap is full, a group whose rows are *all* strictly + // worse than the boundary changes neither the heap nor the + // ties. Bail before the gather below — at high partition + // cardinality this is the common case. + if let Some(max_row) = state.heap.max() { + let boundary = max_row.row(); + if indices + .iter() + .all(|&i| self.scratch_rows.row(i as usize).as_ref() > boundary) + { + continue; + } + } + + // Gather this partition's rows into their own batch, as + // `PartitionedTopK` does. Registering the whole input batch + // instead would pin it — and charge for it — once per + // partition key present in the batch, so a batch spanning P + // partitions would be counted P times over. + let indices_arr = UInt32Array::from(indices); + let sub_batch = take_record_batch(batch, &indices_arr)?; + + // Indices *into `sub_batch`* of rows from this batch that tied + // at the boundary. Coalesced into a single tie entry at the end + // of the partition's loop. Discarded if the boundary moves up + // mid-loop (those rows were tied to the old boundary, which is + // now strictly worse than the new K-th). let mut equal_indices: Vec = Vec::new(); // Lazy-registered: only attached if at least one row reaches // the heap from this batch in this partition. - let mut entry: Option = None; + let mut heap_entry: Option = None; - for &orig_idx in &indices { + for (sub_idx, &orig_idx) in indices_arr.values().iter().enumerate() { let row = self.scratch_rows.row(orig_idx as usize); // Classify against the current K-th-best (the heap top). @@ -1671,21 +1756,21 @@ impl PartitionedTopKRank { match classification { Some(Ordering::Equal) => { - equal_indices.push(orig_idx); + equal_indices.push(sub_idx as u32); continue; } Some(Ordering::Greater) => continue, Some(Ordering::Less) | None => { // Heap path: heap not yet full, or row strictly // better than the current boundary. - let entry_ref = entry.get_or_insert_with(|| { - state.heap.register_batch(batch.clone()) + let entry_ref = heap_entry.get_or_insert_with(|| { + state.heap.register_batch(sub_batch.clone()) }); if let Some(EvictedRow { batch: evicted_batch, index: evicted_index, row_bytes: evicted_bytes, - }) = state.heap.add(entry_ref, row, orig_idx as usize) + }) = state.heap.add(entry_ref, row, sub_idx) { // Compare the new boundary (post-eviction heap // top) against the evicted row's bytes — both @@ -1705,16 +1790,14 @@ impl PartitionedTopKRank { state.ties.clear(); equal_indices.clear(); } else { - // Boundary unchanged — evicted row is tied - // at the (unchanged) boundary; push as a - // single-row entry. - let batch_bytes = - get_record_batch_memory_size(&evicted_batch); - state.ties.push(TieEntry { - batch: evicted_batch, - row_indices: vec![evicted_index as u32], - batch_bytes, - }); + // Boundary unchanged — the evicted row is + // still tied at the boundary. Gather just + // that row: holding `evicted_batch` would + // keep a whole heap batch alive for one row, + // and one such entry per input batch would + // again make memory grow with the input. + let one = UInt32Array::from(vec![evicted_index as u32]); + state.push_ties(take_record_batch(&evicted_batch, &one)?); } } replacements += 1; @@ -1722,21 +1805,34 @@ impl PartitionedTopKRank { } } - if let Some(e) = entry { + let registered_with_heap = heap_entry.is_some(); + if let Some(e) = heap_entry { state.heap.insert_batch_entry(e); state.heap.maybe_compact()?; } // Commit this batch's ties as a single entry. if !equal_indices.is_empty() { - state.ties.push(TieEntry { - batch: batch.clone(), - row_indices: equal_indices, - batch_bytes: input_batch_bytes, - }); + // No row of this group reached the heap, so `sub_batch` is + // not registered there and reusing it here cannot + // double-charge it. Combined with every row having tied, + // `sub_batch` already *is* exactly the tie rows — the + // gather below would just copy it. + let tie_batch = if !registered_with_heap + && equal_indices.len() == sub_batch.num_rows() + { + sub_batch + } else { + take_record_batch(&sub_batch, &UInt32Array::from(equal_indices))? + }; + state.push_ties(tie_batch); } } + // 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); } @@ -1761,12 +1857,13 @@ impl PartitionedTopKRank { partition_converter: _, partition_scratch_rows: _, mut states, + partition_groups: _, k: _, batch_size, } = self; let _timer = metrics.baseline.elapsed_compute().timer(); - let mut sorted_pks: Vec = states.keys().cloned().collect(); + let mut sorted_pks: Vec> = states.keys().cloned().collect(); sorted_pks.sort(); let mut coalescer = BatchCoalescer::new(Arc::clone(&schema), batch_size); @@ -1779,10 +1876,8 @@ impl PartitionedTopKRank { coalescer.push_batch(batch)?; } for tie in ties { - let indices = UInt32Array::from(tie.row_indices); - let tie_batch = take_record_batch(&tie.batch, &indices)?; - (&tie_batch).record_output(&metrics.baseline); - coalescer.push_batch(tie_batch)?; + (&tie.batch).record_output(&metrics.baseline); + coalescer.push_batch(tie.batch)?; } } coalescer.finish_buffered_batch()?; @@ -1800,14 +1895,24 @@ impl PartitionedTopKRank { /// Total memory currently held, including all per-partition states. fn size(&self) -> usize { + // Per partition: the state plus the encoded partition key owned by + // the map. The key bytes are a heap allocation the table slot + // doesn't cover. + let states_contents: usize = self + .states + .iter() + .map(|(pk, state)| pk.capacity() + state.size()) + .sum(); 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::()) + * (size_of::>() + size_of::()) + + self.partition_groups.capacity() + * (size_of::>() + size_of::>()) } } @@ -2074,8 +2179,26 @@ impl PartitionedTopKDenseRank { // value don't re-allocate. let mut runs = std::mem::take(&mut self.ob_runs); runs.clear(); + // Once the partition holds its full K distinct ob values, the + // largest of them is the bar a new value must beat, and that bar + // only ever gets stricter: with K values held there is no free + // slot left for a new (possibly larger) one, and the only way in + // from here is to evict that largest and put something strictly + // smaller in its place. So a row worse than today's bar is worse + // than every later bar too — it is not held now and can never be + // admitted — and can be dropped before it costs a bucket. Rows + // *equal* to the bar must still go through: that value is one of + // the K being held, so its rows belong to a live group. + let boundary: Option<&[u8]> = if state.groups.len() == k { + state.keys.peek().map(Vec::as_slice) + } else { + None + }; for &orig_idx in &indices { let ob_row = self.scratch_rows.row(orig_idx as usize); + if boundary.is_some_and(|b| ob_row.as_ref() > b) { + continue; + } runs.entry_ref(ob_row.as_ref()).or_default().push(orig_idx); } @@ -2950,6 +3073,260 @@ mod tests { Ok((schema, state)) } + /// `PartitionedTopKRank` must not charge (or pin) a whole input batch + /// once per partition key that batch touches. + /// + /// Regression: the heap used to be handed `batch.clone()` — the full + /// input batch — rather than a gather of just that partition's rows, + /// so a batch spanning P partitions was counted P times over. With 500 + /// partitions in one batch that reported ~516x the batch size, against + /// ~15x for the equivalent `ROW_NUMBER` state on identical input. + #[tokio::test] + async fn test_partitioned_topk_rank_size_is_not_per_partition_batch() -> Result<()> { + // The over-count factor was exactly the number of partitions sharing + // a batch, so P is what makes the bug visible at all. 500 turns the + // rank/row_number size ratio from ~1x into ~35x — far outside the + // range any accounting change could drift into. + const P: i32 = 500; + // 8 rows per partition: enough for every partition's k=2 heap to fill + // and then evict, so the heaps hold real state rather than sitting + // half-empty in the fill phase. + const ROWS: i32 = 4000; + + // Distinct values throughout: no ties, so RANK retains exactly what + // ROW_NUMBER does and the two sizes are directly comparable. k=2 is + // the smallest k with a distinct fill phase before eviction begins. + let pks: Vec = (0..ROWS).map(|i| i % P).collect(); + let vals: Vec = (0..ROWS).map(|i| i / P).collect(); + + let (schema, mut rn) = build_partitioned_topk(2)?; + rn.insert_batch(&pk_val_batch(&schema, pks.clone(), vals.clone())?)?; + let rn_size = rn.size(); + + let (schema, mut rk) = build_partitioned_topk_rank(2)?; + rk.insert_batch(&pk_val_batch(&schema, pks, vals)?)?; + let rk_size = rk.size(); + + // Measured ratio here is ~1.04 (18.1x vs 17.4x of the input batch); + // the bug produced ~35x. A 2x bar sits more than an order of magnitude + // clear of both, so this neither flakes when the size calculation is + // legitimately adjusted nor misses a reintroduced per-partition factor. + assert!( + rk_size <= rn_size * 2, + "RANK reported {rk_size} bytes for the same retained rows that \ + ROW_NUMBER reported {rn_size} bytes for; a per-partition factor \ + has crept back in" + ); + Ok(()) + } + + /// A boundary tie must retain only the tied rows, not the batch they + /// arrived in. + /// + /// Regression: `TieEntry` used to hold the source `RecordBatch` plus + /// indices, so one tied row pinned — and was charged for — a whole + /// input batch, for as long as the boundary held. Across a stream of + /// batches that made retained memory grow with the *input* size rather + /// than with `K + ties`. + #[tokio::test] + async fn test_partitioned_topk_rank_ties_do_not_pin_input_batches() -> Result<()> { + // k=1 makes every retained row after the very first one a boundary + // tie, which is the state under test. + let (schema, mut rk) = build_partitioned_topk_rank(1)?; + + // Each batch carries exactly one row tied at the boundary and 999 + // rows that are strictly worse and must be dropped. The 1000:1 ratio + // is the point: holding the tied row costs a few bytes, holding the + // batch it arrived in costs ~8 KB, so the two outcomes cannot be + // confused. + let mut vals = vec![100; 1000]; + vals[0] = 7; + let batch = pk_val_batch(&schema, vec![1; 1000], vals)?; + + rk.insert_batch(&batch)?; + let after_first = rk.size(); + for _ in 0..7 { + rk.insert_batch(&batch)?; + } + let after_eight = rk.size(); + + // Seven more batches retain seven more single rows. Correct growth is + // ~8 bytes per batch (~56 total, just the tie-list slots); the bug grew + // by a whole ~8 KB batch each time (~56 KB total). A 2000-byte bar sits + // between the two with more than an order of magnitude of clearance on + // each side. + let growth = after_eight - after_first; + assert!( + growth < 2000, + "tie list grew {growth} bytes over 7 batches that contributed \ + 1 row each; it is holding the source batches" + ); + Ok(()) + } + + /// One deterministic pseudo-random input shape for the differential + /// tests below: a `k`, and the `(pks, vals)` batches to feed in. + /// + /// Shapes are deliberately tiny. These bugs live in how ties, partitions + /// and batch boundaries interleave, not in volume, so many small shapes + /// cover far more of that space than a few large ones — and they keep the + /// O(n^2) brute-force reference cheap. + struct DiffShape { + seed: u64, + k: usize, + n_partitions: i32, + n_values: i32, + batches: Vec<(Vec, Vec)>, + } + + impl std::fmt::Display for DiffShape { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "seed={} k={} partitions={} values={} batches={}", + self.seed, + self.k, + self.n_partitions, + self.n_values, + self.batches.len() + ) + } + } + + impl DiffShape { + /// `max_values` caps the ORDER BY value domain. Keep it near `k` so + /// rows tie above, at, and below the boundary frequently rather than + /// by chance; raise it for an operator that bounds *distinct* values + /// and so needs more than `k` of them before it will evict anything. + fn new(seed: u64, max_values: i32) -> Self { + use rand::rngs::StdRng; + use rand::{Rng, SeedableRng}; + + let mut rng = StdRng::seed_from_u64(seed); + let k = rng.random_range(1..5usize); + let n_partitions = rng.random_range(1..4i32); + let n_values = rng.random_range(1..max_values); + let n_batches = rng.random_range(1..5usize); + + let batches = (0..n_batches) + .map(|_| { + let rows = rng.random_range(1..12usize); + let pks = (0..rows) + .map(|_| rng.random_range(0..n_partitions)) + .collect(); + let vals = (0..rows).map(|_| rng.random_range(0..n_values)).collect(); + (pks, vals) + }) + .collect(); + + Self { + seed, + k, + n_partitions, + n_values, + batches, + } + } + + /// Every `(pk, val)` fed in, across all batches. + fn all_rows(&self) -> Vec<(i32, i32)> { + self.batches + .iter() + .flat_map(|(pks, vals)| pks.iter().copied().zip(vals.iter().copied())) + .collect() + } + + /// Brute-force reference: the sorted rows a `<= k` filter keeps, given + /// `rank_of(all_rows, pk, val)` for the ranking function under test. + fn expected( + &self, + rank_of: impl Fn(&[(i32, i32)], i32, i32) -> usize, + ) -> Vec<(i32, i32)> { + let all = self.all_rows(); + let mut kept: Vec<(i32, i32)> = all + .iter() + .copied() + .filter(|&(pk, val)| rank_of(&all, pk, val) <= self.k) + .collect(); + kept.sort_unstable(); + kept + } + } + + /// Drain an operator's output into sorted `(pk, val)` pairs, ready to + /// compare against [`DiffShape::expected`]. + async fn sorted_pk_val(stream: SendableRecordBatchStream) -> Result> { + let batches: Vec = stream.try_collect().await?; + let mut rows: Vec<(i32, i32)> = Vec::new(); + for b in &batches { + let pk = b.column(0).as_primitive::(); + let val = b.column(1).as_primitive::(); + for i in 0..b.num_rows() { + rows.push((pk.value(i), val.value(i))); + } + } + rows.sort_unstable(); + Ok(rows) + } + + /// Randomized differential test for `PartitionedTopKRank`. + /// + /// The retention rule is subtle — a K-bounded heap plus a boundary-tie + /// list that must be discarded the moment the K-th-best ORDER BY value + /// improves — and it turns on how partitions, ties and batch boundaries + /// interleave. 64 seeds runs in ~10 ms; deleting the tie-clear on a + /// boundary shift is caught by seed 0. + #[tokio::test] + async fn test_partitioned_topk_rank_matches_bruteforce() -> Result<()> { + for seed in 0..64u64 { + let shape = DiffShape::new(seed, 6); + let (schema, mut state) = build_partitioned_topk_rank(shape.k)?; + for (pks, vals) in &shape.batches { + state.insert_batch(&pk_val_batch(&schema, pks.clone(), vals.clone())?)?; + } + + // RANK: 1 + the number of strictly smaller rows. + let expected = shape.expected(|rows, pk, val| { + 1 + rows.iter().filter(|&&(p, v)| p == pk && v < val).count() + }); + + assert_eq!(sorted_pk_val(state.emit()?).await?, expected, "{shape}"); + } + Ok(()) + } + + /// Randomized differential test for `PartitionedTopKDenseRank`. + /// + /// Same harness as [`test_partitioned_topk_rank_matches_bruteforce`], + /// differing only in the ranking formula and a wider value domain: + /// DENSE_RANK bounds *distinct* values, so a partition needs more than + /// `k` of them before it evicts anything, and eviction is what the + /// admission pre-filter guards. A pre-filter that wrongly rejects + /// boundary-equal rows is caught by seed 6. + #[tokio::test] + async fn test_partitioned_topk_dense_rank_matches_bruteforce() -> Result<()> { + for seed in 0..64u64 { + let shape = DiffShape::new(seed, 8); + let (schema, mut state) = build_partitioned_topk_dense_rank(shape.k)?; + for (pks, vals) in &shape.batches { + state.insert_batch(&pk_val_batch(&schema, pks.clone(), vals.clone())?)?; + } + + // DENSE_RANK: 1 + the number of *distinct* strictly smaller values. + let expected = shape.expected(|rows, pk, val| { + 1 + rows + .iter() + .filter(|&&(p, v)| p == pk && v < val) + .map(|&(_, v)| v) + .collect::>() + .len() + }); + + assert_eq!(sorted_pk_val(state.emit()?).await?, expected, "{shape}"); + } + Ok(()) + } + fn pk_val_batch( schema: &Arc, pks: Vec,