lance_core/
deepsize.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4pub use lance_derive::DeepSizeOf;
5
6use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet};
7use std::mem::{size_of, size_of_val};
8use std::sync::atomic::{AtomicU64, AtomicUsize};
9use std::sync::{Arc, Mutex, RwLock};
10
11use arrow_array::{Array, RecordBatch};
12use arrow_buffer::ArrowNativeType;
13use arrow_data::ArrayData;
14
15pub struct Context {
16    seen: HashSet<usize>,
17}
18
19impl Default for Context {
20    fn default() -> Self {
21        Self::new()
22    }
23}
24
25impl Context {
26    pub fn new() -> Self {
27        Self {
28            seen: HashSet::new(),
29        }
30    }
31
32    /// Returns true if this pointer was NOT previously seen (i.e., it's new).
33    pub fn mark_seen(&mut self, ptr: usize) -> bool {
34        self.seen.insert(ptr)
35    }
36}
37
38pub trait DeepSizeOf {
39    fn deep_size_of(&self) -> usize {
40        size_of_val(self) + self.deep_size_of_children(&mut Context::new())
41    }
42
43    fn deep_size_of_children(&self, context: &mut Context) -> usize;
44}
45
46// Primitives — no heap children
47macro_rules! impl_deep_size_primitive {
48    ($($t:ty),*) => {
49        $(
50            impl DeepSizeOf for $t {
51                fn deep_size_of_children(&self, _context: &mut Context) -> usize {
52                    0
53                }
54            }
55        )*
56    };
57}
58
59impl_deep_size_primitive!(
60    u8,
61    u16,
62    u32,
63    u64,
64    u128,
65    usize,
66    i8,
67    i16,
68    i32,
69    i64,
70    i128,
71    isize,
72    f32,
73    f64,
74    bool,
75    ()
76);
77
78impl DeepSizeOf for str {
79    fn deep_size_of_children(&self, _context: &mut Context) -> usize {
80        0
81    }
82}
83
84impl DeepSizeOf for String {
85    fn deep_size_of_children(&self, _context: &mut Context) -> usize {
86        self.capacity()
87    }
88}
89
90impl DeepSizeOf for bytes::Bytes {
91    fn deep_size_of_children(&self, context: &mut Context) -> usize {
92        if context.mark_seen(self.as_ptr() as usize) {
93            self.len()
94        } else {
95            0
96        }
97    }
98}
99
100impl DeepSizeOf for AtomicU64 {
101    fn deep_size_of_children(&self, _context: &mut Context) -> usize {
102        0
103    }
104}
105
106impl DeepSizeOf for AtomicUsize {
107    fn deep_size_of_children(&self, _context: &mut Context) -> usize {
108        0
109    }
110}
111
112impl<T: DeepSizeOf, const N: usize> DeepSizeOf for [T; N] {
113    fn deep_size_of_children(&self, context: &mut Context) -> usize {
114        self.iter()
115            .map(|item| item.deep_size_of_children(context))
116            .sum()
117    }
118}
119
120impl<T: DeepSizeOf> DeepSizeOf for [T] {
121    fn deep_size_of_children(&self, context: &mut Context) -> usize {
122        // The slice's own element bytes are accounted for by the owner (e.g. the
123        // `size_of_val` in the `Arc`/`Box` impls); here we only sum the heap
124        // children of each element.
125        self.iter()
126            .map(|item| item.deep_size_of_children(context))
127            .sum()
128    }
129}
130
131impl<T: DeepSizeOf> DeepSizeOf for RwLock<T> {
132    fn deep_size_of_children(&self, context: &mut Context) -> usize {
133        self.read()
134            .map(|val| val.deep_size_of_children(context))
135            .unwrap_or(0)
136    }
137}
138
139impl<T: DeepSizeOf> DeepSizeOf for Mutex<T> {
140    fn deep_size_of_children(&self, context: &mut Context) -> usize {
141        self.lock()
142            .map(|val| val.deep_size_of_children(context))
143            .unwrap_or(0)
144    }
145}
146
147// Tuples
148macro_rules! impl_deep_size_tuple {
149    ($($name:ident),+) => {
150        impl<$($name: DeepSizeOf),+> DeepSizeOf for ($($name,)+) {
151            #[allow(non_snake_case)]
152            fn deep_size_of_children(&self, context: &mut Context) -> usize {
153                let ($($name,)+) = self;
154                0 $(+ $name.deep_size_of_children(context))+
155            }
156        }
157    };
158}
159
160impl_deep_size_tuple!(A, B);
161impl_deep_size_tuple!(A, B, C);
162impl_deep_size_tuple!(A, B, C, D);
163impl_deep_size_tuple!(A, B, C, D, E);
164impl_deep_size_tuple!(A, B, C, D, E, F);
165
166impl<T: DeepSizeOf> DeepSizeOf for Vec<T> {
167    fn deep_size_of_children(&self, context: &mut Context) -> usize {
168        self.capacity() * size_of::<T>()
169            + self
170                .iter()
171                .map(|item| item.deep_size_of_children(context))
172                .sum::<usize>()
173    }
174}
175
176impl<T: DeepSizeOf + ?Sized> DeepSizeOf for Box<T> {
177    fn deep_size_of_children(&self, context: &mut Context) -> usize {
178        size_of_val(&**self) + (**self).deep_size_of_children(context)
179    }
180}
181
182impl<T: DeepSizeOf + ?Sized> DeepSizeOf for Arc<T> {
183    fn deep_size_of_children(&self, context: &mut Context) -> usize {
184        if context.mark_seen(Self::as_ptr(self) as *const () as usize) {
185            size_of_val(&**self) + (**self).deep_size_of_children(context)
186        } else {
187            0
188        }
189    }
190}
191
192impl<T: DeepSizeOf> DeepSizeOf for Option<T> {
193    fn deep_size_of_children(&self, context: &mut Context) -> usize {
194        match self {
195            Some(val) => val.deep_size_of_children(context),
196            None => 0,
197        }
198    }
199}
200
201impl<K: DeepSizeOf, V: DeepSizeOf> DeepSizeOf for HashMap<K, V> {
202    fn deep_size_of_children(&self, context: &mut Context) -> usize {
203        // Each bucket holds a key-value pair plus hash metadata (~1 byte control per bucket).
204        // Robin hood / Swiss table capacity is always a power of 2.
205        let capacity_bytes = self.capacity() * (size_of::<K>() + size_of::<V>() + 1);
206        let children: usize = self
207            .iter()
208            .map(|(k, v)| k.deep_size_of_children(context) + v.deep_size_of_children(context))
209            .sum();
210        capacity_bytes + children
211    }
212}
213
214impl<K: DeepSizeOf> DeepSizeOf for HashSet<K> {
215    fn deep_size_of_children(&self, context: &mut Context) -> usize {
216        let capacity_bytes = self.capacity() * (size_of::<K>() + 1);
217        let children: usize = self.iter().map(|k| k.deep_size_of_children(context)).sum();
218        capacity_bytes + children
219    }
220}
221
222impl<K: DeepSizeOf, V: DeepSizeOf> DeepSizeOf for BTreeMap<K, V> {
223    fn deep_size_of_children(&self, context: &mut Context) -> usize {
224        // BTreeMap nodes have ~11 entries each. Rough estimate: per-entry overhead ~3 pointers.
225        let per_entry = size_of::<K>() + size_of::<V>() + 3 * size_of::<usize>();
226        let overhead = self.len() * per_entry;
227        let children: usize = self
228            .iter()
229            .map(|(k, v)| k.deep_size_of_children(context) + v.deep_size_of_children(context))
230            .sum();
231        overhead + children
232    }
233}
234
235impl<K: DeepSizeOf> DeepSizeOf for BTreeSet<K> {
236    fn deep_size_of_children(&self, context: &mut Context) -> usize {
237        let per_entry = size_of::<K>() + 3 * size_of::<usize>();
238        let overhead = self.len() * per_entry;
239        let children: usize = self.iter().map(|k| k.deep_size_of_children(context)).sum();
240        overhead + children
241    }
242}
243
244// Arrow types
245
246fn record_array_data(context: &mut Context, data: &ArrayData) -> usize {
247    let mut total = 0;
248    for buffer in data.buffers() {
249        if context.mark_seen(buffer.as_ptr() as usize) {
250            total += buffer.capacity();
251        }
252    }
253    if let Some(nulls) = data.nulls() {
254        let null_buf = nulls.inner().inner();
255        if context.mark_seen(null_buf.as_ptr() as usize) {
256            total += null_buf.capacity();
257        }
258    }
259    for child in data.child_data() {
260        total += record_array_data(context, child);
261    }
262    total
263}
264
265impl DeepSizeOf for dyn Array {
266    fn deep_size_of_children(&self, context: &mut Context) -> usize {
267        // `to_data()` only clones Arc refs (no data copy) and allocates a small
268        // ArrayData metadata struct. This lets us walk buffer pointers for dedup.
269        // Cost is O(number_of_buffers), not O(data_size).
270        let data = self.to_data();
271        record_array_data(context, &data)
272    }
273}
274
275impl DeepSizeOf for RecordBatch {
276    fn deep_size_of_children(&self, context: &mut Context) -> usize {
277        self.columns()
278            .iter()
279            .map(|col| col.deep_size_of_children(context))
280            .sum()
281    }
282}
283
284impl<T> DeepSizeOf for arrow_buffer::ScalarBuffer<T>
285where
286    T: ArrowNativeType,
287{
288    fn deep_size_of_children(&self, context: &mut Context) -> usize {
289        // Track the underlying buffer pointer to avoid double-counting shared allocations.
290        // Use capacity() rather than len() * size_of::<T>() because sliced buffers retain
291        // their full original allocation.
292        let buf = self.inner();
293        if context.mark_seen(buf.as_ptr() as usize) {
294            buf.capacity()
295        } else {
296            0
297        }
298    }
299}
300
301#[cfg(test)]
302mod tests {
303    use super::*;
304    use arrow_array::{Int32Array, StringArray, StructArray};
305    use arrow_schema::{DataType, Field, Fields, Schema};
306
307    #[test]
308    fn test_basic_record_batch() {
309        let batch = RecordBatch::try_new(
310            Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])),
311            vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
312        )
313        .unwrap();
314
315        let size = batch.deep_size_of();
316        // Should at least include the buffer for 3 i32s
317        assert!(size >= 3 * size_of::<i32>());
318    }
319
320    #[test]
321    fn test_same_batch_dedup() {
322        let batch = RecordBatch::try_new(
323            Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])),
324            vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5]))],
325        )
326        .unwrap();
327
328        let mut ctx = Context::new();
329        let size_a = batch.deep_size_of_children(&mut ctx);
330        let size_b = batch.deep_size_of_children(&mut ctx);
331
332        // First measurement should report buffer sizes
333        assert!(size_a > 0);
334        // Second measurement of the same batch should add nothing (buffers already seen)
335        assert_eq!(size_b, 0);
336    }
337
338    #[test]
339    fn test_arc_dedup() {
340        let batch = Arc::new(
341            RecordBatch::try_new(
342                Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])),
343                vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
344            )
345            .unwrap(),
346        );
347        let clone = Arc::clone(&batch);
348
349        let mut ctx = Context::new();
350        let size_a = batch.deep_size_of_children(&mut ctx);
351        let size_b = clone.deep_size_of_children(&mut ctx);
352
353        assert!(size_a > 0);
354        assert_eq!(size_b, 0);
355    }
356
357    #[test]
358    fn test_multi_column_shared_array() {
359        // Two columns pointing to the same Arc<dyn Array>
360        let array: Arc<dyn Array> = Arc::new(Int32Array::from(vec![10, 20, 30]));
361        let schema = Arc::new(Schema::new(vec![
362            Field::new("a", DataType::Int32, false),
363            Field::new("b", DataType::Int32, false),
364        ]));
365
366        // Single-column batch for reference
367        let one_col = RecordBatch::try_new(
368            Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])),
369            vec![array.clone()],
370        )
371        .unwrap();
372
373        // Two-column batch with the same Arc shared
374        let two_col = RecordBatch::try_new(schema, vec![array.clone(), array]).unwrap();
375
376        let mut ctx1 = Context::new();
377        let size_one = one_col.deep_size_of_children(&mut ctx1);
378
379        let mut ctx2 = Context::new();
380        let size_two = two_col.deep_size_of_children(&mut ctx2);
381
382        // Both should report the same size since the second column's Arc is
383        // already seen and contributes nothing
384        assert_eq!(size_one, size_two);
385    }
386
387    #[test]
388    fn test_nested_struct_array() {
389        let int_array = Int32Array::from(vec![1, 2, 3]);
390        let str_array = StringArray::from(vec!["a", "b", "c"]);
391        let struct_array = StructArray::from(vec![
392            (
393                Arc::new(Field::new("x", DataType::Int32, false)),
394                Arc::new(int_array) as Arc<dyn Array>,
395            ),
396            (
397                Arc::new(Field::new("y", DataType::Utf8, false)),
398                Arc::new(str_array) as Arc<dyn Array>,
399            ),
400        ]);
401
402        let batch = RecordBatch::try_new(
403            Arc::new(Schema::new(vec![Field::new(
404                "s",
405                DataType::Struct(Fields::from(vec![
406                    Field::new("x", DataType::Int32, false),
407                    Field::new("y", DataType::Utf8, false),
408                ])),
409                false,
410            )])),
411            vec![Arc::new(struct_array)],
412        )
413        .unwrap();
414
415        let size = batch.deep_size_of();
416        // Should include buffers for both child arrays
417        assert!(size > 3 * size_of::<i32>());
418    }
419
420    #[test]
421    fn test_std_types() {
422        assert_eq!(42u32.deep_size_of(), size_of::<u32>());
423
424        let s = String::from("hello");
425        assert!(s.deep_size_of() >= size_of::<String>() + 5);
426
427        let v = vec![1u32, 2, 3];
428        assert!(v.deep_size_of() >= size_of::<Vec<u32>>() + 3 * size_of::<u32>());
429
430        let a = Arc::new(42u32);
431        let b = Arc::clone(&a);
432        let mut ctx = Context::new();
433        let size_a = a.deep_size_of_children(&mut ctx);
434        let size_b = b.deep_size_of_children(&mut ctx);
435        assert_eq!(size_a, size_of::<u32>());
436        assert_eq!(size_b, 0);
437    }
438
439    #[test]
440    fn test_derive_macro() {
441        use lance_derive::DeepSizeOf;
442
443        #[derive(DeepSizeOf)]
444        struct Outer {
445            count: u64,
446            label: String,
447            inner: Inner,
448        }
449
450        #[derive(DeepSizeOf)]
451        struct Inner {
452            values: Vec<u32>,
453        }
454
455        let val = Outer {
456            count: 7,
457            label: String::from("hello"),
458            inner: Inner {
459                values: vec![1, 2, 3],
460            },
461        };
462
463        let size = val.deep_size_of();
464        // Must be at least the stack size + heap allocations for label + values
465        assert!(size >= std::mem::size_of::<Outer>() + 5 + 3 * std::mem::size_of::<u32>());
466    }
467}
lance_core/deepsize.rs

lance_core/
deepsize.rs
