make distinct values calc lazy + refactor typed stats

Signed-off-by: Connor Tsui <connor.tsui20@gmail.com>

Author: connortsui20

vortex-btrblocks/src/canonical_compressor.rs

@@ -16,6 +16,7 @@ use vortex_array::arrays::ExtensionArray;
 use vortex_array::arrays::FixedSizeListArray;
 use vortex_array::arrays::ListArray;
 use vortex_array::arrays::ListViewArray;
+use vortex_array::arrays::Primitive;
 use vortex_array::arrays::StructArray;
 use vortex_array::arrays::TemporalArray;
 use vortex_array::arrays::listview::list_from_list_view;
@@ -31,11 +32,16 @@ use vortex_error::VortexResult;
 
 use crate::BtrBlocksCompressorBuilder;
 use crate::CompressorContext;
+use crate::CompressorStats;
+use crate::GenerateStatsOptions;
 use crate::Scheme;
 use crate::SchemeId;
 use crate::StatsCache;
 use crate::compressor::decimal::compress_decimal;
+use crate::compressor::float::FloatStats;
 use crate::compressor::integer::DictScheme as IntDictScheme;
+use crate::compressor::integer::IntegerStats;
+use crate::compressor::string::StringStats;
 use crate::compressor::temporal::compress_temporal;
 
 /// The main compressor type implementing BtrBlocks-inspired compression.
@@ -217,8 +223,44 @@ impl BtrBlocksCompressor {
         }
 
         let before_nbytes = array.nbytes();
+        let needs_distinct = eligible.iter().any(|s| s.needs_distinct_values());
         let mut cache = StatsCache::new();
 
+        // Pre-populate the stats cache with the right `count_distinct_values` setting.
+        // This matches the old `gen_stats` behavior where distinct values were only computed
+        // when Dict was in the scheme list.
+        if let Some(prim) = array.as_opt::<Primitive>() {
+            let prim = prim.to_primitive();
+            if prim.ptype().is_int() {
+                cache.get_or_insert_with::<IntegerStats>(|| {
+                    IntegerStats::generate_opts(
+                        &prim,
+                        GenerateStatsOptions {
+                            count_distinct_values: needs_distinct,
+                        },
+                    )
+                });
+            } else {
+                cache.get_or_insert_with::<FloatStats>(|| {
+                    FloatStats::generate_opts(
+                        &prim,
+                        GenerateStatsOptions {
+                            count_distinct_values: needs_distinct,
+                        },
+                    )
+                });
+            }
+        } else if array.as_opt::<vortex_array::arrays::VarBinView>().is_some() {
+            cache.get_or_insert_with::<StringStats>(|| {
+                StringStats::generate_opts(
+                    &array.to_varbinview(),
+                    GenerateStatsOptions {
+                        count_distinct_values: needs_distinct,
+                    },
+                )
+            });
+        }
+
         if let Some(winner) = self.choose_scheme(&eligible, &array, ctx, &mut cache, excludes)? {
             let compressed = winner.compress(self, &array, ctx, &mut cache, excludes)?;
             if compressed.nbytes() < before_nbytes {

vortex-btrblocks/src/compressor/float/dictionary.rs

@@ -12,13 +12,18 @@ use vortex_array::dtype::half::f16;
 use vortex_array::validity::Validity;
 use vortex_array::vtable::ValidityHelper;
 use vortex_buffer::Buffer;
+use vortex_error::VortexExpect;
 
-use super::stats::ErasedDistinctValues;
+use super::stats::ErasedStats;
 use super::stats::FloatStats;
 
 macro_rules! typed_encode {
     ($stats:ident, $typed:ident, $validity:ident, $typ:ty) => {{
-        let values: Buffer<$typ> = $typed.values.iter().map(|x| x.0).collect();
+        let distinct = $typed.distinct.as_ref().vortex_expect(
+            "this must be present since `DictScheme` declared that we need distinct values",
+        );
+
+        let values: Buffer<$typ> = distinct.distinct_values.iter().map(|x| x.0).collect();
 
         let max_code = values.len();
         let codes = if max_code <= u8::MAX as usize {
@@ -49,10 +54,10 @@ macro_rules! typed_encode {
 /// Compresses a floating-point array into a dictionary arrays according to attached stats.
 pub fn dictionary_encode(stats: &FloatStats) -> DictArray {
     let validity = stats.src.validity();
-    match &stats.distinct_values {
-        ErasedDistinctValues::F16(typed) => typed_encode!(stats, typed, validity, f16),
-        ErasedDistinctValues::F32(typed) => typed_encode!(stats, typed, validity, f32),
-        ErasedDistinctValues::F64(typed) => typed_encode!(stats, typed, validity, f64),
+    match &stats.erased {
+        ErasedStats::F16(typed) => typed_encode!(stats, typed, validity, f16),
+        ErasedStats::F32(typed) => typed_encode!(stats, typed, validity, f32),
+        ErasedStats::F64(typed) => typed_encode!(stats, typed, validity, f64),
     }
 }
 
@@ -118,7 +123,12 @@ mod tests {
             Validity::Array(BoolArray::from_iter([true, true, true, false, true]).into_array());
         let array = PrimitiveArray::new(values, validity);
 
-        let stats = FloatStats::generate(&array);
+        let stats = FloatStats::generate_opts(
+            &array,
+            crate::GenerateStatsOptions {
+                count_distinct_values: true,
+            },
+        );
         let dict_array = dictionary_encode(&stats);
         assert_eq!(dict_array.values().len(), 2);
         assert_eq!(dict_array.codes().len(), 5);

vortex-btrblocks/src/compressor/float/mod.rs

@@ -173,11 +173,11 @@ impl Scheme for ConstantScheme {
         }
 
         // Can only have 1 distinct value.
-        if stats.distinct_values_count != 1 {
-            return Ok(0.0);
+        if stats.distinct_count().is_some_and(|count| count == 1) {
+            return Ok(stats.value_count as f64);
         }
 
-        Ok(stats.value_count as f64)
+        Ok(0.0)
     }
 
     fn compress(
@@ -350,6 +350,10 @@ impl Scheme for DictScheme {
         is_float_primitive(canonical)
     }
 
+    fn needs_distinct_values(&self) -> bool {
+        true
+    }
+
     fn expected_compression_ratio(
         &self,
         compressor: &BtrBlocksCompressor,
@@ -371,13 +375,19 @@ impl Scheme for DictScheme {
             return Ok(0.0);
         }
 
-        // If the array is high cardinality (>50% unique values) skip.
-        if stats.distinct_values_count > stats.value_count / 2 {
-            return Ok(0.0);
+        // If the array is high cardinality (>50% unique values), we do not want to compress as a
+        // dictionary.
+        if stats
+            .distinct_count()
+            .is_some_and(|count| count <= stats.value_count / 2)
+        {
+            // Take a sample and run compression on the sample to determine before/after size.
+            return estimate_compression_ratio_with_sampling(
+                self, compressor, array, ctx, excludes,
+            );
         }
 
-        // Take a sample and run compression on the sample to determine before/after size.
-        estimate_compression_ratio_with_sampling(self, compressor, array, ctx, excludes)
+        Ok(0.0)
     }
 
     fn compress(

vortex-btrblocks/src/compressor/float/stats.rs

@@ -27,43 +27,58 @@ use crate::compressor::rle::RLEStats;
 use crate::sample::sample;
 
 #[derive(Debug, Clone)]
-pub struct DistinctValues<T> {
-    pub values: HashSet<NativeValue<T>, FxBuildHasher>,
+pub struct DistinctInfo<T> {
+    pub(super) distinct_values: HashSet<NativeValue<T>, FxBuildHasher>,
+    distinct_count: u32,
 }
 
 #[derive(Debug, Clone)]
-pub enum ErasedDistinctValues {
-    F16(DistinctValues<f16>),
-    F32(DistinctValues<f32>),
-    F64(DistinctValues<f64>),
+pub struct TypedStats<T> {
+    pub(super) distinct: Option<DistinctInfo<T>>,
+}
+
+#[derive(Debug, Clone)]
+pub enum ErasedStats {
+    F16(TypedStats<f16>),
+    F32(TypedStats<f32>),
+    F64(TypedStats<f64>),
+}
+
+impl ErasedStats {
+    /// Get the count of distinct values, if we have computed it already.
+    fn distinct_count(&self) -> Option<u32> {
+        match self {
+            ErasedStats::F16(x) => x.distinct.as_ref().map(|d| d.distinct_count),
+            ErasedStats::F32(x) => x.distinct.as_ref().map(|d| d.distinct_count),
+            ErasedStats::F64(x) => x.distinct.as_ref().map(|d| d.distinct_count),
+        }
+    }
 }
 
 macro_rules! impl_from_typed {
-    ($typ:ty, $variant:path) => {
-        impl From<DistinctValues<$typ>> for ErasedDistinctValues {
-            fn from(value: DistinctValues<$typ>) -> Self {
-                $variant(value)
+    ($T:ty, $variant:path) => {
+        impl From<TypedStats<$T>> for ErasedStats {
+            fn from(typed: TypedStats<$T>) -> Self {
+                $variant(typed)
             }
         }
     };
 }
 
-impl_from_typed!(f16, ErasedDistinctValues::F16);
-impl_from_typed!(f32, ErasedDistinctValues::F32);
-impl_from_typed!(f64, ErasedDistinctValues::F64);
+impl_from_typed!(f16, ErasedStats::F16);
+impl_from_typed!(f32, ErasedStats::F32);
+impl_from_typed!(f64, ErasedStats::F64);
 
 /// Array of floating-point numbers and relevant stats for compression.
 #[derive(Debug, Clone)]
 pub struct FloatStats {
-    pub(crate) src: PrimitiveArray,
+    pub(super) src: PrimitiveArray,
     // cache for validity.false_count()
-    pub(crate) null_count: u32,
+    pub(super) null_count: u32,
     // cache for validity.true_count()
-    pub(crate) value_count: u32,
-    #[allow(dead_code)]
-    pub(crate) average_run_length: u32,
-    pub(crate) distinct_values: ErasedDistinctValues,
-    pub(crate) distinct_values_count: u32,
+    pub(super) value_count: u32,
+    pub(super) average_run_length: u32,
+    pub(super) erased: ErasedStats,
 }
 
 impl FloatStats {
@@ -78,6 +93,11 @@ impl FloatStats {
             _ => vortex_panic!("cannot generate FloatStats from ptype {}", input.ptype()),
         }
     }
+
+    /// Get the count of distinct values, if we have computed it already.
+    pub fn distinct_count(&self) -> Option<u32> {
+        self.erased.distinct_count()
+    }
 }
 
 impl CompressorStats for FloatStats {
@@ -119,8 +139,8 @@ fn typed_float_stats<T: NativePType + Float>(
     count_distinct_values: bool,
 ) -> VortexResult<FloatStats>
 where
-    DistinctValues<T>: Into<ErasedDistinctValues>,
     NativeValue<T>: Hash + Eq,
+    TypedStats<T>: Into<ErasedStats>,
 {
     // Special case: empty array
     if array.is_empty() {
@@ -129,23 +149,15 @@ where
             null_count: 0,
             value_count: 0,
             average_run_length: 0,
-            distinct_values_count: 0,
-            distinct_values: DistinctValues {
-                values: HashSet::<NativeValue<T>, FxBuildHasher>::with_hasher(FxBuildHasher),
-            }
-            .into(),
+            erased: TypedStats { distinct: None }.into(),
         });
     } else if array.all_invalid()? {
         return Ok(FloatStats {
             src: array.clone(),
             null_count: u32::try_from(array.len())?,
             value_count: 0,
             average_run_length: 0,
-            distinct_values_count: 0,
-            distinct_values: DistinctValues {
-                values: HashSet::<NativeValue<T>, FxBuildHasher>::with_hasher(FxBuildHasher),
-            }
-            .into(),
+            erased: TypedStats { distinct: None }.into(),
         });
     }
 
@@ -208,7 +220,7 @@ where
 
     let null_count = u32::try_from(null_count)?;
     let value_count = u32::try_from(value_count)?;
-    let distinct_values_count = if count_distinct_values {
+    let distinct_count = if count_distinct_values {
         u32::try_from(distinct_values.len())?
     } else {
         u32::MAX
@@ -217,11 +229,13 @@ where
     Ok(FloatStats {
         null_count,
         value_count,
-        distinct_values_count,
         src: array.clone(),
         average_run_length: value_count / runs,
-        distinct_values: DistinctValues {
-            values: distinct_values,
+        erased: TypedStats {
+            distinct: Some(DistinctInfo {
+                distinct_values,
+                distinct_count,
+            }),
         }
         .into(),
     })
@@ -243,12 +257,17 @@ mod tests {
         let floats = buffer![0.0f32, 1.0f32, 2.0f32].into_array();
         let floats = floats.to_primitive();
 
-        let stats = FloatStats::generate(&floats);
+        let stats = FloatStats::generate_opts(
+            &floats,
+            crate::GenerateStatsOptions {
+                count_distinct_values: true,
+            },
+        );
 
         assert_eq!(stats.value_count, 3);
         assert_eq!(stats.null_count, 0);
         assert_eq!(stats.average_run_length, 1);
-        assert_eq!(stats.distinct_values_count, 3);
+        assert_eq!(stats.distinct_count().unwrap(), 3);
     }
 
     #[test]
@@ -258,11 +277,16 @@ mod tests {
             Validity::from_iter([false, true, true]),
         );
 
-        let stats = FloatStats::generate(&floats);
+        let stats = FloatStats::generate_opts(
+            &floats,
+            crate::GenerateStatsOptions {
+                count_distinct_values: true,
+            },
+        );
 
         assert_eq!(stats.value_count, 2);
         assert_eq!(stats.null_count, 1);
         assert_eq!(stats.average_run_length, 1);
-        assert_eq!(stats.distinct_values_count, 2);
+        assert_eq!(stats.distinct_count().unwrap(), 2);
     }
 }