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

vortex_array/stats/
flatbuffers.rs
