datafusion_common/
heap_size.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! Estimating the heap-allocated memory owned by a value.
19//!
20//! The [`DFHeapSize`] trait reports the number of bytes a value owns on the
21//! heap, **excluding** the stack size of the value itself.
22//!
23//! Implementations need to use [`DFHeapSizeCtx`] that is pushed through every
24//! nested call. The context records which allocations have already been measured
25//! so they are only counted once.
26//!
27//! # Example
28//!
29//! ```
30//! use datafusion_common::heap_size::{DFHeapSize, DFHeapSizeCtx};
31//! use std::sync::Arc;
32//!
33//! let shared: Arc<String> = Arc::new("hello".to_string());
34//! let alias = Arc::clone(&shared);
35//!
36//! let mut ctx = DFHeapSizeCtx::default();
37//! // The shared allocation is counted once even when reached twice.
38//! let total = shared.heap_size(&mut ctx) + alias.heap_size(&mut ctx);
39//! assert_eq!(total, shared.heap_size(&mut DFHeapSizeCtx::default()));
40//! ```
41
42use crate::stats::Precision;
43use crate::{ColumnStatistics, ScalarValue, Statistics, TableReference};
44use arrow::array::{
45    Array, FixedSizeListArray, LargeListArray, LargeListViewArray, ListArray,
46    ListViewArray, MapArray, StructArray,
47};
48use arrow::datatypes::{
49    DataType, Field, Fields, IntervalDayTime, IntervalMonthDayNano, IntervalUnit,
50    TimeUnit, UnionFields, UnionMode, i256,
51};
52use chrono::{DateTime, Utc};
53use half::f16;
54use hashbrown::HashSet;
55use std::collections::HashMap;
56use std::fmt::Debug;
57use std::sync::Arc;
58
59/// Trait for computing how many bytes a value has allocated on the heap.
60///
61/// Implementations need to use [`DFHeapSizeCtx`] that is pushed through every
62/// nested call. The context records which allocations have already been measured
63/// so they are only counted once.
64///
65pub trait DFHeapSize {
66    /// Return the number of bytes this value has allocated on the heap,
67    /// including heap memory owned transitively by nested values.
68    ///
69    /// Note that the size of the type itself is not included in the result --
70    /// instead, that size is added by the caller (e.g. container).
71    fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize;
72}
73
74#[derive(Default)]
75pub struct DFHeapSizeCtx {
76    seen: HashSet<usize>,
77}
78
79impl DFHeapSizeCtx {
80    fn count_allocation_once(&mut self, ptr: usize) -> bool {
81        self.seen.insert(ptr)
82    }
83}
84
85impl DFHeapSize for Statistics {
86    fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize {
87        self.num_rows.heap_size(ctx)
88            + self.total_byte_size.heap_size(ctx)
89            + self.column_statistics.heap_size(ctx)
90    }
91}
92
93impl DFHeapSize for TableReference {
94    fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize {
95        match self {
96            TableReference::Bare { table } => table.heap_size(ctx),
97            TableReference::Partial { schema, table } => {
98                schema.heap_size(ctx) + table.heap_size(ctx)
99            }
100            TableReference::Full {
101                catalog,
102                schema,
103                table,
104            } => catalog.heap_size(ctx) + schema.heap_size(ctx) + table.heap_size(ctx),
105        }
106    }
107}
108
109impl<T: Debug + Clone + PartialEq + Eq + PartialOrd + DFHeapSize> DFHeapSize
110    for Precision<T>
111{
112    fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize {
113        self.get_value().map_or_else(|| 0, |v| v.heap_size(ctx))
114    }
115}
116
117impl DFHeapSize for ColumnStatistics {
118    fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize {
119        self.null_count.heap_size(ctx)
120            + self.max_value.heap_size(ctx)
121            + self.min_value.heap_size(ctx)
122            + self.sum_value.heap_size(ctx)
123            + self.distinct_count.heap_size(ctx)
124            + self.byte_size.heap_size(ctx)
125    }
126}
127
128impl DFHeapSize for ScalarValue {
129    fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize {
130        use crate::scalar::ScalarValue::*;
131        match self {
132            Null => 0,
133            Boolean(b) => b.heap_size(ctx),
134            Float16(f) => f.heap_size(ctx),
135            Float32(f) => f.heap_size(ctx),
136            Float64(f) => f.heap_size(ctx),
137            Decimal32(a, b, c) => a.heap_size(ctx) + b.heap_size(ctx) + c.heap_size(ctx),
138            Decimal64(a, b, c) => a.heap_size(ctx) + b.heap_size(ctx) + c.heap_size(ctx),
139            Decimal128(a, b, c) => a.heap_size(ctx) + b.heap_size(ctx) + c.heap_size(ctx),
140            Decimal256(a, b, c) => a.heap_size(ctx) + b.heap_size(ctx) + c.heap_size(ctx),
141            Int8(i) => i.heap_size(ctx),
142            Int16(i) => i.heap_size(ctx),
143            Int32(i) => i.heap_size(ctx),
144            Int64(i) => i.heap_size(ctx),
145            UInt8(u) => u.heap_size(ctx),
146            UInt16(u) => u.heap_size(ctx),
147            UInt32(u) => u.heap_size(ctx),
148            UInt64(u) => u.heap_size(ctx),
149            Utf8(u) => u.heap_size(ctx),
150            Utf8View(u) => u.heap_size(ctx),
151            LargeUtf8(l) => l.heap_size(ctx),
152            Binary(b) => b.heap_size(ctx),
153            BinaryView(b) => b.heap_size(ctx),
154            FixedSizeBinary(a, b) => a.heap_size(ctx) + b.heap_size(ctx),
155            LargeBinary(l) => l.heap_size(ctx),
156            FixedSizeList(f) => f.heap_size(ctx),
157            List(l) => l.heap_size(ctx),
158            LargeList(l) => l.heap_size(ctx),
159            Struct(s) => s.heap_size(ctx),
160            Map(m) => m.heap_size(ctx),
161            Date32(d) => d.heap_size(ctx),
162            Date64(d) => d.heap_size(ctx),
163            Time32Second(t) => t.heap_size(ctx),
164            Time32Millisecond(t) => t.heap_size(ctx),
165            Time64Microsecond(t) => t.heap_size(ctx),
166            Time64Nanosecond(t) => t.heap_size(ctx),
167            TimestampSecond(a, b) => a.heap_size(ctx) + b.heap_size(ctx),
168            TimestampMillisecond(a, b) => a.heap_size(ctx) + b.heap_size(ctx),
169            TimestampMicrosecond(a, b) => a.heap_size(ctx) + b.heap_size(ctx),
170            TimestampNanosecond(a, b) => a.heap_size(ctx) + b.heap_size(ctx),
171            IntervalYearMonth(i) => i.heap_size(ctx),
172            IntervalDayTime(i) => i.heap_size(ctx),
173            IntervalMonthDayNano(i) => i.heap_size(ctx),
174            DurationSecond(d) => d.heap_size(ctx),
175            DurationMillisecond(d) => d.heap_size(ctx),
176            DurationMicrosecond(d) => d.heap_size(ctx),
177            DurationNanosecond(d) => d.heap_size(ctx),
178            Union(a, b, c) => a.heap_size(ctx) + b.heap_size(ctx) + c.heap_size(ctx),
179            Dictionary(a, b) => a.heap_size(ctx) + b.heap_size(ctx),
180            RunEndEncoded(a, b, c) => {
181                a.heap_size(ctx) + b.heap_size(ctx) + c.heap_size(ctx)
182            }
183            ListView(a) => a.heap_size(ctx),
184            LargeListView(a) => a.heap_size(ctx),
185        }
186    }
187}
188
189impl DFHeapSize for DataType {
190    fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize {
191        use DataType::*;
192        match self {
193            Null => 0,
194            Boolean => 0,
195            Int8 => 0,
196            Int16 => 0,
197            Int32 => 0,
198            Int64 => 0,
199            UInt8 => 0,
200            UInt16 => 0,
201            UInt32 => 0,
202            UInt64 => 0,
203            Float16 => 0,
204            Float32 => 0,
205            Float64 => 0,
206            Timestamp(t, s) => t.heap_size(ctx) + s.heap_size(ctx),
207            Date32 => 0,
208            Date64 => 0,
209            Time32(t) => t.heap_size(ctx),
210            Time64(t) => t.heap_size(ctx),
211            Duration(t) => t.heap_size(ctx),
212            Interval(i) => i.heap_size(ctx),
213            Binary => 0,
214            FixedSizeBinary(i) => i.heap_size(ctx),
215            LargeBinary => 0,
216            BinaryView => 0,
217            Utf8 => 0,
218            LargeUtf8 => 0,
219            Utf8View => 0,
220            List(v) => v.heap_size(ctx),
221            ListView(v) => v.heap_size(ctx),
222            FixedSizeList(f, i) => f.heap_size(ctx) + i.heap_size(ctx),
223            LargeList(l) => l.heap_size(ctx),
224            LargeListView(l) => l.heap_size(ctx),
225            Struct(s) => s.heap_size(ctx),
226            Union(u, m) => u.heap_size(ctx) + m.heap_size(ctx),
227            Dictionary(a, b) => a.heap_size(ctx) + b.heap_size(ctx),
228            Decimal32(p, s) => p.heap_size(ctx) + s.heap_size(ctx),
229            Decimal64(p, s) => p.heap_size(ctx) + s.heap_size(ctx),
230            Decimal128(p, s) => p.heap_size(ctx) + s.heap_size(ctx),
231            Decimal256(p, s) => p.heap_size(ctx) + s.heap_size(ctx),
232            Map(m, b) => m.heap_size(ctx) + b.heap_size(ctx),
233            RunEndEncoded(a, b) => a.heap_size(ctx) + b.heap_size(ctx),
234        }
235    }
236}
237
238impl<T: DFHeapSize> DFHeapSize for Vec<T> {
239    fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize {
240        let item_size = size_of::<T>();
241        // account for the contents of the Vec
242        (self.capacity() * item_size) +
243            // add any heap allocations by contents
244            self.iter().map(|t| t.heap_size(ctx)).sum::<usize>()
245    }
246}
247
248impl<K: DFHeapSize, V: DFHeapSize> DFHeapSize for HashMap<K, V> {
249    fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize {
250        let capacity = self.capacity();
251        if capacity == 0 {
252            return 0;
253        }
254
255        // HashMap doesn't provide a way to get its heap size, so this is an approximation based on
256        // the behavior of hashbrown::HashMap as at version 0.16.0, and may become inaccurate
257        // if the implementation changes.
258        let key_val_size = size_of::<(K, V)>();
259        // Overhead for the control tags group, which may be smaller depending on architecture
260        let group_size = 16;
261        // 1 byte of metadata stored per bucket.
262        let metadata_size = 1;
263
264        // Compute the number of buckets for the capacity. Based on hashbrown's capacity_to_buckets
265        let buckets = if capacity < 15 {
266            let min_cap = match key_val_size {
267                0..=1 => 14,
268                2..=3 => 7,
269                _ => 3,
270            };
271            let cap = min_cap.max(capacity);
272            if cap < 4 {
273                4
274            } else if cap < 8 {
275                8
276            } else {
277                16
278            }
279        } else {
280            (capacity.saturating_mul(8) / 7).next_power_of_two()
281        };
282
283        group_size
284            + (buckets * (key_val_size + metadata_size))
285            + self.keys().map(|k| k.heap_size(ctx)).sum::<usize>()
286            + self.values().map(|v| v.heap_size(ctx)).sum::<usize>()
287    }
288}
289
290fn arc_ptr<T>(arc: &Arc<T>) -> usize {
291    Arc::as_ptr(arc) as usize
292}
293
294/// For unsized types, `Arc::as_ptr` returns the data address + metadata - we only need the thin address
295/// Casting through `*const i32` gets us the thin pointer
296fn arc_unsized_ptr<T: ?Sized>(arc: &Arc<T>) -> usize {
297    Arc::as_ptr(arc) as *const i32 as usize
298}
299
300impl<T: DFHeapSize> DFHeapSize for Arc<T> {
301    fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize {
302        let ptr = arc_ptr(self);
303
304        if !ctx.count_allocation_once(ptr) {
305            return 0;
306        }
307
308        // Arc stores weak and strong counts on the heap alongside an instance of T
309        2 * size_of::<usize>() + size_of::<T>() + self.as_ref().heap_size(ctx)
310    }
311}
312
313impl DFHeapSize for Arc<str> {
314    fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize {
315        let ptr = arc_unsized_ptr(self);
316
317        if !ctx.count_allocation_once(ptr) {
318            return 0;
319        }
320
321        // Arc stores weak and strong counts on the heap alongside an instance of T
322        2 * size_of::<usize>() + self.as_ref().heap_size(ctx)
323    }
324}
325
326impl DFHeapSize for Arc<dyn DFHeapSize> {
327    fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize {
328        let ptr = arc_unsized_ptr(self);
329
330        if !ctx.count_allocation_once(ptr) {
331            return 0;
332        }
333
334        // Arc stores weak and strong counts on the heap alongside an instance of T
335        2 * size_of::<usize>() + size_of_val(self.as_ref()) + self.as_ref().heap_size(ctx)
336    }
337}
338
339impl DFHeapSize for Fields {
340    fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize {
341        self.into_iter().map(|f| f.heap_size(ctx)).sum::<usize>()
342    }
343}
344
345impl<T: DFHeapSize> DFHeapSize for Box<T> {
346    fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize {
347        size_of::<T>() + self.as_ref().heap_size(ctx)
348    }
349}
350
351impl<T: DFHeapSize> DFHeapSize for Option<T> {
352    fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize {
353        self.as_ref().map(|inner| inner.heap_size(ctx)).unwrap_or(0)
354    }
355}
356
357impl<A, B> DFHeapSize for (A, B)
358where
359    A: DFHeapSize,
360    B: DFHeapSize,
361{
362    fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize {
363        self.0.heap_size(ctx) + self.1.heap_size(ctx)
364    }
365}
366
367impl<A, B, C> DFHeapSize for (A, B, C)
368where
369    A: DFHeapSize,
370    B: DFHeapSize,
371    C: DFHeapSize,
372{
373    fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize {
374        self.0.heap_size(ctx) + self.1.heap_size(ctx) + self.2.heap_size(ctx)
375    }
376}
377
378impl DFHeapSize for String {
379    fn heap_size(&self, _: &mut DFHeapSizeCtx) -> usize {
380        self.capacity()
381    }
382}
383
384impl DFHeapSize for str {
385    fn heap_size(&self, _: &mut DFHeapSizeCtx) -> usize {
386        // Internal accounting helper for owners like Arc<str>
387        self.len()
388    }
389}
390
391impl DFHeapSize for UnionFields {
392    fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize {
393        self.iter()
394            .map(|f| f.0.heap_size(ctx) + f.1.heap_size(ctx))
395            .sum()
396    }
397}
398
399impl DFHeapSize for Field {
400    fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize {
401        self.name().heap_size(ctx)
402            + self.data_type().heap_size(ctx)
403            + self.is_nullable().heap_size(ctx)
404            + self.dict_is_ordered().heap_size(ctx)
405            + self.metadata().heap_size(ctx)
406    }
407}
408
409impl DFHeapSize for IntervalMonthDayNano {
410    fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize {
411        self.days.heap_size(ctx)
412            + self.months.heap_size(ctx)
413            + self.nanoseconds.heap_size(ctx)
414    }
415}
416
417impl DFHeapSize for IntervalDayTime {
418    fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize {
419        self.days.heap_size(ctx) + self.milliseconds.heap_size(ctx)
420    }
421}
422
423/// Implement [`DFHeapSize`] for types that own no heap allocations.
424macro_rules! impl_zero_heap_size {
425    ($($t:ty),+ $(,)?) => {
426        $(
427            impl DFHeapSize for $t {
428                fn heap_size(&self, _: &mut DFHeapSizeCtx) -> usize {
429                    0 // no heap allocations
430                }
431            }
432        )+
433    };
434}
435
436impl_zero_heap_size!(
437    bool,
438    u8,
439    u16,
440    u32,
441    u64,
442    usize,
443    i8,
444    i16,
445    i32,
446    i64,
447    i128,
448    i256,
449    f16,
450    f32,
451    f64,
452    UnionMode,
453    TimeUnit,
454    IntervalUnit,
455    DateTime<Utc>,
456);
457
458/// Implement [`DFHeapSize`] for Arrow arrays types.
459macro_rules! impl_array_heap_size {
460    ($($t:ty),+ $(,)?) => {
461        $(
462            impl DFHeapSize for $t {
463                fn heap_size(&self, _: &mut DFHeapSizeCtx) -> usize {
464                    self.get_array_memory_size()
465                }
466            }
467        )+
468    };
469}
470
471impl_array_heap_size!(
472    StructArray,
473    LargeListArray,
474    LargeListViewArray,
475    ListArray,
476    ListViewArray,
477    FixedSizeListArray,
478    MapArray,
479);
480
481#[cfg(test)]
482mod tests {
483    use super::*;
484
485    fn size<T: DFHeapSize + ?Sized>(v: &T) -> usize {
486        v.heap_size(&mut DFHeapSizeCtx::default())
487    }
488
489    #[test]
490    fn test_heap_size_arc_avoid_double_accounting() {
491        let a1 = Arc::new(vec![1, 2, 3]);
492        let mut ctx = DFHeapSizeCtx::default();
493        let heap_size = a1.heap_size(&mut ctx);
494
495        let a2 = Arc::clone(&a1);
496        let a3 = Arc::clone(&a1);
497        let a4 = Arc::clone(&a3);
498
499        let mut ctx = DFHeapSizeCtx::default();
500        let heap_size_with_clones = a1.heap_size(&mut ctx)
501            + a2.heap_size(&mut ctx)
502            + a3.heap_size(&mut ctx)
503            + a4.heap_size(&mut ctx);
504
505        assert_eq!(heap_size, heap_size_with_clones);
506    }
507
508    #[test]
509    fn test_heap_size_arc_str_avoid_double_accounting() {
510        let a1: Arc<str> = Arc::from("Hello");
511        let mut ctx = DFHeapSizeCtx::default();
512        let heap_size = a1.heap_size(&mut ctx);
513
514        let a2 = Arc::clone(&a1);
515        let a3 = Arc::clone(&a1);
516        let a4 = Arc::clone(&a3);
517
518        let mut ctx = DFHeapSizeCtx::default();
519        let heap_size_with_clones = a1.heap_size(&mut ctx)
520            + a2.heap_size(&mut ctx)
521            + a3.heap_size(&mut ctx)
522            + a4.heap_size(&mut ctx);
523
524        assert_eq!(heap_size, heap_size_with_clones);
525    }
526
527    #[test]
528    fn test_arc_dyn() {
529        let a1: Arc<dyn DFHeapSize> = Arc::new(String::from("hello"));
530        let baseline = size(&a1);
531
532        let a2 = Arc::clone(&a1);
533        let mut ctx = DFHeapSizeCtx::default();
534        let with_clones = a1.heap_size(&mut ctx) + a2.heap_size(&mut ctx);
535        assert_eq!(baseline, with_clones);
536    }
537
538    #[test]
539    fn test_primitives() {
540        assert_eq!(size(&true), 0);
541        assert_eq!(size(&0u8), 0);
542        assert_eq!(size(&0u16), 0);
543        assert_eq!(size(&0u32), 0);
544        assert_eq!(size(&0u64), 0);
545        assert_eq!(size(&0usize), 0);
546        assert_eq!(size(&0i8), 0);
547        assert_eq!(size(&0i16), 0);
548        assert_eq!(size(&0i32), 0);
549        assert_eq!(size(&0i64), 0);
550        assert_eq!(size(&0i128), 0);
551        assert_eq!(size(&i256::ZERO), 0);
552        assert_eq!(size(&0f32), 0);
553        assert_eq!(size(&0f64), 0);
554        assert_eq!(size(&f16::from_f32(0.0)), 0);
555    }
556
557    #[test]
558    fn test_heap_size_union_mode() {
559        assert_eq!(size(&UnionMode::Sparse), 0);
560        assert_eq!(size(&UnionMode::Dense), 0);
561    }
562
563    #[test]
564    fn test_heap_size_time_units() {
565        assert_eq!(size(&TimeUnit::Second), 0);
566        assert_eq!(size(&IntervalUnit::YearMonth), 0);
567        assert_eq!(size(&DateTime::<Utc>::UNIX_EPOCH), 0);
568        assert_eq!(size(&Utc::now()), 0);
569    }
570
571    #[test]
572    fn test_string() {
573        let mut s = String::with_capacity(32);
574        s.push_str("hello");
575        assert_eq!(size(&s), 32);
576
577        let empty = String::new();
578        assert_eq!(size(&empty), 0);
579    }
580
581    #[test]
582    fn test_owned_str() {
583        let a: Arc<str> = Arc::from("Hello");
584        assert!(size(&a) > 0);
585    }
586
587    #[test]
588    fn test_option() {
589        let some: Option<String> = Some(String::from("hi"));
590        assert_eq!(size(&some), some.as_ref().unwrap().capacity());
591
592        let none: Option<String> = None;
593        assert_eq!(size(&none), 0);
594    }
595
596    #[test]
597    fn test_vec() {
598        let v: Vec<i32> = vec![1, 2, 3];
599        assert!(size(&v) > 0);
600
601        let strings = vec![String::from("ab"), String::from("cdef")];
602        assert!(size(&strings) > 0);
603
604        let empty: Vec<i32> = Vec::new();
605        assert_eq!(size(&empty), 0);
606    }
607
608    #[test]
609    fn test_box() {
610        let b: Box<i32> = Box::new(42);
611        assert!(size(&b) > 0);
612
613        let b: Box<String> = Box::new(String::from("hello"));
614        assert!(size(&b) > 0);
615    }
616
617    #[test]
618    fn test_tuple() {
619        let zero = (1i32, 2i64);
620        assert_eq!(size(&zero), 0);
621
622        let t = (String::from("hello"), String::from("world"));
623        assert!(size(&t) > 0);
624    }
625
626    #[test]
627    fn test_hashmap() {
628        let m: HashMap<i32, i32> = HashMap::new();
629        assert_eq!(size(&m), 0);
630
631        let mut m: HashMap<String, String> = HashMap::new();
632        m.insert("key".into(), "value".into());
633
634        assert!(size(&m) > 0);
635    }
636
637    #[test]
638    fn test_precision() {
639        let exact: Precision<usize> = Precision::Exact(42);
640        assert_eq!(size(&exact), 0);
641
642        let inexact: Precision<usize> = Precision::Inexact(99);
643        assert_eq!(size(&inexact), 0);
644
645        let absent: Precision<usize> = Precision::Absent;
646        assert_eq!(size(&absent), 0);
647    }
648
649    #[test]
650    fn test_scalar_values() {
651        assert_eq!(size(&ScalarValue::Null), 0);
652        assert_eq!(size(&ScalarValue::Int32(Some(42))), 0);
653        assert_eq!(size(&ScalarValue::Boolean(Some(true))), 0);
654        assert_eq!(size(&ScalarValue::Float64(None)), 0);
655
656        let sv = ScalarValue::Utf8(Some(String::from("hello")));
657        assert_eq!(size(&sv), "hello".len());
658
659        let sv = ScalarValue::Utf8(None);
660        assert_eq!(size(&sv), 0);
661    }
662
663    #[test]
664    fn test_data_type_primitives() {
665        assert_eq!(size(&DataType::Int32), 0);
666        assert_eq!(size(&DataType::Utf8), 0);
667        assert_eq!(size(&DataType::Boolean), 0);
668        assert_eq!(size(&DataType::Null), 0);
669    }
670
671    #[test]
672    fn test_data_type_with_field() {
673        let list = DataType::List(Arc::new(Field::new("item", DataType::Int32, true)));
674        assert!(size(&list) > 0);
675    }
676
677    #[test]
678    fn test_table_references() {
679        let tr = TableReference::bare("users");
680        // Arc<str> overhead (two usize counts) plus the bytes of "users".
681        assert!(size(&tr) > 0);
682        let tr = TableReference::full("cat", "schema", "users");
683        assert!(size(&tr) > 0);
684    }
685
686    #[test]
687    fn test_column_statistics() {
688        let mut col = ColumnStatistics::new_unknown();
689        col.max_value = Precision::Exact(ScalarValue::Utf8(Some("hello".into())));
690        col.min_value = Precision::Exact(ScalarValue::Utf8(Some("ab".into())));
691        assert_eq!(size(&col), "hello".len() + "ab".len());
692
693        let mut col = ColumnStatistics::new_unknown();
694        col.max_value = Precision::Exact(ScalarValue::Utf8(Some("hello".into())));
695        let stats = Statistics {
696            num_rows: Precision::Exact(10),
697            total_byte_size: Precision::Absent,
698            column_statistics: vec![col],
699        };
700        assert!(size(&stats) > 0);
701    }
702
703    #[test]
704    fn test_field() {
705        let field = Field::new("temperature", DataType::Float64, true);
706        assert!(size(&field) > 0);
707    }
708
709    #[test]
710    fn test_list_array() {
711        use arrow::array::types::Int32Type;
712
713        let array = ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
714            Some(vec![Some(1), Some(2), Some(3)]),
715            Some(vec![Some(4)]),
716        ]);
717        assert_eq!(size(&array), array.get_array_memory_size());
718        assert!(size(&array) > 0);
719
720        let large =
721            LargeListArray::from_iter_primitive::<Int32Type, _, _>(vec![Some(vec![
722                Some(1),
723                Some(2),
724            ])]);
725        assert_eq!(size(&large), large.get_array_memory_size());
726        assert!(size(&large) > 0);
727    }
728
729    #[test]
730    fn test_struct_array() {
731        use arrow::array::Int32Array;
732
733        let array = StructArray::from(vec![(
734            Arc::new(Field::new("a", DataType::Int32, true)),
735            Arc::new(Int32Array::from(vec![1, 2, 3])) as _,
736        )]);
737        assert_eq!(size(&array), array.get_array_memory_size());
738        assert!(size(&array) > 0);
739    }
740
741    #[test]
742    fn test_fixed_size_list_array() {
743        use arrow::array::Int32Array;
744
745        let values = Arc::new(Int32Array::from(vec![1, 2, 3, 4]));
746        let field = Arc::new(Field::new("item", DataType::Int32, true));
747        let array = FixedSizeListArray::new(field, 2, values, None);
748        assert_eq!(size(&array), array.get_array_memory_size());
749        assert!(size(&array) > 0);
750    }
751
752    #[test]
753    fn test_map_array() {
754        use arrow::array::{Int32Builder, MapBuilder, StringBuilder};
755
756        let mut builder =
757            MapBuilder::new(None, StringBuilder::new(), Int32Builder::new());
758        builder.keys().append_value("key");
759        builder.values().append_value(1);
760        builder.append(true).unwrap();
761        let array = builder.finish();
762        assert_eq!(size(&array), array.get_array_memory_size());
763        assert!(size(&array) > 0);
764    }
765}
datafusion_common/heap_size.rs

datafusion_common/
heap_size.rs
