fix: support stats for ExtensionArray in vortex-file (#1547)

By using ArrayBuilder to accumulate stats

---------

Co-authored-by: Will Manning <will@willmanning.io>

Author: gatesn

encodings/datetime-parts/src/array.rs

@@ -4,7 +4,7 @@ use serde::{Deserialize, Serialize};
 use vortex_array::array::StructArray;
 use vortex_array::compute::try_cast;
 use vortex_array::encoding::ids;
-use vortex_array::stats::{Stat, StatisticsVTable, StatsSet};
+use vortex_array::stats::StatsSet;
 use vortex_array::validity::{ArrayValidity, LogicalValidity, Validity, ValidityVTable};
 use vortex_array::variants::{ExtensionArrayTrait, VariantsVTable};
 use vortex_array::visitor::{ArrayVisitor, VisitorVTable};
@@ -14,7 +14,6 @@ use vortex_array::{
 };
 use vortex_dtype::{DType, PType};
 use vortex_error::{vortex_bail, VortexExpect as _, VortexResult, VortexUnwrap};
-use vortex_scalar::Scalar;
 
 use crate::compute::decode_to_temporal;
 
@@ -166,18 +165,3 @@ impl VisitorVTable<DateTimePartsArray> for DateTimePartsEncoding {
         visitor.visit_child("subsecond", &array.subsecond())
     }
 }
-
-impl StatisticsVTable<DateTimePartsArray> for DateTimePartsEncoding {
-    fn compute_statistics(&self, array: &DateTimePartsArray, stat: Stat) -> VortexResult<StatsSet> {
-        let maybe_stat = match stat {
-            Stat::NullCount => Some(Scalar::from(array.validity().null_count(array.len())?)),
-            _ => None,
-        };
-
-        let mut stats = StatsSet::default();
-        if let Some(value) = maybe_stat {
-            stats.set(stat, value);
-        }
-        Ok(stats)
-    }
-}

encodings/datetime-parts/src/lib.rs

@@ -4,3 +4,4 @@ pub use compress::*;
 mod array;
 mod compress;
 mod compute;
+mod stats;

encodings/datetime-parts/src/stats.rs

@@ -0,0 +1,21 @@
+use vortex_array::stats::{Stat, StatisticsVTable, StatsSet};
+use vortex_array::ArrayLen;
+use vortex_error::VortexResult;
+use vortex_scalar::Scalar;
+
+use crate::{DateTimePartsArray, DateTimePartsEncoding};
+
+impl StatisticsVTable<DateTimePartsArray> for DateTimePartsEncoding {
+    fn compute_statistics(&self, array: &DateTimePartsArray, stat: Stat) -> VortexResult<StatsSet> {
+        let maybe_stat = match stat {
+            Stat::NullCount => Some(Scalar::from(array.validity().null_count(array.len())?)),
+            _ => None,
+        };
+
+        let mut stats = StatsSet::default();
+        if let Some(value) = maybe_stat {
+            stats.set(stat, value);
+        }
+        Ok(stats)
+    }
+}

vortex-array/src/array/varbin/stats.rs

@@ -37,7 +37,7 @@ pub fn compute_varbin_statistics<T: ArrayTrait + ArrayAccessor<[u8]>>(
 
     Ok(match stat {
         Stat::NullCount => {
-            let null_count = array.logical_validity().null_count(array.len())?;
+            let null_count = array.logical_validity().null_count()?;
             if null_count == array.len() {
                 return Ok(StatsSet::nulls(array.len(), array.dtype()));
             }

vortex-array/src/builders/binary.rs

@@ -2,7 +2,7 @@ use std::any::Any;
 
 use arrow_array::builder::{ArrayBuilder as _, BinaryViewBuilder};
 use arrow_array::Array;
-use vortex_dtype::Nullability;
+use vortex_dtype::{DType, Nullability};
 use vortex_error::{vortex_bail, VortexResult};
 
 use crate::arrow::FromArrowArray;
@@ -12,13 +12,15 @@ use crate::ArrayData;
 pub struct BinaryBuilder {
     inner: BinaryViewBuilder,
     nullability: Nullability,
+    dtype: DType,
 }
 
 impl BinaryBuilder {
     pub fn with_capacity(nullability: Nullability, capacity: usize) -> Self {
         Self {
             inner: BinaryViewBuilder::with_capacity(capacity),
             nullability,
+            dtype: DType::Binary(nullability),
         }
     }
 
@@ -40,6 +42,10 @@ impl ArrayBuilder for BinaryBuilder {
         self
     }
 
+    fn dtype(&self) -> &DType {
+        &self.dtype
+    }
+
     fn len(&self) -> usize {
         self.inner.len()
     }
@@ -59,7 +65,7 @@ impl ArrayBuilder for BinaryBuilder {
     fn finish(&mut self) -> VortexResult<ArrayData> {
         let arrow = self.inner.finish();
 
-        if self.nullability == Nullability::NonNullable && arrow.null_count() > 0 {
+        if !self.dtype().is_nullable() && arrow.null_count() > 0 {
             vortex_bail!("Non-nullable builder has null values");
         }
 

vortex-array/src/builders/bool.rs

@@ -2,7 +2,7 @@ use std::any::Any;
 
 use arrow_array::builder::{ArrayBuilder as _, BooleanBuilder as ArrowBooleanBuilder};
 use arrow_array::Array;
-use vortex_dtype::Nullability;
+use vortex_dtype::{DType, Nullability};
 use vortex_error::{vortex_bail, VortexResult};
 
 use crate::arrow::FromArrowArray;
@@ -12,7 +12,7 @@ use crate::ArrayData;
 pub struct BoolBuilder {
     inner: ArrowBooleanBuilder,
     nullability: Nullability,
-    // TODO(ngates): track stats?
+    dtype: DType,
 }
 
 impl BoolBuilder {
@@ -24,6 +24,7 @@ impl BoolBuilder {
         Self {
             inner: ArrowBooleanBuilder::with_capacity(capacity),
             nullability,
+            dtype: DType::Bool(nullability),
         }
     }
 
@@ -52,6 +53,10 @@ impl ArrayBuilder for BoolBuilder {
         self
     }
 
+    fn dtype(&self) -> &DType {
+        &self.dtype
+    }
+
     fn len(&self) -> usize {
         self.inner.len()
     }
@@ -67,7 +72,7 @@ impl ArrayBuilder for BoolBuilder {
     fn finish(&mut self) -> VortexResult<ArrayData> {
         let arrow = self.inner.finish();
 
-        if self.nullability == Nullability::NonNullable && arrow.null_count() > 0 {
+        if !self.dtype().is_nullable() && arrow.null_count() > 0 {
             vortex_bail!("Non-nullable builder has null values");
         }
 

vortex-array/src/builders/extension.rs

@@ -3,10 +3,10 @@ use std::sync::Arc;
 
 use vortex_dtype::{DType, ExtDType};
 use vortex_error::VortexResult;
-use vortex_scalar::{ExtScalar, Scalar};
+use vortex_scalar::ExtScalar;
 
 use crate::array::ExtensionArray;
-use crate::builders::{builder_with_capacity, ArrayBuilder};
+use crate::builders::{builder_with_capacity, ArrayBuilder, ArrayBuilderExt};
 use crate::{ArrayData, IntoArrayData};
 
 pub struct ExtensionBuilder {
@@ -27,8 +27,7 @@ impl ExtensionBuilder {
     }
 
     pub fn append_value(&mut self, value: ExtScalar) -> VortexResult<()> {
-        self.storage
-            .append_scalar(&Scalar::extension(self.ext_dtype(), value.storage()))
+        self.storage.append_scalar(&value.storage())
     }
 
     pub fn append_option(&mut self, value: Option<ExtScalar>) -> VortexResult<()> {
@@ -59,6 +58,10 @@ impl ArrayBuilder for ExtensionBuilder {
         self
     }
 
+    fn dtype(&self) -> &DType {
+        &self.dtype
+    }
+
     fn len(&self) -> usize {
         self.storage.len()
     }

vortex-array/src/builders/mod.rs

@@ -15,19 +15,21 @@ pub use null::*;
 pub use primitive::*;
 pub use utf8::*;
 use vortex_dtype::{match_each_native_ptype, DType};
-use vortex_error::{vortex_err, VortexResult};
+use vortex_error::{vortex_bail, vortex_err, VortexResult};
 use vortex_scalar::{
     BinaryScalar, BoolScalar, ExtScalar, PrimitiveScalar, Scalar, StructScalar, Utf8Scalar,
 };
 
 use crate::builders::struct_::StructBuilder;
 use crate::ArrayData;
 
-pub trait ArrayBuilder {
+pub trait ArrayBuilder: Send {
     fn as_any(&self) -> &dyn Any;
 
     fn as_any_mut(&mut self) -> &mut dyn Any;
 
+    fn dtype(&self) -> &DType;
+
     fn len(&self) -> usize;
 
     fn is_empty(&self) -> bool {
@@ -78,9 +80,16 @@ pub fn builder_with_capacity(dtype: &DType, capacity: usize) -> Box<dyn ArrayBui
     }
 }
 
-impl dyn ArrayBuilder + '_ {
+pub trait ArrayBuilderExt: ArrayBuilder {
     /// A generic function to append a scalar to the builder.
-    pub fn append_scalar(&mut self, scalar: &Scalar) -> VortexResult<()> {
+    fn append_scalar(&mut self, scalar: &Scalar) -> VortexResult<()> {
+        if !scalar.dtype().eq_ignore_nullability(self.dtype()) {
+            vortex_bail!(
+                "Builder has dtype {:?}, scalar has {:?}",
+                self.dtype(),
+                scalar.dtype()
+            )
+        }
         match scalar.dtype() {
             DType::Null => self
                 .as_any_mut()
@@ -130,3 +139,5 @@ impl dyn ArrayBuilder + '_ {
         Ok(())
     }
 }
+
+impl<T: ?Sized + ArrayBuilder> ArrayBuilderExt for T {}

vortex-array/src/builders/null.rs

@@ -1,5 +1,6 @@
 use std::any::Any;
 
+use vortex_dtype::DType;
 use vortex_error::VortexResult;
 
 use crate::array::NullArray;
@@ -31,6 +32,10 @@ impl ArrayBuilder for NullBuilder {
         self
     }
 
+    fn dtype(&self) -> &DType {
+        &DType::Null
+    }
+
     fn len(&self) -> usize {
         self.length
     }

vortex-array/src/builders/primitive.rs

@@ -14,7 +14,6 @@ use crate::ArrayData;
 pub struct PrimitiveBuilder<T: NativePType> {
     inner: ArrowPrimitiveBuilder<T::ArrowPrimitiveType>,
     dtype: DType,
-    // TODO(ngates): track stats?
 }
 
 impl<T: NativePType> PrimitiveBuilder<T>
@@ -64,6 +63,10 @@ where
         self
     }
 
+    fn dtype(&self) -> &DType {
+        &self.dtype
+    }
+
     fn len(&self) -> usize {
         self.inner.len()
     }
@@ -82,8 +85,8 @@ where
     fn finish(&mut self) -> VortexResult<ArrayData> {
         let arrow = self.inner.finish();
 
-        if self.dtype.is_nullable() && arrow.null_count() > 0 {
-            vortex_bail!("Non-nullable builder has null values");
+        if !self.dtype().is_nullable() && arrow.null_count() > 0 {
+            vortex_bail!("Non-nullable builder {} has null values", self.dtype());
         }
 
         Ok(ArrayData::from_arrow(&arrow, self.dtype.is_nullable()))