From 2a56ad5fa05cad4ba51d187b9a1eb80157756b2f Mon Sep 17 00:00:00 2001 From: Michal Piatkowski <291740709+MassivePizza@users.noreply.github.com> Date: Tue, 14 Jul 2026 11:14:10 +0200 Subject: [PATCH 01/17] store concrete type in TopK structures --- .../src/aggregates/topk/hash_table.rs | 23 ++++++++----------- .../physical-plan/src/aggregates/topk/heap.rs | 17 ++++++-------- 2 files changed, 16 insertions(+), 24 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/topk/hash_table.rs b/datafusion/physical-plan/src/aggregates/topk/hash_table.rs index 694780f08547f..66fb8f802c0a8 100644 --- a/datafusion/physical-plan/src/aggregates/topk/hash_table.rs +++ b/datafusion/physical-plan/src/aggregates/topk/hash_table.rs @@ -99,7 +99,7 @@ struct PrimitiveHashTable where Option<::Native>: Comparable, { - owned: ArrayRef, + owned: PrimitiveArray, map: TopKHashTable>, rnd: RandomState, kt: DataType, @@ -194,15 +194,12 @@ 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.clone()) + .finish(); Self { owned, map: TopKHashTable::new(limit, limit * 10), @@ -214,11 +211,10 @@ where 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 = PrimitiveArray::from(ids.to_data()); } fn len(&self) -> usize { @@ -248,11 +244,10 @@ where } fn find_or_insert(&mut self, row_idx: usize, replace_idx: usize) -> (usize, bool) { - let ids = self.owned.as_primitive::(); - let id: Option = if ids.is_null(row_idx) { + let id: Option = if self.owned.is_null(row_idx) { None } else { - Some(ids.value(row_idx)) + Some(self.owned.value(row_idx)) }; // Compute hash and create equality closure for hash table lookup. let hash: u64 = id.hash(&self.rnd); diff --git a/datafusion/physical-plan/src/aggregates/topk/heap.rs b/datafusion/physical-plan/src/aggregates/topk/heap.rs index ca321cdf99784..ad33a39addae3 100644 --- a/datafusion/physical-plan/src/aggregates/topk/heap.rs +++ b/datafusion/physical-plan/src/aggregates/topk/heap.rs @@ -95,7 +95,7 @@ pub struct PrimitiveHeap where ::Native: Comparable, { - batch: ArrayRef, + batch: PrimitiveArray, heap: TopKHeap, desc: bool, data_type: DataType, @@ -106,9 +106,9 @@ where ::Native: Comparable, { pub fn new(limit: usize, desc: bool, data_type: DataType) -> Self { - let owned: ArrayRef = Arc::new(PrimitiveArray::::builder(0).finish()); + let batch = PrimitiveArray::::builder(0).finish(); Self { - batch: owned, + batch, heap: TopKHeap::new(limit, desc), desc, data_type, @@ -121,15 +121,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 +138,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 +148,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); } From 14a5c02289ef6699a28b52d5a6e78aa271ea82ba Mon Sep 17 00:00:00 2001 From: Michal Piatkowski <291740709+MassivePizza@users.noreply.github.com> Date: Tue, 14 Jul 2026 11:36:37 +0200 Subject: [PATCH 02/17] store concrete string type in TopK hash-table --- .../src/aggregates/topk/hash_table.rs | 104 +++++++++--------- 1 file changed, 53 insertions(+), 51 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/topk/hash_table.rs b/datafusion/physical-plan/src/aggregates/topk/hash_table.rs index 66fb8f802c0a8..ba09927d05d38 100644 --- a/datafusion/physical-plan/src/aggregates/topk/hash_table.rs +++ b/datafusion/physical-plan/src/aggregates/topk/hash_table.rs @@ -73,25 +73,64 @@ pub trait ArrowHashTable { fn find_or_insert(&mut self, row_idx: usize, replace_idx: usize) -> (usize, bool); } +enum StringArrayType { + Utf8(StringArray), + Utf8View(StringViewArray), + LargeUtf8(LargeStringArray), +} +impl StringArrayType { + /// 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)`. + fn value(&self, row_idx: usize) -> Option<&str> { + let (is_null, value) = match self { + StringArrayType::Utf8(arr) => (arr.is_null(row_idx), arr.value(row_idx)), + StringArrayType::LargeUtf8(arr) => (arr.is_null(row_idx), arr.value(row_idx)), + StringArrayType::Utf8View(arr) => (arr.is_null(row_idx), arr.value(row_idx)), + }; + if is_null { None } else { Some(value) } + } +} +impl<'a> TryFrom<&'a DataType> for StringArrayType { + type Error = (); + + fn try_from(data_type: &'a DataType) -> std::result::Result { + let vals: Vec<&str> = Vec::new(); + Ok(match data_type { + DataType::Utf8 => StringArrayType::Utf8(vals.into()), + DataType::Utf8View => StringArrayType::Utf8View(vals.into()), + DataType::LargeUtf8 => StringArrayType::LargeUtf8(vals.into()), + _ => return Err(()), + }) + } +} +impl TryFrom for StringArrayType { + type Error = DataType; + + fn try_from(arr: ArrayRef) -> std::result::Result { + Ok(match arr.data_type() { + DataType::Utf8 => StringArrayType::Utf8(arr.as_string().clone()), + DataType::LargeUtf8 => StringArrayType::LargeUtf8(arr.as_string().clone()), + DataType::Utf8View => StringArrayType::Utf8View(arr.as_string_view().clone()), + ty => return Err(ty.clone()), + }) + } +} + /// Returns true if the given data type can be used as a top-K aggregation hash key. /// /// Supported types include Arrow primitives (integers, floats, decimals, intervals) /// and UTF-8 strings (`Utf8`, `LargeUtf8`, `Utf8View`). This is used internally by /// `PriorityMap::supports()` to validate grouping key type compatibility. pub fn is_supported_hash_key_type(kt: &DataType) -> bool { - kt.is_primitive() - || matches!( - kt, - DataType::Utf8 | DataType::Utf8View | DataType::LargeUtf8 - ) + kt.is_primitive() || StringArrayType::try_from(kt).is_ok() } // An implementation of ArrowHashTable for String keys pub struct StringHashTable { - owned: ArrayRef, + owned: StringArrayType, map: TopKHashTable>, rnd: RandomState, - data_type: DataType, } // An implementation of ArrowHashTable for any `ArrowPrimitiveType` key @@ -107,54 +146,18 @@ where 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"), - }; - + let owned = StringArrayType::try_from(&data_type).expect("Unsupported data type"); 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()) } } } impl ArrowHashTable for StringHashTable { fn set_batch(&mut self, ids: ArrayRef) { - self.owned = ids; + self.owned = StringArrayType::try_from(ids).expect("Unsupported data type"); } fn len(&self) -> usize { @@ -171,16 +174,15 @@ 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!(), + match &self.owned { + StringArrayType::Utf8(_) => Arc::new(StringArray::from(ids)), + StringArrayType::LargeUtf8(_) => Arc::new(LargeStringArray::from(ids)), + StringArrayType::Utf8View(_) => Arc::new(StringViewArray::from(ids)), } } fn find_or_insert(&mut self, row_idx: usize, replace_idx: usize) -> (usize, bool) { - let id = self.extract_string_value(row_idx); + let id = self.owned.value(row_idx); // Compute hash and create equality closure for hash table lookup. let hash = self.rnd.hash_one(id.as_deref()); From 293a0e86111b07710b34f11a3af9f08f255a5064 Mon Sep 17 00:00:00 2001 From: Michal Piatkowski <291740709+MassivePizza@users.noreply.github.com> Date: Tue, 14 Jul 2026 11:39:35 +0200 Subject: [PATCH 03/17] swap in-place in TopK heap --- datafusion/physical-plan/src/aggregates/topk/heap.rs | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/topk/heap.rs b/datafusion/physical-plan/src/aggregates/topk/heap.rs index ad33a39addae3..aef6bb5596c2e 100644 --- a/datafusion/physical-plan/src/aggregates/topk/heap.rs +++ b/datafusion/physical-plan/src/aggregates/topk/heap.rs @@ -440,14 +440,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)>) { From 3b28922a0284e51326abab79b95b9fc0b61780ba Mon Sep 17 00:00:00 2001 From: Michal Piatkowski <291740709+MassivePizza@users.noreply.github.com> Date: Wed, 15 Jul 2026 17:08:16 +0200 Subject: [PATCH 04/17] fix error from cherry-pick --- datafusion/physical-plan/src/aggregates/topk/hash_table.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datafusion/physical-plan/src/aggregates/topk/hash_table.rs b/datafusion/physical-plan/src/aggregates/topk/hash_table.rs index ba09927d05d38..0fd8045c7cef8 100644 --- a/datafusion/physical-plan/src/aggregates/topk/hash_table.rs +++ b/datafusion/physical-plan/src/aggregates/topk/hash_table.rs @@ -190,7 +190,7 @@ impl ArrowHashTable for StringHashTable { let eq = move |mi: &Option| id_for_eq.as_deref() == mi.as_deref(); // Use entry API to avoid double lookup - self.map.find_or_insert(hash, id, replace_idx, eq) + self.map.find_or_insert(hash, id.map(ToOwned::to_owned), replace_idx, eq) } } From 6214dfa659cf970eac37ab0a254b8b3683c6c307 Mon Sep 17 00:00:00 2001 From: Michal Piatkowski <291740709+MassivePizza@users.noreply.github.com> Date: Thu, 16 Jul 2026 17:40:35 +0200 Subject: [PATCH 05/17] fmt and clippy --- .../src/aggregates/topk/hash_table.rs | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/topk/hash_table.rs b/datafusion/physical-plan/src/aggregates/topk/hash_table.rs index 0fd8045c7cef8..28c17b60658ed 100644 --- a/datafusion/physical-plan/src/aggregates/topk/hash_table.rs +++ b/datafusion/physical-plan/src/aggregates/topk/hash_table.rs @@ -145,8 +145,8 @@ where } impl StringHashTable { - pub fn new(limit: usize, data_type: DataType) -> Self { - let owned = StringArrayType::try_from(&data_type).expect("Unsupported data type"); + pub fn new(limit: usize, data_type: &DataType) -> Self { + let owned = StringArrayType::try_from(data_type).expect("Unsupported data type"); Self { owned, map: TopKHashTable::new(limit, limit * 10), @@ -185,12 +185,12 @@ impl ArrowHashTable for StringHashTable { let id = self.owned.value(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); + let eq = move |mi: &Option| id == mi.as_deref(); // Use entry API to avoid double lookup - self.map.find_or_insert(hash, id.map(ToOwned::to_owned), replace_idx, eq) + self.map + .find_or_insert(hash, id.map(ToOwned::to_owned), replace_idx, eq) } } @@ -408,9 +408,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::Utf8))), + DataType::LargeUtf8 => return Ok(Box::new(StringHashTable::new(limit, &DataType::LargeUtf8))), + DataType::Utf8View => return Ok(Box::new(StringHashTable::new(limit, &DataType::Utf8View))), _ => {} } From a82f5fae712b019c167a94cd4fb0630df652c5da Mon Sep 17 00:00:00 2001 From: Michal Piatkowski <291740709+MassivePizza@users.noreply.github.com> Date: Mon, 27 Jul 2026 11:43:45 +0200 Subject: [PATCH 06/17] remove redundant data_type in PrimitiveHashTable --- .../physical-plan/src/aggregates/topk/hash_table.rs | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/topk/hash_table.rs b/datafusion/physical-plan/src/aggregates/topk/hash_table.rs index 28c17b60658ed..dffdf40607e98 100644 --- a/datafusion/physical-plan/src/aggregates/topk/hash_table.rs +++ b/datafusion/physical-plan/src/aggregates/topk/hash_table.rs @@ -141,7 +141,6 @@ where owned: PrimitiveArray, map: TopKHashTable>, rnd: RandomState, - kt: DataType, } impl StringHashTable { @@ -200,13 +199,12 @@ where { pub fn new(limit: usize, kt: DataType) -> Self { let owned = PrimitiveArray::::builder(0) - .with_data_type(kt.clone()) + .with_data_type(kt) .finish(); Self { owned, map: TopKHashTable::new(limit, limit * 10), rnd: RandomState::default(), - kt, } } } @@ -233,8 +231,8 @@ 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(), From dfbf3f9c25bb713b2b8e96acd58878fc5d9084f8 Mon Sep 17 00:00:00 2001 From: Michal Piatkowski <291740709+MassivePizza@users.noreply.github.com> Date: Mon, 27 Jul 2026 11:59:36 +0200 Subject: [PATCH 07/17] Undo StringArrayType convenience check, avoiding from(Vec) into drop. --- datafusion/physical-plan/src/aggregates/topk/hash_table.rs | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/datafusion/physical-plan/src/aggregates/topk/hash_table.rs b/datafusion/physical-plan/src/aggregates/topk/hash_table.rs index dffdf40607e98..5117ba380d558 100644 --- a/datafusion/physical-plan/src/aggregates/topk/hash_table.rs +++ b/datafusion/physical-plan/src/aggregates/topk/hash_table.rs @@ -123,7 +123,11 @@ impl TryFrom for StringArrayType { /// and UTF-8 strings (`Utf8`, `LargeUtf8`, `Utf8View`). This is used internally by /// `PriorityMap::supports()` to validate grouping key type compatibility. pub fn is_supported_hash_key_type(kt: &DataType) -> bool { - kt.is_primitive() || StringArrayType::try_from(kt).is_ok() + kt.is_primitive() + || matches!( + kt, + DataType::Utf8 | DataType::Utf8View | DataType::LargeUtf8 + ) } // An implementation of ArrowHashTable for String keys From 5e07c8db7dd4278cf2a83d99b1248a1e289f4f94 Mon Sep 17 00:00:00 2001 From: Michal Piatkowski <291740709+MassivePizza@users.noreply.github.com> Date: Mon, 27 Jul 2026 12:55:59 +0200 Subject: [PATCH 08/17] simplify set_batch in PrimitiveHashTable --- datafusion/physical-plan/src/aggregates/topk/hash_table.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datafusion/physical-plan/src/aggregates/topk/hash_table.rs b/datafusion/physical-plan/src/aggregates/topk/hash_table.rs index 5117ba380d558..b60484d806ec3 100644 --- a/datafusion/physical-plan/src/aggregates/topk/hash_table.rs +++ b/datafusion/physical-plan/src/aggregates/topk/hash_table.rs @@ -218,7 +218,7 @@ where Option<::Native>: Comparable + HashValue, { fn set_batch(&mut self, ids: ArrayRef) { - self.owned = PrimitiveArray::from(ids.to_data()); + self.owned = ids.as_primitive().clone(); } fn len(&self) -> usize { From 55a2e334a8de3bd4bd7a78f2ff650abfe6233841 Mon Sep 17 00:00:00 2001 From: Michal Piatkowski <291740709+MassivePizza@users.noreply.github.com> Date: Mon, 3 Aug 2026 13:47:43 +0200 Subject: [PATCH 09/17] replace StringArrayType enum with arrow trait --- .../src/aggregates/topk/hash_table.rs | 92 +++++++------------ 1 file changed, 32 insertions(+), 60 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/topk/hash_table.rs b/datafusion/physical-plan/src/aggregates/topk/hash_table.rs index b60484d806ec3..2d48451d1ff0c 100644 --- a/datafusion/physical-plan/src/aggregates/topk/hash_table.rs +++ b/datafusion/physical-plan/src/aggregates/topk/hash_table.rs @@ -24,6 +24,7 @@ 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; @@ -73,50 +74,6 @@ pub trait ArrowHashTable { fn find_or_insert(&mut self, row_idx: usize, replace_idx: usize) -> (usize, bool); } -enum StringArrayType { - Utf8(StringArray), - Utf8View(StringViewArray), - LargeUtf8(LargeStringArray), -} -impl StringArrayType { - /// 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)`. - fn value(&self, row_idx: usize) -> Option<&str> { - let (is_null, value) = match self { - StringArrayType::Utf8(arr) => (arr.is_null(row_idx), arr.value(row_idx)), - StringArrayType::LargeUtf8(arr) => (arr.is_null(row_idx), arr.value(row_idx)), - StringArrayType::Utf8View(arr) => (arr.is_null(row_idx), arr.value(row_idx)), - }; - if is_null { None } else { Some(value) } - } -} -impl<'a> TryFrom<&'a DataType> for StringArrayType { - type Error = (); - - fn try_from(data_type: &'a DataType) -> std::result::Result { - let vals: Vec<&str> = Vec::new(); - Ok(match data_type { - DataType::Utf8 => StringArrayType::Utf8(vals.into()), - DataType::Utf8View => StringArrayType::Utf8View(vals.into()), - DataType::LargeUtf8 => StringArrayType::LargeUtf8(vals.into()), - _ => return Err(()), - }) - } -} -impl TryFrom for StringArrayType { - type Error = DataType; - - fn try_from(arr: ArrayRef) -> std::result::Result { - Ok(match arr.data_type() { - DataType::Utf8 => StringArrayType::Utf8(arr.as_string().clone()), - DataType::LargeUtf8 => StringArrayType::LargeUtf8(arr.as_string().clone()), - DataType::Utf8View => StringArrayType::Utf8View(arr.as_string_view().clone()), - ty => return Err(ty.clone()), - }) - } -} - /// Returns true if the given data type can be used as a top-K aggregation hash key. /// /// Supported types include Arrow primitives (integers, floats, decimals, intervals) @@ -131,8 +88,11 @@ pub fn is_supported_hash_key_type(kt: &DataType) -> bool { } // An implementation of ArrowHashTable for String keys -pub struct StringHashTable { - owned: StringArrayType, +pub struct StringHashTable +where + for<'a> &'a S: StringArrayType<'a>, +{ + owned: S, map: TopKHashTable>, rnd: RandomState, } @@ -147,9 +107,13 @@ where rnd: RandomState, } -impl StringHashTable { - pub fn new(limit: usize, data_type: &DataType) -> Self { - let owned = StringArrayType::try_from(data_type).expect("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), @@ -158,9 +122,17 @@ impl StringHashTable { } } -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 = StringArrayType::try_from(ids).expect("Unsupported data type"); + self.owned = ids + .as_any() + .downcast_ref::() + .expect("Unsupported data type") + .clone(); } fn len(&self) -> usize { @@ -177,15 +149,15 @@ impl ArrowHashTable for StringHashTable { fn take_all(&mut self, indexes: Vec) -> ArrayRef { let ids = self.map.take_all(indexes); - match &self.owned { - StringArrayType::Utf8(_) => Arc::new(StringArray::from(ids)), - StringArrayType::LargeUtf8(_) => Arc::new(LargeStringArray::from(ids)), - StringArrayType::Utf8View(_) => Arc::new(StringViewArray::from(ids)), - } + Arc::new(S::from(ids)) } fn find_or_insert(&mut self, row_idx: usize, replace_idx: usize) -> (usize, bool) { - let id = self.owned.value(row_idx); + let id = if self.owned.is_null(row_idx) { + None + } else { + Some((&self.owned).value(row_idx)) + }; // Compute hash and create equality closure for hash table lookup. let hash = self.rnd.hash_one(id); @@ -410,9 +382,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))), _ => {} } From e5861d247fa75bf8aa0a8ec430a637e1dba31fd4 Mon Sep 17 00:00:00 2001 From: Michal Piatkowski <291740709+MassivePizza@users.noreply.github.com> Date: Mon, 3 Aug 2026 14:15:08 +0200 Subject: [PATCH 10/17] use some_value helper --- datafusion-testing | 2 +- .../src/aggregates/topk/hash_table.rs | 43 ++++++------------- 2 files changed, 15 insertions(+), 30 deletions(-) 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 e476dce3a8d27..aaf9fac369622 100644 --- a/datafusion/physical-plan/src/aggregates/topk/hash_table.rs +++ b/datafusion/physical-plan/src/aggregates/topk/hash_table.rs @@ -190,12 +190,7 @@ where row_idx: usize, replace_idx: usize, ) -> (usize, InsertKind) { - let id = if self.owned.is_null(row_idx) { - None - } else { - Some((&self.owned).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); let eq = move |mi: &Option| id == mi.as_deref(); @@ -206,22 +201,14 @@ where } fn insert_null(&mut self, row_idx: usize) -> bool { - let id = if self.owned.is_null(row_idx) { - None - } else { - Some((&self.owned).value(row_idx)) - }; + let id = some_value(&self.owned, row_idx); let hash = self.rnd.hash_one(id); let eq = move |mi: &Option| id == mi.as_deref(); self.map.insert_null(hash, id.map(ToOwned::to_owned), eq) } fn remove_if_null(&mut self, row_idx: usize) -> bool { - let id = if self.owned.is_null(row_idx) { - None - } else { - Some((&self.owned).value(row_idx)) - }; + let id = some_value(&self.owned, row_idx); let hash = self.rnd.hash_one(id); let eq = move |mi: &Option| id == mi.as_deref(); self.map.remove_if_null(hash, eq) @@ -249,11 +236,7 @@ where /// 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: Option = if self.owned.is_null(row_idx) { - None - } else { - Some(self.owned.value(row_idx)) - }; + let id: Option = some_value(&self.owned, row_idx); let hash: u64 = id.hash(&self.rnd); (id, hash) } @@ -298,15 +281,8 @@ where row_idx: usize, replace_idx: usize, ) -> (usize, InsertKind) { - let id: Option = if self.owned.is_null(row_idx) { - None - } else { - Some(self.owned.value(row_idx)) - }; - // Compute hash and create equality closure for hash table lookup. - let hash: u64 = id.hash(&self.rnd); + let (id, hash) = self.id_and_hash(row_idx); let eq = |mi: &Option| id == *mi; - // Use entry API to avoid double lookup self.map.find_or_insert(hash, id, replace_idx, eq) } @@ -556,6 +532,15 @@ has_integer!(u8, u16, u32, u64); has_integer!(IntervalDayTime, IntervalMonthDayNano); hash_float!(f16, f32, f64); +#[inline] +fn some_value(array: A, index: usize) -> Option { + if array.is_null(index) { + None + } else { + Some(array.value(index)) + } +} + pub fn new_hash_table( limit: usize, kt: DataType, From c5131bd18843701d6a6d0abeca8904f094a5c2c0 Mon Sep 17 00:00:00 2001 From: Michal Piatkowski <291740709+MassivePizza@users.noreply.github.com> Date: Tue, 4 Aug 2026 14:55:36 +0200 Subject: [PATCH 11/17] fix clippy --- datafusion/physical-plan/src/aggregates/topk/hash_table.rs | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/datafusion/physical-plan/src/aggregates/topk/hash_table.rs b/datafusion/physical-plan/src/aggregates/topk/hash_table.rs index aaf9fac369622..9d4dfc298141a 100644 --- a/datafusion/physical-plan/src/aggregates/topk/hash_table.rs +++ b/datafusion/physical-plan/src/aggregates/topk/hash_table.rs @@ -533,7 +533,10 @@ has_integer!(IntervalDayTime, IntervalMonthDayNano); hash_float!(f16, f32, f64); #[inline] -fn some_value(array: A, index: usize) -> Option { +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 { From 71ec3ccd998371276adc3bd9eceb8dfebbb0b6ba Mon Sep 17 00:00:00 2001 From: Michal Piatkowski <291740709+MassivePizza@users.noreply.github.com> Date: Tue, 4 Aug 2026 15:12:46 +0200 Subject: [PATCH 12/17] share equality closures --- .../src/aggregates/topk/hash_table.rs | 37 ++++++++++++------- 1 file changed, 24 insertions(+), 13 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/topk/hash_table.rs b/datafusion/physical-plan/src/aggregates/topk/hash_table.rs index 9d4dfc298141a..61410893df865 100644 --- a/datafusion/physical-plan/src/aggregates/topk/hash_table.rs +++ b/datafusion/physical-plan/src/aggregates/topk/hash_table.rs @@ -153,6 +153,11 @@ where rnd: RandomState::default(), } } + + #[inline] + fn eq_fn(id: Option<&str>) -> impl Fn(&Option) -> bool { + move |mi| id == mi.as_deref() + } } impl ArrowHashTable for StringHashTable @@ -193,25 +198,27 @@ where 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); - let eq = move |mi: &Option| id == mi.as_deref(); // Use entry API to avoid double lookup - self.map - .find_or_insert(hash, id.map(ToOwned::to_owned), replace_idx, eq) + self.map.find_or_insert( + hash, + id.map(ToOwned::to_owned), + 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); - let eq = move |mi: &Option| id == mi.as_deref(); - self.map.insert_null(hash, id.map(ToOwned::to_owned), eq) + self.map + .insert_null(hash, id.map(ToOwned::to_owned), Self::eq_fn(id)) } fn remove_if_null(&mut self, row_idx: usize) -> bool { let id = some_value(&self.owned, row_idx); let hash = self.rnd.hash_one(id); - let eq = move |mi: &Option| id == mi.as_deref(); - self.map.remove_if_null(hash, eq) + self.map.remove_if_null(hash, Self::eq_fn(id)) } fn null_map_idxs(&self) -> Vec { @@ -235,11 +242,17 @@ where } /// 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 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 @@ -282,21 +295,19 @@ where replace_idx: usize, ) -> (usize, InsertKind) { let (id, hash) = self.id_and_hash(row_idx); - let eq = |mi: &Option| id == *mi; // 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 eq = move |mi: &Option| id == *mi; - self.map.insert_null(hash, id, eq) + 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 == *mi; - self.map.remove_if_null(hash, eq) + self.map.remove_if_null(hash, Self::eq_fn(id)) } fn null_map_idxs(&self) -> Vec { From 4ba140b1d0f6ae7b3ea35acfddb4163be33daccb Mon Sep 17 00:00:00 2001 From: Michal Piatkowski <291740709+MassivePizza@users.noreply.github.com> Date: Mon, 3 Aug 2026 16:36:36 +0200 Subject: [PATCH 13/17] move ID nullability into HashTableItem --- .../src/aggregates/topk/hash_table.rs | 87 +++++++++---------- 1 file changed, 41 insertions(+), 46 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/topk/hash_table.rs b/datafusion/physical-plan/src/aggregates/topk/hash_table.rs index 61410893df865..1c7083d01fa69 100644 --- a/datafusion/physical-plan/src/aggregates/topk/hash_table.rs +++ b/datafusion/physical-plan/src/aggregates/topk/hash_table.rs @@ -35,11 +35,6 @@ 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. @@ -49,9 +44,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, } @@ -59,10 +54,10 @@ 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, // The maximum number of entries allowed @@ -126,7 +121,7 @@ where for<'a> &'a S: StringArrayType<'a>, { owned: S, - map: TopKHashTable>, + map: TopKHashTable, rnd: RandomState, } @@ -136,7 +131,7 @@ where Option<::Native>: Comparable, { owned: PrimitiveArray, - map: TopKHashTable>, + map: TopKHashTable, rnd: RandomState, } @@ -316,7 +311,7 @@ where } use hashbrown::hash_table::Entry; -impl TopKHashTable { +impl TopKHashTable { pub fn new(limit: usize, capacity: usize) -> Self { Self { map: HashTable::with_capacity(capacity), @@ -328,21 +323,21 @@ impl TopKHashTable { } 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; + self.store[removed_idx].id.take(); self.free_indices.push(removed_idx); } Entry::Vacant(_) => unreachable!(), @@ -363,7 +358,7 @@ 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; } } @@ -374,16 +369,16 @@ impl TopKHashTable { pub fn find_or_insert( &mut self, hash: u64, - id: ID, + id: Option, replace_idx: usize, - mut eq: impl FnMut(&ID) -> bool, + mut eq: impl FnMut(&Option) -> bool, ) -> (usize, InsertKind) { // 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); + 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].as_ref().unwrap().heap_idx == NULL_HEAP_IDX { + if self.store[map_idx].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); @@ -399,15 +394,15 @@ impl TopKHashTable { let heap_idx = self.remove_if_full(replace_idx); let mi = HashTableItem::new(hash, id, heap_idx); let store_idx = if let Some(idx) = self.free_indices.pop() { - self.store[idx] = Some(mi); + 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); } @@ -430,10 +425,10 @@ impl TopKHashTable { pub fn insert_null( &mut self, hash: u64, - id: ID, - mut eq: impl FnMut(&ID) -> bool, + id: Option, + mut eq: impl FnMut(&Option) -> bool, ) -> bool { - let eq_fn = |idx: &usize| eq(&self.store[*idx].as_ref().unwrap().id); + let eq_fn = |idx: &usize| eq(&self.store[*idx].id); if self.map.find(hash, eq_fn).is_some() { return false; } @@ -444,14 +439,14 @@ impl TopKHashTable { 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); + self.store[idx] = mi; idx } else { - self.store.push(Some(mi)); + self.store.push(mi); self.store.len() - 1 }; - 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); } @@ -464,10 +459,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].heap_idx == NULL_HEAP_IDX { self.remove_at(map_idx); self.null_count -= 1; @@ -481,11 +480,7 @@ impl TopKHashTable { self.store .iter() .enumerate() - .filter_map(|(idx, item)| { - item.as_ref() - .filter(|item| item.heap_idx == NULL_HEAP_IDX) - .map(|_| idx) - }) + .filter_map(|(idx, item)| (item.heap_idx == NULL_HEAP_IDX).then_some(idx)) .collect() } @@ -493,10 +488,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(); @@ -506,8 +501,8 @@ 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 } } } @@ -603,7 +598,7 @@ 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() { @@ -636,7 +631,7 @@ 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()); @@ -672,7 +667,7 @@ 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()); From 3c248a53fd63dd8e8656bdca27893c7fcd78d86e Mon Sep 17 00:00:00 2001 From: Michal Piatkowski <291740709+MassivePizza@users.noreply.github.com> Date: Tue, 4 Aug 2026 15:57:10 +0200 Subject: [PATCH 14/17] respect null properly in TopKHashTable::null_map_idxs --- .../src/aggregates/topk/hash_table.rs | 17 ++++++++++++----- 1 file changed, 12 insertions(+), 5 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/topk/hash_table.rs b/datafusion/physical-plan/src/aggregates/topk/hash_table.rs index 1c7083d01fa69..b13c748cf4182 100644 --- a/datafusion/physical-plan/src/aggregates/topk/hash_table.rs +++ b/datafusion/physical-plan/src/aggregates/topk/hash_table.rs @@ -30,7 +30,7 @@ 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::fmt::Debug; use std::hash::BuildHasher; use std::sync::Arc; @@ -310,7 +310,12 @@ where } } -use hashbrown::hash_table::Entry; +impl HashTableItem { + #[inline] + pub fn is_null(&self) -> bool { + self.heap_idx == NULL_HEAP_IDX + } +} impl TopKHashTable { pub fn new(limit: usize, capacity: usize) -> Self { Self { @@ -378,7 +383,7 @@ impl TopKHashTable { { 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].heap_idx == NULL_HEAP_IDX { + 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); @@ -466,7 +471,7 @@ impl TopKHashTable { ) -> 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].heap_idx == NULL_HEAP_IDX + && self.store[map_idx].is_null() { self.remove_at(map_idx); self.null_count -= 1; @@ -480,7 +485,9 @@ impl TopKHashTable { self.store .iter() .enumerate() - .filter_map(|(idx, item)| (item.heap_idx == NULL_HEAP_IDX).then_some(idx)) + .filter_map(|(idx, item)| { + (item.id.is_some() && item.is_null()).then_some(idx) + }) .collect() } From 6d2db94ea5e55dffc94f6af6e3063311bcdab8ea Mon Sep 17 00:00:00 2001 From: Michal Piatkowski <291740709+MassivePizza@users.noreply.github.com> Date: Wed, 5 Aug 2026 15:06:05 +0200 Subject: [PATCH 15/17] concretely typed TopK ArrowHeap storage --- .../physical-plan/src/aggregates/topk/heap.rs | 122 ++++++------------ 1 file changed, 36 insertions(+), 86 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/topk/heap.rs b/datafusion/physical-plan/src/aggregates/topk/heap.rs index aef6bb5596c2e..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}; @@ -98,20 +96,18 @@ where 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 { + pub fn new(limit: usize, desc: bool) -> Self { let batch = PrimitiveArray::::builder(0).finish(); Self { batch, heap: TopKHeap::new(limit, desc), desc, - data_type, } } } @@ -156,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) } } @@ -168,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 { @@ -229,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, @@ -249,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); } @@ -260,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"); @@ -289,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) } } @@ -615,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))), _ => {} } From 31b45de8725e600d452e5e0b7500de1bf38725ab Mon Sep 17 00:00:00 2001 From: Michal Piatkowski <291740709+MassivePizza@users.noreply.github.com> Date: Wed, 5 Aug 2026 15:14:05 +0200 Subject: [PATCH 16/17] reuse (String) allocs in TopKHashTable --- .../src/aggregates/topk/hash_table.rs | 158 ++++++++++-------- 1 file changed, 89 insertions(+), 69 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/topk/hash_table.rs b/datafusion/physical-plan/src/aggregates/topk/hash_table.rs index b13c748cf4182..21d52b7fd6c70 100644 --- a/datafusion/physical-plan/src/aggregates/topk/hash_table.rs +++ b/datafusion/physical-plan/src/aggregates/topk/hash_table.rs @@ -31,6 +31,7 @@ use datafusion_common::exec_datafusion_err; use datafusion_common::hash_utils::RandomState; use half::f16; use hashbrown::hash_table::{Entry, HashTable}; +use std::borrow::BorrowMut; use std::fmt::Debug; use std::hash::BuildHasher; use std::sync::Arc; @@ -60,6 +61,8 @@ struct TopKHashTable { 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) @@ -195,19 +198,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 { @@ -275,10 +273,7 @@ where 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) @@ -292,12 +287,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,18 +305,13 @@ where } } -impl HashTableItem { - #[inline] - pub fn is_null(&self) -> bool { - self.heap_idx == NULL_HEAP_IDX - } -} 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, } @@ -342,7 +332,12 @@ impl TopKHashTable { match self.map.entry(hash, eq, hasher) { Entry::Occupied(entry) => { let (removed_idx, _) = entry.remove(); - self.store[removed_idx].id.take(); + 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!(), @@ -367,38 +362,72 @@ impl TopKHashTable { } } + 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: Option, + id: Option<&Q>, replace_idx: usize, mut eq: impl FnMut(&Option) -> bool, - ) -> (usize, InsertKind) { + ) -> (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].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); - } + + 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() { + debug_assert!(self.store[idx].id.is_none(), "slot should be empty"); self.store[idx] = mi; idx } else { @@ -414,12 +443,7 @@ impl TopKHashTable { // 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,12 +451,16 @@ 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: Option, + id: Option<&Q>, mut eq: impl FnMut(&Option) -> bool, - ) -> 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 +470,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] = mi; - idx - } else { - self.store.push(mi); - self.store.len() - 1 - }; - - let hasher = |idx: &usize| self.store[*idx].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 } @@ -512,6 +527,11 @@ 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 { @@ -612,7 +632,7 @@ mod tests { 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); } @@ -645,16 +665,16 @@ mod tests { 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); @@ -680,18 +700,18 @@ mod tests { 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(()) From 6ed1b35f9b8b1da2d16ae2da466d44952bd4dbd6 Mon Sep 17 00:00:00 2001 From: Michal Piatkowski <291740709+MassivePizza@users.noreply.github.com> Date: Wed, 5 Aug 2026 16:23:20 +0200 Subject: [PATCH 17/17] add use_free_slots comment --- datafusion/physical-plan/src/aggregates/topk/hash_table.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/datafusion/physical-plan/src/aggregates/topk/hash_table.rs b/datafusion/physical-plan/src/aggregates/topk/hash_table.rs index 21d52b7fd6c70..2e2936a1dcec0 100644 --- a/datafusion/physical-plan/src/aggregates/topk/hash_table.rs +++ b/datafusion/physical-plan/src/aggregates/topk/hash_table.rs @@ -362,6 +362,7 @@ impl TopKHashTable { } } + /// Used to avoid pushing pointless copies of primitives to the `free_slots` pool. const fn use_free_slots() -> bool { std::mem::needs_drop::() }