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 1/4] 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>>()
     }
 }
 

From 6a80d0c8ffcb2b549cc7bca554b0ef393e7eec84 Mon Sep 17 00:00:00 2001
From: linfeng <33561138+lyne7-sc@users.noreply.github.com>
Date: Thu, 9 Jul 2026 13:58:23 +0800
Subject: [PATCH 2/4] apply suggestion

Co-authored-by: Oleks V <comphead@users.noreply.github.com>
---
 datafusion/functions-aggregate/src/average.rs | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/datafusion/functions-aggregate/src/average.rs b/datafusion/functions-aggregate/src/average.rs
index ddd544a42a9c1..16f74281b0fc5 100644
--- a/datafusion/functions-aggregate/src/average.rs
+++ b/datafusion/functions-aggregate/src/average.rs
@@ -970,7 +970,11 @@ where
     }
 
     fn size(&self) -> usize {
-        self.counts.capacity() * size_of::<u64>()
-            + self.sums.capacity() * size_of::<T::Native>()
+    // Heap buffers
+    self.counts.capacity() * size_of::<u64>()
+        + self.sums.capacity() * size_of::<T::Native>()
+        // Vec struct overhead (ptr, len, cap) for each field
+        + size_of::<Vec<u64>>()
+        + size_of::<Vec<T::Native>>()
     }
 }

From 4a9647e728e3fec1b192dfd44903f482e50c4aec Mon Sep 17 00:00:00 2001
From: linfeng <33561138+lyne7-sc@users.noreply.github.com>
Date: Thu, 9 Jul 2026 14:20:33 +0800
Subject: [PATCH 3/4] fmt and account for avg null_state memory

---
 datafusion/functions-aggregate/src/average.rs | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/datafusion/functions-aggregate/src/average.rs b/datafusion/functions-aggregate/src/average.rs
index 16f74281b0fc5..278a861de2024 100644
--- a/datafusion/functions-aggregate/src/average.rs
+++ b/datafusion/functions-aggregate/src/average.rs
@@ -970,11 +970,13 @@ where
     }
 
     fn size(&self) -> usize {
-    // Heap buffers
-    self.counts.capacity() * size_of::<u64>()
+        // Heap buffers
+        self.counts.capacity() * size_of::<u64>()
         + self.sums.capacity() * size_of::<T::Native>()
         // Vec struct overhead (ptr, len, cap) for each field
         + size_of::<Vec<u64>>()
         + size_of::<Vec<T::Native>>()
+        // Null tracking buffers
+        + self.null_state.size()
     }
 }

From feb4f145958504447c059622db76bfd3ba76974a Mon Sep 17 00:00:00 2001
From: linfeng <33561138+lyne7-sc@users.noreply.github.com>
Date: Thu, 9 Jul 2026 14:27:53 +0800
Subject: [PATCH 4/4] account for median vec overhead

---
 datafusion/functions-aggregate/src/median.rs | 1 +
 1 file changed, 1 insertion(+)

diff --git a/datafusion/functions-aggregate/src/median.rs b/datafusion/functions-aggregate/src/median.rs
index 96410d652b352..9a6ef3e7e5fc5 100644
--- a/datafusion/functions-aggregate/src/median.rs
+++ b/datafusion/functions-aggregate/src/median.rs
@@ -547,6 +547,7 @@ impl<T: ArrowNumericType + Send> GroupsAccumulator for MedianGroupsAccumulator<T
             .sum::<usize>()
             // account for size of self.group_values too
             + self.group_values.capacity() * size_of::<Vec<T::Native>>()
+            + size_of::<Vec<Vec<T::Native>>>()
     }
 }