diff --git a/vortex-btrblocks/src/builder.rs b/vortex-btrblocks/src/builder.rs index 26632e70860..b1aa20401b6 100644 --- a/vortex-btrblocks/src/builder.rs +++ b/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, } 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) -> 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 = 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 { + 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 { + 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(()) + } } diff --git a/vortex-btrblocks/src/lib.rs b/vortex-btrblocks/src/lib.rs index 37915b427b6..3c791bf4072 100644 --- a/vortex-btrblocks/src/lib.rs +++ b/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; diff --git a/vortex-btrblocks/src/schemes/integer/delta.rs b/vortex-btrblocks/src/schemes/integer/delta.rs index 229e8224b67..98dde74c25e 100644 --- a/vortex-btrblocks/src/schemes/integer/delta.rs +++ b/vortex-btrblocks/src/schemes/integer/delta.rs @@ -35,26 +35,39 @@ use crate::SchemeExt; /// stride), so a later cascade layer (FoR / BitPacking) packs the smaller residuals. It only /// pays off when those residuals span meaningfully fewer bits than the values themselves. /// -/// The minimum penalized compression ratio required for Delta to be selected is configurable via -/// [`DeltaScheme::new`]; [`DeltaScheme::default`] uses a ratio of `1.25`. +/// The scheme reports its raw estimated ratio (`full_width / delta_bits`). Delta's selection +/// policy — the "delta tax" multiplier and the minimum-win floor — lives on the compressor's +/// cost model as a scheme prior (registered by the btrblocks default cost model), not in the +/// scheme. #[derive(Debug, Copy, Clone, PartialEq)] pub struct DeltaScheme { - min_ratio: f64, + /// Deprecated scheme-level minimum-ratio floor override. `None` defers entirely to the + /// cost model's prior. + min_ratio: Option, } +/// The `ALL_SCHEMES` Delta instance: all selection policy comes from the cost model. +pub(crate) const DELTA_SCHEME: DeltaScheme = DeltaScheme { min_ratio: None }; + impl DeltaScheme { - /// Creates a Delta scheme requiring `min_ratio` after the delta penalty before it wins. - /// - /// Pass a higher ratio to make Delta more conservative, or a lower one to select it more - /// eagerly. [`DeltaScheme::default`] uses a ratio of `1.25`. + /// Creates a Delta scheme with a scheme-level minimum-ratio floor: Delta is skipped + /// unless its penalized ratio strictly exceeds `min_ratio`. + #[deprecated( + since = "0.1.0", + note = "configure the floor on the compressor's cost model via \ + `SizeCost::with_scheme_prior` instead; a scheme-level floor only tightens \ + policy (the model's prior floor still applies)" + )] pub const fn new(min_ratio: f64) -> Self { - Self { min_ratio } + Self { + min_ratio: Some(min_ratio), + } } } impl Default for DeltaScheme { fn default() -> Self { - Self::new(1.25) + DELTA_SCHEME } } @@ -64,7 +77,21 @@ impl Default for DeltaScheme { /// carries a structural sign bit on its residuals. We therefore require Delta to be meaningfully /// (~5%) smaller than the best alternative before it wins, rather than picking it for a /// single-bit gain. This factor encodes that "delta tax". -const DELTA_PENALTY: f64 = 0.95; +/// +/// The penalty is applied by the cost model (registered as Delta's [`SchemePrior`] multiplier +/// in [`default_cost_model`]); the scheme itself only uses it to honor the deprecated +/// scheme-level floor from [`DeltaScheme::new`]. +/// +/// [`SchemePrior`]: vortex_compressor::cost::SchemePrior +/// [`default_cost_model`]: crate::builder +pub(crate) const DELTA_PENALTY: f64 = 0.95; + +/// Minimum effective (post-penalty) compression ratio Delta must strictly exceed before it is +/// selected, registered as Delta's [`SchemePrior`] floor in [`default_cost_model`]. +/// +/// [`SchemePrior`]: vortex_compressor::cost::SchemePrior +/// [`default_cost_model`]: crate::builder +pub(crate) const DELTA_MIN_RATIO: f64 = 1.25; /// Minimum length before Delta is worth considering (one FastLanes chunk). const MIN_DELTA_LEN: usize = 1024; @@ -131,7 +158,7 @@ impl Scheme for DeltaScheme { // Estimating Delta needs the real transposed-delta span, so defer to a callback that // delta-encodes the array and measures the residual range. - let min_ratio = self.min_ratio; + let min_ratio_override = self.min_ratio; SchemeEvaluation::Deferred(DeferredEvaluation::Callback(Box::new( move |_compressor, data, _ctx, exec_ctx| { let primitive = data.array().clone().execute::(exec_ctx)?; @@ -152,10 +179,15 @@ impl Scheme for DeltaScheme { None => return Ok(ResolvedEvaluation::Skip), }; - let ratio = full_width / delta_bits * DELTA_PENALTY; - if ratio <= min_ratio { + let ratio = full_width / delta_bits; + + // Deprecated scheme-level floor (`DeltaScheme::new`): reproduces the historical + // post-penalty comparison while the constructor remains. The model's prior floor + // also applies, so this override can only tighten policy. + if min_ratio_override.is_some_and(|min_ratio| ratio * DELTA_PENALTY <= min_ratio) { return Ok(ResolvedEvaluation::Skip); } + Ok(ResolvedEvaluation::Candidate( CandidateEstimate::from_compression_ratio(ratio), )) @@ -192,3 +224,135 @@ impl Scheme for DeltaScheme { Delta::try_new(compressed_bases, compressed_deltas, 0, len).map(IntoArray::into_array) } } + +#[cfg(test)] +mod tests { + use std::sync::LazyLock; + + use rand::RngExt; + use rand::SeedableRng; + use rand::rngs::StdRng; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::Constant; + use vortex_array::arrays::ConstantArray; + use vortex_array::dtype::Nullability; + use vortex_array::scalar::Scalar; + use vortex_array::validity::Validity; + use vortex_buffer::Buffer; + use vortex_session::VortexSession; + + use super::*; + use crate::builder::default_cost_model; + + static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); + + /// A compressor over `schemes` with the btrblocks default cost model (which carries + /// Delta's prior). + fn compressor_with_default_model(schemes: Vec<&'static dyn Scheme>) -> CascadingCompressor { + CascadingCompressor::new(schemes).with_cost_model(std::sync::Arc::new(default_cost_model())) + } + + /// Immediate fixed-ratio competitor; its `compress` emits a tiny constant array so the + /// winner is observable from the output encoding. + #[derive(Debug)] + struct FixedRatioScheme { + ratio: f64, + } + + impl Scheme for FixedRatioScheme { + fn scheme_name(&self) -> &'static str { + "test.fixed_ratio" + } + + 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 { + Ok(ConstantArray::new( + Scalar::primitive(0u64, Nullability::NonNullable), + data.array_len(), + ) + .into_array()) + } + } + + /// Near-monotone u64 data: Delta's residual span is a few bits, so its effective + /// (post-prior) ratio is comfortably above 2.0 but nowhere near its 64-bit best case. + fn monotone_jitter_u64() -> ArrayRef { + let mut rng = StdRng::seed_from_u64(42); + let mut value = 1_700_000_000_000u64; + let values: Buffer = (0..4096) + .map(|_| { + value += 900 + rng.random_range(0..200); + value + }) + .collect(); + PrimitiveArray::new(values, Validity::NonNullable).into_array() + } + + /// With the incumbent below Delta's achievable effective ratio, Delta must win under + /// the default model's prior. + #[test] + fn delta_wins_over_low_incumbent() -> VortexResult<()> { + static COMPETITOR: FixedRatioScheme = FixedRatioScheme { ratio: 2.0 }; + let compressor = compressor_with_default_model(vec![&COMPETITOR, &DELTA_SCHEME]); + + let mut exec_ctx = SESSION.create_execution_ctx(); + let compressed = compressor.compress(&monotone_jitter_u64(), &mut exec_ctx)?; + + assert!(compressed.is::()); + Ok(()) + } + + /// With the incumbent at Delta's penalized best case (`full_width * DELTA_PENALTY`), + /// Delta must not be chosen: its measured effective ratio is far below the incumbent — + /// the same decision the historical scheme-side `max_ratio <= best` skip produced on + /// this input. + #[test] + fn delta_loses_at_best_case_tie() -> VortexResult<()> { + static COMPETITOR: FixedRatioScheme = FixedRatioScheme { + ratio: 64.0 * DELTA_PENALTY, + }; + let compressor = compressor_with_default_model(vec![&COMPETITOR, &DELTA_SCHEME]); + + let mut exec_ctx = SESSION.create_execution_ctx(); + let compressed = compressor.compress(&monotone_jitter_u64(), &mut exec_ctx)?; + + assert!(compressed.is::()); + Ok(()) + } + + /// Wide random residuals put Delta's effective ratio below the prior's floor, so with + /// no competitor the array must stay canonical. + #[test] + fn delta_skips_below_min_ratio() -> VortexResult<()> { + let compressor = compressor_with_default_model(vec![&DELTA_SCHEME]); + + let mut rng = StdRng::seed_from_u64(43); + let values: Buffer = (0..4096).map(|_| rng.random::()).collect(); + let array = PrimitiveArray::new(values, Validity::NonNullable).into_array(); + + let mut exec_ctx = SESSION.create_execution_ctx(); + let compressed = compressor.compress(&array, &mut exec_ctx)?; + + assert!(!compressed.is::()); + Ok(()) + } +} diff --git a/vortex-btrblocks/src/schemes/integer/mod.rs b/vortex-btrblocks/src/schemes/integer/mod.rs index 3aae2ae5601..44f42b42ad1 100644 --- a/vortex-btrblocks/src/schemes/integer/mod.rs +++ b/vortex-btrblocks/src/schemes/integer/mod.rs @@ -18,6 +18,12 @@ mod pco; pub use bitpacking::BitPackingScheme; #[cfg(feature = "unstable_encodings")] +pub(crate) use delta::DELTA_MIN_RATIO; +#[cfg(feature = "unstable_encodings")] +pub(crate) use delta::DELTA_PENALTY; +#[cfg(feature = "unstable_encodings")] +pub(crate) use delta::DELTA_SCHEME; +#[cfg(feature = "unstable_encodings")] pub use delta::DeltaScheme; pub use for_::FoRScheme; #[cfg(feature = "pco")] diff --git a/vortex-compressor/src/compressor/mod.rs b/vortex-compressor/src/compressor/mod.rs index fbbb838d592..c6f611cf7aa 100644 --- a/vortex-compressor/src/compressor/mod.rs +++ b/vortex-compressor/src/compressor/mod.rs @@ -70,7 +70,7 @@ impl CascadingCompressor { Self { schemes, root_exclusions, - cost_model: Arc::new(SizeCost), + cost_model: Arc::new(SizeCost::default()), } } diff --git a/vortex-compressor/src/compressor/tests.rs b/vortex-compressor/src/compressor/tests.rs index ff5e3a94132..9289b4a6d91 100644 --- a/vortex-compressor/src/compressor/tests.rs +++ b/vortex-compressor/src/compressor/tests.rs @@ -542,11 +542,11 @@ struct PruneAllDeferredModel; impl CostModel for PruneAllDeferredModel { fn cost(&self, candidate: &Candidate<'_>) -> Option { - SizeCost.cost(candidate) + SizeCost::default().cost(candidate) } fn canonical_cost(&self, data: &ArrayAndStats, n_values: u64) -> Cost { - SizeCost.canonical_cost(data, n_values) + SizeCost::default().canonical_cost(data, n_values) } fn lower_bound(&self, _scheme: SchemeId, _data: &ArrayAndStats) -> Option { diff --git a/vortex-compressor/src/cost/mod.rs b/vortex-compressor/src/cost/mod.rs index fda8c2c0c86..68d04c27048 100644 --- a/vortex-compressor/src/cost/mod.rs +++ b/vortex-compressor/src/cost/mod.rs @@ -48,4 +48,5 @@ pub use model::Cost; pub use model::CostModel; mod size; +pub use size::SchemePrior; pub use size::SizeCost; diff --git a/vortex-compressor/src/cost/size.rs b/vortex-compressor/src/cost/size.rs index 34117e2b557..0a2a5d46f02 100644 --- a/vortex-compressor/src/cost/size.rs +++ b/vortex-compressor/src/cost/size.rs @@ -3,11 +3,39 @@ //! The default, size-only cost model. +use vortex_utils::aliases::hash_map::HashMap; + use crate::cost::Candidate; use crate::cost::Cost; use crate::cost::CostModel; +use crate::scheme::SchemeId; use crate::stats::ArrayAndStats; +/// A per-scheme selection prior applied by [`SizeCost`]: **policy, not measurement**. +/// +/// Priors let the model demand that a scheme win by a margin (`multiplier < 1.0`) or clear a +/// floor (`min_ratio`) before it is selected, without the scheme hardcoding that judgment +/// into its own estimate. +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct SchemePrior { + /// Multiplier applied to the scheme's raw ratio signal before any comparison. Values + /// below `1.0` handicap the scheme, demanding it win by a margin. + pub multiplier: f64, + + /// The effective (post-multiplier) ratio the scheme must strictly exceed to be + /// considered at all. + pub min_ratio: f64, +} + +impl Default for SchemePrior { + fn default() -> Self { + Self { + multiplier: 1.0, + min_ratio: 1.0, + } + } +} + /// The default cost model: maximize estimated compression ratio. /// /// `SizeCost` preserves the compressor's historical ratio-argmax ordering: @@ -21,12 +49,43 @@ use crate::stats::ArrayAndStats; /// historical `ratio > 1.0` validity gate. /// - Zero-byte sample results and non-finite or subnormal ratios are rejected (`None`), /// reproducing the historical estimate-validity checks. -#[derive(Debug, Default, Clone, Copy)] -pub struct SizeCost; +/// +/// Per-scheme [`SchemePrior`]s registered via [`with_scheme_prior`] additionally handicap a +/// scheme's raw ratio (`effective = raw * multiplier`) and gate it behind a floor +/// (`effective > min_ratio`). Schemes without a registered prior are priced off their raw +/// ratio unchanged. +/// +/// [`with_scheme_prior`]: SizeCost::with_scheme_prior +#[derive(Debug, Default, Clone)] +pub struct SizeCost { + /// Per-scheme selection priors; schemes without an entry are priced unmodified. + priors: HashMap, +} + +impl SizeCost { + /// Registers a selection prior for `scheme`, replacing any existing entry. + pub fn with_scheme_prior(mut self, scheme: SchemeId, prior: SchemePrior) -> Self { + self.priors.insert(scheme, prior); + self + } +} impl CostModel for SizeCost { fn cost(&self, candidate: &Candidate<'_>) -> Option { - let ratio = candidate.estimate().estimated_compression_ratio()?; + let raw = candidate.estimate().estimated_compression_ratio()?; + + let ratio = match self.priors.get(&candidate.scheme_id()) { + Some(prior) => { + let effective = raw * prior.multiplier; + // Non-finite `effective` falls through to the validity check below. + if effective <= prior.min_ratio { + return None; + } + effective + } + None => raw, + }; + (ratio.is_finite() && !ratio.is_subnormal()).then(|| Cost::new(-ratio)) } @@ -53,6 +112,7 @@ mod tests { use crate::scheme::CompressorContext; use crate::scheme::Scheme; use crate::scheme::SchemeEvaluation; + use crate::scheme::SchemeExt; use crate::stats::GenerateStatsOptions; #[derive(Debug)] @@ -98,12 +158,12 @@ mod tests { CandidateEstimate::zero_bytes, CandidateEstimate::from_compression_ratio, ); - SizeCost.cost(&Candidate::new(&TestScheme, estimate, &data, None, &[])) + SizeCost::default().cost(&Candidate::new(&TestScheme, estimate, &data, None, &[])) } fn canonical_cost() -> Cost { let data = test_data(); - SizeCost.canonical_cost(&data, 4) + SizeCost::default().canonical_cost(&data, 4) } /// The historical estimate-validity rule that `SizeCost` must reproduce: finite, @@ -190,4 +250,43 @@ mod tests { fn zero_bytes_is_rejected() { assert!(cost(None).is_none()); } + + /// A prior must reproduce the historical scheme-side policy arithmetic bit-for-bit: + /// `effective = raw * multiplier` (the exact expression Delta used), gated at + /// `effective > min_ratio`, with the effective ratio driving the cost. + #[test] + fn prior_reproduces_post_penalty_ratios_bit_for_bit() { + let model = SizeCost::default().with_scheme_prior( + TestScheme.id(), + SchemePrior { + multiplier: 0.95, + min_ratio: 1.25, + }, + ); + + // Raw ratios spanning both sides of the floor, plus the exact boundary. + let raws = [64.0, 64.0 / 9.0, 8.0, 2.0, 1.4, 1.25 / 0.95, 1.32, 1.3, 1.0]; + let data = test_data(); + for raw in raws { + let expected = raw * 0.95; + let candidate = Candidate::new( + &TestScheme, + CandidateEstimate::from_compression_ratio(raw), + &data, + None, + &[], + ); + let cost = model.cost(&candidate); + if expected > 1.25 { + let cost = cost.expect("above the floor must be priceable"); + // Bit-exact: the cost is the negated product, not a re-derived value. + assert_eq!(cost.value(), -expected, "raw {raw}"); + } else { + assert!(cost.is_none(), "raw {raw} must be gated by the floor"); + } + } + + // Schemes without an entry are priced off the raw ratio unchanged. + assert_eq!(cost(Some(1.1)).map(Cost::value), Some(-1.1)); + } } diff --git a/vortex-file/src/strategy.rs b/vortex-file/src/strategy.rs index 7db2b436a38..56a99861d91 100644 --- a/vortex-file/src/strategy.rs +++ b/vortex-file/src/strategy.rs @@ -218,7 +218,8 @@ impl WriteStrategyBuilder { /// Override the default [`BtrBlocksCompressorBuilder`] used for compression. /// /// The builder is finalized during [`build`](Self::build), producing two compressors: one for - /// data (with `IntDictScheme` excluded) and one for stats. + /// data (with `IntDictScheme` excluded) and one for stats. Any cost model configured on the + /// builder (via `BtrBlocksCompressorBuilder::with_cost_model`) is carried into both. pub fn with_btrblocks_builder(mut self, builder: BtrBlocksCompressorBuilder) -> Self { self.compressor = CompressorConfig::BtrBlocks(builder); self diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index ea7f8cbeda5..9aac9ee37ed 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -58,8 +58,13 @@ use vortex_array::stats::PRUNING_STATS; use vortex_array::stream::ArrayStreamAdapter; use vortex_array::stream::ArrayStreamExt; use vortex_array::validity::Validity; +use vortex_btrblocks::ArrayAndStats; use vortex_btrblocks::BtrBlocksCompressorBuilder; use vortex_btrblocks::SchemeExt; +use vortex_btrblocks::cost::Candidate; +use vortex_btrblocks::cost::Cost; +use vortex_btrblocks::cost::CostModel; +use vortex_btrblocks::cost::SizeCost; use vortex_btrblocks::schemes::string::StringDictScheme; use vortex_buffer::Buffer; use vortex_buffer::ByteBufferMut; @@ -67,6 +72,7 @@ use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_io::session::RuntimeSession; use vortex_layout::Layout; +use vortex_layout::LayoutStrategy; use vortex_layout::layouts::flat::writer::FlatLayoutStrategy; use vortex_layout::layouts::zoned::LegacyStats; use vortex_layout::layouts::zoned::Zoned; @@ -2292,3 +2298,65 @@ async fn repro_8166_binary_gt_all_ff_max() -> VortexResult<()> { assert_eq!(result.len(), 1); Ok(()) } + +/// Inverts the default size model's preferences, so the *worst*-priced candidate wins. Only +/// useful for proving that a model injected through the write strategy actually drives +/// selection inside the file writer. +#[derive(Debug)] +struct InvertedSizeCost; + +impl CostModel for InvertedSizeCost { + fn cost(&self, candidate: &Candidate<'_>) -> Option { + SizeCost::default() + .cost(candidate) + .map(|cost| Cost::new(-cost.value())) + } + + fn canonical_cost(&self, _data: &ArrayAndStats, _n_values: u64) -> Cost { + Cost::new(f64::MAX) + } +} + +/// A cost model injected via `WriteStrategyBuilder::with_btrblocks_builder` must reach the +/// writer's compressors: on the same highly compressible data, the inverted model must +/// produce a strictly larger file than the default model. +#[tokio::test] +async fn cost_model_reaches_file_compressors() -> VortexResult<()> { + async fn write_len(array: &ArrayRef, strategy: Arc) -> VortexResult { + let mut buf = ByteBufferMut::empty(); + SESSION + .write_options() + .with_strategy(strategy) + .write(&mut buf, array.clone().to_array_stream()) + .await?; + Ok(buf.len()) + } + + let values: Buffer = (0..65_536i64).map(|i| (i % 7) * 123_456).collect(); + let array = StructArray::from_fields(&[( + "col", + PrimitiveArray::new(values, Validity::NonNullable).into_array(), + )])? + .into_array(); + + let default_len = write_len( + &array, + crate::strategy::WriteStrategyBuilder::default().build(), + ) + .await?; + let inverted_len = write_len( + &array, + crate::strategy::WriteStrategyBuilder::default() + .with_btrblocks_builder( + BtrBlocksCompressorBuilder::default().with_cost_model(Arc::new(InvertedSizeCost)), + ) + .build(), + ) + .await?; + + assert!( + inverted_len > default_len, + "a live cost model must change the written file: inverted {inverted_len} vs default {default_len}" + ); + Ok(()) +}