Rework Scheme estimation in compressor (#7230)

## Summary

Tracking Issue: #7216

Adds a new `CompressionEstimate` type in
`vortex-compressor/src/estimate.rs` that the
`expected_compression_ratio` method now returns. Additionally moves some
things around for clarity.

Note that this is not just a refactor, there is subtle logic that has
changed in a few places (that I think is better, not actually sure). I'm
happy to split out some stuff into other PRs if that helps.

### Future Work

- I also would like to add a variant called `Exact` that returns the
fully compressed array in the case that we can only determine if a
scheme is a candidate by compressing the whole thing without any errors;
the only case where we want to do this is `SequenceArray` (and maybe
there's an argument to do this for `ConstantArray` too, but the
semantics around `ConstantArray` should be even more special regardless,
imo).
- This might be in a `ResolvedEstimate` enum instead.
- There are also a bunch of TODOs littered everywhere that are easily
fixed, but I want to do those in a followup.
- We probably want to hardcode the `ConstantScheme` logic into the
compressor since I cannot think of any reason why you would not want to
have a `ConstantScheme` (except when you have a very small array, and at
that point you don't care about perf regardless).

## API Changes

`expected_compression_ratio` now only takes stats and compressor context
(it does not take the compressor at all) and returns a
`CompressionEstimate`. This method must be super quick, and any sampling
or expensive operations are now deferred to later by the compressor.

## Testing

Just a few extra tests, am relying on the existing test suite as it's
not like completely new logic is happening.

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

Author: connortsui20

Cargo.lock

@@ -65,15 +65,5 @@ name = "compress_listview"
 harness = false
 test = false
 
-[[bench]]
-name = "dict_encode"
-harness = false
-test = false
-
-[[bench]]
-name = "stats_calc"
-harness = false
-test = false
-
 [package.metadata.cargo-machete]
 ignored = ["getrandom_v03"]

vortex-btrblocks/Cargo.toml

@@ -66,13 +66,11 @@ pub use builder::BtrBlocksCompressorBuilder;
 pub use canonical_compressor::BtrBlocksCompressor;
 pub use schemes::patches::compress_patches;
 pub use vortex_compressor::CascadingCompressor;
-pub use vortex_compressor::builtins::integer_dictionary_encode;
 pub use vortex_compressor::ctx::CompressorContext;
 pub use vortex_compressor::ctx::MAX_CASCADE;
 pub use vortex_compressor::scheme::Scheme;
 pub use vortex_compressor::scheme::SchemeExt;
 pub use vortex_compressor::scheme::SchemeId;
-pub use vortex_compressor::scheme::estimate_compression_ratio_with_sampling;
 pub use vortex_compressor::stats::ArrayAndStats;
 pub use vortex_compressor::stats::BoolStats;
 pub use vortex_compressor::stats::FloatStats;

vortex-btrblocks/public-api.lock

@@ -10,6 +10,7 @@ use vortex_array::ToCanonical;
 use vortex_array::arrays::PrimitiveArray;
 use vortex_array::arrays::decimal::narrowed_decimal;
 use vortex_array::dtype::DecimalType;
+use vortex_compressor::estimate::CompressionEstimate;
 use vortex_decimal_byte_parts::DecimalByteParts;
 use vortex_error::VortexResult;
 
@@ -42,12 +43,11 @@ impl Scheme for DecimalScheme {
 
     fn expected_compression_ratio(
         &self,
-        _compressor: &CascadingCompressor,
         _data: &mut ArrayAndStats,
         _ctx: CompressorContext,
-    ) -> VortexResult<f64> {
+    ) -> CompressionEstimate {
         // Decimal compression is almost always beneficial (narrowing + primitive compression).
-        Ok(f64::MAX)
+        CompressionEstimate::AlwaysUse
     }
 
     fn compress(

vortex-btrblocks/src/lib.rs

@@ -15,6 +15,7 @@ use vortex_array::IntoArray;
 use vortex_array::ToCanonical;
 use vortex_array::arrays::primitive::PrimitiveArrayExt;
 use vortex_array::dtype::PType;
+use vortex_compressor::estimate::CompressionEstimate;
 use vortex_compressor::scheme::ChildSelection;
 use vortex_compressor::scheme::DescendantExclusion;
 use vortex_error::VortexResult;
@@ -28,7 +29,6 @@ use crate::CompressorContext;
 use crate::Scheme;
 use crate::SchemeExt;
 use crate::compress_patches;
-use crate::estimate_compression_ratio_with_sampling;
 
 /// ALP (Adaptive Lossless floating-Point) encoding.
 #[derive(Debug, Copy, Clone, PartialEq, Eq)]
@@ -73,22 +73,21 @@ impl Scheme for ALPScheme {
 
     fn expected_compression_ratio(
         &self,
-        compressor: &CascadingCompressor,
         data: &mut ArrayAndStats,
         ctx: CompressorContext,
-    ) -> VortexResult<f64> {
+    ) -> CompressionEstimate {
         // ALP encodes floats as integers. Without integer compression afterward, the encoded ints
         // are the same size.
         if ctx.finished_cascading() {
-            return Ok(0.0);
+            return CompressionEstimate::Skip;
         }
 
         // We don't support ALP for f16.
-        if data.float_stats().source().ptype() == PType::F16 {
-            return Ok(0.0);
+        if data.array_as_primitive().ptype() == PType::F16 {
+            return CompressionEstimate::Skip;
         }
 
-        estimate_compression_ratio_with_sampling(self, compressor, data.array(), ctx)
+        CompressionEstimate::Sample
     }
 
     fn compress(
@@ -97,9 +96,7 @@ impl Scheme for ALPScheme {
         data: &mut ArrayAndStats,
         ctx: CompressorContext,
     ) -> VortexResult<ArrayRef> {
-        let stats = data.float_stats();
-
-        let alp_encoded = alp_encode(stats.source(), None)?;
+        let alp_encoded = alp_encode(&data.array_as_primitive(), None)?;
 
         // Compress the ALP ints.
         let compressed_alp_ints =
@@ -124,15 +121,15 @@ impl Scheme for ALPRDScheme {
 
     fn expected_compression_ratio(
         &self,
-        compressor: &CascadingCompressor,
         data: &mut ArrayAndStats,
-        ctx: CompressorContext,
-    ) -> VortexResult<f64> {
-        if data.float_stats().source().ptype() == PType::F16 {
-            return Ok(0.0);
+        _ctx: CompressorContext,
+    ) -> CompressionEstimate {
+        // We don't support ALPRD for f16.
+        if data.array_as_primitive().ptype() == PType::F16 {
+            return CompressionEstimate::Skip;
         }
 
-        estimate_compression_ratio_with_sampling(self, compressor, data.array(), ctx)
+        CompressionEstimate::Sample
     }
 
     fn compress(
@@ -141,15 +138,15 @@ impl Scheme for ALPRDScheme {
         data: &mut ArrayAndStats,
         _ctx: CompressorContext,
     ) -> VortexResult<ArrayRef> {
-        let stats = data.float_stats();
+        let primitive_array = data.array_as_primitive();
 
-        let encoder = match stats.source().ptype() {
-            PType::F32 => RDEncoder::new(stats.source().as_slice::<f32>()),
-            PType::F64 => RDEncoder::new(stats.source().as_slice::<f64>()),
+        let encoder = match primitive_array.ptype() {
+            PType::F32 => RDEncoder::new(primitive_array.as_slice::<f32>()),
+            PType::F64 => RDEncoder::new(primitive_array.as_slice::<f64>()),
             ptype => vortex_panic!("cannot ALPRD compress ptype {ptype}"),
         };
 
-        let alp_rd = encoder.encode(stats.source());
+        let alp_rd = encoder.encode(&primitive_array);
         let dtype = alp_rd.dtype().clone();
         let right_bit_width = alp_rd.right_bit_width();
         let mut parts = ALPRDArrayOwnedExt::into_data_parts(alp_rd);
@@ -191,24 +188,25 @@ impl Scheme for NullDominatedSparseScheme {
 
     fn expected_compression_ratio(
         &self,
-        _compressor: &CascadingCompressor,
         data: &mut ArrayAndStats,
         _ctx: CompressorContext,
-    ) -> VortexResult<f64> {
+    ) -> CompressionEstimate {
+        let len = data.array_len() as f64;
         let stats = data.float_stats();
+        let value_count = stats.value_count();
 
-        if stats.value_count() == 0 {
-            // All nulls should use ConstantScheme instead of this.
-            return Ok(0.0);
+        // All-null arrays should be compressed as constant instead anyways.
+        if value_count == 0 {
+            return CompressionEstimate::Skip;
         }
 
         // If the majority (90%) of values is null, this will compress well.
-        if stats.null_count() as f64 / stats.source().len() as f64 > 0.9 {
-            return Ok(stats.source().len() as f64 / stats.value_count() as f64);
+        if stats.null_count() as f64 / len > 0.9 {
+            return CompressionEstimate::Ratio(len / value_count as f64);
         }
 
         // Otherwise we don't go this route.
-        Ok(0.0)
+        CompressionEstimate::Skip
     }
 
     fn compress(
@@ -217,10 +215,8 @@ impl Scheme for NullDominatedSparseScheme {
         data: &mut ArrayAndStats,
         ctx: CompressorContext,
     ) -> VortexResult<ArrayRef> {
-        let stats = data.float_stats();
-
         // We pass None as we only run this pathway for NULL-dominated float arrays.
-        let sparse_encoded = Sparse::encode(&stats.source().clone().into_array(), None)?;
+        let sparse_encoded = Sparse::encode(data.array(), None)?;
 
         if let Some(sparse) = sparse_encoded.as_opt::<Sparse>() {
             let indices = sparse.patches().indices().to_primitive().narrow()?;
@@ -250,17 +246,26 @@ impl Scheme for PcoScheme {
         is_float_primitive(canonical)
     }
 
+    fn expected_compression_ratio(
+        &self,
+        _data: &mut ArrayAndStats,
+        _ctx: CompressorContext,
+    ) -> CompressionEstimate {
+        CompressionEstimate::Sample
+    }
+
     fn compress(
         &self,
         _compressor: &CascadingCompressor,
         data: &mut ArrayAndStats,
         _ctx: CompressorContext,
     ) -> VortexResult<ArrayRef> {
-        let stats = data.float_stats();
-        Ok(
-            vortex_pco::Pco::from_primitive(stats.source(), pco::DEFAULT_COMPRESSION_LEVEL, 8192)?
-                .into_array(),
-        )
+        Ok(vortex_pco::Pco::from_primitive(
+            &data.array_as_primitive(),
+            pco::DEFAULT_COMPRESSION_LEVEL,
+            8192,
+        )?
+        .into_array())
     }
 }
 
@@ -406,7 +411,8 @@ mod scheme_selection_tests {
         let array = PrimitiveArray::new(Buffer::copy_from(&values), Validity::NonNullable);
         let btr = BtrBlocksCompressor::default();
         let compressed = btr.compress(&array.into_array())?;
-        assert!(compressed.is::<Dict>());
+        assert!(compressed.is::<ALP>());
+        assert!(compressed.children()[0].is::<Dict>());
         Ok(())
     }