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 half::f16;
7use vortex_error::VortexResult;
8use vortex_error::vortex_bail;
9use vortex_session::VortexSession;
10
11use crate::dtype::DType;
12use crate::dtype::Nullability;
13use crate::dtype::PType;
14use crate::expr::stats::Precision;
15use crate::expr::stats::Stat;
16use crate::flatbuffers::WriteFlatBuffer;
17use crate::flatbuffers::array as fba;
18use crate::scalar::PValue;
19use crate::scalar::ScalarValue;
20use crate::stats::StatsSet;
21use crate::stats::StatsSetRef;
22
23impl WriteFlatBuffer for StatsSetRef<'_> {
24    type Target<'t> = fba::ArrayStats<'t>;
25
26    /// All statistics written must be exact
27    fn write_flatbuffer<'fb>(
28        &self,
29        fbb: &mut FlatBufferBuilder<'fb>,
30    ) -> VortexResult<WIPOffset<Self::Target<'fb>>> {
31        self.with_typed_stats_set(|stats_set| stats_set.values.write_flatbuffer(fbb))
32    }
33}
34
35impl WriteFlatBuffer for StatsSet {
36    type Target<'t> = fba::ArrayStats<'t>;
37
38    /// All statistics written must be exact
39    fn write_flatbuffer<'fb>(
40        &self,
41        fbb: &mut FlatBufferBuilder<'fb>,
42    ) -> VortexResult<WIPOffset<Self::Target<'fb>>> {
43        let (min_precision, min) = match self.get(Stat::Min) {
44            Precision::Exact(min) => (
45                fba::Precision::Exact,
46                Some(fbb.create_vector(&ScalarValue::to_proto_bytes::<Vec<u8>>(Some(&min)))),
47            ),
48            Precision::Inexact(min) => (
49                fba::Precision::Inexact,
50                Some(fbb.create_vector(&ScalarValue::to_proto_bytes::<Vec<u8>>(Some(&min)))),
51            ),
52            Precision::Absent => (fba::Precision::Inexact, None),
53        };
54
55        let (max_precision, max) = match self.get(Stat::Max) {
56            Precision::Exact(max) => (
57                fba::Precision::Exact,
58                Some(fbb.create_vector(&ScalarValue::to_proto_bytes::<Vec<u8>>(Some(&max)))),
59            ),
60            Precision::Inexact(max) => (
61                fba::Precision::Inexact,
62                Some(fbb.create_vector(&ScalarValue::to_proto_bytes::<Vec<u8>>(Some(&max)))),
63            ),
64            Precision::Absent => (fba::Precision::Inexact, None),
65        };
66
67        let sum = self.get(Stat::Sum).as_exact().map(|sum| {
68            let sum = canonicalize_nan_sum(sum);
69            fbb.create_vector(&ScalarValue::to_proto_bytes::<Vec<u8>>(Some(&sum)))
70        });
71
72        let stat_args = &fba::ArrayStatsArgs {
73            min,
74            min_precision,
75            max,
76            max_precision,
77            sum,
78            is_sorted: self
79                .get_as::<bool>(Stat::IsSorted, &DType::Bool(Nullability::NonNullable))
80                .as_exact(),
81            is_strict_sorted: self
82                .get_as::<bool>(Stat::IsStrictSorted, &DType::Bool(Nullability::NonNullable))
83                .as_exact(),
84            is_constant: self
85                .get_as::<bool>(Stat::IsConstant, &DType::Bool(Nullability::NonNullable))
86                .as_exact(),
87            null_count: self
88                .get_as::<u64>(Stat::NullCount, &PType::U64.into())
89                .as_exact(),
90            uncompressed_size_in_bytes: self
91                .get_as::<u64>(Stat::UncompressedSizeInBytes, &PType::U64.into())
92                .as_exact(),
93            nan_count: self
94                .get_as::<u64>(Stat::NaNCount, &PType::U64.into())
95                .as_exact(),
96        };
97
98        Ok(fba::ArrayStats::create(fbb, stat_args))
99    }
100}
101
102/// Replace a NaN float sum with the canonical quiet NaN of its width.
103///
104/// A float sum is NaN when the data holds both `+inf` and `-inf`, because `inf + -inf` is an IEEE 754
105/// invalid operation. IEEE 754 does not specify the bits of that result, and targets disagree:
106/// x86_64 sets the sign bit and aarch64 does not. Without this step, the bytes of a written file
107/// depend on the architecture that wrote it. A NaN sum has no payload information, so nothing is
108/// lost. Stats that are already written are read back unchanged.
109fn canonicalize_nan_sum(sum: ScalarValue) -> ScalarValue {
110    match sum {
111        ScalarValue::Primitive(PValue::F16(v)) if v.is_nan() => {
112            ScalarValue::Primitive(PValue::F16(f16::from_bits(0x7e00)))
113        }
114        ScalarValue::Primitive(PValue::F32(v)) if v.is_nan() => {
115            ScalarValue::Primitive(PValue::F32(f32::from_bits(0x7fc0_0000)))
116        }
117        ScalarValue::Primitive(PValue::F64(v)) if v.is_nan() => {
118            ScalarValue::Primitive(PValue::F64(f64::from_bits(0x7ff8_0000_0000_0000)))
119        }
120        sum => sum,
121    }
122}
123
124impl StatsSet {
125    /// Creates a [`StatsSet`] from a flatbuffers array [`fba::ArrayStats<'a>`].
126    pub fn from_flatbuffer<'a>(
127        fb: &fba::ArrayStats<'a>,
128        array_dtype: &DType,
129        session: &VortexSession,
130    ) -> VortexResult<Self> {
131        let mut stats_set = StatsSet::default();
132
133        for stat in Stat::all() {
134            match stat {
135                Stat::IsConstant => {
136                    if let Some(is_constant) = fb.is_constant() {
137                        stats_set.set(Stat::IsConstant, Precision::Exact(is_constant.into()));
138                    }
139                }
140                Stat::IsSorted => {
141                    if let Some(is_sorted) = fb.is_sorted() {
142                        stats_set.set(Stat::IsSorted, Precision::Exact(is_sorted.into()));
143                    }
144                }
145                Stat::IsStrictSorted => {
146                    if let Some(is_strict_sorted) = fb.is_strict_sorted() {
147                        stats_set.set(
148                            Stat::IsStrictSorted,
149                            Precision::Exact(is_strict_sorted.into()),
150                        );
151                    }
152                }
153                Stat::Max => {
154                    if let Some(max) = fb.max()
155                        && let Some(stat_dtype) = stat.dtype(array_dtype)
156                    {
157                        let value =
158                            ScalarValue::from_proto_bytes(max.bytes(), &stat_dtype, session)?;
159                        let Some(value) = value else {
160                            continue;
161                        };
162
163                        stats_set.set(
164                            Stat::Max,
165                            match fb.max_precision() {
166                                fba::Precision::Exact => Precision::Exact(value),
167                                fba::Precision::Inexact => Precision::Inexact(value),
168                                other => vortex_bail!("Corrupted max_precision field: {other:?}"),
169                            },
170                        );
171                    }
172                }
173                Stat::Min => {
174                    if let Some(min) = fb.min()
175                        && let Some(stat_dtype) = stat.dtype(array_dtype)
176                    {
177                        let value =
178                            ScalarValue::from_proto_bytes(min.bytes(), &stat_dtype, session)?;
179                        let Some(value) = value else {
180                            continue;
181                        };
182
183                        stats_set.set(
184                            Stat::Min,
185                            match fb.min_precision() {
186                                fba::Precision::Exact => Precision::Exact(value),
187                                fba::Precision::Inexact => Precision::Inexact(value),
188                                other => vortex_bail!("Corrupted min_precision field: {other:?}"),
189                            },
190                        );
191                    }
192                }
193                Stat::NullCount => {
194                    if let Some(null_count) = fb.null_count() {
195                        stats_set.set(Stat::NullCount, Precision::Exact(null_count.into()));
196                    }
197                }
198                Stat::UncompressedSizeInBytes => {
199                    if let Some(uncompressed_size_in_bytes) = fb.uncompressed_size_in_bytes() {
200                        stats_set.set(
201                            Stat::UncompressedSizeInBytes,
202                            Precision::Exact(uncompressed_size_in_bytes.into()),
203                        );
204                    }
205                }
206                Stat::Sum => {
207                    if let Some(sum) = fb.sum()
208                        && let Some(stat_dtype) = stat.dtype(array_dtype)
209                    {
210                        let value =
211                            ScalarValue::from_proto_bytes(sum.bytes(), &stat_dtype, session)?;
212                        let Some(value) = value else {
213                            continue;
214                        };
215
216                        stats_set.set(Stat::Sum, Precision::Exact(value));
217                    }
218                }
219                Stat::NaNCount => {
220                    if let Some(nan_count) = fb.nan_count() {
221                        stats_set.set(
222                            Stat::NaNCount,
223                            Precision::Exact(ScalarValue::from(nan_count)),
224                        );
225                    }
226                }
227            }
228        }
229
230        Ok(stats_set)
231    }
232}
233
234#[cfg(test)]
235mod tests {
236    use flatbuffers::FlatBufferBuilder;
237    use flatbuffers::root;
238    use vortex_error::VortexResult;
239    use vortex_error::vortex_err;
240
241    use crate::array_session;
242    use crate::dtype::DType;
243    use crate::dtype::Nullability;
244    use crate::dtype::PType;
245    use crate::expr::stats::Precision;
246    use crate::expr::stats::Stat;
247    use crate::flatbuffers::WriteFlatBuffer;
248    use crate::flatbuffers::array as fba;
249    use crate::scalar::PValue;
250    use crate::scalar::ScalarValue;
251    use crate::stats::StatsSet;
252
253    fn roundtrip_sum(sum: f64, expected_bits: u64) -> VortexResult<u64> {
254        let stats = StatsSet::of(
255            Stat::Sum,
256            Precision::exact(ScalarValue::Primitive(PValue::F64(sum))),
257        );
258
259        let mut fbb = FlatBufferBuilder::new();
260        let offset = stats.write_flatbuffer(&mut fbb)?;
261        fbb.finish_minimal(offset);
262        let fb = root::<fba::ArrayStats>(fbb.finished_data())?;
263        // ScalarValue.f64_value is protobuf field 6, wire type fixed64.
264        let mut expected_bytes = vec![0x31];
265        expected_bytes.extend_from_slice(&expected_bits.to_le_bytes());
266        assert_eq!(fb.sum().expect("serialized sum").bytes(), expected_bytes);
267
268        let dtype = DType::Primitive(PType::F64, Nullability::NonNullable);
269        let read = StatsSet::from_flatbuffer(&fb, &dtype, &array_session())?;
270        Ok(read
271            .get_as::<f64>(
272                Stat::Sum,
273                &DType::Primitive(PType::F64, Nullability::Nullable),
274            )
275            .as_exact()
276            .ok_or_else(|| vortex_err!("sum must be exact after a round trip"))?
277            .to_bits())
278    }
279
280    #[test]
281    fn nan_sum_is_written_canonical() -> VortexResult<()> {
282        // x86_64 gives this NaN for `inf + -inf`.
283        assert_eq!(
284            roundtrip_sum(f64::from_bits(0xfff8_0000_0000_0000), 0x7ff8_0000_0000_0000)?,
285            0x7ff8_0000_0000_0000
286        );
287        assert_eq!(
288            roundtrip_sum(f64::from_bits(0x7ff8_0000_dead_beef), 0x7ff8_0000_0000_0000)?,
289            0x7ff8_0000_0000_0000
290        );
291
292        assert_eq!(roundtrip_sum(3.0, 3.0f64.to_bits())?, 3.0f64.to_bits());
293        assert_eq!(
294            roundtrip_sum(f64::NEG_INFINITY, f64::NEG_INFINITY.to_bits())?,
295            f64::NEG_INFINITY.to_bits()
296        );
297
298        Ok(())
299    }
300}
vortex_array/stats/flatbuffers.rs

vortex_array/stats/
flatbuffers.rs
