re_byte_size/
arrow_sizes.rs

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