From f231437b0d457bc9e3fd1cb506921799bd7c276f 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 01/12] store concrete string type in TopK hash-table

---
 datafusion/physical-plan/src/aggregates/topk/hash_table.rs | 6 +-----
 1 file changed, 1 insertion(+), 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 61410893df865..1ef6179936e36 100644
--- a/datafusion/physical-plan/src/aggregates/topk/hash_table.rs
+++ b/datafusion/physical-plan/src/aggregates/topk/hash_table.rs
@@ -113,11 +113,7 @@ pub trait ArrowHashTable {
 /// 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

From ab035f773ef600687efec0bb453f054c07f2391c 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 02/12] 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 1ef6179936e36..61410893df865 100644
--- a/datafusion/physical-plan/src/aggregates/topk/hash_table.rs
+++ b/datafusion/physical-plan/src/aggregates/topk/hash_table.rs
@@ -113,7 +113,11 @@ pub trait ArrowHashTable {
 /// 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 edc5cc03f5c7a37962748b3638169f197564db18 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 03/12] 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<T> 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<ID: KeyType> {
+pub struct HashTableItem<ID> {
     hash: u64,
-    pub id: ID,
+    pub id: Option<ID>,
     pub heap_idx: usize,
 }
 
@@ -59,10 +54,10 @@ pub struct HashTableItem<ID: KeyType> {
 /// 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<ID: KeyType> {
+struct TopKHashTable<ID> {
     map: HashTable<usize>,
     // Store the actual items separately to allow for index-based access
-    store: Vec<Option<HashTableItem<ID>>>,
+    store: Vec<HashTableItem<ID>>,
     // Free indexes in the store for reuse
     free_indices: Vec<usize>,
     // The maximum number of entries allowed
@@ -126,7 +121,7 @@ where
     for<'a> &'a S: StringArrayType<'a>,
 {
     owned: S,
-    map: TopKHashTable<Option<String>>,
+    map: TopKHashTable<String>,
     rnd: RandomState,
 }
 
@@ -136,7 +131,7 @@ where
     Option<<VAL as ArrowPrimitiveType>::Native>: Comparable,
 {
     owned: PrimitiveArray<VAL>,
-    map: TopKHashTable<Option<VAL::Native>>,
+    map: TopKHashTable<VAL::Native>,
     rnd: RandomState,
 }
 
@@ -316,7 +311,7 @@ where
 }
 
 use hashbrown::hash_table::Entry;
-impl<ID: KeyType + PartialEq> TopKHashTable<ID> {
+impl<ID: PartialEq> TopKHashTable<ID> {
     pub fn new(limit: usize, capacity: usize) -> Self {
         Self {
             map: HashTable::with_capacity(capacity),
@@ -328,21 +323,21 @@ impl<ID: KeyType + PartialEq> TopKHashTable<ID> {
     }
 
     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<ID: KeyType + PartialEq> TopKHashTable<ID> {
 
     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<ID: KeyType + PartialEq> TopKHashTable<ID> {
     pub fn find_or_insert(
         &mut self,
         hash: u64,
-        id: ID,
+        id: Option<ID>,
         replace_idx: usize,
-        mut eq: impl FnMut(&ID) -> bool,
+        mut eq: impl FnMut(&Option<ID>) -> 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<ID: KeyType + PartialEq> TopKHashTable<ID> {
         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<ID: KeyType + PartialEq> TopKHashTable<ID> {
     pub fn insert_null(
         &mut self,
         hash: u64,
-        id: ID,
-        mut eq: impl FnMut(&ID) -> bool,
+        id: Option<ID>,
+        mut eq: impl FnMut(&Option<ID>) -> 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<ID: KeyType + PartialEq> TopKHashTable<ID> {
 
         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<ID: KeyType + PartialEq> TopKHashTable<ID> {
     /// 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<ID>) -> 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<ID: KeyType + PartialEq> TopKHashTable<ID> {
         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<ID: KeyType + PartialEq> TopKHashTable<ID> {
         self.map.len()
     }
 
-    pub fn take_all(&mut self, idxs: Vec<usize>) -> Vec<ID> {
+    pub fn take_all(&mut self, idxs: Vec<usize>) -> Vec<Option<ID>> {
         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<ID: KeyType + PartialEq> TopKHashTable<ID> {
     }
 }
 
-impl<ID: KeyType> HashTableItem<ID> {
-    pub fn new(hash: u64, id: ID, heap_idx: usize) -> Self {
+impl<ID> HashTableItem<ID> {
+    pub fn new(hash: u64, id: Option<ID>, 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::<usize, usize>::new();
         // Create TopKHashTable with limit=5 and capacity=3 to force resizing
-        let mut map = TopKHashTable::<Option<String>>::new(5, 3);
+        let mut map = TopKHashTable::<String>::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::<Option<String>>::new(2, 10);
+        let mut map = TopKHashTable::<String>::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::<Option<String>>::new(1, 10);
+        let mut map = TopKHashTable::<String>::new(1, 10);
 
         let a = Some("a".to_string());
         let b = Some("b".to_string());

From 1e3089bf7500c3417b4a7efad5ee349fc4c45d77 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 04/12] 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<ID> HashTableItem<ID> {
+    #[inline]
+    pub fn is_null(&self) -> bool {
+        self.heap_idx == NULL_HEAP_IDX
+    }
+}
 impl<ID: PartialEq> TopKHashTable<ID> {
     pub fn new(limit: usize, capacity: usize) -> Self {
         Self {
@@ -378,7 +383,7 @@ impl<ID: PartialEq> TopKHashTable<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].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<ID: PartialEq> TopKHashTable<ID> {
     ) -> 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<ID: PartialEq> TopKHashTable<ID> {
         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 24f3d43a3a8706a4a631bfb0615858a459de69f5 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 05/12] 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<VAL>,
     heap: TopKHeap<VAL::Native>,
     desc: bool,
-    data_type: DataType,
 }
 
 impl<VAL: ArrowPrimitiveType> PrimitiveHeap<VAL>
 where
     <VAL as ArrowPrimitiveType>::Native: Comparable,
 {
-    pub fn new(limit: usize, desc: bool, data_type: DataType) -> Self {
+    pub fn new(limit: usize, desc: bool) -> Self {
         let batch = PrimitiveArray::<VAL>::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::<VAL>::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<S: Array>
+where
+    for<'a> &'a S: StringArrayType<'a>,
+{
+    batch: S,
     heap: TopKHeap<Option<String>>,
     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<S> StringHeap<S>
+where
+    S: Array + From<Vec<Option<String>>>,
+    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::<i32>().value(idx),
-        DataType::LargeUtf8 => batch.as_string::<i64>().value(idx),
-        DataType::Utf8View => batch.as_string_view().value(idx),
-        _ => unreachable!("Unsupported string type: {data_type}"),
-    }
-}
-
-impl ArrowHeap for StringHeap {
+impl<S> ArrowHeap for StringHeap<S>
+where
+    S: Array + Clone + From<Vec<Option<String>>> + 'static,
+    for<'a> &'a S: StringArrayType<'a>,
+{
     fn set_batch(&mut self, vals: ArrayRef) {
-        self.batch = vals;
+        self.batch = vals
+            .as_any()
+            .downcast_ref::<S>()
+            .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<usize>) {
         let (vals, map_idxs) = self.heap.drain();
-        // Use Arrow builders to safely construct arrays from the owned
-        // `Option<String>` 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<Box<dyn ArrowHeap + Send>> {
-    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::<StringArray>::new(limit, desc))),
+        DataType::LargeUtf8 => return Ok(Box::new(StringHeap::<LargeStringArray>::new(limit, desc))),
+        DataType::Utf8View => return Ok(Box::new(StringHeap::<StringViewArray>::new(limit, desc))),
         _ => {}
     }
 

From d56173254ed5b9cf1d4a6b28856249b3996e730f 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 06/12] 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<ID> {
     store: Vec<HashTableItem<ID>>,
     // Free indexes in the store for reuse
     free_indices: Vec<usize>,
+    // Pool of reusable value locations, usually Strings
+    free_slots: Vec<ID>,
     // 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<VAL> = 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<ID> HashTableItem<ID> {
-    #[inline]
-    pub fn is_null(&self) -> bool {
-        self.heap_idx == NULL_HEAP_IDX
-    }
-}
 impl<ID: PartialEq> TopKHashTable<ID> {
     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<ID: PartialEq> TopKHashTable<ID> {
         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<ID: PartialEq> TopKHashTable<ID> {
         }
     }
 
+    const fn use_free_slots() -> bool {
+        std::mem::needs_drop::<ID>()
+    }
+
     /// 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<Q>(
         &mut self,
         hash: u64,
-        id: Option<ID>,
+        id: Option<&Q>,
         replace_idx: usize,
         mut eq: impl FnMut(&Option<ID>) -> bool,
-    ) -> (usize, InsertKind) {
+    ) -> (usize, InsertKind)
+    where
+        Q: ToOwned<Owned = ID> + ?Sized,
+        ID: BorrowMut<Q::Owned>,
+    {
         // 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<Q>(&mut self, hash: u64, id: Option<&Q>, heap_idx: usize) -> usize
+    where
+        Q: ToOwned<Owned = ID> + ?Sized,
+        ID: BorrowMut<Q::Owned>,
+    {
+        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<ID: PartialEq> TopKHashTable<ID> {
 
         // 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<ID: PartialEq> TopKHashTable<ID> {
     /// 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<Q>(
         &mut self,
         hash: u64,
-        id: Option<ID>,
+        id: Option<&Q>,
         mut eq: impl FnMut(&Option<ID>) -> bool,
-    ) -> bool {
+    ) -> bool
+    where
+        Q: ToOwned<Owned = ID> + ?Sized,
+        ID: BorrowMut<Q::Owned>,
+    {
         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<ID: PartialEq> TopKHashTable<ID> {
             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<ID> HashTableItem<ID> {
     pub fn new(hash: u64, id: Option<ID>, 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<String> {
@@ -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 6702dac4adc2f7bc0fdf907c23e33bab7f52006e 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 07/12] 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<ID: PartialEq> TopKHashTable<ID> {
         }
     }
 
+    /// Used to avoid pushing pointless copies of primitives to the `free_slots` pool.
     const fn use_free_slots() -> bool {
         std::mem::needs_drop::<ID>()
     }

From c9fb728d1337d6be1455826c3251ba64a9e92b3e Mon Sep 17 00:00:00 2001
From: Michal Piatkowski <291740709+MassivePizza@users.noreply.github.com>
Date: Fri, 21 Aug 2026 15:52:53 +0200
Subject: [PATCH 08/12] fix clippy

---
 datafusion/physical-plan/src/aggregates/topk/heap.rs         | 2 +-
 datafusion/physical-plan/src/aggregates/topk/priority_map.rs | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/datafusion/physical-plan/src/aggregates/topk/heap.rs b/datafusion/physical-plan/src/aggregates/topk/heap.rs
index 819041f76b45f..378dc1acad85b 100644
--- a/datafusion/physical-plan/src/aggregates/topk/heap.rs
+++ b/datafusion/physical-plan/src/aggregates/topk/heap.rs
@@ -567,7 +567,7 @@ pub fn is_supported_heap_type(vt: &DataType) -> bool {
 pub fn new_heap(
     limit: usize,
     desc: bool,
-    vt: DataType,
+    vt: &DataType,
 ) -> Result<Box<dyn ArrowHeap + Send>> {
     macro_rules! downcast_helper {
         ($vt:ty, $d:ident) => {
diff --git a/datafusion/physical-plan/src/aggregates/topk/priority_map.rs b/datafusion/physical-plan/src/aggregates/topk/priority_map.rs
index f46cb22a7a63c..b359b3a6b3395 100644
--- a/datafusion/physical-plan/src/aggregates/topk/priority_map.rs
+++ b/datafusion/physical-plan/src/aggregates/topk/priority_map.rs
@@ -46,7 +46,7 @@ impl PriorityMap {
     ) -> Result<Self> {
         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,

From 4d8b2c5201052bf8ef299ec979ef7a71365c59fb Mon Sep 17 00:00:00 2001
From: Michal Piatkowski <291740709+MassivePizza@users.noreply.github.com>
Date: Fri, 21 Aug 2026 15:57:31 +0200
Subject: [PATCH 09/12] use cast instead of from(to_data)

---
 datafusion/physical-plan/src/aggregates/topk/heap.rs | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/datafusion/physical-plan/src/aggregates/topk/heap.rs b/datafusion/physical-plan/src/aggregates/topk/heap.rs
index 378dc1acad85b..2498b59fe1e0d 100644
--- a/datafusion/physical-plan/src/aggregates/topk/heap.rs
+++ b/datafusion/physical-plan/src/aggregates/topk/heap.rs
@@ -25,8 +25,7 @@
 
 use arrow::array::types::{IntervalDayTime, IntervalMonthDayNano};
 use arrow::array::{
-    Array, ArrayAccessor, ArrayRef, ArrowPrimitiveType, LargeStringArray, PrimitiveArray,
-    StringArray, StringArrayType, StringViewArray, downcast_primitive,
+    Array, ArrayAccessor, ArrayRef, ArrowPrimitiveType, AsArray, LargeStringArray, PrimitiveArray, StringArray, StringArrayType, StringViewArray, downcast_primitive,
 };
 use arrow::buffer::ScalarBuffer;
 use arrow::datatypes::{DataType, i256};
@@ -103,7 +102,7 @@ where
     <VAL as ArrowPrimitiveType>::Native: Comparable,
 {
     pub fn new(limit: usize, desc: bool) -> Self {
-        let batch = PrimitiveArray::<VAL>::builder(0).finish();
+        let batch = PrimitiveArray::<VAL>::new_null(0);
         Self {
             batch,
             heap: TopKHeap::new(limit, desc),
@@ -117,7 +116,7 @@ where
     <VAL as ArrowPrimitiveType>::Native: Comparable,
 {
     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 {

From 7d2aedfccfa188bae209ca2885a03a9e7df23d85 Mon Sep 17 00:00:00 2001
From: Michal Piatkowski <291740709+MassivePizza@users.noreply.github.com>
Date: Fri, 21 Aug 2026 16:03:16 +0200
Subject: [PATCH 10/12] fix fmt

---
 datafusion/physical-plan/src/aggregates/topk/heap.rs | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/datafusion/physical-plan/src/aggregates/topk/heap.rs b/datafusion/physical-plan/src/aggregates/topk/heap.rs
index 2498b59fe1e0d..2799857739437 100644
--- a/datafusion/physical-plan/src/aggregates/topk/heap.rs
+++ b/datafusion/physical-plan/src/aggregates/topk/heap.rs
@@ -25,7 +25,8 @@
 
 use arrow::array::types::{IntervalDayTime, IntervalMonthDayNano};
 use arrow::array::{
-    Array, ArrayAccessor, ArrayRef, ArrowPrimitiveType, AsArray, LargeStringArray, PrimitiveArray, StringArray, StringArrayType, StringViewArray, downcast_primitive,
+    Array, ArrayAccessor, ArrayRef, ArrowPrimitiveType, AsArray, LargeStringArray,
+    PrimitiveArray, StringArray, StringArrayType, StringViewArray, downcast_primitive,
 };
 use arrow::buffer::ScalarBuffer;
 use arrow::datatypes::{DataType, i256};

From cb02950b696ed16092e55ae204de25b1fd8a311c Mon Sep 17 00:00:00 2001
From: Michal Piatkowski <291740709+MassivePizza@users.noreply.github.com>
Date: Fri, 21 Aug 2026 17:25:13 +0200
Subject: [PATCH 11/12] Store exact type in TopKHeap and TopKHashTable. This
 restores behavior from before
 https://github.com/apache/datafusion/pull/23609.

---
 .../src/aggregates/topk/hash_table.rs         | 16 ++++++-------
 .../physical-plan/src/aggregates/topk/heap.rs | 23 ++++++++++++++-----
 .../src/aggregates/topk/priority_map.rs       |  6 ++---
 3 files changed, 26 insertions(+), 19 deletions(-)

diff --git a/datafusion/physical-plan/src/aggregates/topk/hash_table.rs b/datafusion/physical-plan/src/aggregates/topk/hash_table.rs
index 2e2936a1dcec0..f6ea8651a6e09 100644
--- a/datafusion/physical-plan/src/aggregates/topk/hash_table.rs
+++ b/datafusion/physical-plan/src/aggregates/topk/hash_table.rs
@@ -136,6 +136,7 @@ where
     owned: PrimitiveArray<VAL>,
     map: TopKHashTable<VAL::Native>,
     rnd: RandomState,
+    value_type: DataType,
 }
 
 impl<S> StringHashTable<S>
@@ -144,9 +145,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(),
         }
@@ -223,14 +223,12 @@ impl<VAL: ArrowPrimitiveType> PrimitiveHashTable<VAL>
 where
     Option<<VAL as ArrowPrimitiveType>::Native>: Comparable + HashValue,
 {
-    pub fn new(limit: usize, kt: DataType) -> Self {
-        let owned = PrimitiveArray::<VAL>::builder(0)
-            .with_data_type(kt)
-            .finish();
+    pub fn new(limit: usize, value_type: DataType) -> Self {
         Self {
-            owned,
+            owned: PrimitiveArray::<VAL>::new_null(0).with_data_type(value_type.clone()),
             map: TopKHashTable::new(limit, limit * 10),
             rnd: RandomState::default(),
+            value_type,
         }
     }
 
@@ -270,8 +268,8 @@ where
 
     fn take_all(&mut self, indexes: Vec<usize>) -> ArrayRef {
         let ids = self.map.take_all(indexes);
-        let mut builder: PrimitiveBuilder<VAL> = PrimitiveArray::builder(ids.len())
-            .with_data_type(self.owned.data_type().clone());
+        let mut builder: PrimitiveBuilder<VAL> =
+            PrimitiveArray::builder(ids.len()).with_data_type(self.value_type.clone());
         for id in ids.into_iter() {
             builder.append_option(id);
         }
diff --git a/datafusion/physical-plan/src/aggregates/topk/heap.rs b/datafusion/physical-plan/src/aggregates/topk/heap.rs
index 2799857739437..71d30d27b2d34 100644
--- a/datafusion/physical-plan/src/aggregates/topk/heap.rs
+++ b/datafusion/physical-plan/src/aggregates/topk/heap.rs
@@ -75,6 +75,7 @@ struct TopKHeap<VAL: ValueType> {
 
 /// 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;
@@ -96,18 +97,19 @@ where
     batch: PrimitiveArray<VAL>,
     heap: TopKHeap<VAL::Native>,
     desc: bool,
+    value_type: DataType,
 }
 
 impl<VAL: ArrowPrimitiveType> PrimitiveHeap<VAL>
 where
     <VAL as ArrowPrimitiveType>::Native: Comparable,
 {
-    pub fn new(limit: usize, desc: bool) -> Self {
-        let batch = PrimitiveArray::<VAL>::new_null(0);
+    pub fn new(limit: usize, desc: bool, value_type: DataType) -> Self {
         Self {
-            batch,
+            batch: PrimitiveArray::<VAL>::new_null(0).with_data_type(value_type.clone()),
             heap: TopKHeap::new(limit, desc),
             desc,
+            value_type,
         }
     }
 }
@@ -116,6 +118,10 @@ impl<VAL: ArrowPrimitiveType> ArrowHeap for PrimitiveHeap<VAL>
 where
     <VAL as ArrowPrimitiveType>::Native: Comparable,
 {
+    fn value_type(&self) -> &DataType {
+        &self.value_type
+    }
+
     fn set_batch(&mut self, vals: ArrayRef) {
         self.batch = vals.as_primitive().clone();
     }
@@ -152,7 +158,7 @@ where
         let nulls = None;
         let (vals, map_idxs) = self.heap.drain();
         let arr = PrimitiveArray::<VAL>::new(ScalarBuffer::from(vals), nulls)
-            .with_data_type(self.batch.data_type().clone());
+            .with_data_type(self.value_type.clone());
         (Arc::new(arr), map_idxs)
     }
 }
@@ -193,6 +199,11 @@ where
     S: Array + Clone + From<Vec<Option<String>>> + 'static,
     for<'a> &'a S: StringArrayType<'a>,
 {
+    fn value_type(&self) -> &DataType {
+        // Strings don't store any metadata.
+        self.batch.data_type()
+    }
+
     fn set_batch(&mut self, vals: ArrayRef) {
         self.batch = vals
             .as_any()
@@ -567,11 +578,11 @@ pub fn is_supported_heap_type(vt: &DataType) -> bool {
 pub fn new_heap(
     limit: usize,
     desc: bool,
-    vt: &DataType,
+    vt: DataType,
 ) -> Result<Box<dyn ArrowHeap + Send>> {
     macro_rules! downcast_helper {
         ($vt:ty, $d:ident) => {
-            return Ok(Box::new(PrimitiveHeap::<$vt>::new(limit, desc)))
+            return Ok(Box::new(PrimitiveHeap::<$vt>::new(limit, desc, vt)))
         };
     }
 
diff --git a/datafusion/physical-plan/src/aggregates/topk/priority_map.rs b/datafusion/physical-plan/src/aggregates/topk/priority_map.rs
index b359b3a6b3395..fd4b0eb93ab8c 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<dyn ArrowHeap + Send>,
     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<Self> {
         Ok(Self {
             map: new_hash_table(capacity, key_type)?,
-            heap: new_heap(capacity, descending, &val_type)?,
+            heap: new_heap(capacity, descending, val_type)?,
             capacity,
             mapper: Vec::with_capacity(capacity),
-            val_type,
             null_count: 0,
         })
     }
@@ -140,7 +138,7 @@ impl PriorityMap {
             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_idxs.len());
             concat(&[vals.as_ref(), nulls.as_ref()])?
         };
         let ids = self.map.take_all(map_idxs);

From a374445fcb08d6e18aa4374023165158fa9ef2ba Mon Sep 17 00:00:00 2001
From: Michal Piatkowski <291740709+MassivePizza@users.noreply.github.com>
Date: Fri, 21 Aug 2026 17:44:17 +0200
Subject: [PATCH 12/12] inline TopKHeap swap

---
 datafusion/physical-plan/src/aggregates/topk/heap.rs | 1 +
 1 file changed, 1 insertion(+)

diff --git a/datafusion/physical-plan/src/aggregates/topk/heap.rs b/datafusion/physical-plan/src/aggregates/topk/heap.rs
index 71d30d27b2d34..23c51a5858770 100644
--- a/datafusion/physical-plan/src/aggregates/topk/heap.rs
+++ b/datafusion/physical-plan/src/aggregates/topk/heap.rs
@@ -404,6 +404,7 @@ impl<VAL: ValueType> TopKHeap<VAL> {
         }
     }
 
+    #[inline]
     fn swap(&mut self, a_idx: usize, b_idx: usize, mapper: &mut Vec<(usize, usize)>) {
         self.heap.swap(a_idx, b_idx);