Relocate Delta's selection policy into cost-model priors; plumb the model through the builders

vortex-btrblocks/src/builder.rs

@@ -3,6 +3,12 @@
 
 //! Builder for configuring `BtrBlocksCompressor` instances.
 
+use std::sync::Arc;
+
+use vortex_compressor::cost::CostModel;
+#[cfg(feature = "unstable_encodings")]
+use vortex_compressor::cost::SchemePrior;
+use vortex_compressor::cost::SizeCost;
 use vortex_utils::aliases::hash_set::HashSet;
 
 use crate::BtrBlocksCompressor;
@@ -35,9 +41,9 @@ pub const ALL_SCHEMES: &[&dyn Scheme] = &[
     &integer::RunEndScheme,
     &integer::SequenceScheme,
     &integer::IntRLEScheme,
-    // Prefer all other schemes above delta, for now (since its slower to decompress).
+    // Delta's selection policy (penalty + floor) lives in `default_cost_model`.
     #[cfg(feature = "unstable_encodings")]
-    &integer::DeltaScheme::new(1.25),
+    &integer::DELTA_SCHEME,
     ////////////////////////////////////////////////////////////////////////////////////////////////
     // Float schemes.
     ////////////////////////////////////////////////////////////////////////////////////////////////
@@ -90,12 +96,14 @@ pub const ALL_SCHEMES: &[&dyn Scheme] = &[
 #[derive(Debug, Clone)]
 pub struct BtrBlocksCompressorBuilder {
     schemes: Vec<&'static dyn Scheme>,
+    cost_model: Arc<dyn CostModel>,
 }
 
 impl Default for BtrBlocksCompressorBuilder {
     fn default() -> Self {
         Self {
             schemes: ALL_SCHEMES.to_vec(),
+            cost_model: Arc::new(default_cost_model()),
         }
     }
 }
@@ -107,9 +115,20 @@ impl BtrBlocksCompressorBuilder {
     pub fn empty() -> Self {
         Self {
             schemes: Vec::new(),
+            cost_model: Arc::new(default_cost_model()),
         }
     }
 
+    /// Replaces the cost model used to price candidates during scheme selection.
+    ///
+    /// The default is the btrblocks default model: [`SizeCost`] with Delta's selection
+    /// prior. The model is attached to the compressor at [`build`](Self::build); the file
+    /// writer's strategy builder carries it to both its data and stats compressors.
+    pub fn with_cost_model(mut self, cost_model: Arc<dyn CostModel>) -> Self {
+        self.cost_model = cost_model;
+        self
+    }
+
     /// Adds an external compression scheme not in [`ALL_SCHEMES`].
     ///
     /// This allows encoding crates outside of `vortex-btrblocks` to register their own schemes
@@ -204,13 +223,58 @@ impl BtrBlocksCompressorBuilder {
 
     /// Builds the configured [`BtrBlocksCompressor`].
     pub fn build(self) -> BtrBlocksCompressor {
-        BtrBlocksCompressor(CascadingCompressor::new(self.schemes))
+        BtrBlocksCompressor(CascadingCompressor::new(self.schemes).with_cost_model(self.cost_model))
     }
 }
 
+/// The default cost model for btrblocks compressors: ratio-argmax ([`SizeCost`]) with Delta's
+/// selection prior.
+///
+/// Prefer all other schemes above Delta unless it wins by a real margin: Delta is slower to
+/// decompress (it breaks random access and adds a prefix-sum decode pass), so its raw ratio
+/// is handicapped by the "delta tax" multiplier and gated behind a minimum effective ratio.
+/// This prior is **policy, not measurement** — the scheme reports raw ratios, and this is
+/// the one place the judgment lives.
+pub(crate) fn default_cost_model() -> SizeCost {
+    let model = SizeCost::default();
+    #[cfg(feature = "unstable_encodings")]
+    let model = model.with_scheme_prior(
+        integer::DELTA_SCHEME.id(),
+        SchemePrior {
+            multiplier: integer::DELTA_PENALTY,
+            min_ratio: integer::DELTA_MIN_RATIO,
+        },
+    );
+    model
+}
+
 #[cfg(test)]
 mod tests {
+    use std::sync::LazyLock;
+
+    use vortex_array::ArrayRef;
+    use vortex_array::Canonical;
+    use vortex_array::ExecutionCtx;
+    use vortex_array::IntoArray;
+    use vortex_array::VortexSessionExecute;
+    use vortex_array::arrays::ConstantArray;
+    use vortex_array::arrays::PrimitiveArray;
+    use vortex_array::dtype::Nullability;
+    use vortex_array::scalar::Scalar;
+    use vortex_array::validity::Validity;
+    use vortex_buffer::buffer;
+    use vortex_compressor::cost::Candidate;
+    use vortex_compressor::cost::Cost;
+    use vortex_compressor::scheme::CandidateEstimate;
+    use vortex_compressor::scheme::SchemeEvaluation;
+    use vortex_error::VortexResult;
+    use vortex_session::VortexSession;
+
     use super::*;
+    use crate::ArrayAndStats;
+    use crate::CompressorContext;
+
+    static SESSION: LazyLock<VortexSession> = LazyLock::new(vortex_array::array_session);
 
     #[test]
     fn empty_starts_with_no_schemes() {
@@ -223,4 +287,103 @@ mod tests {
         let builder = BtrBlocksCompressorBuilder::default();
         assert_eq!(builder.schemes.len(), ALL_SCHEMES.len());
     }
+
+    /// Inverts [`SizeCost`]'s preferences, so the *worst*-priced candidate wins. Only useful
+    /// for proving that an injected model actually drives selection.
+    #[derive(Debug)]
+    struct InvertedSizeCost;
+
+    impl CostModel for InvertedSizeCost {
+        fn cost(&self, candidate: &Candidate<'_>) -> Option<Cost> {
+            SizeCost::default()
+                .cost(candidate)
+                .map(|cost| Cost::new(-cost.value()))
+        }
+
+        fn canonical_cost(&self, _data: &ArrayAndStats, _n_values: u64) -> Cost {
+            Cost::new(f64::MAX)
+        }
+    }
+
+    /// Fixed-ratio scheme whose output is a constant array carrying `marker`, so the winner
+    /// is observable from the compressed output.
+    #[derive(Debug)]
+    struct MarkerScheme {
+        name: &'static str,
+        ratio: f64,
+        marker: i32,
+    }
+
+    impl Scheme for MarkerScheme {
+        fn scheme_name(&self) -> &'static str {
+            self.name
+        }
+
+        fn matches(&self, canonical: &Canonical) -> bool {
+            canonical.dtype().is_int()
+        }
+
+        fn evaluate(
+            &self,
+            _data: &ArrayAndStats,
+            _compress_ctx: CompressorContext,
+            _exec_ctx: &mut ExecutionCtx,
+        ) -> SchemeEvaluation {
+            SchemeEvaluation::Candidate(CandidateEstimate::from_compression_ratio(self.ratio))
+        }
+
+        fn compress(
+            &self,
+            _compressor: &CascadingCompressor,
+            data: &ArrayAndStats,
+            _compress_ctx: CompressorContext,
+            _exec_ctx: &mut ExecutionCtx,
+        ) -> VortexResult<ArrayRef> {
+            Ok(ConstantArray::new(
+                Scalar::primitive(self.marker, Nullability::NonNullable),
+                data.array_len(),
+            )
+            .into_array())
+        }
+    }
+
+    static HIGH_RATIO: MarkerScheme = MarkerScheme {
+        name: "test.high_ratio",
+        ratio: 3.0,
+        marker: 111,
+    };
+    static LOW_RATIO: MarkerScheme = MarkerScheme {
+        name: "test.low_ratio",
+        ratio: 2.0,
+        marker: 222,
+    };
+
+    /// A model injected via [`BtrBlocksCompressorBuilder::with_cost_model`] must actually
+    /// drive selection: under the default model the higher ratio wins, under the inverted
+    /// model the lower ratio wins.
+    #[test]
+    fn cost_model_plumbing_is_live() -> VortexResult<()> {
+        let array =
+            PrimitiveArray::new(buffer![5i32, 6, 7, 8, 9, 10], Validity::NonNullable).into_array();
+        let marker = |scheme: &MarkerScheme| {
+            Some(Scalar::primitive(scheme.marker, Nullability::NonNullable))
+        };
+
+        let compressed = BtrBlocksCompressorBuilder::empty()
+            .with_new_scheme(&HIGH_RATIO)
+            .with_new_scheme(&LOW_RATIO)
+            .build()
+            .compress(&array, &mut SESSION.create_execution_ctx())?;
+        assert_eq!(compressed.as_constant(), marker(&HIGH_RATIO));
+
+        let compressed = BtrBlocksCompressorBuilder::empty()
+            .with_new_scheme(&HIGH_RATIO)
+            .with_new_scheme(&LOW_RATIO)
+            .with_cost_model(Arc::new(InvertedSizeCost))
+            .build()
+            .compress(&array, &mut SESSION.create_execution_ctx())?;
+        assert_eq!(compressed.as_constant(), marker(&LOW_RATIO));
+
+        Ok(())
+    }
 }

vortex-btrblocks/src/lib.rs

@@ -78,6 +78,7 @@ pub use builder::BtrBlocksCompressorBuilder;
 pub use canonical_compressor::BtrBlocksCompressor;
 pub use schemes::patches::compress_patches;
 pub use vortex_compressor::CascadingCompressor;
+pub use vortex_compressor::cost;
 pub use vortex_compressor::scheme::CompressorContext;
 pub use vortex_compressor::scheme::MAX_CASCADE;
 pub use vortex_compressor::scheme::Scheme;