vortex_array/stats/
flatbuffers.rs

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

vortex_array/stats/
flatbuffers.rs
