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

vortex_array/stats/
flatbuffers.rs
