Add more non copying stats operations; replace stats replace with inherit_from (#4279)

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

Author: robert3005

vortex-array/src/array/mod.rs

@@ -22,7 +22,7 @@ use crate::arrays::{
 use crate::builders::ArrayBuilder;
 use crate::compute::{ComputeFn, Cost, InvocationArgs, IsConstantOpts, Output, is_constant_opts};
 use crate::serde::ArrayChildren;
-use crate::stats::{Precision, Stat, StatsSetRef};
+use crate::stats::{Precision, Stat, StatsProviderExt, StatsSetRef};
 use crate::vtable::{
     ArrayVTable, CanonicalVTable, ComputeVTable, OperationsVTable, SerdeVTable, VTable,
     ValidityVTable, VisitorVTable,
@@ -516,7 +516,7 @@ impl<V: VTable> Array for ArrayAdapter<V> {
         canonical
             .as_ref()
             .statistics()
-            .replace(self.statistics().to_owned());
+            .inherit_from(self.statistics());
         Ok(canonical)
     }
 

vortex-array/src/compute/is_constant.rs

@@ -12,7 +12,7 @@ use vortex_scalar::Scalar;
 use crate::Array;
 use crate::arrays::{ConstantVTable, NullVTable};
 use crate::compute::{ComputeFn, ComputeFnVTable, InvocationArgs, Kernel, Options, Output};
-use crate::stats::{Precision, Stat, StatsProvider};
+use crate::stats::{Precision, Stat, StatsProvider, StatsProviderExt};
 use crate::vtable::VTable;
 
 static IS_CONSTANT_FN: LazyLock<ComputeFn> = LazyLock::new(|| {

vortex-array/src/compute/is_sorted.rs

@@ -12,7 +12,7 @@ use vortex_scalar::Scalar;
 use crate::Array;
 use crate::arrays::{ConstantVTable, NullVTable};
 use crate::compute::{ComputeFn, ComputeFnVTable, InvocationArgs, Kernel, Options, Output};
-use crate::stats::{Precision, Stat};
+use crate::stats::{Precision, Stat, StatsProviderExt};
 use crate::vtable::VTable;
 
 static IS_SORTED_FN: LazyLock<ComputeFn> = LazyLock::new(|| {

vortex-array/src/compute/nan_count.rs

@@ -10,7 +10,7 @@ use vortex_scalar::{Scalar, ScalarValue};
 
 use crate::Array;
 use crate::compute::{ComputeFn, ComputeFnVTable, InvocationArgs, Kernel, Output, UnaryArgs};
-use crate::stats::{Precision, Stat};
+use crate::stats::{Precision, Stat, StatsProviderExt};
 use crate::vtable::VTable;
 
 static NAN_COUNT_FN: LazyLock<ComputeFn> = LazyLock::new(|| {

vortex-array/src/compute/take.rs

@@ -10,7 +10,7 @@ use vortex_scalar::Scalar;
 
 use crate::arrays::{ConstantArray, ConstantVTable};
 use crate::compute::{ComputeFn, ComputeFnVTable, InvocationArgs, Kernel, Output};
-use crate::stats::{Precision, Stat, StatsSet};
+use crate::stats::{Precision, Stat, StatsProviderExt, StatsSet};
 use crate::vtable::VTable;
 use crate::{Array, ArrayRef, Canonical, IntoArray};
 

vortex-array/src/serde.rs

@@ -208,7 +208,7 @@ impl WriteFlatBuffer for ArrayNodeFlatBuffer<'_> {
         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));
+        let stats = Some(self.array.statistics().write_flatbuffer(fbb));
 
         fba::ArrayNode::create(
             fbb,

vortex-array/src/stats/array.rs

@@ -9,7 +9,7 @@ use parking_lot::RwLock;
 use vortex_error::{VortexError, VortexResult, vortex_panic};
 use vortex_scalar::{Scalar, ScalarValue};
 
-use super::{Precision, Stat, StatType, StatsProvider, StatsSet, StatsSetIntoIter};
+use super::{Precision, Stat, StatsProvider, StatsSet, StatsSetIntoIter, TypedStatsSetRef};
 use crate::Array;
 use crate::compute::{
     MinMaxResult, is_constant, is_sorted, is_strict_sorted, min_max, nan_count, sum,
@@ -75,19 +75,32 @@ impl StatsSetRef<'_> {
     }
 
     pub fn inherit_from(&self, stats: StatsSetRef<'_>) {
-        stats.with_iter(|iter| self.inherit(iter));
+        // Only inherit if the underlying stats are different
+        if !Arc::ptr_eq(&self.array_stats.inner, &stats.array_stats.inner) {
+            stats.with_iter(|iter| self.inherit(iter));
+        }
     }
 
     pub fn inherit<'a>(&self, iter: impl Iterator<Item = &'a (Stat, Precision<ScalarValue>)>) {
-        // TODO(ngates): depending on statistic, this should choose the more precise one
         let mut guard = self.array_stats.inner.write();
         for (stat, value) in iter {
-            guard.set(*stat, value.clone());
+            if !value.is_exact() {
+                if !guard.get(*stat).is_some_and(|v| v.is_exact()) {
+                    guard.set(*stat, value.clone());
+                }
+            } else {
+                guard.set(*stat, value.clone());
+            }
         }
     }
 
-    pub fn replace(&self, stats: StatsSet) {
-        *self.array_stats.inner.write() = stats;
+    pub fn with_typed_stats_set<U, F: FnOnce(TypedStatsSetRef) -> U>(&self, apply: F) -> U {
+        apply(
+            self.array_stats
+                .inner
+                .read()
+                .as_typed_ref(self.dyn_array_ref.dtype()),
+        )
     }
 
     pub fn to_owned(&self) -> StatsSet {
@@ -165,32 +178,6 @@ impl StatsSetRef<'_> {
 }
 
 impl StatsSetRef<'_> {
-    pub fn get_as<U: for<'a> TryFrom<&'a Scalar, Error = VortexError>>(
-        &self,
-        stat: Stat,
-    ) -> Option<Precision<U>> {
-        self.get(stat).map(|v| {
-            v.map(|v| {
-                U::try_from(&v).unwrap_or_else(|err| {
-                    vortex_panic!(
-                        err,
-                        "Failed to get stat {} as {}",
-                        stat,
-                        std::any::type_name::<U>()
-                    )
-                })
-            })
-        })
-    }
-
-    pub fn get_as_bound<S, U>(&self) -> Option<S::Bound>
-    where
-        S: StatType<U>,
-        U: for<'a> TryFrom<&'a Scalar, Error = VortexError>,
-    {
-        self.get_as::<U>(S::STAT).map(|v| v.bound::<S>())
-    }
-
     pub fn compute_as<U: for<'a> TryFrom<&'a Scalar, Error = VortexError>>(
         &self,
         stat: Stat,

vortex-array/src/stats/flatbuffers.rs

@@ -7,7 +7,19 @@ use vortex_error::{VortexError, vortex_bail};
 use vortex_flatbuffers::{ReadFlatBuffer, WriteFlatBuffer, array as fba};
 use vortex_scalar::ScalarValue;
 
-use crate::stats::{Precision, Stat, StatsSet};
+use crate::stats::{Precision, Stat, StatsSet, StatsSetRef};
+
+impl WriteFlatBuffer for StatsSetRef<'_> {
+    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>> {
+        self.with_typed_stats_set(|stats_set| stats_set.values.write_flatbuffer(fbb))
+    }
+}
 
 impl WriteFlatBuffer for StatsSet {
     type Target<'t> = fba::ArrayStats<'t>;

vortex-file/src/pruning.rs

@@ -5,12 +5,11 @@ use std::sync::Arc;
 
 use vortex_array::ArrayRef;
 use vortex_array::arrays::{ConstantArray, StructArray};
-use vortex_array::stats::{Stat, StatsSet};
+use vortex_array::stats::{Stat, StatsProvider, StatsSet};
 use vortex_array::validity::Validity;
 use vortex_dtype::{Field, FieldName, FieldNames, FieldPath, StructFields};
 use vortex_error::{VortexExpect, VortexResult, vortex_bail, vortex_err};
 use vortex_expr::pruning::field_path_stat_field_name;
-use vortex_scalar::Scalar;
 use vortex_utils::aliases::hash_map::HashMap;
 use vortex_utils::aliases::hash_set::HashSet;
 
@@ -43,19 +42,17 @@ pub fn extract_relevant_file_stats_as_struct_row(
         let Some(stat_set) = stats_sets.get(field_idx) else {
             vortex_bail!("missing stat field {} from stats set", field)
         };
+        let typed_stats = stat_set.as_typed_ref(&field_dtype);
 
         for stat in stats {
-            let Some(stat_value) = stat_set.get(*stat).and_then(|p| p.as_exact()) else {
-                vortex_bail!("missing stat {}, {} from stats set", field, stat)
-            };
             if stat == &Stat::Max || stat == &Stat::Min || stat == &Stat::NaNCount {
-                if let Some(stat_dtype) = stat.dtype(&field_dtype) {
-                    columns.push((
-                        field_path_stat_field_name(field_path, *stat),
-                        ConstantArray::new(Scalar::new(stat_dtype, stat_value.clone()), 1)
-                            .to_array(),
-                    ));
-                }
+                let Some(stat_value) = typed_stats.get(*stat).and_then(|p| p.as_exact()) else {
+                    vortex_bail!("missing stat {}, {} from stats set", field, stat)
+                };
+                columns.push((
+                    field_path_stat_field_name(field_path, *stat),
+                    ConstantArray::new(stat_value, 1).to_array(),
+                ));
             } else {
                 todo!("unsupported file prune stat {stat}")
             }

vortex-layout/src/layouts/flat/writer.rs

@@ -140,7 +140,7 @@ mod tests {
     use futures::stream;
     use vortex_array::arrays::{BoolArray, PrimitiveArray, StructArray};
     use vortex_array::builders::{ArrayBuilder, VarBinViewBuilder};
-    use vortex_array::stats::{Precision, Stat};
+    use vortex_array::stats::{Precision, Stat, StatsProviderExt};
     use vortex_array::validity::Validity;
     use vortex_array::{Array, ArrayContext, ArrayRef, IntoArray, ToCanonical};
     use vortex_buffer::buffer;