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

vortex_array/stats/
flatbuffers.rs
