Truncate stats on write (#3272)

When truncating min we take the prefix and for max we increment the
value until it doesn't overflow. Needs some tests

Signed-off-by: Robert Kruszewski <[email protected]>

Author: robert3005

Cargo.lock

@@ -139,7 +139,9 @@ struct ArrayStreamSink {
 
 	~ArrayStreamSink() {
 		// "should dctor a sink, before closing it
-		D_ASSERT(sink == nullptr);
+		// If you throw during writes then the stack will be unwound and the destructor is going to be called before the
+		// close method is invoked thus triggering following assertion failure and will clobber the exception printing
+		// D_ASSERT(sink == nullptr);
 	}
 
 	vx_array_sink *sink;

duckdb-vortex/src/include/vortex_common.hpp

@@ -107,7 +107,7 @@ impl<T: NativeDecimalType> DecimalBuilder<T> {
     }
 }
 
-impl<T: NativeDecimalType + Default + Send + 'static> ArrayBuilder for DecimalBuilder<T> {
+impl<T: NativeDecimalType> ArrayBuilder for DecimalBuilder<T> {
     fn as_any(&self) -> &dyn Any {
         self
     }

vortex-array/src/builders/decimal.rs

@@ -43,6 +43,7 @@ pub use extension::*;
 pub use list::*;
 pub use null::*;
 pub use primitive::*;
+pub use struct_::*;
 pub use varbinview::*;
 use vortex_dtype::{DType, match_each_native_ptype};
 use vortex_error::{VortexResult, vortex_bail, vortex_err};
@@ -53,7 +54,6 @@ use vortex_scalar::{
 };
 
 use crate::arrays::precision_to_storage_size;
-use crate::builders::struct_::StructBuilder;
 use crate::{Array, ArrayRef};
 
 pub trait ArrayBuilder: Send {

vortex-array/src/builders/mod.rs

@@ -3,7 +3,6 @@ use std::iter;
 use std::sync::Arc;
 
 use flatbuffers::{FlatBufferBuilder, Follow, WIPOffset, root};
-use itertools::Itertools;
 use vortex_buffer::{Alignment, ByteBuffer};
 use vortex_dtype::{DType, TryFromBytes};
 use vortex_error::{
@@ -44,12 +43,10 @@ impl dyn Array + '_ {
         options: &SerializeOptions,
     ) -> VortexResult<Vec<ByteBuffer>> {
         // Collect all array buffers
-        let mut array_buffers = vec![];
-        for a in self.depth_first_traversal() {
-            for buffer in a.buffers() {
-                array_buffers.push(buffer);
-            }
-        }
+        let array_buffers = self
+            .depth_first_traversal()
+            .flat_map(|f| f.buffers())
+            .collect::<Vec<_>>();
 
         // Allocate result buffers, including a possible padding buffer for each.
         let mut buffers = vec![];
@@ -183,28 +180,28 @@ impl WriteFlatBuffer for ArrayNodeFlatBuffer<'_> {
         // Assign buffer indices for all child arrays.
         let nbuffers = u16::try_from(self.array.nbuffers())
             .vortex_expect("Array can have at most u16::MAX buffers");
-        let child_buffer_idx = self.buffer_idx + nbuffers;
+        let mut child_buffer_idx = self.buffer_idx + nbuffers;
 
-        let children = self
+        let children = &self
             .array
             .children()
             .iter()
-            .scan(child_buffer_idx, |buffer_idx, child| {
+            .map(|child| {
                 // Update the number of buffers required.
                 let msg = ArrayNodeFlatBuffer {
                     ctx: self.ctx,
                     array: child,
-                    buffer_idx: *buffer_idx,
+                    buffer_idx: child_buffer_idx,
                 }
                 .write_flatbuffer(fbb);
-                *buffer_idx = u16::try_from(child.nbuffers_recursive())
+                child_buffer_idx = u16::try_from(child.nbuffers_recursive())
                     .ok()
-                    .and_then(|nbuffers| nbuffers.checked_add(*buffer_idx))
+                    .and_then(|nbuffers| nbuffers.checked_add(child_buffer_idx))
                     .vortex_expect("Too many buffers (u16) for Array");
-                Some(msg)
+                msg
             })
-            .collect_vec();
-        let children = Some(fbb.create_vector(&children));
+            .collect::<Vec<_>>();
+        let children = Some(fbb.create_vector(children));
 
         let buffers = Some(fbb.create_vector_from_iter((0..nbuffers).map(|i| i + self.buffer_idx)));
         let stats = Some(self.array.statistics().to_owned().write_flatbuffer(fbb));

vortex-array/src/serde.rs

@@ -156,9 +156,9 @@ impl StatsSetRef<'_> {
 
     pub fn compute_all(&self, stats: &[Stat]) -> VortexResult<StatsSet> {
         let mut stats_set = StatsSet::default();
-        for stat in stats {
-            if let Some(s) = self.compute_stat(*stat)? {
-                stats_set.set(*stat, Precision::exact(s))
+        for &stat in stats {
+            if let Some(s) = self.compute_stat(stat)? {
+                stats_set.set(stat, Precision::exact(s))
             }
         }
         Ok(stats_set)

vortex-array/src/stats/array.rs

@@ -1,38 +1,57 @@
-use enum_iterator::all;
 use flatbuffers::{FlatBufferBuilder, Follow, WIPOffset};
-use vortex_error::VortexError;
-use vortex_flatbuffers::{ReadFlatBuffer, WriteFlatBuffer};
+use vortex_error::{VortexError, vortex_bail};
+use vortex_flatbuffers::{ReadFlatBuffer, WriteFlatBuffer, array as fba};
 use vortex_scalar::ScalarValue;
 
 use super::traits::{StatsProvider, StatsProviderExt};
 use crate::stats::{Precision, Stat, StatsSet};
 
 impl WriteFlatBuffer for StatsSet {
-    type Target<'t> = crate::flatbuffers::ArrayStats<'t>;
+    type Target<'t> = fba::ArrayStats<'t>;
 
     /// All statistics written must be exact
     fn write_flatbuffer<'fb>(
         &self,
         fbb: &mut FlatBufferBuilder<'fb>,
     ) -> WIPOffset<Self::Target<'fb>> {
-        let min = self
+        let (min_precision, min) = self
             .get(Stat::Min)
-            .and_then(Precision::as_exact)
-            .map(|min| fbb.create_vector(&min.to_protobytes::<Vec<u8>>()));
+            .map(|sum| {
+                (
+                    if sum.is_exact() {
+                        fba::Precision::Exact
+                    } else {
+                        fba::Precision::Inexact
+                    },
+                    Some(fbb.create_vector(&sum.into_inner().to_protobytes::<Vec<u8>>())),
+                )
+            })
+            .unwrap_or_else(|| (fba::Precision::Inexact, None));
 
-        let max = self
+        let (max_precision, max) = self
             .get(Stat::Max)
-            .and_then(Precision::as_exact)
-            .map(|max| fbb.create_vector(&max.to_protobytes::<Vec<u8>>()));
+            .map(|sum| {
+                (
+                    if sum.is_exact() {
+                        fba::Precision::Exact
+                    } else {
+                        fba::Precision::Inexact
+                    },
+                    Some(fbb.create_vector(&sum.into_inner().to_protobytes::<Vec<u8>>())),
+                )
+            })
+            .unwrap_or_else(|| (fba::Precision::Inexact, None));
 
         let sum = self
             .get(Stat::Sum)
             .and_then(Precision::as_exact)
             .map(|sum| fbb.create_vector(&sum.to_protobytes::<Vec<u8>>()));
 
-        let stat_args = &crate::flatbuffers::ArrayStatsArgs {
+        let stat_args = &fba::ArrayStatsArgs {
             min,
+            min_precision,
             max,
+            max_precision,
             sum,
             is_sorted: self
                 .get_as::<bool>(Stat::IsSorted)
@@ -54,20 +73,20 @@ impl WriteFlatBuffer for StatsSet {
                 .and_then(Precision::as_exact),
         };
 
-        crate::flatbuffers::ArrayStats::create(fbb, stat_args)
+        fba::ArrayStats::create(fbb, stat_args)
     }
 }
 
 impl ReadFlatBuffer for StatsSet {
-    type Source<'a> = crate::flatbuffers::ArrayStats<'a>;
+    type Source<'a> = fba::ArrayStats<'a>;
     type Error = VortexError;
 
     fn read_flatbuffer<'buf>(
         fb: &<Self::Source<'buf> as Follow<'buf>>::Inner,
     ) -> Result<Self, Self::Error> {
         let mut stats_set = StatsSet::default();
 
-        for stat in all::<Stat>() {
+        for stat in Stat::all() {
             match stat {
                 Stat::IsConstant => {
                     if let Some(is_constant) = fb.is_constant() {
@@ -89,17 +108,27 @@ impl ReadFlatBuffer for StatsSet {
                 }
                 Stat::Max => {
                     if let Some(max) = fb.max() {
+                        let value = ScalarValue::from_protobytes(max.bytes())?;
                         stats_set.set(
                             Stat::Max,
-                            Precision::Exact(ScalarValue::from_protobytes(max.bytes())?),
+                            match fb.max_precision() {
+                                fba::Precision::Exact => Precision::Exact(value),
+                                fba::Precision::Inexact => Precision::Inexact(value),
+                                _ => vortex_bail!("Corrupted max_precision field"),
+                            },
                         );
                     }
                 }
                 Stat::Min => {
                     if let Some(min) = fb.min() {
+                        let value = ScalarValue::from_protobytes(min.bytes())?;
                         stats_set.set(
                             Stat::Min,
-                            Precision::Exact(ScalarValue::from_protobytes(min.bytes())?),
+                            match fb.min_precision() {
+                                fba::Precision::Exact => Precision::Exact(value),
+                                fba::Precision::Inexact => Precision::Inexact(value),
+                                _ => vortex_bail!("Corrupted min_precision field"),
+                            },
                         );
                     }
                 }

vortex-array/src/stats/flatbuffers.rs

@@ -5,7 +5,7 @@ use std::hash::Hash;
 
 use arrow_buffer::bit_iterator::BitIterator;
 use arrow_buffer::{BooleanBufferBuilder, MutableBuffer};
-use enum_iterator::{Sequence, last};
+use enum_iterator::{Sequence, all, last};
 use log::debug;
 use num_enum::{IntoPrimitive, TryFromPrimitive};
 pub use stats_set::*;
@@ -37,19 +37,6 @@ pub const PRUNING_STATS: &[Stat] = &[
     Stat::NaNCount,
 ];
 
-/// Stats to keep when serializing arrays to layouts
-pub const STATS_TO_WRITE: &[Stat] = &[
-    Stat::Min,
-    Stat::Max,
-    Stat::NullCount,
-    Stat::NaNCount,
-    Stat::Sum,
-    Stat::IsConstant,
-    Stat::IsSorted,
-    Stat::IsStrictSorted,
-    Stat::UncompressedSizeInBytes,
-];
-
 #[derive(
     Debug,
     Clone,
@@ -236,6 +223,10 @@ impl Stat {
             Self::NaNCount => "nan_count",
         }
     }
+
+    pub fn all() -> impl Iterator<Item = Stat> {
+        all::<Self>()
+    }
 }
 
 pub fn as_stat_bitset_bytes(stats: &[Stat]) -> Vec<u8> {

vortex-array/src/stats/mod.rs

@@ -25,8 +25,8 @@ where
 {
     fn clone(&self) -> Self {
         match self {
-            Self::Exact(e) => Self::Exact(e.clone()),
-            Self::Inexact(ie) => Self::Inexact(ie.clone()),
+            Exact(e) => Exact(e.clone()),
+            Inexact(ie) => Inexact(ie.clone()),
         }
     }
 }
@@ -129,7 +129,8 @@ impl<T> Precision<T> {
         Ok(precision)
     }
 
-    pub(crate) fn into_inner(self) -> T {
+    /// Unwrap the underlying value
+    pub fn into_inner(self) -> T {
         match self {
             Exact(val) | Inexact(val) => val,
         }

vortex-array/src/stats/precision.rs

@@ -27,6 +27,11 @@ impl BufferString {
     pub fn into_inner(self) -> ByteBuffer {
         self.0
     }
+
+    /// Returns reference to the inner [`ByteBuffer`].
+    pub fn inner(&self) -> &ByteBuffer {
+        &self.0
+    }
 }
 
 impl Debug for BufferString {