re_byte_size/
arrow_sizes.rs

1use arrow::{
2    array::{Array, ArrayRef, ListArray, RecordBatch},
3    buffer::{Buffer, ScalarBuffer},
4    datatypes::{ArrowNativeType, DataType, Field, Fields, Schema, UnionFields},
5};
6
7use super::SizeBytes;
8
9impl SizeBytes for Buffer {
10    #[inline]
11    fn heap_size_bytes(&self) -> u64 {
12        self.capacity() as u64
13    }
14}
15
16impl SizeBytes for dyn Array {
17    #[inline]
18    fn heap_size_bytes(&self) -> u64 {
19        self.get_buffer_memory_size() as u64
20    }
21}
22
23impl<T: Array> SizeBytes for &T {
24    #[inline]
25    fn heap_size_bytes(&self) -> u64 {
26        self.get_buffer_memory_size() as u64
27    }
28}
29
30impl SizeBytes for ArrayRef {
31    #[inline]
32    fn heap_size_bytes(&self) -> u64 {
33        self.get_buffer_memory_size() as u64
34    }
35}
36
37impl<T: ArrowNativeType> SizeBytes for ScalarBuffer<T> {
38    #[inline]
39    fn heap_size_bytes(&self) -> u64 {
40        self.inner().capacity() as _
41    }
42}
43
44impl SizeBytes for ListArray {
45    #[inline]
46    fn heap_size_bytes(&self) -> u64 {
47        self.get_buffer_memory_size() as u64
48    }
49}
50
51impl SizeBytes for RecordBatch {
52    #[inline]
53    fn heap_size_bytes(&self) -> u64 {
54        self.schema().heap_size_bytes()
55            + self
56                .columns()
57                .iter()
58                .map(|array| array.heap_size_bytes())
59                .sum::<u64>()
60    }
61}
62
63impl SizeBytes for Schema {
64    #[inline]
65    fn heap_size_bytes(&self) -> u64 {
66        let Self { fields, metadata } = self;
67        fields.heap_size_bytes() + metadata.heap_size_bytes()
68    }
69}
70
71impl SizeBytes for Fields {
72    #[inline]
73    fn heap_size_bytes(&self) -> u64 {
74        self.iter().map(|field| field.heap_size_bytes()).sum()
75    }
76}
77
78impl SizeBytes for Field {
79    #[inline]
80    fn heap_size_bytes(&self) -> u64 {
81        self.name().heap_size_bytes() + self.data_type().heap_size_bytes()
82    }
83}
84
85impl SizeBytes for DataType {
86    #[inline]
87    fn heap_size_bytes(&self) -> u64 {
88        match self {
89            Self::Null
90            | Self::Boolean
91            | Self::Int8
92            | Self::Int16
93            | Self::Int32
94            | Self::Int64
95            | Self::UInt8
96            | Self::UInt16
97            | Self::UInt32
98            | Self::UInt64
99            | Self::Float16
100            | Self::Float32
101            | Self::Float64
102            | Self::Date32
103            | Self::Date64
104            | Self::Binary
105            | Self::LargeBinary
106            | Self::Utf8
107            | Self::LargeUtf8
108            | Self::BinaryView
109            | Self::Decimal128(_, _)
110            | Self::Decimal256(_, _)
111            | Self::FixedSizeBinary(_)
112            | Self::Utf8View => 0,
113            Self::Timestamp(_time_unit, _tz) => 0,
114
115            Self::Time32(_time_unit) | Self::Time64(_time_unit) | Self::Duration(_time_unit) => 0,
116
117            Self::Interval(_interval_unit) => 0,
118
119            Self::List(field)
120            | Self::ListView(field)
121            | Self::FixedSizeList(field, _)
122            | Self::LargeList(field)
123            | Self::Map(field, _)
124            | Self::LargeListView(field) => field.heap_size_bytes(),
125
126            Self::Union(fields, _) => fields.heap_size_bytes(),
127            Self::Struct(fields) => fields.heap_size_bytes(),
128
129            Self::Dictionary(key, value) => key.heap_size_bytes() + value.heap_size_bytes(),
130
131            Self::RunEndEncoded(field, field1) => {
132                field.heap_size_bytes() + field1.heap_size_bytes()
133            }
134        }
135    }
136}
137
138impl SizeBytes for UnionFields {
139    fn heap_size_bytes(&self) -> u64 {
140        self.iter().map(|(_, field)| field.heap_size_bytes()).sum()
141    }
142}