vortex_array/stats/
flatbuffers.rs

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

vortex_array/stats/
flatbuffers.rs
