diff --git a/datafusion/common/src/heap_size.rs b/datafusion/common/src/heap_size.rs index 405736dbf9c9b..037f807fce9d8 100644 --- a/datafusion/common/src/heap_size.rs +++ b/datafusion/common/src/heap_size.rs @@ -76,6 +76,12 @@ pub struct DFHeapSizeCtx { seen: HashSet, } +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,11 +287,21 @@ impl DFHeapSize for HashMap { } } +fn arc_ptr(arc: &Arc) -> 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(arc: &Arc) -> usize { + Arc::as_ptr(arc) as *const i32 as usize +} + impl DFHeapSize for Arc { 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 !ctx.count_allocation_once(ptr) { return 0; } @@ -296,9 +312,9 @@ impl DFHeapSize for Arc { impl DFHeapSize for Arc { 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 !ctx.count_allocation_once(ptr) { return 0; } @@ -309,9 +325,9 @@ impl DFHeapSize for Arc { impl DFHeapSize for Arc { 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 !ctx.count_allocation_once(ptr) { return 0; }