diff --git a/datafusion/physical-plan/src/aggregates/topk/hash_table.rs b/datafusion/physical-plan/src/aggregates/topk/hash_table.rs index 61410893df865..117cf9b378900 100644 --- a/datafusion/physical-plan/src/aggregates/topk/hash_table.rs +++ b/datafusion/physical-plan/src/aggregates/topk/hash_table.rs @@ -30,28 +30,27 @@ use datafusion_common::Result; use datafusion_common::exec_datafusion_err; use datafusion_common::hash_utils::RandomState; use half::f16; -use hashbrown::hash_table::HashTable; +use hashbrown::hash_table::{Entry, HashTable}; +use std::borrow::BorrowMut; use std::fmt::Debug; use std::hash::BuildHasher; use std::sync::Arc; -/// A "type alias" for Keys which are stored in our map -pub trait KeyType: Clone + Comparable + Debug {} +/// `heap_idx` assigned to groups whose aggregate values are all NULL. +/// Such groups are tracked in the hash table only, without entering heap, +/// so they can be emitted with a NULL aggregate value at the end. +const NULL_HEAP_IDX: usize = usize::MAX - 1; -impl KeyType for T where T: Clone + Comparable + Debug {} - -/// `heap_idx` assigned to groups whose aggregate values are all NULL. Such -/// groups are tracked in the hash table only (they never enter the heap), so -/// they can be emitted with a NULL aggregate value at the end. -const NULL_HEAP_IDX: usize = usize::MAX; +/// `heap_idx` needs to distinguish vacant slots from items that have a NULL key and value. +const VACANT_HEAP_IDX: usize = usize::MAX; /// An entry in our hash table that: /// 1. memoizes the hash /// 2. contains the key (ID) /// 3. contains the value (heap_idx - an index into the corresponding heap) -pub struct HashTableItem { +pub struct HashTableItem { hash: u64, - pub id: ID, + pub id: Option, pub heap_idx: usize, } @@ -59,12 +58,14 @@ pub struct HashTableItem { /// 1. limits the number of entries to the top K /// 2. Allocates a capacity greater than top K to maintain a low-fill factor and prevent resizing /// 3. Tracks indexes to allow corresponding heap to refer to entries by index vs hash -struct TopKHashTable { +struct TopKHashTable { map: HashTable, // Store the actual items separately to allow for index-based access - store: Vec>>, + store: Vec>, // Free indexes in the store for reuse free_indices: Vec, + // Pool of reusable value locations, usually Strings + free_slots: Vec, // The maximum number of entries allowed limit: usize, // Number of entries registered as all-NULL (heap_idx == NULL_HEAP_IDX) @@ -103,8 +104,8 @@ pub trait ArrowHashTable { /// Remove the group at `row_idx` if it is registered as all-NULL. Returns /// true if a NULL registration was removed. fn remove_if_null(&mut self, row_idx: usize) -> bool; - /// Store indexes of all groups registered as all-NULL - fn null_map_idxs(&self) -> Vec; + /// Store indexes of all groups registered as all-NULL. Returns null count. + fn null_map_idxs(&self, out: &mut Vec) -> usize; } /// Returns true if the given data type can be used as a top-K aggregation hash key. @@ -126,7 +127,7 @@ where for<'a> &'a S: StringArrayType<'a>, { owned: S, - map: TopKHashTable>, + map: TopKHashTable, rnd: RandomState, } @@ -136,8 +137,9 @@ where Option<::Native>: Comparable, { owned: PrimitiveArray, - map: TopKHashTable>, + map: TopKHashTable, rnd: RandomState, + value_type: DataType, } impl StringHashTable @@ -146,9 +148,8 @@ where for<'a> &'a S: StringArrayType<'a>, { pub fn new(limit: usize) -> Self { - let owned = S::from(Vec::new()); Self { - owned, + owned: S::from(Vec::new()), map: TopKHashTable::new(limit, limit * 10), rnd: RandomState::default(), } @@ -200,19 +201,14 @@ where let hash = self.rnd.hash_one(id); // Use entry API to avoid double lookup - self.map.find_or_insert( - hash, - id.map(ToOwned::to_owned), - replace_idx, - Self::eq_fn(id), - ) + self.map + .find_or_insert(hash, id, replace_idx, Self::eq_fn(id)) } fn insert_null(&mut self, row_idx: usize) -> bool { let id = some_value(&self.owned, row_idx); let hash = self.rnd.hash_one(id); - self.map - .insert_null(hash, id.map(ToOwned::to_owned), Self::eq_fn(id)) + self.map.insert_null(hash, id, Self::eq_fn(id)) } fn remove_if_null(&mut self, row_idx: usize) -> bool { @@ -221,8 +217,8 @@ where self.map.remove_if_null(hash, Self::eq_fn(id)) } - fn null_map_idxs(&self) -> Vec { - self.map.null_map_idxs() + fn null_map_idxs(&self, out: &mut Vec) -> usize { + self.map.null_map_idxs(out) } } @@ -230,14 +226,12 @@ impl PrimitiveHashTable where Option<::Native>: Comparable + HashValue, { - pub fn new(limit: usize, kt: DataType) -> Self { - let owned = PrimitiveArray::::builder(0) - .with_data_type(kt) - .finish(); + pub fn new(limit: usize, value_type: DataType) -> Self { Self { - owned, + owned: PrimitiveArray::::new_null(0).with_data_type(value_type.clone()), map: TopKHashTable::new(limit, limit * 10), rnd: RandomState::default(), + value_type, } } @@ -277,13 +271,10 @@ where fn take_all(&mut self, indexes: Vec) -> ArrayRef { let ids = self.map.take_all(indexes); - let mut builder: PrimitiveBuilder = PrimitiveArray::builder(ids.len()) - .with_data_type(self.owned.data_type().clone()); + let mut builder: PrimitiveBuilder = + PrimitiveArray::builder(ids.len()).with_data_type(self.value_type.clone()); for id in ids.into_iter() { - match id { - None => builder.append_null(), - Some(id) => builder.append_value(id), - } + builder.append_option(id); } let ids = builder.finish(); Arc::new(ids) @@ -297,12 +288,12 @@ where let (id, hash) = self.id_and_hash(row_idx); // Use entry API to avoid double lookup self.map - .find_or_insert(hash, id, replace_idx, Self::eq_fn(id)) + .find_or_insert(hash, id.as_ref(), replace_idx, Self::eq_fn(id)) } fn insert_null(&mut self, row_idx: usize) -> bool { let (id, hash) = self.id_and_hash(row_idx); - self.map.insert_null(hash, id, Self::eq_fn(id)) + self.map.insert_null(hash, id.as_ref(), Self::eq_fn(id)) } fn remove_if_null(&mut self, row_idx: usize) -> bool { @@ -310,39 +301,47 @@ where self.map.remove_if_null(hash, Self::eq_fn(id)) } - fn null_map_idxs(&self) -> Vec { - self.map.null_map_idxs() + fn null_map_idxs(&self, out: &mut Vec) -> usize { + self.map.null_map_idxs(out) } } -use hashbrown::hash_table::Entry; -impl TopKHashTable { +impl TopKHashTable { pub fn new(limit: usize, capacity: usize) -> Self { Self { map: HashTable::with_capacity(capacity), store: Vec::with_capacity(capacity), free_indices: Vec::new(), + free_slots: Vec::new(), limit, null_count: 0, } } pub fn heap_idx_at(&self, map_idx: usize) -> usize { - self.store[map_idx].as_ref().unwrap().heap_idx + self.store[map_idx].heap_idx } /// Remove the entry stored at `map_idx`, freeing its store slot for reuse fn remove_at(&mut self, map_idx: usize) { - let item_to_remove = self.store[map_idx].as_ref().unwrap(); + let item_to_remove = &self.store[map_idx]; let hash = item_to_remove.hash; let id_to_remove = &item_to_remove.id; - let eq = |&idx: &usize| self.store[idx].as_ref().unwrap().id == *id_to_remove; - let hasher = |idx: &usize| self.store[*idx].as_ref().unwrap().hash; + let eq = |idx: &usize| self.store[*idx].id == *id_to_remove; + let hasher = |idx: &usize| self.store[*idx].hash; match self.map.entry(hash, eq, hasher) { Entry::Occupied(entry) => { let (removed_idx, _) = entry.remove(); - self.store[removed_idx] = None; + let mi = &mut self.store[removed_idx]; + debug_assert!(mi.is_occupied()); + mi.heap_idx = VACANT_HEAP_IDX; + match mi.id.take() { + Some(slot) if Self::use_free_slots() => { + self.free_slots.push(slot); + } + _ => (), + } self.free_indices.push(removed_idx); } Entry::Vacant(_) => unreachable!(), @@ -363,63 +362,91 @@ impl TopKHashTable { fn update_heap_idx(&mut self, mapper: &[(usize, usize)]) { for (m, h) in mapper { - self.store[*m].as_mut().unwrap().heap_idx = *h; + self.store[*m].heap_idx = *h; } } + /// Used to avoid pushing pointless copies of primitives to the `free_slots` pool. + const fn use_free_slots() -> bool { + std::mem::needs_drop::() + } + /// Find an existing entry or insert a new one, avoiding double hash table lookup. /// Returns (map_idx, kind) where kind describes whether the group already /// existed, was newly inserted, or was converted from an all-NULL group. /// If inserting a new entry and the table is full, replaces the entry at replace_idx. - pub fn find_or_insert( + pub fn find_or_insert( &mut self, hash: u64, - id: ID, + id: Option<&Q>, replace_idx: usize, - mut eq: impl FnMut(&ID) -> bool, - ) -> (usize, InsertKind) { + mut eq: impl FnMut(&Option) -> bool, + ) -> (usize, InsertKind) + where + Q: ToOwned + ?Sized, + { // Check if entry exists - this is the only hash table lookup let mut replaced_null = false; - { - let eq_fn = |idx: &usize| eq(&self.store[*idx].as_ref().unwrap().id); - if let Some(&map_idx) = self.map.find(hash, eq_fn) { - if self.store[map_idx].as_ref().unwrap().heap_idx == NULL_HEAP_IDX { - // This group was registered as all-NULL but now produced a - // value: unregister it so it is inserted as a valued group - self.remove_at(map_idx); - self.null_count -= 1; - replaced_null = true; - } else { - return (map_idx, InsertKind::Existing); - } + + let eq_fn = |idx: &usize| eq(&self.store[*idx].id); + if let Some(&map_idx) = self.map.find(hash, eq_fn) { + if self.store[map_idx].is_null() { + // This group was registered as all-NULL but now produced a + // value: unregister it so it is inserted as a valued group + self.remove_at(map_idx); + self.null_count -= 1; + replaced_null = true; + } else { + return (map_idx, InsertKind::Existing); } } // Entry doesn't exist - compute heap_idx and prepare item let heap_idx = self.remove_if_full(replace_idx); + let store_idx = self.push_store_item(hash, id, heap_idx); + let kind = if replaced_null { + InsertKind::ReplacedNull + } else { + InsertKind::New + }; + (store_idx, kind) + } + + fn push_store_item(&mut self, hash: u64, id: Option<&Q>, heap_idx: usize) -> usize + where + Q: ToOwned + ?Sized, + { + let id = if Self::use_free_slots() { + id.map(|id| match self.free_slots.pop() { + Some(mut slot) => { + id.clone_into(&mut slot); + slot + } + _ => id.to_owned(), + }) + } else { + debug_assert!(self.free_slots.is_empty(), "primitives should not pool"); + id.map(ToOwned::to_owned) + }; let mi = HashTableItem::new(hash, id, heap_idx); let store_idx = if let Some(idx) = self.free_indices.pop() { - self.store[idx] = Some(mi); + debug_assert!(self.store[idx].is_vacant(), "slot should be empty"); + self.store[idx] = mi; idx } else { - self.store.push(Some(mi)); + self.store.push(mi); self.store.len() - 1 }; // Reserve space if needed - let hasher = |idx: &usize| self.store[*idx].as_ref().unwrap().hash; + let hasher = |idx: &usize| self.store[*idx].hash; if self.map.len() == self.map.capacity() { self.map.reserve(self.limit, hasher); } // Insert without checking again since we already confirmed it doesn't exist self.map.insert_unique(hash, store_idx, hasher); - let kind = if replaced_null { - InsertKind::ReplacedNull - } else { - InsertKind::New - }; - (store_idx, kind) + store_idx } /// Register a group whose aggregate values are all NULL, unless it is @@ -427,13 +454,17 @@ impl TopKHashTable { /// never enter the heap. At most `limit` NULL groups are tracked: they all /// tie on the sort key, so any `limit` of them is a valid top-k superset. /// Returns true if the group was newly registered. - pub fn insert_null( + pub fn insert_null( &mut self, hash: u64, - id: ID, - mut eq: impl FnMut(&ID) -> bool, - ) -> bool { - let eq_fn = |idx: &usize| eq(&self.store[*idx].as_ref().unwrap().id); + id: Option<&Q>, + mut eq: impl FnMut(&Option) -> bool, + ) -> bool + where + Q: ToOwned + ?Sized, + ID: BorrowMut, + { + let eq_fn = |idx: &usize| eq(&self.store[*idx].id); if self.map.find(hash, eq_fn).is_some() { return false; } @@ -442,20 +473,7 @@ impl TopKHashTable { return false; } - let mi = HashTableItem::new(hash, id, NULL_HEAP_IDX); - let store_idx = if let Some(idx) = self.free_indices.pop() { - self.store[idx] = Some(mi); - idx - } else { - self.store.push(Some(mi)); - self.store.len() - 1 - }; - - let hasher = |idx: &usize| self.store[*idx].as_ref().unwrap().hash; - if self.map.len() == self.map.capacity() { - self.map.reserve(self.limit, hasher); - } - self.map.insert_unique(hash, store_idx, hasher); + _ = self.push_store_item(hash, id, NULL_HEAP_IDX); self.null_count += 1; true } @@ -464,10 +482,14 @@ impl TopKHashTable { /// all-NULL group produces a value that loses to the current top-k: the /// group can no longer reach the top-k, but it must not be emitted with a /// NULL value either. Returns true if a NULL registration was removed. - pub fn remove_if_null(&mut self, hash: u64, mut eq: impl FnMut(&ID) -> bool) -> bool { - let eq_fn = |idx: &usize| eq(&self.store[*idx].as_ref().unwrap().id); + pub fn remove_if_null( + &mut self, + hash: u64, + mut eq: impl FnMut(&Option) -> bool, + ) -> bool { + let eq_fn = |idx: &usize| eq(&self.store[*idx].id); if let Some(&map_idx) = self.map.find(hash, eq_fn) - && self.store[map_idx].as_ref().unwrap().heap_idx == NULL_HEAP_IDX + && self.store[map_idx].is_null() { self.remove_at(map_idx); self.null_count -= 1; @@ -477,39 +499,57 @@ impl TopKHashTable { } /// Store indexes of all groups registered as all-NULL - pub fn null_map_idxs(&self) -> Vec { - self.store + pub fn null_map_idxs(&self, out: &mut Vec) -> usize { + let iter = self + .store .iter() .enumerate() - .filter_map(|(idx, item)| { - item.as_ref() - .filter(|item| item.heap_idx == NULL_HEAP_IDX) - .map(|_| idx) - }) - .collect() + .filter_map(|(idx, item)| item.is_null().then_some(idx)); + + let prev_len = out.len(); + out.extend(iter); + + debug_assert_eq!(out.len(), self.null_count + prev_len); + self.null_count } pub fn len(&self) -> usize { self.map.len() } - pub fn take_all(&mut self, idxs: Vec) -> Vec { + pub fn take_all(&mut self, idxs: Vec) -> Vec> { let ids = idxs .into_iter() - .map(|idx| self.store[idx].take().unwrap().id) + .map(|idx| self.store[idx].id.take()) .collect(); self.map.clear(); self.store.clear(); self.free_indices.clear(); + // No need to clear free_slots. self.null_count = 0; ids } } -impl HashTableItem { - pub fn new(hash: u64, id: ID, heap_idx: usize) -> Self { +impl HashTableItem { + pub fn new(hash: u64, id: Option, heap_idx: usize) -> Self { Self { hash, id, heap_idx } } + + #[inline] + pub fn is_null(&self) -> bool { + self.heap_idx == NULL_HEAP_IDX + } + + #[inline] + pub fn is_vacant(&self) -> bool { + self.heap_idx == VACANT_HEAP_IDX + } + + #[inline] + pub fn is_occupied(&self) -> bool { + self.heap_idx != VACANT_HEAP_IDX + } } impl HashValue for Option { @@ -585,6 +625,13 @@ mod tests { use arrow_schema::TimeUnit; use std::collections::BTreeMap; + fn assert_nulls(map: &TopKHashTable, expected: &[usize]) { + assert_eq!(map.null_count, expected.len()); + let mut actual = Vec::new(); + assert_eq!(map.null_map_idxs(&mut actual), expected.len()); + assert_eq!(actual, expected); + } + #[test] fn should_emit_correct_type() -> Result<()> { let ids = @@ -603,14 +650,14 @@ mod tests { fn should_resize_properly() -> Result<()> { let mut heap_to_map = BTreeMap::::new(); // Create TopKHashTable with limit=5 and capacity=3 to force resizing - let mut map = TopKHashTable::>::new(5, 3); + let mut map = TopKHashTable::::new(5, 3); // Insert 5 entries, tracking the heap-to-map index mapping for (heap_idx, id) in ["1", "2", "3", "4", "5"].iter().enumerate() { let value = Some(id.to_string()); let hash = heap_idx as u64; let (map_idx, kind) = - map.find_or_insert(hash, value.clone(), heap_idx, |v| *v == value); + map.find_or_insert(hash, value.as_ref(), heap_idx, |v| *v == value); assert_eq!(kind, InsertKind::New, "Entry should be new"); heap_to_map.insert(heap_idx, map_idx); } @@ -636,32 +683,29 @@ mod tests { #[test] fn should_track_null_groups() -> Result<()> { - let mut map = TopKHashTable::>::new(2, 10); + let mut map = TopKHashTable::::new(2, 10); let a = Some("a".to_string()); let b = Some("b".to_string()); let c = Some("c".to_string()); // register two all-NULL groups; the third exceeds the NULL group limit - assert!(map.insert_null(100, a.clone(), |v| *v == a)); - assert!(map.insert_null(200, b.clone(), |v| *v == b)); - assert!(!map.insert_null(300, c.clone(), |v| *v == c)); + assert!(map.insert_null(100, a.as_ref(), |v| *v == a)); + assert!(map.insert_null(200, b.as_ref(), |v| *v == b)); + assert!(!map.insert_null(300, c.as_ref(), |v| *v == c)); // re-registering an existing NULL group is a no-op - assert!(!map.insert_null(100, a.clone(), |v| *v == a)); - assert_eq!(map.null_count, 2); - assert_eq!(map.null_map_idxs(), vec![0, 1]); + assert!(!map.insert_null(100, a.as_ref(), |v| *v == a)); + assert_nulls(&map, &[0, 1]); // a valued insert for a NULL group converts it to a valued group - let (map_idx, kind) = map.find_or_insert(200, b.clone(), 0, |v| *v == b); + let (map_idx, kind) = map.find_or_insert(200, b.as_ref(), 0, |v| *v == b); assert_eq!(kind, InsertKind::ReplacedNull, "NULL group should convert"); assert_eq!(map.heap_idx_at(map_idx), 0, "Heap should append at 0"); - assert_eq!(map.null_count, 1); - assert_eq!(map.null_map_idxs(), vec![0]); + assert_nulls(&map, &[0]); // remove the remaining NULL group; removing twice is a no-op map.remove_if_null(100, |v| *v == a); - assert_eq!(map.null_count, 0); - assert!(map.null_map_idxs().is_empty()); + assert_nulls(&map, &[]); map.remove_if_null(100, |v| *v == a); // removing a valued group via remove_if_null is a no-op map.remove_if_null(200, |v| *v == b); @@ -670,26 +714,43 @@ mod tests { Ok(()) } + #[test] + fn should_track_null_id() { + let mut map = TopKHashTable::::new(2, 10); + + let hash = 0; + let none = Option::<&str>::None; + + assert!(map.insert_null(hash, none, Option::is_none)); + assert_nulls(&map, &[0]); + + assert!(!map.insert_null(hash, none, Option::is_none)); + assert_nulls(&map, &[0]); + + assert!(map.remove_if_null(hash, Option::is_none)); + assert_nulls(&map, &[]); + } + #[test] fn should_reuse_all_freed_store_slots() -> Result<()> { - let mut map = TopKHashTable::>::new(1, 10); + let mut map = TopKHashTable::::new(1, 10); let a = Some("a".to_string()); let b = Some("b".to_string()); let c = Some("c".to_string()); - let (b_idx, kind) = map.find_or_insert(100, b.clone(), 0, |v| *v == b); + let (b_idx, kind) = map.find_or_insert(100, b.as_ref(), 0, |v| *v == b); assert_eq!(kind, InsertKind::New); - assert!(map.insert_null(200, a.clone(), |v| *v == a)); + assert!(map.insert_null(200, a.as_ref(), |v| *v == a)); // Converting a NULL group while the valued heap is full frees two // slots: the NULL registration and the evicted valued group. - let (_, kind) = map.find_or_insert(200, a.clone(), b_idx, |v| *v == a); + let (_, kind) = map.find_or_insert(200, a.as_ref(), b_idx, |v| *v == a); assert_eq!(kind, InsertKind::ReplacedNull); // Both freed slots must remain reusable. Otherwise repeated // conversions make the backing store grow without bound. - assert!(map.insert_null(300, c.clone(), |v| *v == c)); + assert!(map.insert_null(300, c.as_ref(), |v| *v == c)); assert_eq!(map.store.len(), 2); Ok(()) diff --git a/datafusion/physical-plan/src/aggregates/topk/heap.rs b/datafusion/physical-plan/src/aggregates/topk/heap.rs index 7982d94d39121..89ddaaa0becd7 100644 --- a/datafusion/physical-plan/src/aggregates/topk/heap.rs +++ b/datafusion/physical-plan/src/aggregates/topk/heap.rs @@ -23,12 +23,10 @@ //! Supported value types include Arrow primitives (integers, floats, decimals, intervals) //! and UTF-8 strings (`Utf8`, `LargeUtf8`, `Utf8View`) using lexicographic ordering. -use arrow::array::{ArrayRef, ArrowPrimitiveType, PrimitiveArray, downcast_primitive}; -use arrow::array::{LargeStringBuilder, StringBuilder, StringViewBuilder}; +use arrow::array::types::{IntervalDayTime, IntervalMonthDayNano}; use arrow::array::{ - StringArray, - cast::AsArray, - types::{IntervalDayTime, IntervalMonthDayNano}, + Array, ArrayAccessor, ArrayRef, ArrowPrimitiveType, AsArray, LargeStringArray, + PrimitiveArray, StringArray, StringArrayType, StringViewArray, downcast_primitive, }; use arrow::buffer::ScalarBuffer; use arrow::datatypes::{DataType, i256}; @@ -42,8 +40,20 @@ use std::fmt::{Debug, Display, Formatter}; use std::sync::Arc; /// A custom version of `Ord` that only exists to we can implement it for the Values in our heap -pub trait Comparable { - fn comp(&self, other: &Self) -> Ordering; +pub trait Comparable { + fn comp(&self, other: &Rhs) -> Ordering; +} + +impl Comparable for str { + fn comp(&self, other: &String) -> Ordering { + self.cmp(other.as_str()) + } +} + +impl Comparable for String { + fn comp(&self, other: &Self) -> Ordering { + self.cmp(other) + } } impl Comparable for Option { @@ -53,15 +63,52 @@ impl Comparable for Option { } /// A "type alias" for Values which are stored in our heap -pub trait ValueType: Comparable + Clone + Debug {} +pub trait ValueType: Comparable + Clone + Debug + Default {} + +impl ValueType for T where T: Comparable + Clone + Debug + Default {} -impl ValueType for T where T: Comparable + Clone + Debug {} +const VACANT_MAP_IDX: usize = usize::MAX; /// An entry in our heap, which contains both the value and a index into an external HashTable struct HeapItem { val: VAL, map_idx: usize, } +impl HeapItem { + fn is_vacant(&self) -> bool { + self.map_idx == VACANT_MAP_IDX + } + + fn is_occupied(&self) -> bool { + self.map_idx != VACANT_MAP_IDX + } + + fn map_idx(&self) -> Option { + self.is_occupied().then_some(self.map_idx) + } + + fn take(&mut self) -> Option { + self.is_occupied().then(|| std::mem::take(self)) + } + + fn as_ref(&self) -> &Self { + debug_assert!(self.is_occupied()); + self + } + + fn as_mut(&mut self) -> &mut Self { + debug_assert!(self.is_occupied()); + self + } +} +impl Default for HeapItem { + fn default() -> Self { + Self { + val: VAL::default(), + map_idx: VACANT_MAP_IDX, + } + } +} /// A custom heap implementation that allows several things that couldn't be achieved with /// `collections::BinaryHeap`: @@ -73,11 +120,12 @@ struct TopKHeap { desc: bool, len: usize, capacity: usize, - heap: Vec>>, + heap: Vec>, } /// An interface to hide the generic type signature of TopKHeap behind arrow arrays pub trait ArrowHeap { + fn value_type(&self) -> &DataType; fn set_batch(&mut self, vals: ArrayRef); fn is_worse(&self, idx: usize) -> bool; fn worst_map_idx(&self) -> usize; @@ -99,20 +147,19 @@ where batch: PrimitiveArray, heap: TopKHeap, desc: bool, - data_type: DataType, + value_type: DataType, } impl PrimitiveHeap where ::Native: Comparable, { - pub fn new(limit: usize, desc: bool, data_type: DataType) -> Self { - let batch = PrimitiveArray::::builder(0).finish(); + pub fn new(limit: usize, desc: bool, value_type: DataType) -> Self { Self { - batch, + batch: PrimitiveArray::::new_null(0).with_data_type(value_type.clone()), heap: TopKHeap::new(limit, desc), desc, - data_type, + value_type, } } } @@ -121,8 +168,12 @@ impl ArrowHeap for PrimitiveHeap where ::Native: Comparable, { + fn value_type(&self) -> &DataType { + &self.value_type + } + fn set_batch(&mut self, vals: ArrayRef) { - self.batch = PrimitiveArray::from(vals.to_data()); + self.batch = vals.as_primitive().clone(); } fn is_worse(&self, row_idx: usize) -> bool { @@ -140,7 +191,7 @@ where fn insert(&mut self, row_idx: usize, map_idx: usize, map: &mut Vec<(usize, usize)>) { let new_val = self.batch.value(row_idx); - self.heap.append_or_replace(new_val, map_idx, map); + self.heap.append_or_replace(&new_val, map_idx, map); } fn replace_if_better( @@ -150,14 +201,14 @@ where map: &mut Vec<(usize, usize)>, ) { let new_val = self.batch.value(row_idx); - self.heap.replace_if_better(heap_idx, new_val, map); + self.heap.replace_if_better(&new_val, heap_idx, map); } fn drain(&mut self) -> (ArrayRef, Vec) { let nulls = None; let (vals, map_idxs) = self.heap.drain(); let arr = PrimitiveArray::::new(ScalarBuffer::from(vals), nulls) - .with_data_type(self.data_type.clone()); + .with_data_type(self.value_type.clone()); (Arc::new(arr), map_idxs) } } @@ -169,58 +220,46 @@ where /// borrowed strings are compared before allocation, and only allocated when the /// heap confirms they improve the top-K set. /// -pub struct StringHeap { - batch: ArrayRef, - heap: TopKHeap>, +pub struct StringHeap +where + for<'a> &'a S: StringArrayType<'a>, +{ + batch: S, + heap: TopKHeap, desc: bool, - data_type: DataType, } -impl StringHeap { - pub fn new(limit: usize, desc: bool, data_type: DataType) -> Self { - let batch: ArrayRef = Arc::new(StringArray::from(Vec::<&str>::new())); +impl StringHeap +where + S: Array + From>>, + for<'a> &'a S: StringArrayType<'a>, +{ + pub fn new(limit: usize, desc: bool) -> Self { + let batch = S::from(Vec::new()); Self { batch, heap: TopKHeap::new(limit, desc), desc, - data_type, } } - - /// Extracts a string value from the current batch at the given row index. - /// - /// Panics if the row index is out of bounds or if the data type is not one of - /// the supported UTF-8 string types. - /// - /// Note: Null values should not appear in the input; the aggregation layer - /// ensures nulls are filtered before reaching this code. - fn value(&self, row_idx: usize) -> &str { - extract_string_value(&self.batch, &self.data_type, row_idx) - } } -/// Helper to extract a string value from an ArrayRef at a given index. -/// -/// Supports `Utf8`, `LargeUtf8`, and `Utf8View` data types. -/// -/// # Panics -/// Panics if the index is out of bounds or if the data type is unsupported. -fn extract_string_value<'a>( - batch: &'a ArrayRef, - data_type: &DataType, - idx: usize, -) -> &'a str { - match data_type { - DataType::Utf8 => batch.as_string::().value(idx), - DataType::LargeUtf8 => batch.as_string::().value(idx), - DataType::Utf8View => batch.as_string_view().value(idx), - _ => unreachable!("Unsupported string type: {data_type}"), +impl ArrowHeap for StringHeap +where + S: Array + Clone + From> + 'static, + for<'a> &'a S: StringArrayType<'a>, +{ + fn value_type(&self) -> &DataType { + // Strings don't store any metadata. + self.batch.data_type() } -} -impl ArrowHeap for StringHeap { fn set_batch(&mut self, vals: ArrayRef) { - self.batch = vals; + self.batch = vals + .as_any() + .downcast_ref::() + .expect("Unsupported data type") + .clone(); } fn is_worse(&self, row_idx: usize) -> bool { @@ -230,15 +269,9 @@ impl ArrowHeap for StringHeap { // Compare borrowed `&str` against the worst heap value first to avoid // allocating a `String` unless this row would actually replace an // existing heap entry. - let new_val = self.value(row_idx); - let worst_val = self.heap.worst_val().expect("Missing root"); - match worst_val { - None => false, - Some(worst_str) => { - (!self.desc && new_val > worst_str.as_str()) - || (self.desc && new_val < worst_str.as_str()) - } - } + let new_val = (&self.batch).value(row_idx); + let worst_val = self.heap.worst_val().expect("Missing root").as_str(); + (!self.desc && new_val > worst_val) || (self.desc && new_val < worst_val) } fn worst_map_idx(&self) -> usize { @@ -246,12 +279,7 @@ impl ArrowHeap for StringHeap { } fn insert(&mut self, row_idx: usize, map_idx: usize, map: &mut Vec<(usize, usize)>) { - // When appending (heap not full) we must allocate to own the string - // because it will be stored in the heap. For replacements we avoid - // allocation until `replace_if_better` confirms a replacement is - // necessary. - let new_str = self.value(row_idx).to_string(); - let new_val = Some(new_str); + let new_val = (&self.batch).value(row_idx); self.heap.append_or_replace(new_val, map_idx, map); } @@ -261,62 +289,14 @@ impl ArrowHeap for StringHeap { row_idx: usize, map: &mut Vec<(usize, usize)>, ) { - let new_str = self.value(row_idx); - let existing = self.heap.heap[heap_idx] - .as_ref() - .expect("Missing heap item"); - - // Compare borrowed reference first—no allocation yet. - // We compare the borrowed `&str` with the stored `Option` and - // only allocate (`to_string()`) when a replacement is required. - match &existing.val { - None => { - // Existing is null; new value always wins - let new_val = Some(new_str.to_string()); - self.heap.replace_if_better(heap_idx, new_val, map); - } - Some(existing_str) => { - // Compare borrowed strings first - if (!self.desc && new_str < existing_str.as_str()) - || (self.desc && new_str > existing_str.as_str()) - { - let new_val = Some(new_str.to_string()); - self.heap.replace_if_better(heap_idx, new_val, map); - } - // Else: no improvement, no allocation - } - } + let new_val = (&self.batch).value(row_idx); + self.heap.replace_if_better(new_val, heap_idx, map); } fn drain(&mut self) -> (ArrayRef, Vec) { let (vals, map_idxs) = self.heap.drain(); - // Use Arrow builders to safely construct arrays from the owned - // `Option` values. Builders avoid needing to maintain - // references to temporary storage. - - // Macro to eliminate duplication across string builder types. - // All three builders share the same interface for append_value, - // append_null, and finish, differing only in their concrete types. - macro_rules! build_string_array { - ($builder_type:ty) => {{ - let mut builder = <$builder_type>::new(); - for val in vals { - match val { - Some(s) => builder.append_value(&s), - None => builder.append_null(), - } - } - Arc::new(builder.finish()) - }}; - } - - let arr: ArrayRef = match self.data_type { - DataType::Utf8 => build_string_array!(StringBuilder), - DataType::LargeUtf8 => build_string_array!(LargeStringBuilder), - DataType::Utf8View => build_string_array!(StringViewBuilder), - _ => unreachable!("Unsupported string type: {}", self.data_type), - }; - (arr, map_idxs) + let vals = Arc::new(S::from(vals)); + (vals, map_idxs) } } @@ -326,18 +306,17 @@ impl TopKHeap { desc, capacity: limit, len: 0, - heap: (0..=limit).map(|_| None).collect::>(), + heap: (0..=limit).map(|_| HeapItem::default()).collect::>(), } } pub fn worst_val(&self) -> Option<&VAL> { let root = self.heap.first()?; - let hi = root.as_ref()?; - Some(&hi.val) + root.is_occupied().then_some(&root.val) } pub fn worst_map_idx(&self) -> usize { - self.heap[0].as_ref().map(|hi| hi.map_idx).unwrap_or(0) + self.heap[0].map_idx().unwrap_or(0) } pub fn is_full(&self) -> bool { @@ -348,12 +327,14 @@ impl TopKHeap { self.len } - pub fn append_or_replace( + pub fn append_or_replace( &mut self, - new_val: VAL, + new_val: &Q, map_idx: usize, map: &mut Vec<(usize, usize)>, - ) { + ) where + Q: ToOwned + ?Sized, + { if self.is_full() { self.replace_root(new_val, map_idx, map); } else { @@ -361,9 +342,14 @@ impl TopKHeap { } } - fn append(&mut self, new_val: VAL, map_idx: usize, mapper: &mut Vec<(usize, usize)>) { - let hi = HeapItem::new(new_val, map_idx); - self.heap[self.len] = Some(hi); + fn append(&mut self, new_val: &Q, map_idx: usize, mapper: &mut Vec<(usize, usize)>) + where + Q: ToOwned + ?Sized, + { + let hi = &mut self.heap[self.len]; + debug_assert!(hi.is_vacant()); + hi.map_idx = map_idx; + new_val.clone_into(&mut hi.val); self.heapify_up(self.len, mapper); self.len += 1; } @@ -384,43 +370,48 @@ impl TopKHeap { } pub fn drain(&mut self) -> (Vec, Vec) { - let mut map = Vec::with_capacity(self.len); + let mut map = Vec::new(); let mut vals = Vec::with_capacity(self.len); let mut map_idxs = Vec::with_capacity(self.len); while let Some(worst_hi) = self.pop(&mut map) { vals.push(worst_hi.val); map_idxs.push(worst_hi.map_idx); + map.clear(); } vals.reverse(); map_idxs.reverse(); (vals, map_idxs) } - fn replace_root( + fn replace_root( &mut self, - new_val: VAL, + new_val: &Q, map_idx: usize, mapper: &mut Vec<(usize, usize)>, - ) { - let hi = self.heap[0].as_mut().expect("No root"); - hi.val = new_val; + ) where + Q: ToOwned + ?Sized, + { + let hi = self.heap[0].as_mut(); + new_val.clone_into(&mut hi.val); hi.map_idx = map_idx; self.heapify_down(0, mapper); } - pub fn replace_if_better( + pub fn replace_if_better( &mut self, + new_val: &Q, heap_idx: usize, - new_val: VAL, mapper: &mut Vec<(usize, usize)>, - ) { - let existing = self.heap[heap_idx].as_mut().expect("Missing heap item"); + ) where + Q: ToOwned + Comparable + ?Sized, + { + let existing = self.heap[heap_idx].as_mut(); if (!self.desc && new_val.comp(&existing.val) != Ordering::Less) || (self.desc && new_val.comp(&existing.val) != Ordering::Greater) { return; } - existing.val = new_val; + new_val.clone_into(&mut existing.val); self.heapify_down(heap_idx, mapper); } @@ -428,8 +419,8 @@ impl TopKHeap { let desc = self.desc; while idx != 0 { let parent_idx = (idx - 1) / 2; - let node = self.heap[idx].as_ref().expect("No heap item"); - let parent = self.heap[parent_idx].as_ref().expect("No heap item"); + let node = self.heap[idx].as_ref(); + let parent = self.heap[parent_idx].as_ref(); if (!desc && node.val.comp(&parent.val) != Ordering::Greater) || (desc && node.val.comp(&parent.val) != Ordering::Less) { @@ -440,40 +431,45 @@ impl TopKHeap { } } + #[inline] fn swap(&mut self, a_idx: usize, b_idx: usize, mapper: &mut Vec<(usize, usize)>) { self.heap.swap(a_idx, b_idx); - let b_hi = self.heap[b_idx].as_ref().expect("Missing heap entry"); - let a_hi = self.heap[a_idx].as_ref().expect("Missing heap entry"); + let b_hi = self.heap[b_idx].as_ref(); + let a_hi = self.heap[a_idx].as_ref(); - mapper.push((b_hi.map_idx, b_idx)); - mapper.push((a_hi.map_idx, a_idx)); + mapper.extend([(b_hi.map_idx, b_idx), (a_hi.map_idx, a_idx)]); } - fn heapify_down(&mut self, node_idx: usize, mapper: &mut Vec<(usize, usize)>) { - let left_child = node_idx * 2 + 1; + fn heapify_down(&mut self, mut node_idx: usize, mapper: &mut Vec<(usize, usize)>) { let desc = self.desc; - let entry = self.heap.get(node_idx).expect("Missing node!"); - let entry = entry.as_ref().expect("Missing node!"); - let mut best_idx = node_idx; - let mut best_val = &entry.val; - for child_idx in left_child..=left_child + 1 { - if let Some(Some(child)) = self.heap.get(child_idx) - && ((!desc && child.val.comp(best_val) == Ordering::Greater) - || (desc && child.val.comp(best_val) == Ordering::Less)) - { - best_val = &child.val; - best_idx = child_idx; + loop { + let left_child = node_idx * 2 + 1; + let entry = self.heap[node_idx].as_ref(); + let mut best_idx = node_idx; + let mut best_val = &entry.val; + for child_idx in left_child..=left_child + 1 { + if let Some(child) = self.heap.get(child_idx) + && child.is_occupied() + && ((!desc && child.val.comp(best_val) == Ordering::Greater) + || (desc && child.val.comp(best_val) == Ordering::Less)) + { + best_val = &child.val; + best_idx = child_idx; + } + } + if best_val.comp(&entry.val) == Ordering::Equal { + break; } - } - if best_val.comp(&entry.val) != Ordering::Equal { self.swap(best_idx, node_idx, mapper); - self.heapify_down(best_idx, mapper); + node_idx = best_idx; } } fn _tree_print(&self, idx: usize, prefix: &str, is_tail: bool, output: &mut String) { - if let Some(Some(hi)) = self.heap.get(idx) { + if let Some(hi) = self.heap.get(idx) + && hi.is_occupied() + { let connector = if idx != 0 { if is_tail { "└── " } else { "├── " } } else { @@ -514,15 +510,6 @@ impl Display for TopKHeap { } } -impl HeapItem { - pub fn new(val: VAL, buk_idx: usize) -> Self { - Self { - val, - map_idx: buk_idx, - } - } -} - impl Debug for HeapItem { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { f.write_str("bucket=")?; @@ -618,13 +605,6 @@ pub fn new_heap( desc: bool, vt: DataType, ) -> Result> { - if matches!( - vt, - DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View - ) { - return Ok(Box::new(StringHeap::new(limit, desc, vt))); - } - macro_rules! downcast_helper { ($vt:ty, $d:ident) => { return Ok(Box::new(PrimitiveHeap::<$vt>::new(limit, desc, vt))) @@ -633,6 +613,9 @@ pub fn new_heap( downcast_primitive! { vt => (downcast_helper, vt), + DataType::Utf8 => return Ok(Box::new(StringHeap::::new(limit, desc))), + DataType::LargeUtf8 => return Ok(Box::new(StringHeap::::new(limit, desc))), + DataType::Utf8View => return Ok(Box::new(StringHeap::::new(limit, desc))), _ => {} } @@ -651,7 +634,7 @@ mod tests { fn should_append() -> Result<()> { let mut map = vec![]; let mut heap = TopKHeap::new(10, false); - heap.append_or_replace(1, 1, &mut map); + heap.append_or_replace(&1, 1, &mut map); let actual = heap.to_string(); assert_snapshot!(actual, @"val=1 idx=0, bucket=1"); @@ -664,10 +647,10 @@ mod tests { let mut map = vec![]; let mut heap = TopKHeap::new(10, false); - heap.append_or_replace(1, 1, &mut map); + heap.append_or_replace(&1, 1, &mut map); assert_eq!(map, vec![]); - heap.append_or_replace(2, 2, &mut map); + heap.append_or_replace(&2, 2, &mut map); assert_eq!(map, vec![(2, 0), (1, 1)]); let actual = heap.to_string(); @@ -684,9 +667,9 @@ mod tests { let mut map = vec![]; let mut heap = TopKHeap::new(3, false); - heap.append_or_replace(1, 1, &mut map); - heap.append_or_replace(2, 2, &mut map); - heap.append_or_replace(3, 3, &mut map); + heap.append_or_replace(&1, 1, &mut map); + heap.append_or_replace(&2, 2, &mut map); + heap.append_or_replace(&3, 3, &mut map); let actual = heap.to_string(); assert_snapshot!(actual, @r" val=3 idx=0, bucket=3 @@ -695,7 +678,7 @@ mod tests { "); let mut map = vec![]; - heap.append_or_replace(0, 0, &mut map); + heap.append_or_replace(&0, 0, &mut map); let actual = heap.to_string(); assert_snapshot!(actual, @r" val=2 idx=0, bucket=2 @@ -712,10 +695,10 @@ mod tests { let mut map = vec![]; let mut heap = TopKHeap::new(4, false); - heap.append_or_replace(1, 1, &mut map); - heap.append_or_replace(2, 2, &mut map); - heap.append_or_replace(3, 3, &mut map); - heap.append_or_replace(4, 4, &mut map); + heap.append_or_replace(&1, 1, &mut map); + heap.append_or_replace(&2, 2, &mut map); + heap.append_or_replace(&3, 3, &mut map); + heap.append_or_replace(&4, 4, &mut map); let actual = heap.to_string(); assert_snapshot!(actual, @r" val=4 idx=0, bucket=4 @@ -725,7 +708,7 @@ mod tests { "); let mut map = vec![]; - heap.replace_if_better(1, 0, &mut map); + heap.replace_if_better(&0, 1, &mut map); let actual = heap.to_string(); assert_snapshot!(actual, @r" val=4 idx=0, bucket=4 @@ -743,8 +726,8 @@ mod tests { let mut map = vec![]; let mut heap = TopKHeap::new(10, false); - heap.append_or_replace(1, 1, &mut map); - heap.append_or_replace(2, 2, &mut map); + heap.append_or_replace(&1, 1, &mut map); + heap.append_or_replace(&2, 2, &mut map); let actual = heap.to_string(); assert_snapshot!(actual, @r" @@ -763,8 +746,8 @@ mod tests { let mut map = vec![]; let mut heap = TopKHeap::new(10, false); - heap.append_or_replace(1, 1, &mut map); - heap.append_or_replace(2, 2, &mut map); + heap.append_or_replace(&1, 1, &mut map); + heap.append_or_replace(&2, 2, &mut map); let actual = heap.to_string(); assert_snapshot!(actual, @r" @@ -779,4 +762,137 @@ mod tests { Ok(()) } + + #[test] + fn should_append_string() -> Result<()> { + let mut map = vec![]; + let mut heap = TopKHeap::::new(10, false); + heap.append_or_replace("A", 1, &mut map); + + let actual = heap.to_string(); + assert_snapshot!(actual, @r#"val="A" idx=0, bucket=1"#); + + Ok(()) + } + + #[test] + fn should_heapify_up_string() -> Result<()> { + let mut map = vec![]; + let mut heap = TopKHeap::::new(10, false); + + heap.append_or_replace("A", 1, &mut map); + assert_eq!(map, vec![]); + + heap.append_or_replace("B", 2, &mut map); + assert_eq!(map, vec![(2, 0), (1, 1)]); + + let actual = heap.to_string(); + assert_snapshot!(actual, @r#" + val="B" idx=0, bucket=2 + └── val="A" idx=1, bucket=1 + "#); + + Ok(()) + } + + #[test] + fn should_heapify_down_string() -> Result<()> { + let mut map = vec![]; + let mut heap = TopKHeap::::new(3, false); + + heap.append_or_replace("B", 1, &mut map); + heap.append_or_replace("C", 2, &mut map); + heap.append_or_replace("D", 3, &mut map); + let actual = heap.to_string(); + assert_snapshot!(actual, @r#" + val="D" idx=0, bucket=3 + ├── val="B" idx=1, bucket=1 + └── val="C" idx=2, bucket=2 + "#); + + let mut map = vec![]; + heap.append_or_replace("A", 0, &mut map); + let actual = heap.to_string(); + assert_snapshot!(actual, @r#" + val="C" idx=0, bucket=2 + ├── val="B" idx=1, bucket=1 + └── val="A" idx=2, bucket=0 + "#); + assert_eq!(map, vec![(2, 0), (0, 2)]); + + Ok(()) + } + + #[test] + fn should_replace_string() -> Result<()> { + let mut map = vec![]; + let mut heap = TopKHeap::::new(4, false); + + heap.append_or_replace("B", 1, &mut map); + heap.append_or_replace("C", 2, &mut map); + heap.append_or_replace("D", 3, &mut map); + heap.append_or_replace("E", 4, &mut map); + let actual = heap.to_string(); + assert_snapshot!(actual, @r#" + val="E" idx=0, bucket=4 + ├── val="D" idx=1, bucket=3 + │ └── val="B" idx=3, bucket=1 + └── val="C" idx=2, bucket=2 + "#); + + let mut map = vec![]; + heap.replace_if_better("A", 1, &mut map); + let actual = heap.to_string(); + assert_snapshot!(actual, @r#" + val="E" idx=0, bucket=4 + ├── val="B" idx=1, bucket=1 + │ └── val="A" idx=3, bucket=3 + └── val="C" idx=2, bucket=2 + "#); + assert_eq!(map, vec![(1, 1), (3, 3)]); + + Ok(()) + } + + #[test] + fn should_find_worst_string() -> Result<()> { + let mut map = vec![]; + let mut heap = TopKHeap::::new(10, false); + + heap.append_or_replace("A", 1, &mut map); + heap.append_or_replace("B", 2, &mut map); + + let actual = heap.to_string(); + assert_snapshot!(actual, @r#" + val="B" idx=0, bucket=2 + └── val="A" idx=1, bucket=1 + "#); + + assert_eq!(heap.worst_val().map(String::as_str), Some("B")); + assert_eq!(heap.worst_map_idx(), 2); + + Ok(()) + } + + #[test] + fn should_drain_string() -> Result<()> { + let mut map = vec![]; + let mut heap = TopKHeap::::new(10, false); + + heap.append_or_replace("A", 1, &mut map); + heap.append_or_replace("B", 2, &mut map); + + let actual = heap.to_string(); + assert_snapshot!(actual, @r#" + val="B" idx=0, bucket=2 + └── val="A" idx=1, bucket=1 + "#); + + let (vals, map_idxs) = heap.drain(); + assert_eq!(vals, vec!["A", "B"]); + assert_eq!(map_idxs, vec![1, 2]); + assert_eq!(heap.len(), 0); + + Ok(()) + } } diff --git a/datafusion/physical-plan/src/aggregates/topk/priority_map.rs b/datafusion/physical-plan/src/aggregates/topk/priority_map.rs index f36e0f762849c..008580882c798 100644 --- a/datafusion/physical-plan/src/aggregates/topk/priority_map.rs +++ b/datafusion/physical-plan/src/aggregates/topk/priority_map.rs @@ -30,7 +30,6 @@ pub struct PriorityMap { heap: Box, capacity: usize, mapper: Vec<(usize, usize)>, - val_type: DataType, /// Mirror of the map's all-NULL group count, kept as a plain field so the /// per-row `insert` path can check it without a `dyn` call (measured to /// regress the topk_aggregate benchmarks when read through the trait) @@ -46,10 +45,9 @@ impl PriorityMap { ) -> Result { Ok(Self { map: new_hash_table(capacity, key_type)?, - heap: new_heap(capacity, descending, val_type.clone())?, + heap: new_heap(capacity, descending, val_type)?, capacity, mapper: Vec::with_capacity(capacity), - val_type, null_count: 0, }) } @@ -60,13 +58,14 @@ impl PriorityMap { } pub fn insert(&mut self, row_idx: usize) -> Result<()> { - assert!(self.map.len() <= self.capacity, "Overflow"); debug_assert_eq!(self.null_count, 0); // if we're full, and the new val is worse than all our values, just bail if self.heap.is_worse(row_idx) { return Ok(()); } + assert!(self.map.len() <= self.capacity, "Overflow"); + self.insert_eligible(row_idx) } @@ -74,10 +73,6 @@ impl PriorityMap { /// separate from [`Self::insert`] so the common no-NULL path does not pay /// for NULL bookkeeping on every row. pub fn insert_with_null_groups(&mut self, row_idx: usize) -> Result<()> { - // valued groups are capped at `capacity`; up to `capacity` additional - // all-NULL groups may be tracked alongside them - assert!(self.map.len() <= 2 * self.capacity, "Overflow"); - if self.heap.is_worse(row_idx) { // A group that was registered as all-NULL now has a value that // loses to the current top-k: it can no longer reach the top-k, @@ -87,9 +82,19 @@ impl PriorityMap { } return Ok(()); } + self.null_capacity_assert(); + self.insert_eligible(row_idx) } + #[inline] + fn null_capacity_assert(&self) { + // TODO: do debug_assert instead? + // valued groups are capped at `capacity`; up to `capacity` additional + // all-NULL groups may be tracked alongside them + assert!(self.map.len() <= 2 * self.capacity, "Overflow"); + } + fn insert_eligible(&mut self, row_idx: usize) -> Result<()> { let map = &mut self.mapper; @@ -125,8 +130,8 @@ impl PriorityMap { /// must still appear in the aggregation output; such groups all tie on the /// sort key, so tracking up to `capacity` of them preserves top-k semantics. pub fn insert_null(&mut self, row_idx: usize) { - assert!(self.map.len() <= 2 * self.capacity, "Overflow"); if self.map.insert_null(row_idx) { + self.null_capacity_assert(); self.null_count += 1; } } @@ -135,12 +140,11 @@ impl PriorityMap { let (vals, mut map_idxs) = self.heap.drain(); // Groups whose values are all NULL are tracked in the map only; // append them with a NULL value so they are not lost from the output - let null_idxs = self.map.null_map_idxs(); - let vals = if null_idxs.is_empty() { + let null_count = self.map.null_map_idxs(&mut map_idxs); + let vals = if null_count == 0 { vals } else { - map_idxs.extend(null_idxs.iter().copied()); - let nulls = new_null_array(&self.val_type, null_idxs.len()); + let nulls = new_null_array(self.heap.value_type(), null_count); concat(&[vals.as_ref(), nulls.as_ref()])? }; let ids = self.map.take_all(map_idxs); @@ -611,6 +615,31 @@ mod tests { Ok(()) } + #[test] + fn should_emit_null_id_with_all_null_values() -> Result<()> { + let ids: ArrayRef = Arc::new(StringArray::from(vec![None, Some("1")])); + let vals: ArrayRef = Arc::new(Int64Array::from(vec![None, Some(5)])); + let mut agg = PriorityMap::new(DataType::Utf8, DataType::Int64, 2, true)?; + agg.set_batch(ids, vals); + agg.insert_null(0); + agg.insert_with_null_groups(1)?; + + let cols = agg.emit()?; + let batch = RecordBatch::try_new(test_schema(), cols)?; + let actual = format!("{}", pretty_format_batches(&[batch])?); + assert_snapshot!(actual, @r" + +----------+--------------+ + | trace_id | timestamp_ms | + +----------+--------------+ + | 1 | 5 | + | | | + +----------+--------------+ + " + ); + + Ok(()) + } + #[test] fn should_cap_null_groups_at_limit() -> Result<()> { let ids: ArrayRef = Arc::new(StringArray::from(vec!["1", "2", "3", "4", "5"]));