re_byte_size/
arrow_sizes.rs1use 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}