From ce007af719af0dd22dc80fc9d810a362785ab140 Mon Sep 17 00:00:00 2001
From: saadtajwar <saadtajwar19@gmail.com>
Date: Mon, 6 Jul 2026 16:44:25 -0400
Subject: [PATCH 1/2] refactor: centralizing shared-allocation accounting for
 Arc DFHeapSize impls

---
 datafusion/common/src/heap_size.rs | 24 ++++++++++++++++++------
 1 file changed, 18 insertions(+), 6 deletions(-)

diff --git a/datafusion/common/src/heap_size.rs b/datafusion/common/src/heap_size.rs
index 869946d82414f..4376161cff14b 100644
--- a/datafusion/common/src/heap_size.rs
+++ b/datafusion/common/src/heap_size.rs
@@ -281,11 +281,23 @@ impl<K: DFHeapSize, V: DFHeapSize> DFHeapSize for HashMap<K, V> {
     }
 }
 
+fn count_allocation_once(ptr: usize, ctx: &mut DFHeapSizeCtx) -> bool {
+    ctx.seen.insert(ptr)
+}
+
+fn arc_ptr<T>(arc: &Arc<T>) -> usize {
+    Arc::as_ptr(arc) as usize
+}
+
+fn arc_unsized_ptr<T: ?Sized>(arc: &Arc<T>) -> usize {
+    Arc::as_ptr(arc) as *const i32 as usize
+}
+
 impl<T: DFHeapSize> DFHeapSize for Arc<T> {
     fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize {
-        let ptr = Arc::as_ptr(self) as usize;
+        let ptr = arc_ptr(self);
 
-        if !ctx.seen.insert(ptr) {
+        if !count_allocation_once(ptr, ctx) {
             return 0;
         }
 
@@ -296,9 +308,9 @@ impl<T: DFHeapSize> DFHeapSize for Arc<T> {
 
 impl DFHeapSize for Arc<str> {
     fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize {
-        let ptr = Arc::as_ptr(self) as *const i32 as usize;
+        let ptr = arc_unsized_ptr(self);
 
-        if !ctx.seen.insert(ptr) {
+        if !count_allocation_once(ptr, ctx) {
             return 0;
         }
 
@@ -309,9 +321,9 @@ impl DFHeapSize for Arc<str> {
 
 impl DFHeapSize for Arc<dyn DFHeapSize> {
     fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize {
-        let ptr = Arc::as_ptr(self) as *const i32 as usize;
+        let ptr = arc_unsized_ptr(self);
 
-        if !ctx.seen.insert(ptr) {
+        if !count_allocation_once(ptr, ctx) {
             return 0;
         }
 

From 6eb6077d83b1a4e4daf72490dc90d0a72f0e0f9f Mon Sep 17 00:00:00 2001
From: saadtajwar <saadtajwar19@gmail.com>
Date: Tue, 7 Jul 2026 08:33:23 -0400
Subject: [PATCH 2/2] addressing comments

---
 datafusion/common/src/heap_size.rs | 18 +++++++++++-------
 1 file changed, 11 insertions(+), 7 deletions(-)

diff --git a/datafusion/common/src/heap_size.rs b/datafusion/common/src/heap_size.rs
index 4376161cff14b..667d6cf05635d 100644
--- a/datafusion/common/src/heap_size.rs
+++ b/datafusion/common/src/heap_size.rs
@@ -76,6 +76,12 @@ pub struct DFHeapSizeCtx {
     seen: HashSet<usize>,
 }
 
+impl DFHeapSizeCtx {
+    fn count_allocation_once(&mut self, ptr: usize) -> bool {
+        self.seen.insert(ptr)
+    }
+}
+
 impl DFHeapSize for Statistics {
     fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize {
         self.num_rows.heap_size(ctx)
@@ -281,14 +287,12 @@ impl<K: DFHeapSize, V: DFHeapSize> DFHeapSize for HashMap<K, V> {
     }
 }
 
-fn count_allocation_once(ptr: usize, ctx: &mut DFHeapSizeCtx) -> bool {
-    ctx.seen.insert(ptr)
-}
-
 fn arc_ptr<T>(arc: &Arc<T>) -> usize {
     Arc::as_ptr(arc) as usize
 }
 
+/// For unsized types, `Arc::as_ptr` returns the data address + metadata - we only need the thin address
+/// Casting through `*const i32` gets us the thin pointer
 fn arc_unsized_ptr<T: ?Sized>(arc: &Arc<T>) -> usize {
     Arc::as_ptr(arc) as *const i32 as usize
 }
@@ -297,7 +301,7 @@ impl<T: DFHeapSize> DFHeapSize for Arc<T> {
     fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize {
         let ptr = arc_ptr(self);
 
-        if !count_allocation_once(ptr, ctx) {
+        if !ctx.count_allocation_once(ptr) {
             return 0;
         }
 
@@ -310,7 +314,7 @@ impl DFHeapSize for Arc<str> {
     fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize {
         let ptr = arc_unsized_ptr(self);
 
-        if !count_allocation_once(ptr, ctx) {
+        if !ctx.count_allocation_once(ptr) {
             return 0;
         }
 
@@ -323,7 +327,7 @@ impl DFHeapSize for Arc<dyn DFHeapSize> {
     fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize {
         let ptr = arc_unsized_ptr(self);
 
-        if !count_allocation_once(ptr, ctx) {
+        if !ctx.count_allocation_once(ptr) {
             return 0;
         }