vortex_array/stats/
flatbuffers.rs

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

vortex_array/stats/
flatbuffers.rs
