vortex_array/stats/
flatbuffers.rs

1use enum_iterator::all;
2use flatbuffers::{FlatBufferBuilder, Follow, WIPOffset};
3use vortex_error::VortexError;
4use vortex_flatbuffers::{ReadFlatBuffer, WriteFlatBuffer};
5use vortex_scalar::ScalarValue;
6
7use super::traits::{StatsProvider, StatsProviderExt};
8use crate::stats::{Precision, Stat, StatsSet};
9
10impl WriteFlatBuffer for StatsSet {
11    type Target<'t> = crate::flatbuffers::ArrayStats<'t>;
12
13    /// All statistics written must be exact
14    fn write_flatbuffer<'fb>(
15        &self,
16        fbb: &mut FlatBufferBuilder<'fb>,
17    ) -> WIPOffset<Self::Target<'fb>> {
18        let min = self
19            .get(Stat::Min)
20            .and_then(Precision::as_exact)
21            .map(|min| fbb.create_vector(&min.to_protobytes::<Vec<u8>>()));
22
23        let max = self
24            .get(Stat::Max)
25            .and_then(Precision::as_exact)
26            .map(|max| fbb.create_vector(&max.to_protobytes::<Vec<u8>>()));
27
28        let sum = self
29            .get(Stat::Sum)
30            .and_then(Precision::as_exact)
31            .map(|sum| fbb.create_vector(&sum.to_protobytes::<Vec<u8>>()));
32
33        let stat_args = &crate::flatbuffers::ArrayStatsArgs {
34            min,
35            max,
36            sum,
37            is_sorted: self
38                .get_as::<bool>(Stat::IsSorted)
39                .and_then(Precision::as_exact),
40            is_strict_sorted: self
41                .get_as::<bool>(Stat::IsStrictSorted)
42                .and_then(Precision::as_exact),
43            is_constant: self
44                .get_as::<bool>(Stat::IsConstant)
45                .and_then(Precision::as_exact),
46            null_count: self
47                .get_as::<u64>(Stat::NullCount)
48                .and_then(Precision::as_exact),
49            uncompressed_size_in_bytes: self
50                .get_as::<u64>(Stat::UncompressedSizeInBytes)
51                .and_then(Precision::as_exact),
52            nan_count: self
53                .get_as::<u64>(Stat::NaNCount)
54                .and_then(Precision::as_exact),
55        };
56
57        crate::flatbuffers::ArrayStats::create(fbb, stat_args)
58    }
59}
60
61impl ReadFlatBuffer for StatsSet {
62    type Source<'a> = crate::flatbuffers::ArrayStats<'a>;
63    type Error = VortexError;
64
65    fn read_flatbuffer<'buf>(
66        fb: &<Self::Source<'buf> as Follow<'buf>>::Inner,
67    ) -> Result<Self, Self::Error> {
68        let mut stats_set = StatsSet::default();
69
70        for stat in all::<Stat>() {
71            match stat {
72                Stat::IsConstant => {
73                    if let Some(is_constant) = fb.is_constant() {
74                        stats_set.set(Stat::IsConstant, Precision::Exact(is_constant.into()));
75                    }
76                }
77                Stat::IsSorted => {
78                    if let Some(is_sorted) = fb.is_sorted() {
79                        stats_set.set(Stat::IsSorted, Precision::Exact(is_sorted.into()));
80                    }
81                }
82                Stat::IsStrictSorted => {
83                    if let Some(is_strict_sorted) = fb.is_strict_sorted() {
84                        stats_set.set(
85                            Stat::IsStrictSorted,
86                            Precision::Exact(is_strict_sorted.into()),
87                        );
88                    }
89                }
90                Stat::Max => {
91                    if let Some(max) = fb.max() {
92                        stats_set.set(
93                            Stat::Max,
94                            Precision::Exact(ScalarValue::from_protobytes(max.bytes())?),
95                        );
96                    }
97                }
98                Stat::Min => {
99                    if let Some(min) = fb.min() {
100                        stats_set.set(
101                            Stat::Min,
102                            Precision::Exact(ScalarValue::from_protobytes(min.bytes())?),
103                        );
104                    }
105                }
106                Stat::NullCount => {
107                    if let Some(null_count) = fb.null_count() {
108                        stats_set.set(Stat::NullCount, Precision::Exact(null_count.into()));
109                    }
110                }
111                Stat::UncompressedSizeInBytes => {
112                    if let Some(uncompressed_size_in_bytes) = fb.uncompressed_size_in_bytes() {
113                        stats_set.set(
114                            Stat::UncompressedSizeInBytes,
115                            Precision::Exact(uncompressed_size_in_bytes.into()),
116                        );
117                    }
118                }
119                Stat::Sum => {
120                    if let Some(sum) = fb.sum() {
121                        stats_set.set(
122                            Stat::Sum,
123                            Precision::Exact(ScalarValue::from_protobytes(sum.bytes())?),
124                        );
125                    }
126                }
127                Stat::NaNCount => {
128                    if let Some(nan_count) = fb.nan_count() {
129                        stats_set.set(
130                            Stat::NaNCount,
131                            Precision::Exact(ScalarValue::from(nan_count)),
132                        );
133                    }
134                }
135            }
136        }
137
138        Ok(stats_set)
139    }
140}
vortex_array/stats/flatbuffers.rs

vortex_array/stats/
flatbuffers.rs
