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| min.write_flatbuffer(fbb));
22
23        let max = self
24            .get(Stat::Max)
25            .and_then(Precision::as_exact)
26            .map(|max| max.write_flatbuffer(fbb));
27
28        let sum = self
29            .get(Stat::Sum)
30            .and_then(Precision::as_exact)
31            .map(|max| max.write_flatbuffer(fbb));
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        };
53
54        crate::flatbuffers::ArrayStats::create(fbb, stat_args)
55    }
56}
57
58impl ReadFlatBuffer for StatsSet {
59    type Source<'a> = crate::flatbuffers::ArrayStats<'a>;
60    type Error = VortexError;
61
62    fn read_flatbuffer<'buf>(
63        fb: &<Self::Source<'buf> as Follow<'buf>>::Inner,
64    ) -> Result<Self, Self::Error> {
65        let mut stats_set = StatsSet::default();
66
67        for stat in all::<Stat>() {
68            match stat {
69                Stat::IsConstant => {
70                    if let Some(is_constant) = fb.is_constant() {
71                        stats_set.set(Stat::IsConstant, Precision::Exact(is_constant.into()));
72                    }
73                }
74                Stat::IsSorted => {
75                    if let Some(is_sorted) = fb.is_sorted() {
76                        stats_set.set(Stat::IsSorted, Precision::Exact(is_sorted.into()));
77                    }
78                }
79                Stat::IsStrictSorted => {
80                    if let Some(is_strict_sorted) = fb.is_strict_sorted() {
81                        stats_set.set(
82                            Stat::IsStrictSorted,
83                            Precision::Exact(is_strict_sorted.into()),
84                        );
85                    }
86                }
87                Stat::Max => {
88                    if let Some(max) = fb.max() {
89                        stats_set.set(Stat::Max, Precision::Exact(ScalarValue::try_from(max)?));
90                    }
91                }
92                Stat::Min => {
93                    if let Some(min) = fb.min() {
94                        stats_set.set(Stat::Min, Precision::Exact(ScalarValue::try_from(min)?));
95                    }
96                }
97                Stat::NullCount => {
98                    if let Some(null_count) = fb.null_count() {
99                        stats_set.set(Stat::NullCount, Precision::Exact(null_count.into()));
100                    }
101                }
102                Stat::UncompressedSizeInBytes => {
103                    if let Some(uncompressed_size_in_bytes) = fb.uncompressed_size_in_bytes() {
104                        stats_set.set(
105                            Stat::UncompressedSizeInBytes,
106                            Precision::Exact(uncompressed_size_in_bytes.into()),
107                        );
108                    }
109                }
110                Stat::Sum => {
111                    if let Some(sum) = fb.sum() {
112                        stats_set.set(Stat::Sum, Precision::Exact(ScalarValue::try_from(sum)?));
113                    }
114                }
115            }
116        }
117
118        Ok(stats_set)
119    }
120}
vortex_array/stats/flatbuffers.rs

vortex_array/stats/
flatbuffers.rs
