add ArrayAndStats

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

Author: connortsui20

vortex-btrblocks/src/builder.rs

@@ -17,9 +17,10 @@ pub const ALL_SCHEMES: &[&dyn Scheme] = &[
     // Integer schemes.
     &crate::compressor::integer::UncompressedScheme as &dyn Scheme,
     &crate::compressor::integer::ConstantScheme,
+    // NOTE: For must precede BitPacking to avoid unnecessary patches.
     &crate::compressor::integer::FORScheme,
-    &crate::compressor::integer::ZigZagScheme,
     &crate::compressor::integer::BitPackingScheme,
+    &crate::compressor::integer::ZigZagScheme,
     &crate::compressor::integer::SparseScheme,
     &crate::compressor::integer::DictScheme,
     &crate::compressor::integer::RunEndScheme,

vortex-btrblocks/src/canonical_compressor.rs

@@ -16,7 +16,6 @@ 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;
@@ -30,18 +29,14 @@ use vortex_array::scalar::Scalar;
 use vortex_array::vtable::ValidityHelper;
 use vortex_error::VortexResult;
 
+use crate::ArrayAndStats;
 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.
@@ -223,69 +218,41 @@ 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,
-                    },
-                )
+        let merged_opts = eligible
+            .iter()
+            .fold(GenerateStatsOptions::default(), |acc, s| {
+                acc.merge(s.stats_options())
             });
-        }
 
-        if let Some(winner) = self.choose_scheme(&eligible, &array, ctx, &mut cache, excludes)? {
-            let compressed = winner.compress(self, &array, ctx, &mut cache, excludes)?;
+        let mut ctx = ctx;
+        ctx.stats_options = merged_opts;
+
+        let mut data = ArrayAndStats::new(array, merged_opts);
+
+        if let Some(winner) = self.choose_scheme(&eligible, &mut data, ctx, excludes)? {
+            let compressed = winner.compress(self, &mut data, ctx, excludes)?;
             if compressed.nbytes() < before_nbytes {
                 return Ok(compressed);
             }
         }
 
         // No scheme improved on the original.
-        Ok(array)
+        Ok(data.into_array())
     }
 
     /// Evaluates each candidate scheme and returns the one with the best compression ratio
     /// (must be > 1.0).
     fn choose_scheme(
         &self,
         schemes: &[&'static dyn Scheme],
-        array: &ArrayRef,
+        data: &mut ArrayAndStats,
         ctx: CompressorContext,
-        cache: &mut StatsCache,
         excludes: &[SchemeId],
     ) -> VortexResult<Option<&'static dyn Scheme>> {
         let mut best: Option<(&'static dyn Scheme, f64)> = None;
 
         for &scheme in schemes {
-            let ratio = self.evaluate_scheme(scheme, array, ctx, cache, excludes)?;
+            let ratio = self.evaluate_scheme(scheme, data, ctx, excludes)?;
             if is_valid_ratio(ratio) && ratio > 1.0 && best.is_none_or(|(_, r)| ratio > r) {
                 best = Some((scheme, ratio));
             }
@@ -298,12 +265,11 @@ impl BtrBlocksCompressor {
     fn evaluate_scheme(
         &self,
         scheme: &'static dyn Scheme,
-        array: &ArrayRef,
+        data: &mut ArrayAndStats,
         ctx: CompressorContext,
-        cache: &mut StatsCache,
         excludes: &[SchemeId],
     ) -> VortexResult<f64> {
-        let ratio = scheme.expected_compression_ratio(self, array, ctx, cache, excludes)?;
+        let ratio = scheme.expected_compression_ratio(self, data, ctx, excludes)?;
 
         tracing::debug!(
             scheme = %scheme.id(),