From eda70906b4f6ff1e640ca106523b90b9adacd896 Mon Sep 17 00:00:00 2001
From: linfeng <lilinfeng0310@gmail.com>
Date: Mon, 6 Jul 2026 22:56:25 +0800
Subject: [PATCH] Fix memory size accounting for grouped median and avg

---
 datafusion/functions-aggregate/src/average.rs | 3 ++-
 datafusion/functions-aggregate/src/median.rs  | 6 +++---
 2 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/datafusion/functions-aggregate/src/average.rs b/datafusion/functions-aggregate/src/average.rs
index 06c76946343dc..ddd544a42a9c1 100644
--- a/datafusion/functions-aggregate/src/average.rs
+++ b/datafusion/functions-aggregate/src/average.rs
@@ -970,6 +970,7 @@ where
     }
 
     fn size(&self) -> usize {
-        self.counts.capacity() * size_of::<u64>() + self.sums.capacity() * size_of::<T>()
+        self.counts.capacity() * size_of::<u64>()
+            + self.sums.capacity() * size_of::<T::Native>()
     }
 }
diff --git a/datafusion/functions-aggregate/src/median.rs b/datafusion/functions-aggregate/src/median.rs
index 4a0da10f51845..96410d652b352 100644
--- a/datafusion/functions-aggregate/src/median.rs
+++ b/datafusion/functions-aggregate/src/median.rs
@@ -543,10 +543,10 @@ impl<T: ArrowNumericType + Send> GroupsAccumulator for MedianGroupsAccumulator<T
     fn size(&self) -> usize {
         self.group_values
             .iter()
-            .map(|values| values.capacity() * size_of::<T>())
+            .map(|values| values.capacity() * size_of::<T::Native>())
             .sum::<usize>()
-            // account for size of self.grou_values too
-            + self.group_values.capacity() * size_of::<Vec<T>>()
+            // account for size of self.group_values too
+            + self.group_values.capacity() * size_of::<Vec<T::Native>>()
     }
 }