diff --git a/datafusion/physical-plan/src/topk/mod.rs b/datafusion/physical-plan/src/topk/mod.rs index 0ca700cb37655..3d1615b50af57 100644 --- a/datafusion/physical-plan/src/topk/mod.rs +++ b/datafusion/physical-plan/src/topk/mod.rs @@ -1952,29 +1952,75 @@ struct DenseRankPartitionState { /// INVARIANT: `keys` and `groups.keys()` hold the same set. Every /// insertion into / removal from `groups` must mirror into `keys`. keys: BinaryHeap>, + /// Running total of the heap allocations owned by the *contents* of + /// `groups` and `keys`: the key bytes, the per-key `Vec` + /// buffers, and each entry's `row_indices`. Excludes the two + /// containers' own tables, which `capacity()` reports in O(1). + /// + /// INVARIANT: equals `recompute_contents_bytes` (test-only, so not + /// linkable from rustdoc). Every mutation of `groups` or `keys` must + /// adjust it; the + /// `test_partitioned_topk_dense_rank_contents_bytes_tracks_recompute` + /// test checks this against a full recompute after a randomized + /// workload. + contents_bytes: usize, } impl DenseRankPartitionState { fn size(&self) -> usize { let table_overhead = self.groups.capacity() * (size_of::>() + size_of::>()); - let contents: usize = self + // The heap's backing Vec: one `Vec` slot per reserved element. + let keys_overhead = self.keys.capacity() * size_of::>(); + table_overhead + self.contents_bytes + keys_overhead + } + + /// Bytes owned by one entry's `row_indices` buffer. + fn entry_bytes(entry: &GroupEntry) -> usize { + entry.row_indices.capacity() * size_of::() + } + + /// Track a previously unseen ob value, holding `run_indices` as its + /// first (and so far only) entry, and charge everything it allocates. + /// + /// `keys` and `groups` each own a copy of the key bytes; the clone's + /// capacity is read after the fact rather than assumed equal to the + /// original's, so the charge matches what was really allocated. + fn insert_new_group( + &mut self, + ob_key: Vec, + run_indices: Vec, + batch_id: u32, + ) { + let key_copy = ob_key.clone(); + self.contents_bytes += key_copy.capacity() + ob_key.capacity(); + self.keys.push(key_copy); + + let entry = GroupEntry { + row_indices: run_indices, + batch_id, + }; + self.contents_bytes += Self::entry_bytes(&entry); + let entries = vec![entry]; + self.contents_bytes += entries.capacity() * size_of::(); + self.groups.insert(ob_key, entries); + } + + /// The value [`Self::contents_bytes`] must hold, computed the slow + /// way. Used to assert the incremental accounting in tests. + #[cfg(test)] + fn recompute_contents_bytes(&self) -> usize { + let groups: usize = self .groups .iter() .map(|(key, entries)| { key.capacity() + entries.capacity() * size_of::() - + entries - .iter() - .map(|e| e.row_indices.capacity() * size_of::()) - .sum::() + + entries.iter().map(Self::entry_bytes).sum::() }) .sum(); - // `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 + // `keys` duplicates every key's bytes. + groups + self.keys.iter().map(|k| k.capacity()).sum::() } } @@ -2208,24 +2254,23 @@ impl PartitionedTopKDenseRank { // new `GroupEntry` (one entry per contributing batch). if let Some(entries) = state.groups.get_mut(&ob_key) { batch_entry.uses += 1; - entries.push(GroupEntry { + let before = entries.capacity() * size_of::(); + let entry = GroupEntry { row_indices: run_indices, batch_id, - }); + }; + state.contents_bytes += DenseRankPartitionState::entry_bytes(&entry); + entries.push(entry); + // The push may have grown the buffer. + state.contents_bytes += + entries.capacity() * size_of::() - before; 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![GroupEntry { - row_indices: run_indices, - batch_id, - }], - ); + state.insert_new_group(ob_key, run_indices, batch_id); continue; } @@ -2234,12 +2279,23 @@ impl PartitionedTopKDenseRank { 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. + // and its ordered mirror. `keys` and `groups` own + // separate copies of the key bytes, so both are + // uncharged — with their own capacities, which is why + // `remove_entry` is used to recover the map's copy + // rather than assuming it matches the popped one. let evicted_key = state.keys.pop().expect("max key present"); - let evicted = state + let (map_key, evicted) = state .groups - .remove(&evicted_key) + .remove_entry(&evicted_key) .expect("keys mirrors groups"); + state.contents_bytes -= evicted_key.capacity() + + map_key.capacity() + + evicted.capacity() * size_of::() + + evicted + .iter() + .map(DenseRankPartitionState::entry_bytes) + .sum::(); for e in &evicted { replacements += e.row_indices.len(); if e.batch_id == batch_id { @@ -2253,14 +2309,7 @@ impl PartitionedTopKDenseRank { } } batch_entry.uses += 1; - state.keys.push(ob_key.clone()); - state.groups.insert( - ob_key, - vec![GroupEntry { - row_indices: run_indices, - batch_id, - }], - ); + state.insert_new_group(ob_key, run_indices, batch_id); } // else: ob >= max — drop the whole run. } @@ -2312,11 +2361,58 @@ impl PartitionedTopKDenseRank { let mut sorted_pks: Vec> = states.keys().cloned().collect(); sorted_pks.sort(); - let mut coalescer = BatchCoalescer::new(Arc::clone(&schema), batch_size); + let mut out: Vec> = Vec::new(); + // Gather every retained row with a single `interleave_record_batch` + // per output batch rather than one `take_record_batch` per + // `GroupEntry`. A group entry holds only the rows one source batch + // contributed at one ob value, so entries are numerous and tiny — + // with P partitions, K distinct ob values and B contributing + // batches there are up to P × K × B of them, and gathering each + // one separately builds and tears down that many `RecordBatch`es. + // `interleave` takes `(batch_pos, row)` pairs across *different* + // source batches in one call, which is exactly the shape here. + // + // The pairs are pushed in emit order — partitions in sorted key + // order, ob values ascending within a partition, entries in + // insertion order within an ob value — so the interleaved output + // is already ordered and needs no post-sort. + // + // `indices` carries the global `batch_id` rather than a position + // into a fixed `batch_refs` slice: for Dictionary columns, + // `interleave` does work proportional to the *number of input + // arrays*, not just the ones the indices reference, so passing + // every batch in the store on every chunk turns emit into + // O(chunks × total batches) instead of O(chunks × batches the + // chunk actually uses). Each chunk below rebuilds a batch slice + // scoped to just its own `batch_id`s. + let mut indices: Vec<(u32, usize)> = Vec::with_capacity(batch_size); + let mut flush = |indices: &mut Vec<(u32, usize)>| -> Result<()> { + let mut batch_refs = Vec::new(); + let mut local_pos = HashMap::new(); + let mut local_indices = Vec::with_capacity(indices.len()); + for &(batch_id, row) in indices.iter() { + let pos = *local_pos.entry(batch_id).or_insert_with(|| { + batch_refs.push(&store.batches[&batch_id].batch); + batch_refs.len() - 1 + }); + local_indices.push((pos, row)); + } + let b = interleave_record_batch(&batch_refs, &local_indices)?; + (&b).record_output(&metrics.baseline); + out.push(Ok(b)); + indices.clear(); + Ok(()) + }; + + // Chunk at `batch_size` so the operator emits the same batch sizes + // as before and no single `interleave` output exceeds `batch_size`. for pk in sorted_pks { - let DenseRankPartitionState { groups, keys: _ } = - states.remove(&pk).expect("key from states.keys()"); + let DenseRankPartitionState { + groups, + keys: _, + contents_bytes: _, + } = 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)> = @@ -2324,22 +2420,17 @@ impl PartitionedTopKDenseRank { 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(batch, &indices)?; - coalescer.push_batch(sub)?; + for row in entry.row_indices { + indices.push((entry.batch_id, row as usize)); + if indices.len() == batch_size { + flush(&mut indices)?; + } + } } } } - coalescer.finish_buffered_batch()?; - - let mut out: Vec> = Vec::new(); - while let Some(b) = coalescer.next_completed_batch() { - (&b).record_output(&metrics.baseline); - out.push(Ok(b)); + if !indices.is_empty() { + flush(&mut indices)?; } Ok(Box::pin(RecordBatchStreamAdapter::new( @@ -2383,8 +2474,11 @@ impl PartitionedTopKDenseRank { mod tests { use super::*; use crate::metrics::MetricValue; - use arrow::array::{BooleanArray, Float64Array, Int32Array, StringArray}; - use arrow::datatypes::{DataType, Field, Schema}; + use arrow::array::{ + AsArray, BooleanArray, Float64Array, Int32Array, StringArray, + StringDictionaryBuilder, + }; + use arrow::datatypes::{DataType, Field, Int32Type, Schema}; use arrow_schema::SortOptions; use datafusion_common::{assert_batches_eq, exec_datafusion_err}; use datafusion_execution::memory_pool::GreedyMemoryPool; @@ -3393,16 +3487,26 @@ mod tests { /// Drain an operator's output into sorted `(pk, val)` pairs, ready to /// compare against [`DiffShape::expected`]. + /// + /// Asserts the rows arrived already in `(pk ASC, val ASC)` order before + /// sorting them. That order is load-bearing rather than cosmetic — it is + /// what `PartitionedTopKExec::compute_properties` advertises, and so what + /// lets the window above it run `mode=Sorted` with no `SortExec` — and + /// the content comparison this feeds would otherwise sort it away. 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::(); + 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))); } } + assert!( + rows.windows(2).all(|w| w[0] <= w[1]), + "emitted rows are not in (pk, val) order: {rows:?}" + ); rows.sort_unstable(); Ok(rows) } @@ -3465,6 +3569,42 @@ mod tests { Ok(()) } + /// `DenseRankPartitionState::contents_bytes` is maintained + /// incrementally at four mutation sites (append to an existing ob + /// group, insert a new one with room, evict-then-insert, and the + /// key/entry buffer growth each can trigger). Drift there would + /// silently give the reservation a wrong total, so check it against a + /// recompute over the same randomized workload the correctness + /// differential test uses — its shapes are tuned to exercise + /// eviction, which is the case with the most bookkeeping. + #[tokio::test] + async fn test_partitioned_topk_dense_rank_contents_bytes_tracks_recompute() + -> Result<()> { + let mut saw_eviction = false; + 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())?)?; + // Check after every batch, not just at the end: a + // compensating pair of errors within one batch would + // survive an end-only assertion. + for (pk, partition) in &state.states { + assert_eq!( + partition.contents_bytes, + partition.recompute_contents_bytes(), + "seed {seed}, partition {pk:?}: {shape}" + ); + } + } + saw_eviction |= state.metrics.row_replacements.value() > 0; + } + // Guards the guard: if the shapes stopped evicting, case C would + // go unchecked and this test would still pass. + assert!(saw_eviction, "workload never evicted a group"); + Ok(()) + } + /// The `(pk Int32, val Int32)` schema every `PartitionedTopK*` test /// builds against. `val_nullable` is what the null-ordering tests vary. fn pk_val_schema(val_nullable: bool) -> Arc { @@ -4412,6 +4552,85 @@ mod tests { Ok(()) } + /// Regression test for the `emit` rewrite: passing the *entire* store + /// to every `interleave_record_batch` call made Dictionary columns + /// O(chunks × total batches) instead of O(chunks × batches the chunk + /// actually uses), since `interleave_dictionaries` does per-input-array + /// work for every array it's handed. `batch_size` is 8, so 20 retained + /// rows spread across 2 source batches emit in 3 chunks — with several + /// chunks mixing rows from both source batches — the exact shape that + /// exercises the per-chunk batch-slice rebuild. + #[tokio::test] + async fn test_partitioned_topk_dense_rank_emit_dictionary_spans_chunks() -> Result<()> + { + let schema = Arc::new(Schema::new(vec![ + Field::new("pk", DataType::Int32, false), + Field::new( + "val", + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)), + 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, + 20, // k: large enough to retain every distinct value below + 8, // batch_size + &Arc::new(RuntimeEnv::default()), + &ExecutionPlanMetricsSet::new(), + )?; + + let dict_batch = |vals: &[&str]| -> Result { + let mut builder = StringDictionaryBuilder::::new(); + for v in vals { + builder.append_value(v); + } + Ok(RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![1; vals.len()])), + Arc::new(builder.finish()), + ], + )?) + }; + + // Two source batches (two distinct `batch_id`s), 10 distinct values + // each, all under a single partition key so every output chunk + // draws from both. + state.insert_batch(&dict_batch(&[ + "v00", "v01", "v02", "v03", "v04", "v05", "v06", "v07", "v08", "v09", + ])?)?; + state.insert_batch(&dict_batch(&[ + "v10", "v11", "v12", "v13", "v14", "v15", "v16", "v17", "v18", "v19", + ])?)?; + + let results: Vec<_> = state.emit()?.try_collect().await?; + assert_eq!(results.iter().map(|b| b.num_rows()).sum::(), 20); + assert_eq!(results.len(), 3, "expected 3 chunks of batch_size=8"); + let val_col = results[0].column(1).as_dictionary::(); + let dict_values = val_col.values().as_string::(); + let first_val = dict_values.value(val_col.keys().value(0) as usize); + assert_eq!(first_val, "v00", "rows must still emit in ob-sorted order"); + 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.