diff --git a/datafusion-testing b/datafusion-testing index 13bbae38776c2..7833a65d5b08b 160000 --- a/datafusion-testing +++ b/datafusion-testing @@ -1 +1 @@ -Subproject commit 13bbae38776c2bfbc1fab1be7e7220222d4284bf +Subproject commit 7833a65d5b08be2ca484ea938f471cf01df54e18 diff --git a/datafusion/physical-plan/src/aggregates/topk/hash_table.rs b/datafusion/physical-plan/src/aggregates/topk/hash_table.rs index adc8f8c315b32..2e2936a1dcec0 100644 --- a/datafusion/physical-plan/src/aggregates/topk/hash_table.rs +++ b/datafusion/physical-plan/src/aggregates/topk/hash_table.rs @@ -24,21 +24,18 @@ use arrow::array::{ Array, ArrayRef, ArrowPrimitiveType, LargeStringArray, PrimitiveArray, StringArray, StringViewArray, builder::PrimitiveBuilder, cast::AsArray, downcast_primitive, }; +use arrow::array::{ArrayAccessor, StringArrayType}; use arrow::datatypes::{DataType, i256}; 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 {} - -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. @@ -48,9 +45,9 @@ const NULL_HEAP_IDX: usize = usize::MAX; /// 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, } @@ -58,12 +55,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) @@ -120,11 +119,13 @@ pub fn is_supported_hash_key_type(kt: &DataType) -> bool { } // An implementation of ArrowHashTable for String keys -pub struct StringHashTable { - owned: ArrayRef, - map: TopKHashTable>, +pub struct StringHashTable +where + for<'a> &'a S: StringArrayType<'a>, +{ + owned: S, + map: TopKHashTable, rnd: RandomState, - data_type: DataType, } // An implementation of ArrowHashTable for any `ArrowPrimitiveType` key @@ -132,69 +133,42 @@ struct PrimitiveHashTable where Option<::Native>: Comparable, { - owned: ArrayRef, - map: TopKHashTable>, + owned: PrimitiveArray, + map: TopKHashTable, rnd: RandomState, - kt: DataType, } -impl StringHashTable { - pub fn new(limit: usize, data_type: DataType) -> Self { - let vals: Vec<&str> = Vec::new(); - let owned: ArrayRef = match data_type { - DataType::Utf8 => Arc::new(StringArray::from(vals)), - DataType::Utf8View => Arc::new(StringViewArray::from(vals)), - DataType::LargeUtf8 => Arc::new(LargeStringArray::from(vals)), - _ => panic!("Unsupported data type"), - }; - +impl StringHashTable +where + S: Array + From>>, + for<'a> &'a S: StringArrayType<'a>, +{ + pub fn new(limit: usize) -> Self { + let owned = S::from(Vec::new()); Self { owned, map: TopKHashTable::new(limit, limit * 10), rnd: RandomState::default(), - data_type, } } - /// Extracts the string value at the given row index, handling nulls and different string types. - /// - /// Returns `None` if the value is null, otherwise `Some(value.to_string())`. - fn extract_string_value(&self, row_idx: usize) -> Option { - let is_null_and_value = match self.data_type { - DataType::Utf8 => { - let arr = self.owned.as_string::(); - (arr.is_null(row_idx), arr.value(row_idx)) - } - DataType::LargeUtf8 => { - let arr = self.owned.as_string::(); - (arr.is_null(row_idx), arr.value(row_idx)) - } - DataType::Utf8View => { - let arr = self.owned.as_string_view(); - (arr.is_null(row_idx), arr.value(row_idx)) - } - _ => panic!("Unsupported data type"), - }; - - let (is_null, value) = is_null_and_value; - if is_null { - None - } else { - Some(value.to_string()) - } - } - - /// Computes the id and its hash for the given row, for hash table lookups - fn id_and_hash(&self, row_idx: usize) -> (Option, u64) { - let id = self.extract_string_value(row_idx); - let hash = self.rnd.hash_one(id.as_deref()); - (id, hash) + #[inline] + fn eq_fn(id: Option<&str>) -> impl Fn(&Option) -> bool { + move |mi| id == mi.as_deref() } } -impl ArrowHashTable for StringHashTable { +impl ArrowHashTable for StringHashTable +where + S: Array + Clone + From>> + 'static, + for<'a> &'a S: StringArrayType<'a>, +{ fn set_batch(&mut self, ids: ArrayRef) { - self.owned = ids; + self.owned = ids + .as_any() + .downcast_ref::() + .expect("Unsupported data type") + .clone(); } fn len(&self) -> usize { @@ -211,12 +185,7 @@ impl ArrowHashTable for StringHashTable { fn take_all(&mut self, indexes: Vec) -> ArrayRef { let ids = self.map.take_all(indexes); - match self.data_type { - DataType::Utf8 => Arc::new(StringArray::from(ids)), - DataType::LargeUtf8 => Arc::new(LargeStringArray::from(ids)), - DataType::Utf8View => Arc::new(StringViewArray::from(ids)), - _ => unreachable!(), - } + Arc::new(S::from(ids)) } fn find_or_insert( @@ -224,28 +193,25 @@ impl ArrowHashTable for StringHashTable { row_idx: usize, replace_idx: usize, ) -> (usize, InsertKind) { - let id = self.extract_string_value(row_idx); - + let id = some_value(&self.owned, row_idx); // Compute hash and create equality closure for hash table lookup. - let hash = self.rnd.hash_one(id.as_deref()); - let id_for_eq = id.clone(); - let eq = move |mi: &Option| id_for_eq.as_deref() == mi.as_deref(); + let hash = self.rnd.hash_one(id); // Use entry API to avoid double lookup - self.map.find_or_insert(hash, id, replace_idx, eq) + self.map + .find_or_insert(hash, id, replace_idx, Self::eq_fn(id)) } fn insert_null(&mut self, row_idx: usize) -> bool { - let (id, hash) = self.id_and_hash(row_idx); - let id_for_eq = id.clone(); - let eq = move |mi: &Option| id_for_eq.as_deref() == mi.as_deref(); - self.map.insert_null(hash, id, eq) + let id = some_value(&self.owned, row_idx); + let hash = self.rnd.hash_one(id); + self.map.insert_null(hash, id, Self::eq_fn(id)) } fn remove_if_null(&mut self, row_idx: usize) -> bool { - let (id, hash) = self.id_and_hash(row_idx); - let eq = move |mi: &Option| id.as_deref() == mi.as_deref(); - self.map.remove_if_null(hash, eq) + let id = some_value(&self.owned, row_idx); + let hash = self.rnd.hash_one(id); + self.map.remove_if_null(hash, Self::eq_fn(id)) } fn null_map_idxs(&self) -> Vec { @@ -255,43 +221,39 @@ impl ArrowHashTable for StringHashTable { impl PrimitiveHashTable where - Option<::Native>: Comparable, - Option<::Native>: HashValue, + Option<::Native>: Comparable + HashValue, { pub fn new(limit: usize, kt: DataType) -> Self { - let owned = Arc::new( - PrimitiveArray::::builder(0) - .with_data_type(kt.clone()) - .finish(), - ); + let owned = PrimitiveArray::::builder(0) + .with_data_type(kt) + .finish(); Self { owned, map: TopKHashTable::new(limit, limit * 10), rnd: RandomState::default(), - kt, } } /// Computes the id and its hash for the given row, for hash table lookups + #[inline] fn id_and_hash(&self, row_idx: usize) -> (Option, u64) { - let ids = self.owned.as_primitive::(); - let id: Option = if ids.is_null(row_idx) { - None - } else { - Some(ids.value(row_idx)) - }; + let id: Option = some_value(&self.owned, row_idx); let hash: u64 = id.hash(&self.rnd); (id, hash) } + + #[inline] + fn eq_fn(id: Option) -> impl Fn(&Option) -> bool { + move |mi| id == *mi + } } impl ArrowHashTable for PrimitiveHashTable where - Option<::Native>: Comparable, - Option<::Native>: HashValue, + Option<::Native>: Comparable + HashValue, { fn set_batch(&mut self, ids: ArrayRef) { - self.owned = ids; + self.owned = ids.as_primitive().clone(); } fn len(&self) -> usize { @@ -308,13 +270,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.kt.clone()); + let mut builder: PrimitiveBuilder = PrimitiveArray::builder(ids.len()) + .with_data_type(self.owned.data_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) @@ -325,30 +284,20 @@ where row_idx: usize, replace_idx: usize, ) -> (usize, InsertKind) { - let ids = self.owned.as_primitive::(); - let id: Option = if ids.is_null(row_idx) { - None - } else { - Some(ids.value(row_idx)) - }; - // Compute hash and create equality closure for hash table lookup. - let hash: u64 = id.hash(&self.rnd); - let eq = |mi: &Option| id == *mi; - + 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, eq) + self.map + .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); - let eq = move |mi: &Option| id == *mi; - self.map.insert_null(hash, id, eq) + self.map.insert_null(hash, id.as_ref(), Self::eq_fn(id)) } fn remove_if_null(&mut self, row_idx: usize) -> bool { let (id, hash) = self.id_and_hash(row_idx); - let eq = move |mi: &Option| id == *mi; - self.map.remove_if_null(hash, eq) + self.map.remove_if_null(hash, Self::eq_fn(id)) } fn null_map_idxs(&self) -> Vec { @@ -356,34 +305,39 @@ where } } -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; + match self.store[removed_idx].id.take() { + Some(slot) if Self::use_free_slots() => { + self.free_slots.push(slot); + } + _ => (), + } self.free_indices.push(removed_idx); } Entry::Vacant(_) => unreachable!(), @@ -404,63 +358,93 @@ 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, + ID: BorrowMut, + { // 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, + ID: BorrowMut, + { + let id = if Self::use_free_slots() { + id.map(|id| match self.free_slots.pop() { + Some(mut slot) => { + id.clone_into(slot.borrow_mut()); + 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].id.is_none(), "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 @@ -468,36 +452,26 @@ 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); - if self.map.find(hash, eq_fn).is_some() { - return false; - } + 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; } + if self.null_count >= self.limit { 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 } @@ -506,10 +480,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; @@ -524,9 +502,7 @@ impl TopKHashTable { .iter() .enumerate() .filter_map(|(idx, item)| { - item.as_ref() - .filter(|item| item.heap_idx == NULL_HEAP_IDX) - .map(|_| idx) + (item.id.is_some() && item.is_null()).then_some(idx) }) .collect() } @@ -535,10 +511,10 @@ impl TopKHashTable { 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(); @@ -548,10 +524,15 @@ impl TopKHashTable { } } -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 + } } impl HashValue for Option { @@ -585,6 +566,18 @@ has_integer!(u8, u16, u32, u64); has_integer!(IntervalDayTime, IntervalMonthDayNano); hash_float!(f16, f32, f64); +#[inline] +fn some_value<'a, A>(array: &'a A, index: usize) -> Option<<&'a A as ArrayAccessor>::Item> +where + &'a A: ArrayAccessor, +{ + if array.is_null(index) { + None + } else { + Some(array.value(index)) + } +} + pub fn new_hash_table( limit: usize, kt: DataType, @@ -597,9 +590,9 @@ pub fn new_hash_table( downcast_primitive! { kt => (downcast_helper, kt), - DataType::Utf8 => return Ok(Box::new(StringHashTable::new(limit, DataType::Utf8))), - DataType::LargeUtf8 => return Ok(Box::new(StringHashTable::new(limit, DataType::LargeUtf8))), - DataType::Utf8View => return Ok(Box::new(StringHashTable::new(limit, DataType::Utf8View))), + DataType::Utf8 => return Ok(Box::new(StringHashTable::::new(limit))), + DataType::LargeUtf8 => return Ok(Box::new(StringHashTable::::new(limit))), + DataType::Utf8View => return Ok(Box::new(StringHashTable::::new(limit))), _ => {} } @@ -633,14 +626,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); } @@ -666,23 +659,23 @@ 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!(!map.insert_null(100, a.as_ref(), |v| *v == a)); assert_eq!(map.null_count, 2); assert_eq!(map.null_map_idxs(), vec![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); @@ -702,24 +695,24 @@ mod tests { #[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 ca321cdf99784..819041f76b45f 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, LargeStringArray, PrimitiveArray, + StringArray, StringArrayType, StringViewArray, downcast_primitive, }; use arrow::buffer::ScalarBuffer; use arrow::datatypes::{DataType, i256}; @@ -95,23 +93,21 @@ pub struct PrimitiveHeap where ::Native: Comparable, { - batch: ArrayRef, + batch: PrimitiveArray, heap: TopKHeap, desc: bool, - data_type: DataType, } impl PrimitiveHeap where ::Native: Comparable, { - pub fn new(limit: usize, desc: bool, data_type: DataType) -> Self { - let owned: ArrayRef = Arc::new(PrimitiveArray::::builder(0).finish()); + pub fn new(limit: usize, desc: bool) -> Self { + let batch = PrimitiveArray::::builder(0).finish(); Self { - batch: owned, + batch, heap: TopKHeap::new(limit, desc), desc, - data_type, } } } @@ -121,15 +117,14 @@ where ::Native: Comparable, { fn set_batch(&mut self, vals: ArrayRef) { - self.batch = vals; + self.batch = PrimitiveArray::from(vals.to_data()); } fn is_worse(&self, row_idx: usize) -> bool { if !self.heap.is_full() { return false; } - let vals = self.batch.as_primitive::(); - let new_val = vals.value(row_idx); + let new_val = self.batch.value(row_idx); let worst_val = self.heap.worst_val().expect("Missing root"); (!self.desc && new_val > *worst_val) || (self.desc && new_val < *worst_val) } @@ -139,8 +134,7 @@ where } fn insert(&mut self, row_idx: usize, map_idx: usize, map: &mut Vec<(usize, usize)>) { - let vals = self.batch.as_primitive::(); - let new_val = vals.value(row_idx); + let new_val = self.batch.value(row_idx); self.heap.append_or_replace(new_val, map_idx, map); } @@ -150,8 +144,7 @@ where row_idx: usize, map: &mut Vec<(usize, usize)>, ) { - let vals = self.batch.as_primitive::(); - let new_val = vals.value(row_idx); + let new_val = self.batch.value(row_idx); self.heap.replace_if_better(heap_idx, new_val, map); } @@ -159,7 +152,7 @@ where 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.batch.data_type().clone()); (Arc::new(arr), map_idxs) } } @@ -171,58 +164,41 @@ 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, +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 { +impl ArrowHeap for StringHeap +where + S: Array + Clone + From>> + 'static, + for<'a> &'a S: StringArrayType<'a>, +{ 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 { @@ -232,7 +208,7 @@ 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 new_val = (&self.batch).value(row_idx); let worst_val = self.heap.worst_val().expect("Missing root"); match worst_val { None => false, @@ -252,7 +228,7 @@ impl ArrowHeap for StringHeap { // 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_str = (&self.batch).value(row_idx).to_string(); let new_val = Some(new_str); self.heap.append_or_replace(new_val, map_idx, map); } @@ -263,7 +239,7 @@ impl ArrowHeap for StringHeap { row_idx: usize, map: &mut Vec<(usize, usize)>, ) { - let new_str = self.value(row_idx); + let new_str = (&self.batch).value(row_idx); let existing = self.heap.heap[heap_idx] .as_ref() .expect("Missing heap item"); @@ -292,33 +268,8 @@ impl ArrowHeap for StringHeap { 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) } } @@ -443,14 +394,13 @@ impl TopKHeap { } fn swap(&mut self, a_idx: usize, b_idx: usize, mapper: &mut Vec<(usize, usize)>) { - let a_hi = self.heap[a_idx].take().expect("Missing heap entry"); - let b_hi = self.heap[b_idx].take().expect("Missing heap entry"); + self.heap.swap(a_idx, b_idx); - mapper.push((a_hi.map_idx, b_idx)); - mapper.push((b_hi.map_idx, a_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"); - self.heap[a_idx] = Some(b_hi); - self.heap[b_idx] = Some(a_hi); + mapper.push((b_hi.map_idx, b_idx)); + mapper.push((a_hi.map_idx, a_idx)); } fn heapify_down(&mut self, node_idx: usize, mapper: &mut Vec<(usize, usize)>) { @@ -619,21 +569,17 @@ 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))) + return Ok(Box::new(PrimitiveHeap::<$vt>::new(limit, desc))) }; } 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))), _ => {} }