From ab15859609d85c70faf3c6c31eeb7a56421ef752 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 9 Jul 2026 14:28:18 +0000 Subject: [PATCH 1/5] Introduce crate-private Candidate bookkeeping in scheme selection Pure refactor of choose_best_scheme and the sampling estimator: - estimate_compression_ratio_with_sampling now returns a SampledEstimate carrying the compressed sample array alongside its score, instead of measuring nbytes and dropping the array. - Selection bookkeeping switches from (scheme, EstimateScore) tuples to a crate-private Candidate carrying the mechanical facts about each option (scheme, score, input bytes, value count, sampled array, cascade ancestry). Candidates are still compared by ratio, equal ratios retain the existing evaluation-order winner, and the sampled array is dropped at the end of selection, so behavior is unchanged. The not-yet-read Candidate fields are the inputs a pluggable cost model will use in the follow-up commits without requiring a new Scheme API. Co-Authored-By: Claude Fable 5 Signed-off-by: Connor Tsui --- vortex-compressor/src/candidate.rs | 43 ++++++++++++++++++++++ vortex-compressor/src/compressor/sample.rs | 21 ++++++++++- vortex-compressor/src/compressor/select.rs | 43 +++++++++++++++------- vortex-compressor/src/compressor/tests.rs | 4 +- vortex-compressor/src/lib.rs | 1 + 5 files changed, 94 insertions(+), 18 deletions(-) create mode 100644 vortex-compressor/src/candidate.rs diff --git a/vortex-compressor/src/candidate.rs b/vortex-compressor/src/candidate.rs new file mode 100644 index 00000000000..7b3add41e20 --- /dev/null +++ b/vortex-compressor/src/candidate.rs @@ -0,0 +1,43 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Selection-time candidate bookkeeping. + +use vortex_array::ArrayRef; + +use crate::scheme::EstimateScore; +use crate::scheme::Scheme; +use crate::scheme::SchemeId; + +/// The mechanical facts the compressor knows about one (scheme, estimate) option during +/// scheme selection. +/// +/// Selection currently reads only `scheme` and `score`; the remaining facts are carried so +/// that a pluggable cost model can price candidates without any new [`Scheme`] API +/// (see #7697). +pub(crate) struct Candidate { + /// The scheme that produced this estimate. + pub(crate) scheme: &'static dyn Scheme, + + /// The ranked estimate: an estimated or sample-measured compression ratio, or the + /// zero-byte special case. + pub(crate) score: EstimateScore, + + /// Uncompressed size in bytes of the array under selection. + #[expect(dead_code, reason = "consumed by the cost model in a follow-up")] + pub(crate) input_nbytes: u64, + + /// Number of values in the array under selection. + #[expect(dead_code, reason = "consumed by the cost model in a follow-up")] + pub(crate) n_values: u64, + + /// The compressed sample array, when this estimate came from sampling. Its encoding tree + /// is the best available prediction of the full-array encoding tree. Dropped at the end + /// of selection. + #[expect(dead_code, reason = "consumed by the cost model in a follow-up")] + pub(crate) sampled: Option, + + /// The cascade ancestry `(scheme_id, child_index)` of the selection site. + #[expect(dead_code, reason = "consumed by the cost model in a follow-up")] + pub(crate) cascade: Vec<(SchemeId, usize)>, +} diff --git a/vortex-compressor/src/compressor/sample.rs b/vortex-compressor/src/compressor/sample.rs index ba1271d4a6d..a10eebccaaf 100644 --- a/vortex-compressor/src/compressor/sample.rs +++ b/vortex-compressor/src/compressor/sample.rs @@ -141,11 +141,25 @@ fn partition_indices(length: usize, num_partitions: u32) -> Vec<(usize, usize)> .collect() } +/// A sampling-based estimate: the ranked score together with the compressed sample array it +/// was measured on. +pub(crate) struct SampledEstimate { + /// The ranked estimate measured on the sample. + pub(crate) score: EstimateScore, + + /// The compressed sample array. Its encoding tree is the best available prediction of + /// the full-array encoding tree. + pub(crate) sampled: ArrayRef, +} + /// Estimates compression ratio by compressing a ~1% sample of the data. /// /// Creates a new [`ArrayAndStats`] for the sample so that stats are generated from the sample, not /// the full array. /// +/// Returns the compressed sample alongside its score so the selection loop can keep the +/// sample's encoding tree for the winning candidate. +/// /// # Errors /// /// Returns an error if sample compression fails. @@ -155,7 +169,7 @@ pub(super) fn estimate_compression_ratio_with_sampling( array: &ArrayRef, compress_ctx: CompressorContext, exec_ctx: &mut ExecutionCtx, -) -> VortexResult { +) -> VortexResult { let sample_array = if compress_ctx.is_sample() { array.clone() } else { @@ -186,7 +200,10 @@ pub(super) fn estimate_compression_ratio_with_sampling( trace::zero_byte_sample_result(scheme.id(), before); } - Ok(score) + Ok(SampledEstimate { + score, + sampled: compressed, + }) } #[cfg(test)] diff --git a/vortex-compressor/src/compressor/select.rs b/vortex-compressor/src/compressor/select.rs index 3c73d2d4cdb..33f1a2184d1 100644 --- a/vortex-compressor/src/compressor/select.rs +++ b/vortex-compressor/src/compressor/select.rs @@ -9,6 +9,7 @@ use vortex_error::VortexResult; use super::ROOT_SCHEME_ID; use super::sample::estimate_compression_ratio_with_sampling; use crate::CascadingCompressor; +use crate::candidate::Candidate; use crate::scheme::CompressionEstimate; use crate::scheme::CompressorContext; use crate::scheme::DeferredEstimate; @@ -38,12 +39,9 @@ impl WinnerEstimate { } } -/// Returns `true` if `score` beats the current best estimate. -fn is_better_score( - score: EstimateScore, - best: Option<&(&'static dyn Scheme, EstimateScore)>, -) -> bool { - score.is_valid() && best.is_none_or(|(_, best_score)| score.beats(*best_score)) +/// Returns `true` if `score` beats the current best candidate. +fn is_better_score(score: EstimateScore, best: Option<&Candidate>) -> bool { + score.is_valid() && best.is_none_or(|candidate| score.beats(candidate.score)) } impl CascadingCompressor { @@ -69,9 +67,20 @@ impl CascadingCompressor { compress_ctx: CompressorContext, exec_ctx: &mut ExecutionCtx, ) -> VortexResult> { - let mut best: Option<(&'static dyn Scheme, EstimateScore)> = None; + let mut best: Option = None; let mut deferred: Vec<(&'static dyn Scheme, DeferredEstimate)> = Vec::new(); + let input_nbytes = data.array().nbytes(); + let n_values = data.array().len() as u64; + let new_candidate = |scheme, score, sampled| Candidate { + scheme, + score, + input_nbytes, + n_values, + sampled, + cascade: compress_ctx.cascade_history().to_vec(), + }; + // Pass 1: evaluate every immediate verdict. Stash deferred work for pass 2. { let _verdict_pass = trace::verdict_pass_span().entered(); @@ -85,7 +94,7 @@ impl CascadingCompressor { let score = EstimateScore::FiniteCompression(ratio); if is_better_score(score, best.as_ref()) { - best = Some((scheme, score)); + best = Some(new_candidate(scheme, score, None)); } } CompressionEstimate::Deferred(deferred_estimate) => { @@ -99,10 +108,10 @@ impl CascadingCompressor { // short-circuit with `Skip` when they cannot beat it. for (scheme, deferred_estimate) in deferred { let _span = trace::scheme_eval_span(scheme.id()).entered(); - let threshold: Option = best.map(|(_, score)| score); + let threshold: Option = best.as_ref().map(|candidate| candidate.score); match deferred_estimate { DeferredEstimate::Sample => { - let score = estimate_compression_ratio_with_sampling( + let estimate = estimate_compression_ratio_with_sampling( self, scheme, data.array(), @@ -110,8 +119,12 @@ impl CascadingCompressor { exec_ctx, )?; - if is_better_score(score, best.as_ref()) { - best = Some((scheme, score)); + if is_better_score(estimate.score, best.as_ref()) { + best = Some(new_candidate( + scheme, + estimate.score, + Some(estimate.sampled), + )); } } DeferredEstimate::Callback(callback) => { @@ -124,7 +137,7 @@ impl CascadingCompressor { let score = EstimateScore::FiniteCompression(ratio); if is_better_score(score, best.as_ref()) { - best = Some((scheme, score)); + best = Some(new_candidate(scheme, score, None)); } } } @@ -132,7 +145,9 @@ impl CascadingCompressor { } } - Ok(best.map(|(scheme, score)| (scheme, WinnerEstimate::Score(score)))) + // The sampled arrays carried by candidates are dropped here; only the winning scheme + // and its score survive selection. + Ok(best.map(|candidate| (candidate.scheme, WinnerEstimate::Score(candidate.score)))) } // TODO(connor): Lots of room for optimization here. diff --git a/vortex-compressor/src/compressor/tests.rs b/vortex-compressor/src/compressor/tests.rs index 08eafa369de..a89c01040f0 100644 --- a/vortex-compressor/src/compressor/tests.rs +++ b/vortex-compressor/src/compressor/tests.rs @@ -694,13 +694,13 @@ fn sampling_uses_scheme_stats_options() -> VortexResult<()> { // Before the fix this panicked with: // "this must be present since `DictScheme` declared that we need distinct values" let mut exec_ctx = SESSION.create_execution_ctx(); - let score = estimate_compression_ratio_with_sampling( + let estimate = estimate_compression_ratio_with_sampling( &compressor, &FloatDictScheme, &array, ctx, &mut exec_ctx, )?; - assert!(matches!(score, EstimateScore::FiniteCompression(ratio) if ratio.is_finite())); + assert!(matches!(estimate.score, EstimateScore::FiniteCompression(ratio) if ratio.is_finite())); Ok(()) } diff --git a/vortex-compressor/src/lib.rs b/vortex-compressor/src/lib.rs index 55bb9b188f6..f7445c070db 100644 --- a/vortex-compressor/src/lib.rs +++ b/vortex-compressor/src/lib.rs @@ -66,6 +66,7 @@ pub mod builtins; pub mod scheme; pub mod stats; +mod candidate; mod compressor; pub use compressor::CascadingCompressor; From b0075cc206db24b4d6d467f0349dfb1f31a3ff0e Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 9 Jul 2026 14:33:23 +0000 Subject: [PATCH 2/5] Add cost module: Cost, CostModel trait, and the SizeCost default Introduce the first cost-model slice tracked by #8701: - Cost is an opaque finite f64 newtype where lower is better. Units belong to each model because costs are compared only within one model. - CostModel computes a Candidate's cost (None rejects it) and the cost of the canonical baseline that every selected candidate must strictly beat. - SizeCost assigns Cost(-ratio) to candidates and Cost(-1.0) to canonical, reproducing the historical ratio-argmax selection. SizeCost uses the exact ratio signal rather than recomputing input/ratio. That avoids IEEE division collapsing distinct ratios to the same byte estimate and changing the winner through evaluation-order tie-breaking. Candidate becomes public because it is the trait's argument and is re-exported through the cost module. Unit tests pin the validity gate, pairwise ordering, and tie behavior against the historical rules. The selection loop is wired to the new model in the next commit. Co-Authored-By: Claude Fable 5 Signed-off-by: Connor Tsui --- Cargo.lock | 1 + vortex-compressor/Cargo.toml | 1 + vortex-compressor/src/candidate.rs | 30 ++--- vortex-compressor/src/cost/mod.rs | 98 ++++++++++++++ vortex-compressor/src/cost/size.rs | 208 +++++++++++++++++++++++++++++ vortex-compressor/src/lib.rs | 1 + 6 files changed, 323 insertions(+), 16 deletions(-) create mode 100644 vortex-compressor/src/cost/mod.rs create mode 100644 vortex-compressor/src/cost/size.rs diff --git a/Cargo.lock b/Cargo.lock index 74236675f9f..871b73a08c2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9900,6 +9900,7 @@ dependencies = [ "num-traits", "parking_lot", "rand 0.10.2", + "rstest", "rustc-hash", "tracing", "vortex-array", diff --git a/vortex-compressor/Cargo.toml b/vortex-compressor/Cargo.toml index 24a2d8e40d5..efabccd5759 100644 --- a/vortex-compressor/Cargo.toml +++ b/vortex-compressor/Cargo.toml @@ -28,6 +28,7 @@ vortex-utils = { workspace = true } [dev-dependencies] divan = { workspace = true } +rstest = { workspace = true } vortex-array = { workspace = true, features = ["_test-harness"] } vortex-session = { workspace = true } diff --git a/vortex-compressor/src/candidate.rs b/vortex-compressor/src/candidate.rs index 7b3add41e20..8806e75b9a6 100644 --- a/vortex-compressor/src/candidate.rs +++ b/vortex-compressor/src/candidate.rs @@ -12,32 +12,30 @@ use crate::scheme::SchemeId; /// The mechanical facts the compressor knows about one (scheme, estimate) option during /// scheme selection. /// -/// Selection currently reads only `scheme` and `score`; the remaining facts are carried so -/// that a pluggable cost model can price candidates without any new [`Scheme`] API -/// (see #7697). -pub(crate) struct Candidate { +/// A `Candidate` is the argument to [`CostModel::cost`]: it carries everything a model can +/// price without any new [`Scheme`] API. Candidates are built by the compressor during +/// selection and dropped when selection ends. +/// +/// [`CostModel::cost`]: crate::cost::CostModel::cost +#[derive(Debug)] +pub struct Candidate { /// The scheme that produced this estimate. - pub(crate) scheme: &'static dyn Scheme, + pub scheme: &'static dyn Scheme, /// The ranked estimate: an estimated or sample-measured compression ratio, or the /// zero-byte special case. - pub(crate) score: EstimateScore, + pub score: EstimateScore, /// Uncompressed size in bytes of the array under selection. - #[expect(dead_code, reason = "consumed by the cost model in a follow-up")] - pub(crate) input_nbytes: u64, + pub input_nbytes: u64, /// Number of values in the array under selection. - #[expect(dead_code, reason = "consumed by the cost model in a follow-up")] - pub(crate) n_values: u64, + pub n_values: u64, /// The compressed sample array, when this estimate came from sampling. Its encoding tree - /// is the best available prediction of the full-array encoding tree. Dropped at the end - /// of selection. - #[expect(dead_code, reason = "consumed by the cost model in a follow-up")] - pub(crate) sampled: Option, + /// is the best available prediction of the full-array encoding tree. + pub sampled: Option, /// The cascade ancestry `(scheme_id, child_index)` of the selection site. - #[expect(dead_code, reason = "consumed by the cost model in a follow-up")] - pub(crate) cascade: Vec<(SchemeId, usize)>, + pub cascade: Vec<(SchemeId, usize)>, } diff --git a/vortex-compressor/src/cost/mod.rs b/vortex-compressor/src/cost/mod.rs new file mode 100644 index 00000000000..a3d596e8abc --- /dev/null +++ b/vortex-compressor/src/cost/mod.rs @@ -0,0 +1,98 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Pluggable cost models for scheme selection. +//! +//! A [`CostModel`] is the policy half of scheme selection: [`Scheme`]s produce mechanical +//! signals (estimated or sample-measured compression ratios, collected into a [`Candidate`]), +//! and the model prices each candidate. The compressor picks the candidate with the lowest +//! cost that (strictly) beats [`CostModel::canonical_cost`] — the price of leaving the array +//! in its canonical encoding. +//! +//! The default model is [`SizeCost`], which reproduces the compressor's historical +//! ratio-argmax selection bit-exactly. +//! +//! # What sits outside the model +//! +//! Two selection mechanisms are deliberately *not* routed through the cost model: +//! +//! - [`EstimateVerdict::AlwaysUse`] short-circuits selection entirely. It expresses semantic +//! normalization (e.g. decimal byte-parts, temporal decomposition), not a priced trade-off. +//! - The byte-acceptance gate: after the winning scheme compresses the full array, the result +//! is kept only if it is byte-wise smaller than its input. This is an axiom for **all** +//! models — compression never grows bytes. A model that prefers canonical (e.g. for speed) +//! expresses that by pricing every candidate at or above `canonical_cost`, so selection +//! returns no winner and the array stays canonical; the gate never forces a *bad* encoding, +//! only "no encoding". (The gate's `AnyScalarFn` carve-out is likewise semantic +//! denormalization, not cost.) +//! +//! # Determinism +//! +//! Cost models must be pure functions of the candidate and their own configuration: no +//! timing measurements, no I/O, no global state. The compressor's output must be a +//! deterministic function of its input and configuration. +//! +//! [`Scheme`]: crate::scheme::Scheme +//! [`EstimateVerdict::AlwaysUse`]: crate::scheme::EstimateVerdict::AlwaysUse + +mod size; + +use std::fmt::Debug; + +pub use size::SizeCost; + +pub use crate::candidate::Candidate; +use crate::stats::ArrayAndStats; + +/// An opaque, totally ordered cost. Lower is better. +/// +/// Units are defined by each [`CostModel`], not by the framework: costs are only ever +/// compared *within* one model on one compressor instance, so cross-model comparability is +/// intentionally not provided. +/// +/// Values must be finite. `NaN`/infinite prices are not representable orderings; a model +/// that cannot price a candidate rejects it by returning `None` from [`CostModel::cost`]. +#[derive(Debug, Clone, Copy, PartialEq, PartialOrd)] +pub struct Cost(f64); + +impl Cost { + /// Creates a cost from a finite value. + /// + /// # Panics + /// + /// Panics in debug builds if `value` is not finite. + pub fn new(value: f64) -> Self { + debug_assert!(value.is_finite(), "Cost must be finite, got {value}"); + Self(value) + } + + /// Returns the raw cost value. + pub fn value(self) -> f64 { + self.0 + } +} + +/// Prices [`Candidate`]s during scheme selection. +/// +/// Implementations must be deterministic pure functions of their inputs and cheap: `cost` is +/// called once per scored candidate per selection site, i.e. `O(schemes × cascade levels)` +/// times per compressed chunk. +/// +/// Selection semantics: the winner is the candidate with the minimum cost, requiring +/// strictly lower cost both to displace the running best (ties break by scheme registration +/// order) and to beat [`canonical_cost`]. Returning `None` from [`cost`] rejects the +/// candidate outright. +/// +/// See the [module docs](self) for the acceptance axioms that apply to all models. +/// +/// [`canonical_cost`]: CostModel::canonical_cost +/// [`cost`]: CostModel::cost +pub trait CostModel: Debug + Send + Sync + 'static { + /// Estimated cost of choosing this candidate. Lower is better; `None` rejects the + /// candidate. + fn cost(&self, candidate: &Candidate) -> Option; + + /// Cost of leaving the array canonical — the baseline every candidate must strictly + /// beat to be selected. + fn canonical_cost(&self, data: &ArrayAndStats, n_values: u64) -> Cost; +} diff --git a/vortex-compressor/src/cost/size.rs b/vortex-compressor/src/cost/size.rs new file mode 100644 index 00000000000..4057b628db3 --- /dev/null +++ b/vortex-compressor/src/cost/size.rs @@ -0,0 +1,208 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! The default, size-only cost model. + +use crate::candidate::Candidate; +use crate::cost::Cost; +use crate::cost::CostModel; +use crate::scheme::EstimateScore; +use crate::stats::ArrayAndStats; + +/// The default cost model: maximize estimated compression ratio. +/// +/// `SizeCost` reproduces the compressor's historical ratio-argmax selection **bit-exactly**: +/// +/// - A candidate is priced at `Cost(-ratio)` over the exact ratio signal the candidate +/// carries (an analytic estimate or a sample-measured `before / after`), so argmin cost is +/// argmax ratio with identical tie behavior. Pricing off the ratio rather than recomputing +/// bytes (`input / ratio`) matters: IEEE division could collapse strict ratio inequalities +/// into cost ties and silently flip winners via the registration-order tie-break. +/// - The canonical baseline is `Cost(-1.0)`, so `cost < canonical_cost` reproduces the +/// 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; + +impl CostModel for SizeCost { + fn cost(&self, candidate: &Candidate) -> Option { + match candidate.score { + EstimateScore::FiniteCompression(ratio) => { + (ratio.is_finite() && !ratio.is_subnormal()).then(|| Cost::new(-ratio)) + } + EstimateScore::ZeroBytes => None, + } + } + + fn canonical_cost(&self, _data: &ArrayAndStats, _n_values: u64) -> Cost { + Cost::new(-1.0) + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use vortex_array::ArrayRef; + use vortex_array::Canonical; + use vortex_array::ExecutionCtx; + use vortex_array::IntoArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::validity::Validity; + use vortex_buffer::buffer; + use vortex_error::VortexResult; + + use super::*; + use crate::CascadingCompressor; + use crate::scheme::CompressionEstimate; + use crate::scheme::CompressorContext; + use crate::scheme::EstimateVerdict; + use crate::scheme::Scheme; + use crate::stats::GenerateStatsOptions; + + #[derive(Debug)] + struct TestScheme; + + impl Scheme for TestScheme { + fn scheme_name(&self) -> &'static str { + "test.size_cost" + } + + fn matches(&self, _canonical: &Canonical) -> bool { + true + } + + fn expected_compression_ratio( + &self, + _data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> CompressionEstimate { + CompressionEstimate::Verdict(EstimateVerdict::Skip) + } + + fn compress( + &self, + _compressor: &CascadingCompressor, + _data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> VortexResult { + unreachable!("test helper should never be selected for compression") + } + } + + fn candidate(score: EstimateScore) -> Candidate { + Candidate { + scheme: &TestScheme, + score, + input_nbytes: 1024, + n_values: 256, + sampled: None, + cascade: Vec::new(), + } + } + + fn canonical_cost() -> Cost { + let array = PrimitiveArray::new(buffer![1i32, 2, 3, 4], Validity::NonNullable).into_array(); + let data = ArrayAndStats::new(array, GenerateStatsOptions::default()); + SizeCost.canonical_cost(&data, 4) + } + + /// The historical estimate-validity rule that `SizeCost` must reproduce: finite, + /// non-subnormal, and strictly better than canonical (`ratio > 1.0`). + fn legacy_is_valid(ratio: f64) -> bool { + ratio.is_finite() && !ratio.is_subnormal() && ratio > 1.0 + } + + /// Every ratio regime: clear winners, near-1.0 boundaries, exact ties, sub-1.0, + /// non-positive, and non-finite/subnormal garbage. + const RATIOS: &[f64] = &[ + 100.0, + 3.0, + 2.0, + 1.0 + f64::EPSILON, + 1.0, + 1.0 - f64::EPSILON, + 0.5, + 0.0, + -1.0, + f64::MIN_POSITIVE / 2.0, + f64::INFINITY, + f64::NEG_INFINITY, + f64::NAN, + ]; + + /// `cost < canonical_cost` must reproduce `legacy_is_valid` exactly, with non-priceable + /// ratios rejected as `None`. + #[test] + fn validity_matches_ratio_gate() { + let canonical = canonical_cost(); + for &ratio in RATIOS { + let cost = SizeCost.cost(&candidate(EstimateScore::FiniteCompression(ratio))); + let valid = cost.is_some_and(|cost| cost < canonical); + assert_eq!( + valid, + legacy_is_valid(ratio), + "validity mismatch for ratio {ratio}" + ); + } + } + + /// For every ordered pair of ratios, the cost ordering must equal the (reversed) ratio + /// ordering the selector historically used: a challenger displaces the best iff its + /// ratio is strictly greater. Equal ratios must produce equal costs (ties keep the + /// incumbent via strict `<`). + #[test] + fn cost_ordering_matches_ratio_ordering() { + let valid: Vec = RATIOS + .iter() + .copied() + .filter(|&r| legacy_is_valid(r)) + .collect(); + for &best in &valid { + for &challenger in &valid { + let best_cost = SizeCost + .cost(&candidate(EstimateScore::FiniteCompression(best))) + .expect("valid ratio must be priceable"); + let challenger_cost = SizeCost + .cost(&candidate(EstimateScore::FiniteCompression(challenger))) + .expect("valid ratio must be priceable"); + + let displaces = challenger_cost < best_cost; + assert_eq!( + displaces, + challenger > best, + "ordering mismatch: challenger {challenger} vs best {best}" + ); + assert_eq!( + challenger_cost == best_cost, + challenger == best, + "tie mismatch: challenger {challenger} vs best {best}" + ); + } + } + } + + #[rstest] + #[case::nan(f64::NAN)] + #[case::infinite(f64::INFINITY)] + #[case::neg_infinite(f64::NEG_INFINITY)] + #[case::subnormal(f64::MIN_POSITIVE / 2.0)] + fn unpriceable_ratios_are_rejected(#[case] ratio: f64) { + assert!( + SizeCost + .cost(&candidate(EstimateScore::FiniteCompression(ratio))) + .is_none() + ); + } + + #[test] + fn zero_bytes_is_rejected() { + assert!( + SizeCost + .cost(&candidate(EstimateScore::ZeroBytes)) + .is_none() + ); + } +} diff --git a/vortex-compressor/src/lib.rs b/vortex-compressor/src/lib.rs index f7445c070db..c3bd91f5f92 100644 --- a/vortex-compressor/src/lib.rs +++ b/vortex-compressor/src/lib.rs @@ -63,6 +63,7 @@ //! with a short `jq` query. pub mod builtins; +pub mod cost; pub mod scheme; pub mod stats; From 81364c8ce436b5c9b541fdb87bb052f67ddd5cf5 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 9 Jul 2026 14:37:05 +0000 Subject: [PATCH 3/5] Wire scheme selection to the cost model; selection becomes argmin cost choose_best_scheme now computes every scored candidate's cost through the compressor's CostModel. The winner is the minimum-cost candidate whose cost is strictly below CostModel::canonical_cost. Strict comparisons preserve the existing evaluation-order behavior for equal costs. The legacy comparison helpers whose policy moved into SizeCost and the selection loop are removed. CascadingCompressor gains with_cost_model(Arc) and uses SizeCost by default, so the default output remains unchanged. AlwaysUse short-circuiting and the final byte-acceptance gate remain outside the model boundary. The winner_compress trace span gains estimated_cost alongside estimated_ratio, and WinnerEstimate retains the winning cost for tracing. The golden suite has no snapshot changes under the default SizeCost. Co-Authored-By: Claude Fable 5 Signed-off-by: Connor Tsui --- vortex-compressor/src/compressor/cascade.rs | 3 +- vortex-compressor/src/compressor/mod.rs | 18 ++++- vortex-compressor/src/compressor/select.rs | 89 ++++++++++++++++----- vortex-compressor/src/compressor/tests.rs | 10 +-- vortex-compressor/src/scheme/estimate.rs | 21 ----- vortex-compressor/src/trace.rs | 7 +- 6 files changed, 97 insertions(+), 51 deletions(-) diff --git a/vortex-compressor/src/compressor/cascade.rs b/vortex-compressor/src/compressor/cascade.rs index 9dbbc611448..32ee935add2 100644 --- a/vortex-compressor/src/compressor/cascade.rs +++ b/vortex-compressor/src/compressor/cascade.rs @@ -270,7 +270,7 @@ impl CascadingCompressor { let actual_ratio = (after_nbytes != 0).then(|| before_nbytes as f64 / after_nbytes as f64); let accepted = after_nbytes < before_nbytes; - trace::record_winner_compress_result(after_nbytes, None, actual_ratio, accepted); + trace::record_winner_compress_result(after_nbytes, None, None, actual_ratio, accepted); return if accepted { Ok(compressed) @@ -311,6 +311,7 @@ impl CascadingCompressor { trace::record_winner_compress_result( after_nbytes, winner_estimate.trace_ratio(), + winner_estimate.trace_cost(), actual_ratio, accepted, ); diff --git a/vortex-compressor/src/compressor/mod.rs b/vortex-compressor/src/compressor/mod.rs index a661970950c..63f2bb189db 100644 --- a/vortex-compressor/src/compressor/mod.rs +++ b/vortex-compressor/src/compressor/mod.rs @@ -9,7 +9,11 @@ mod sample; mod select; mod structural; +use std::sync::Arc; + use crate::builtins::IntDictScheme; +use crate::cost::CostModel; +use crate::cost::SizeCost; use crate::scheme::ChildSelection; use crate::scheme::DescendantExclusion; use crate::scheme::Scheme; @@ -46,10 +50,13 @@ pub struct CascadingCompressor { /// Descendant exclusion rules for the compressor's own cascading (e.g. excluding Dict from /// list offsets). root_exclusions: Vec, + + /// The cost model pricing candidates during scheme selection. + cost_model: Arc, } impl CascadingCompressor { - /// Creates a new compressor with the given schemes. + /// Creates a new compressor with the given schemes and the default [`SizeCost`] model. /// /// Root-level exclusion rules (e.g. excluding Dict from list offsets) are built automatically. pub fn new(schemes: Vec<&'static dyn Scheme>) -> Self { @@ -63,8 +70,17 @@ impl CascadingCompressor { Self { schemes, root_exclusions, + cost_model: Arc::new(SizeCost), } } + + /// Replaces the cost model used to price candidates during scheme selection. + /// + /// The default is [`SizeCost`], which maximizes estimated compression ratio. + pub fn with_cost_model(mut self, cost_model: Arc) -> Self { + self.cost_model = cost_model; + self + } } // NB: Cascading compression logic is located in `vortex-compressor/src/compressor/cascade.rs`. diff --git a/vortex-compressor/src/compressor/select.rs b/vortex-compressor/src/compressor/select.rs index 33f1a2184d1..8ebadf29e3e 100644 --- a/vortex-compressor/src/compressor/select.rs +++ b/vortex-compressor/src/compressor/select.rs @@ -10,6 +10,7 @@ use super::ROOT_SCHEME_ID; use super::sample::estimate_compression_ratio_with_sampling; use crate::CascadingCompressor; use crate::candidate::Candidate; +use crate::cost::Cost; use crate::scheme::CompressionEstimate; use crate::scheme::CompressorContext; use crate::scheme::DeferredEstimate; @@ -25,8 +26,13 @@ use crate::trace; pub(super) enum WinnerEstimate { /// The scheme must be used immediately. AlwaysUse, - /// The scheme won by a ranked estimate. - Score(EstimateScore), + /// The scheme won by a ranked estimate, priced by the compressor's cost model. + Score { + /// The winning candidate's ranked estimate. + score: EstimateScore, + /// The winning candidate's cost under the compressor's cost model. + cost: Cost, + }, } impl WinnerEstimate { @@ -34,20 +40,30 @@ impl WinnerEstimate { pub(super) fn trace_ratio(self) -> Option { match self { Self::AlwaysUse => None, - Self::Score(score) => score.finite_ratio(), + Self::Score { score, .. } => score.finite_ratio(), } } -} -/// Returns `true` if `score` beats the current best candidate. -fn is_better_score(score: EstimateScore, best: Option<&Candidate>) -> bool { - score.is_valid() && best.is_none_or(|candidate| score.beats(candidate.score)) + /// Returns the traceable cost for the winning estimate. + pub(super) fn trace_cost(self) -> Option { + match self { + Self::AlwaysUse => None, + Self::Score { cost, .. } => Some(cost.value()), + } + } } impl CascadingCompressor { /// Calls [`expected_compression_ratio`] on each candidate and returns the winning scheme along /// with its resolved winner estimate, or `None` if no scheme beats the canonical encoding. /// + /// Each scored candidate is priced by the compressor's [`CostModel`]; the winner is the + /// candidate with the minimum cost, and only candidates pricing strictly below + /// [`CostModel::canonical_cost`] are eligible. + /// + /// [`CostModel`]: crate::cost::CostModel + /// [`CostModel::canonical_cost`]: crate::cost::CostModel::canonical_cost + /// /// Selection runs in two passes. Pass 1 evaluates every immediate /// [`CompressionEstimate::Verdict`] and tracks the running best. [`Scheme`]s returning /// [`CompressionEstimate::Deferred`] are stashed for pass 2 so that we do not make any @@ -57,7 +73,8 @@ impl CascadingCompressor { /// current best [`EstimateScore`] as an early-exit hint so the callback can return /// [`EstimateVerdict::Skip`] without doing expensive work when it cannot beat the threshold. /// - /// Ties are broken by registration order within each pass. + /// Ties are broken by registration order within each pass (displacing the running best + /// requires a strictly lower cost). /// /// [`expected_compression_ratio`]: Scheme::expected_compression_ratio pub(super) fn choose_best_scheme( @@ -67,9 +84,13 @@ impl CascadingCompressor { compress_ctx: CompressorContext, exec_ctx: &mut ExecutionCtx, ) -> VortexResult> { - let mut best: Option = None; + let mut best: Option<(Cost, Candidate)> = None; let mut deferred: Vec<(&'static dyn Scheme, DeferredEstimate)> = Vec::new(); + let canonical_cost = self + .cost_model + .canonical_cost(data, data.array().len() as u64); + let input_nbytes = data.array().nbytes(); let n_values = data.array().len() as u64; let new_candidate = |scheme, score, sampled| Candidate { @@ -92,9 +113,10 @@ impl CascadingCompressor { } CompressionEstimate::Verdict(EstimateVerdict::Ratio(ratio)) => { let score = EstimateScore::FiniteCompression(ratio); + let candidate = new_candidate(scheme, score, None); - if is_better_score(score, best.as_ref()) { - best = Some(new_candidate(scheme, score, None)); + if let Some(cost) = self.price(&candidate, canonical_cost, best.as_ref()) { + best = Some((cost, candidate)); } } CompressionEstimate::Deferred(deferred_estimate) => { @@ -108,7 +130,8 @@ impl CascadingCompressor { // short-circuit with `Skip` when they cannot beat it. for (scheme, deferred_estimate) in deferred { let _span = trace::scheme_eval_span(scheme.id()).entered(); - let threshold: Option = best.as_ref().map(|candidate| candidate.score); + let threshold: Option = + best.as_ref().map(|(_, candidate)| candidate.score); match deferred_estimate { DeferredEstimate::Sample => { let estimate = estimate_compression_ratio_with_sampling( @@ -118,13 +141,10 @@ impl CascadingCompressor { compress_ctx.clone(), exec_ctx, )?; + let candidate = new_candidate(scheme, estimate.score, Some(estimate.sampled)); - if is_better_score(estimate.score, best.as_ref()) { - best = Some(new_candidate( - scheme, - estimate.score, - Some(estimate.sampled), - )); + if let Some(cost) = self.price(&candidate, canonical_cost, best.as_ref()) { + best = Some((cost, candidate)); } } DeferredEstimate::Callback(callback) => { @@ -135,9 +155,12 @@ impl CascadingCompressor { } EstimateVerdict::Ratio(ratio) => { let score = EstimateScore::FiniteCompression(ratio); + let candidate = new_candidate(scheme, score, None); - if is_better_score(score, best.as_ref()) { - best = Some(new_candidate(scheme, score, None)); + if let Some(cost) = + self.price(&candidate, canonical_cost, best.as_ref()) + { + best = Some((cost, candidate)); } } } @@ -146,8 +169,30 @@ impl CascadingCompressor { } // The sampled arrays carried by candidates are dropped here; only the winning scheme - // and its score survive selection. - Ok(best.map(|candidate| (candidate.scheme, WinnerEstimate::Score(candidate.score)))) + // and its estimate survive selection. + Ok(best.map(|(cost, candidate)| { + ( + candidate.scheme, + WinnerEstimate::Score { + score: candidate.score, + cost, + }, + ) + })) + } + + /// Prices a candidate and returns its cost iff it becomes the new best: the model must + /// price it (`Some`), and the cost must be strictly below both the canonical baseline and + /// the best so far. Strict `<` preserves registration-order tie-breaking. + fn price( + &self, + candidate: &Candidate, + canonical_cost: Cost, + best: Option<&(Cost, Candidate)>, + ) -> Option { + let cost = self.cost_model.cost(candidate)?; + (cost < canonical_cost && best.is_none_or(|&(best_cost, _)| cost < best_cost)) + .then_some(cost) } // TODO(connor): Lots of room for optimization here. diff --git a/vortex-compressor/src/compressor/tests.rs b/vortex-compressor/src/compressor/tests.rs index a89c01040f0..d8ba5b15c73 100644 --- a/vortex-compressor/src/compressor/tests.rs +++ b/vortex-compressor/src/compressor/tests.rs @@ -378,7 +378,7 @@ fn callback_skip_is_ignored() -> VortexResult<()> { assert!(matches!( winner, - Some((scheme, WinnerEstimate::Score(EstimateScore::FiniteCompression(2.0)))) + Some((scheme, WinnerEstimate::Score { score: EstimateScore::FiniteCompression(2.0), .. })) if scheme.id() == DirectRatioScheme.id() )); Ok(()) @@ -396,7 +396,7 @@ fn callback_ratio_competes_numerically() -> VortexResult<()> { assert!(matches!( winner, - Some((scheme, WinnerEstimate::Score(EstimateScore::FiniteCompression(3.0)))) + Some((scheme, WinnerEstimate::Score { score: EstimateScore::FiniteCompression(3.0), .. })) if scheme.id() == CallbackRatioScheme.id() )); Ok(()) @@ -414,7 +414,7 @@ fn zero_byte_sample_loses_to_finite_ratio() -> VortexResult<()> { assert!(matches!( winner, - Some((scheme, WinnerEstimate::Score(EstimateScore::FiniteCompression(100.0)))) + Some((scheme, WinnerEstimate::Score { score: EstimateScore::FiniteCompression(100.0), .. })) if scheme.id() == HugeRatioScheme.id() )); Ok(()) @@ -432,7 +432,7 @@ fn finite_ratio_displaces_zero_byte_sample() -> VortexResult<()> { assert!(matches!( winner, - Some((scheme, WinnerEstimate::Score(EstimateScore::FiniteCompression(100.0)))) + Some((scheme, WinnerEstimate::Score { score: EstimateScore::FiniteCompression(100.0), .. })) if scheme.id() == HugeRatioScheme.id() )); Ok(()) @@ -646,7 +646,7 @@ fn ratio_tie_between_immediate_and_deferred_favors_immediate() -> VortexResult<( assert!(matches!( winner, - Some((scheme, WinnerEstimate::Score(EstimateScore::FiniteCompression(r)))) + Some((scheme, WinnerEstimate::Score { score: EstimateScore::FiniteCompression(r), .. })) if scheme.id() == DirectRatioScheme.id() && r == 2.0 )); Ok(()) diff --git a/vortex-compressor/src/scheme/estimate.rs b/vortex-compressor/src/scheme/estimate.rs index 20b9dc77e9b..937004d70b2 100644 --- a/vortex-compressor/src/scheme/estimate.rs +++ b/vortex-compressor/src/scheme/estimate.rs @@ -130,27 +130,6 @@ impl EstimateScore { Self::ZeroBytes => None, } } - - /// Returns whether this estimate is eligible to compete. - pub(crate) fn is_valid(self) -> bool { - match self { - Self::FiniteCompression(ratio) => { - ratio.is_finite() && !ratio.is_subnormal() && ratio > 1.0 - } - Self::ZeroBytes => false, - } - } - - /// Returns whether this estimate beats another valid estimate. - pub(crate) fn beats(self, other: Self) -> bool { - match (self, other) { - (Self::ZeroBytes, _) => false, - (Self::FiniteCompression(_), Self::ZeroBytes) => true, - (Self::FiniteCompression(ratio), Self::FiniteCompression(best_ratio)) => { - ratio > best_ratio - } - } - } } impl fmt::Debug for DeferredEstimate { diff --git a/vortex-compressor/src/trace.rs b/vortex-compressor/src/trace.rs index ee594621199..da7bc5cd8a9 100644 --- a/vortex-compressor/src/trace.rs +++ b/vortex-compressor/src/trace.rs @@ -87,7 +87,7 @@ pub(super) fn zero_byte_sample_result(scheme: SchemeId, sampled_before: u64) { /// Builds a span covering the winning scheme's full-array compression. /// /// `scheme_chosen` and `input_nbytes` are known up front. `compressed_nbytes`, -/// `estimated_ratio`, `achieved_ratio`, and `accepted` are filled in by +/// `estimated_ratio`, `estimated_cost`, `achieved_ratio`, and `accepted` are filled in by /// [`record_winner_compress_result`] once the encode completes. #[inline] pub(super) fn winner_compress_span(scheme: SchemeId, before_nbytes: u64) -> tracing::Span { @@ -98,6 +98,7 @@ pub(super) fn winner_compress_span(scheme: SchemeId, before_nbytes: u64) -> trac input_nbytes = before_nbytes, compressed_nbytes = tracing::field::Empty, estimated_ratio = tracing::field::Empty, + estimated_cost = tracing::field::Empty, achieved_ratio = tracing::field::Empty, accepted = tracing::field::Empty, ) @@ -108,6 +109,7 @@ pub(super) fn winner_compress_span(scheme: SchemeId, before_nbytes: u64) -> trac pub(super) fn record_winner_compress_result( compressed_nbytes: u64, estimated_ratio: Option, + estimated_cost: Option, achieved_ratio: Option, accepted: bool, ) { @@ -116,6 +118,9 @@ pub(super) fn record_winner_compress_result( if let Some(r) = estimated_ratio { span.record("estimated_ratio", r); } + if let Some(c) = estimated_cost { + span.record("estimated_cost", c); + } if let Some(r) = achieved_ratio { span.record("achieved_ratio", r); } From 33c9399bae8e3586c418d1e0af3d1c25f2513b3f Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Fri, 10 Jul 2026 13:31:44 +0100 Subject: [PATCH 4/5] Make schemes produce candidate evidence for cost models Replace the ratio-named scheme contract with an explicit candidate pipeline: - Scheme::evaluate returns SchemeEvaluation. - Immediate and deferred evaluations produce an opaque CandidateEstimate. - The compressor adds the scheme, input, sample, and cascade context to construct a borrowed Candidate for CostModel::cost. - Selection retains and compares only the model-defined Cost. Remove the current-winner ratio feedback path, CompressionRatioThreshold, and estimated-ratio winner tracing. Deferred callbacks no longer observe selection state, so Delta and Sequence always finish their evaluation before the active model decides whether their candidates are competitive. Any future pruning must be expressed as a cost-model capability rather than a SizeCost special case in scheme evaluation. SizeCost continues to consume the candidate's transient compression-ratio evidence and preserves the historical winners, canonical acceptance rule, and golden outputs. Tests cover custom ordering, canonical rejection, immediate and sampled candidate evidence, and a lower-ratio deferred candidate selected by a custom model. Signed-off-by: Connor Tsui --- vortex-btrblocks/src/schemes/binary/zstd.rs | 10 +- .../src/schemes/binary/zstd_buffers.rs | 10 +- vortex-btrblocks/src/schemes/decimal.rs | 9 +- vortex-btrblocks/src/schemes/float/alp.rs | 15 +- vortex-btrblocks/src/schemes/float/alprd.rs | 13 +- vortex-btrblocks/src/schemes/float/pco.rs | 10 +- vortex-btrblocks/src/schemes/float/rle.rs | 15 +- vortex-btrblocks/src/schemes/float/sparse.rs | 16 +- .../src/schemes/integer/bitpacking.rs | 13 +- vortex-btrblocks/src/schemes/integer/delta.rs | 35 +- vortex-btrblocks/src/schemes/integer/for_.rs | 18 +- vortex-btrblocks/src/schemes/integer/pco.rs | 13 +- vortex-btrblocks/src/schemes/integer/rle.rs | 15 +- .../src/schemes/integer/runend.rs | 13 +- .../src/schemes/integer/sequence.rs | 35 +- .../src/schemes/integer/sparse.rs | 20 +- .../src/schemes/integer/zigzag.rs | 15 +- vortex-btrblocks/src/schemes/string/fsst.rs | 10 +- vortex-btrblocks/src/schemes/string/onpair.rs | 10 +- vortex-btrblocks/src/schemes/string/sparse.rs | 16 +- vortex-btrblocks/src/schemes/string/zstd.rs | 10 +- .../src/schemes/string/zstd_buffers.rs | 10 +- vortex-btrblocks/src/schemes/temporal.rs | 9 +- vortex-compressor/src/builtins/dict/binary.rs | 15 +- vortex-compressor/src/builtins/dict/float.rs | 15 +- .../src/builtins/dict/integer.rs | 14 +- vortex-compressor/src/builtins/dict/string.rs | 15 +- vortex-compressor/src/candidate.rs | 95 +++- vortex-compressor/src/compressor/cascade.rs | 9 +- vortex-compressor/src/compressor/mod.rs | 9 +- vortex-compressor/src/compressor/sample.rs | 34 +- vortex-compressor/src/compressor/select.rs | 205 ++++---- vortex-compressor/src/compressor/tests.rs | 450 +++++++++++------- vortex-compressor/src/cost/mod.rs | 112 ++++- vortex-compressor/src/cost/size.rs | 81 ++-- vortex-compressor/src/lib.rs | 4 +- vortex-compressor/src/scheme/estimate.rs | 143 +++--- vortex-compressor/src/scheme/mod.rs | 54 ++- vortex-compressor/src/trace.rs | 20 +- vortex-tensor/src/encodings/l2_denorm.rs | 9 +- 40 files changed, 882 insertions(+), 742 deletions(-) diff --git a/vortex-btrblocks/src/schemes/binary/zstd.rs b/vortex-btrblocks/src/schemes/binary/zstd.rs index 83c09d24e94..c1b3f688a25 100644 --- a/vortex-btrblocks/src/schemes/binary/zstd.rs +++ b/vortex-btrblocks/src/schemes/binary/zstd.rs @@ -7,8 +7,8 @@ use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; -use vortex_compressor::scheme::CompressionEstimate; -use vortex_compressor::scheme::DeferredEstimate; +use vortex_compressor::scheme::DeferredEvaluation; +use vortex_compressor::scheme::SchemeEvaluation; use vortex_error::VortexResult; use crate::ArrayAndStats; @@ -29,13 +29,13 @@ impl Scheme for ZstdScheme { canonical.dtype().is_binary() } - fn expected_compression_ratio( + fn evaluate( &self, _data: &ArrayAndStats, _compress_ctx: CompressorContext, _exec_ctx: &mut ExecutionCtx, - ) -> CompressionEstimate { - CompressionEstimate::Deferred(DeferredEstimate::Sample) + ) -> SchemeEvaluation { + SchemeEvaluation::Deferred(DeferredEvaluation::Sample) } fn compress( diff --git a/vortex-btrblocks/src/schemes/binary/zstd_buffers.rs b/vortex-btrblocks/src/schemes/binary/zstd_buffers.rs index a84672ef860..e932ce1978a 100644 --- a/vortex-btrblocks/src/schemes/binary/zstd_buffers.rs +++ b/vortex-btrblocks/src/schemes/binary/zstd_buffers.rs @@ -7,8 +7,8 @@ use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; -use vortex_compressor::scheme::CompressionEstimate; -use vortex_compressor::scheme::DeferredEstimate; +use vortex_compressor::scheme::DeferredEvaluation; +use vortex_compressor::scheme::SchemeEvaluation; use vortex_error::VortexResult; use crate::ArrayAndStats; @@ -29,13 +29,13 @@ impl Scheme for ZstdBuffersScheme { canonical.dtype().is_binary() } - fn expected_compression_ratio( + fn evaluate( &self, _data: &ArrayAndStats, _compress_ctx: CompressorContext, _exec_ctx: &mut ExecutionCtx, - ) -> CompressionEstimate { - CompressionEstimate::Deferred(DeferredEstimate::Sample) + ) -> SchemeEvaluation { + SchemeEvaluation::Deferred(DeferredEvaluation::Sample) } fn compress( diff --git a/vortex-btrblocks/src/schemes/decimal.rs b/vortex-btrblocks/src/schemes/decimal.rs index 20c5cd80a3c..dec21ec2e2b 100644 --- a/vortex-btrblocks/src/schemes/decimal.rs +++ b/vortex-btrblocks/src/schemes/decimal.rs @@ -11,8 +11,7 @@ use vortex_array::arrays::DecimalArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::decimal::narrowed_decimal; use vortex_array::dtype::DecimalType; -use vortex_compressor::scheme::CompressionEstimate; -use vortex_compressor::scheme::EstimateVerdict; +use vortex_compressor::scheme::SchemeEvaluation; use vortex_decimal_byte_parts::DecimalByteParts; use vortex_error::VortexResult; @@ -43,14 +42,14 @@ impl Scheme for DecimalScheme { 1 } - fn expected_compression_ratio( + fn evaluate( &self, _data: &ArrayAndStats, _compress_ctx: CompressorContext, _exec_ctx: &mut ExecutionCtx, - ) -> CompressionEstimate { + ) -> SchemeEvaluation { // Decimal compression is almost always beneficial (narrowing + primitive compression). - CompressionEstimate::Verdict(EstimateVerdict::AlwaysUse) + SchemeEvaluation::AlwaysUse } fn compress( diff --git a/vortex-btrblocks/src/schemes/float/alp.rs b/vortex-btrblocks/src/schemes/float/alp.rs index b8a26abf348..20cb78b66cf 100644 --- a/vortex-btrblocks/src/schemes/float/alp.rs +++ b/vortex-btrblocks/src/schemes/float/alp.rs @@ -15,9 +15,8 @@ use vortex_array::arrays::Patched; use vortex_array::arrays::patched::use_experimental_patches; use vortex_array::arrays::primitive::PrimitiveArrayExt; use vortex_array::dtype::PType; -use vortex_compressor::scheme::CompressionEstimate; -use vortex_compressor::scheme::DeferredEstimate; -use vortex_compressor::scheme::EstimateVerdict; +use vortex_compressor::scheme::DeferredEvaluation; +use vortex_compressor::scheme::SchemeEvaluation; use vortex_error::VortexResult; use crate::ArrayAndStats; @@ -45,24 +44,24 @@ impl Scheme for ALPScheme { 1 } - fn expected_compression_ratio( + fn evaluate( &self, data: &ArrayAndStats, compress_ctx: CompressorContext, _exec_ctx: &mut ExecutionCtx, - ) -> CompressionEstimate { + ) -> SchemeEvaluation { // ALP encodes floats as integers. Without integer compression afterward, the encoded ints // are the same size. if compress_ctx.finished_cascading() { - return CompressionEstimate::Verdict(EstimateVerdict::Skip); + return SchemeEvaluation::Skip; } // We don't support ALP for f16. if data.array_as_primitive().ptype() == PType::F16 { - return CompressionEstimate::Verdict(EstimateVerdict::Skip); + return SchemeEvaluation::Skip; } - CompressionEstimate::Deferred(DeferredEstimate::Sample) + SchemeEvaluation::Deferred(DeferredEvaluation::Sample) } fn compress( diff --git a/vortex-btrblocks/src/schemes/float/alprd.rs b/vortex-btrblocks/src/schemes/float/alprd.rs index 662a47a1d9c..34880f391c8 100644 --- a/vortex-btrblocks/src/schemes/float/alprd.rs +++ b/vortex-btrblocks/src/schemes/float/alprd.rs @@ -12,9 +12,8 @@ use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::primitive::PrimitiveArrayExt; use vortex_array::dtype::PType; -use vortex_compressor::scheme::CompressionEstimate; -use vortex_compressor::scheme::DeferredEstimate; -use vortex_compressor::scheme::EstimateVerdict; +use vortex_compressor::scheme::DeferredEvaluation; +use vortex_compressor::scheme::SchemeEvaluation; use vortex_error::VortexResult; use vortex_error::vortex_panic; @@ -37,18 +36,18 @@ impl Scheme for ALPRDScheme { canonical.dtype().is_float() } - fn expected_compression_ratio( + fn evaluate( &self, data: &ArrayAndStats, _compress_ctx: CompressorContext, _exec_ctx: &mut ExecutionCtx, - ) -> CompressionEstimate { + ) -> SchemeEvaluation { // We don't support ALPRD for f16. if data.array_as_primitive().ptype() == PType::F16 { - return CompressionEstimate::Verdict(EstimateVerdict::Skip); + return SchemeEvaluation::Skip; } - CompressionEstimate::Deferred(DeferredEstimate::Sample) + SchemeEvaluation::Deferred(DeferredEvaluation::Sample) } fn compress( diff --git a/vortex-btrblocks/src/schemes/float/pco.rs b/vortex-btrblocks/src/schemes/float/pco.rs index dc9a96133d5..89a3332e709 100644 --- a/vortex-btrblocks/src/schemes/float/pco.rs +++ b/vortex-btrblocks/src/schemes/float/pco.rs @@ -7,8 +7,8 @@ use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; -use vortex_compressor::scheme::CompressionEstimate; -use vortex_compressor::scheme::DeferredEstimate; +use vortex_compressor::scheme::DeferredEvaluation; +use vortex_compressor::scheme::SchemeEvaluation; use vortex_error::VortexResult; use crate::ArrayAndStats; @@ -29,13 +29,13 @@ impl Scheme for PcoScheme { canonical.dtype().is_float() } - fn expected_compression_ratio( + fn evaluate( &self, _data: &ArrayAndStats, _compress_ctx: CompressorContext, _exec_ctx: &mut ExecutionCtx, - ) -> CompressionEstimate { - CompressionEstimate::Deferred(DeferredEstimate::Sample) + ) -> SchemeEvaluation { + SchemeEvaluation::Deferred(DeferredEvaluation::Sample) } fn compress( diff --git a/vortex-btrblocks/src/schemes/float/rle.rs b/vortex-btrblocks/src/schemes/float/rle.rs index c2683ed131c..4e494ee0edf 100644 --- a/vortex-btrblocks/src/schemes/float/rle.rs +++ b/vortex-btrblocks/src/schemes/float/rle.rs @@ -7,10 +7,9 @@ use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_compressor::scheme::AncestorExclusion; -use vortex_compressor::scheme::CompressionEstimate; -use vortex_compressor::scheme::DeferredEstimate; +use vortex_compressor::scheme::DeferredEvaluation; use vortex_compressor::scheme::DescendantExclusion; -use vortex_compressor::scheme::EstimateVerdict; +use vortex_compressor::scheme::SchemeEvaluation; use vortex_error::VortexResult; use crate::ArrayAndStats; @@ -48,22 +47,22 @@ impl Scheme for FloatRLEScheme { rle_ancestor_exclusions() } - fn expected_compression_ratio( + fn evaluate( &self, data: &ArrayAndStats, compress_ctx: CompressorContext, exec_ctx: &mut ExecutionCtx, - ) -> CompressionEstimate { + ) -> SchemeEvaluation { // RLE is only useful when we cascade it with another encoding. if compress_ctx.finished_cascading() { - return CompressionEstimate::Verdict(EstimateVerdict::Skip); + return SchemeEvaluation::Skip; } if data.float_stats(exec_ctx).average_run_length() < RUN_LENGTH_THRESHOLD { - return CompressionEstimate::Verdict(EstimateVerdict::Skip); + return SchemeEvaluation::Skip; } - CompressionEstimate::Deferred(DeferredEstimate::Sample) + SchemeEvaluation::Deferred(DeferredEvaluation::Sample) } fn compress( diff --git a/vortex-btrblocks/src/schemes/float/sparse.rs b/vortex-btrblocks/src/schemes/float/sparse.rs index 42d09e323c5..16012a69b9d 100644 --- a/vortex-btrblocks/src/schemes/float/sparse.rs +++ b/vortex-btrblocks/src/schemes/float/sparse.rs @@ -9,10 +9,10 @@ use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::primitive::PrimitiveArrayExt; +use vortex_compressor::scheme::CandidateEstimate; use vortex_compressor::scheme::ChildSelection; -use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::DescendantExclusion; -use vortex_compressor::scheme::EstimateVerdict; +use vortex_compressor::scheme::SchemeEvaluation; use vortex_error::VortexResult; use vortex_sparse::Sparse; use vortex_sparse::SparseExt as _; @@ -52,28 +52,30 @@ impl Scheme for NullDominatedSparseScheme { }] } - fn expected_compression_ratio( + fn evaluate( &self, data: &ArrayAndStats, _compress_ctx: CompressorContext, exec_ctx: &mut ExecutionCtx, - ) -> CompressionEstimate { + ) -> SchemeEvaluation { let len = data.array_len() as f64; let stats = data.float_stats(exec_ctx); let value_count = stats.value_count(); // All-null arrays should be compressed as constant instead anyways. if value_count == 0 { - return CompressionEstimate::Verdict(EstimateVerdict::Skip); + return SchemeEvaluation::Skip; } // If the majority (90%) of values is null, this will compress well. if stats.null_count() as f64 / len > 0.9 { - return CompressionEstimate::Verdict(EstimateVerdict::Ratio(len / value_count as f64)); + return SchemeEvaluation::Candidate(CandidateEstimate::from_compression_ratio( + len / value_count as f64, + )); } // Otherwise we don't go this route. - CompressionEstimate::Verdict(EstimateVerdict::Skip) + SchemeEvaluation::Skip } fn compress( diff --git a/vortex-btrblocks/src/schemes/integer/bitpacking.rs b/vortex-btrblocks/src/schemes/integer/bitpacking.rs index 2424402649e..1e06ae4469f 100644 --- a/vortex-btrblocks/src/schemes/integer/bitpacking.rs +++ b/vortex-btrblocks/src/schemes/integer/bitpacking.rs @@ -10,9 +10,8 @@ use vortex_array::IntoArray; use vortex_array::arrays::Patched; use vortex_array::arrays::patched::use_experimental_patches; use vortex_array::arrays::primitive::PrimitiveArrayExt; -use vortex_compressor::scheme::CompressionEstimate; -use vortex_compressor::scheme::DeferredEstimate; -use vortex_compressor::scheme::EstimateVerdict; +use vortex_compressor::scheme::DeferredEvaluation; +use vortex_compressor::scheme::SchemeEvaluation; use vortex_error::VortexResult; use vortex_fastlanes::BitPacked; use vortex_fastlanes::bitpack_compress::bit_width_histogram; @@ -38,20 +37,20 @@ impl Scheme for BitPackingScheme { canonical.dtype().is_int() } - fn expected_compression_ratio( + fn evaluate( &self, data: &ArrayAndStats, _compress_ctx: CompressorContext, exec_ctx: &mut ExecutionCtx, - ) -> CompressionEstimate { + ) -> SchemeEvaluation { let stats = data.integer_stats(exec_ctx); // BitPacking only works for non-negative values. if stats.erased().min_is_negative() { - return CompressionEstimate::Verdict(EstimateVerdict::Skip); + return SchemeEvaluation::Skip; } - CompressionEstimate::Deferred(DeferredEstimate::Sample) + SchemeEvaluation::Deferred(DeferredEvaluation::Sample) } fn compress( diff --git a/vortex-btrblocks/src/schemes/integer/delta.rs b/vortex-btrblocks/src/schemes/integer/delta.rs index ed087f5d845..229e8224b67 100644 --- a/vortex-btrblocks/src/schemes/integer/delta.rs +++ b/vortex-btrblocks/src/schemes/integer/delta.rs @@ -13,12 +13,12 @@ use vortex_compressor::builtins::FloatDictScheme; use vortex_compressor::builtins::IntDictScheme; use vortex_compressor::builtins::StringDictScheme; use vortex_compressor::scheme::AncestorExclusion; +use vortex_compressor::scheme::CandidateEstimate; use vortex_compressor::scheme::ChildSelection; -use vortex_compressor::scheme::CompressionEstimate; -use vortex_compressor::scheme::DeferredEstimate; +use vortex_compressor::scheme::DeferredEvaluation; use vortex_compressor::scheme::DescendantExclusion; -use vortex_compressor::scheme::EstimateScore; -use vortex_compressor::scheme::EstimateVerdict; +use vortex_compressor::scheme::ResolvedEvaluation; +use vortex_compressor::scheme::SchemeEvaluation; use vortex_error::VortexResult; use vortex_fastlanes::Delta; @@ -114,36 +114,29 @@ impl Scheme for DeltaScheme { ] } - fn expected_compression_ratio( + fn evaluate( &self, data: &ArrayAndStats, compress_ctx: CompressorContext, _exec_ctx: &mut ExecutionCtx, - ) -> CompressionEstimate { + ) -> SchemeEvaluation { // Delta only pays off if a later cascade layer (FoR/BitPacking) packs the residuals. if compress_ctx.finished_cascading() { - return CompressionEstimate::Verdict(EstimateVerdict::Skip); + return SchemeEvaluation::Skip; } // Too short to transpose into FastLanes chunks meaningfully. if data.array_len() < MIN_DELTA_LEN { - return CompressionEstimate::Verdict(EstimateVerdict::Skip); + return SchemeEvaluation::Skip; } // 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; - CompressionEstimate::Deferred(DeferredEstimate::Callback(Box::new( - move |_compressor, data, best_so_far, _ctx, exec_ctx| { + SchemeEvaluation::Deferred(DeferredEvaluation::Callback(Box::new( + move |_compressor, data, _ctx, exec_ctx| { let primitive = data.array().clone().execute::(exec_ctx)?; let full_width = primitive.ptype().bit_width() as f64; - // Delta's best case is residuals collapsing to a single bit. If even that, after - // the penalty, can't beat the incumbent, skip before doing the encode work. - let threshold = best_so_far.and_then(EstimateScore::finite_ratio); - if threshold.is_some_and(|t| full_width * DELTA_PENALTY <= t) { - return Ok(EstimateVerdict::Skip); - } - // Measure the actual FastLanes transposed-delta span. This is the lane-stride // difference that gets bit-packed, not the lag-1 difference (which the transpose // makes optimistic), so it is what truly drives the compressed size. @@ -156,14 +149,16 @@ impl Scheme for DeltaScheme { // SequenceScheme already captures more cheaply, so defer to it. let delta_bits = match span.checked_ilog2() { Some(l) => (l + 1) as f64, - None => return Ok(EstimateVerdict::Skip), + None => return Ok(ResolvedEvaluation::Skip), }; let ratio = full_width / delta_bits * DELTA_PENALTY; if ratio <= min_ratio { - return Ok(EstimateVerdict::Skip); + return Ok(ResolvedEvaluation::Skip); } - Ok(EstimateVerdict::Ratio(ratio)) + Ok(ResolvedEvaluation::Candidate( + CandidateEstimate::from_compression_ratio(ratio), + )) }, ))) } diff --git a/vortex-btrblocks/src/schemes/integer/for_.rs b/vortex-btrblocks/src/schemes/integer/for_.rs index 832ed4819fd..71c3fa8e441 100644 --- a/vortex-btrblocks/src/schemes/integer/for_.rs +++ b/vortex-btrblocks/src/schemes/integer/for_.rs @@ -13,9 +13,9 @@ use vortex_compressor::builtins::FloatDictScheme; use vortex_compressor::builtins::IntDictScheme; use vortex_compressor::builtins::StringDictScheme; use vortex_compressor::scheme::AncestorExclusion; +use vortex_compressor::scheme::CandidateEstimate; use vortex_compressor::scheme::ChildSelection; -use vortex_compressor::scheme::CompressionEstimate; -use vortex_compressor::scheme::EstimateVerdict; +use vortex_compressor::scheme::SchemeEvaluation; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_fastlanes::FoR; @@ -63,29 +63,29 @@ impl Scheme for FoRScheme { ] } - fn expected_compression_ratio( + fn evaluate( &self, data: &ArrayAndStats, compress_ctx: CompressorContext, exec_ctx: &mut ExecutionCtx, - ) -> CompressionEstimate { + ) -> SchemeEvaluation { // FoR only subtracts the min. Without further compression (e.g. BitPacking), the output is // the same size. if compress_ctx.finished_cascading() { - return CompressionEstimate::Verdict(EstimateVerdict::Skip); + return SchemeEvaluation::Skip; } let stats = data.integer_stats(exec_ctx); // Only apply when the min is not already zero. if stats.erased().min_is_zero() { - return CompressionEstimate::Verdict(EstimateVerdict::Skip); + return SchemeEvaluation::Skip; } // Difference between max and min. let for_bitwidth = match stats.erased().max_minus_min().checked_ilog2() { Some(l) => l + 1, // If max-min == 0, the we should be compressing this as a constant array. - None => return CompressionEstimate::Verdict(EstimateVerdict::Skip), + None => return SchemeEvaluation::Skip, }; // If BitPacking can be applied (only non-negative values) and FoR doesn't reduce bit width @@ -99,7 +99,7 @@ impl Scheme for FoRScheme { { let bitpack_bitwidth = max_log + 1; if for_bitwidth >= bitpack_bitwidth { - return CompressionEstimate::Verdict(EstimateVerdict::Skip); + return SchemeEvaluation::Skip; } } @@ -110,7 +110,7 @@ impl Scheme for FoRScheme { .try_into() .vortex_expect("bit width must fit in u32"); - CompressionEstimate::Verdict(EstimateVerdict::Ratio( + SchemeEvaluation::Candidate(CandidateEstimate::from_compression_ratio( full_width as f64 / for_bitwidth as f64, )) } diff --git a/vortex-btrblocks/src/schemes/integer/pco.rs b/vortex-btrblocks/src/schemes/integer/pco.rs index d7f182f588e..fe0351e1950 100644 --- a/vortex-btrblocks/src/schemes/integer/pco.rs +++ b/vortex-btrblocks/src/schemes/integer/pco.rs @@ -7,9 +7,8 @@ use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; -use vortex_compressor::scheme::CompressionEstimate; -use vortex_compressor::scheme::DeferredEstimate; -use vortex_compressor::scheme::EstimateVerdict; +use vortex_compressor::scheme::DeferredEvaluation; +use vortex_compressor::scheme::SchemeEvaluation; use vortex_error::VortexResult; use crate::ArrayAndStats; @@ -30,20 +29,20 @@ impl Scheme for PcoScheme { canonical.dtype().is_int() } - fn expected_compression_ratio( + fn evaluate( &self, data: &ArrayAndStats, _compress_ctx: CompressorContext, _exec_ctx: &mut ExecutionCtx, - ) -> CompressionEstimate { + ) -> SchemeEvaluation { use vortex_array::dtype::PType; // Pco does not support I8 or U8. if matches!(data.array_as_primitive().ptype(), PType::I8 | PType::U8) { - return CompressionEstimate::Verdict(EstimateVerdict::Skip); + return SchemeEvaluation::Skip; } - CompressionEstimate::Deferred(DeferredEstimate::Sample) + SchemeEvaluation::Deferred(DeferredEvaluation::Sample) } fn compress( diff --git a/vortex-btrblocks/src/schemes/integer/rle.rs b/vortex-btrblocks/src/schemes/integer/rle.rs index 5bc4e5296f1..55577aa13f9 100644 --- a/vortex-btrblocks/src/schemes/integer/rle.rs +++ b/vortex-btrblocks/src/schemes/integer/rle.rs @@ -10,10 +10,9 @@ use vortex_array::IntoArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::primitive::PrimitiveArrayExt; use vortex_compressor::scheme::AncestorExclusion; -use vortex_compressor::scheme::CompressionEstimate; -use vortex_compressor::scheme::DeferredEstimate; +use vortex_compressor::scheme::DeferredEvaluation; use vortex_compressor::scheme::DescendantExclusion; -use vortex_compressor::scheme::EstimateVerdict; +use vortex_compressor::scheme::SchemeEvaluation; #[cfg(feature = "unstable_encodings")] use vortex_compressor::scheme::SchemeId; use vortex_error::VortexResult; @@ -169,21 +168,21 @@ impl Scheme for IntRLEScheme { rle_ancestor_exclusions() } - fn expected_compression_ratio( + fn evaluate( &self, data: &ArrayAndStats, compress_ctx: CompressorContext, exec_ctx: &mut ExecutionCtx, - ) -> CompressionEstimate { + ) -> SchemeEvaluation { // RLE is only useful when we cascade it with another encoding. if compress_ctx.finished_cascading() { - return CompressionEstimate::Verdict(EstimateVerdict::Skip); + return SchemeEvaluation::Skip; } if data.integer_stats(exec_ctx).average_run_length() < RUN_LENGTH_THRESHOLD { - return CompressionEstimate::Verdict(EstimateVerdict::Skip); + return SchemeEvaluation::Skip; } - CompressionEstimate::Deferred(DeferredEstimate::Sample) + SchemeEvaluation::Deferred(DeferredEvaluation::Sample) } fn compress( diff --git a/vortex-btrblocks/src/schemes/integer/runend.rs b/vortex-btrblocks/src/schemes/integer/runend.rs index 175f33e7406..cb6e25f054c 100644 --- a/vortex-btrblocks/src/schemes/integer/runend.rs +++ b/vortex-btrblocks/src/schemes/integer/runend.rs @@ -14,10 +14,9 @@ use vortex_compressor::builtins::IntDictScheme; use vortex_compressor::builtins::StringDictScheme; use vortex_compressor::scheme::AncestorExclusion; use vortex_compressor::scheme::ChildSelection; -use vortex_compressor::scheme::CompressionEstimate; -use vortex_compressor::scheme::DeferredEstimate; +use vortex_compressor::scheme::DeferredEvaluation; use vortex_compressor::scheme::DescendantExclusion; -use vortex_compressor::scheme::EstimateVerdict; +use vortex_compressor::scheme::SchemeEvaluation; use vortex_error::VortexResult; use vortex_runend::RunEnd; use vortex_runend::compress::runend_encode; @@ -97,18 +96,18 @@ impl Scheme for RunEndScheme { ] } - fn expected_compression_ratio( + fn evaluate( &self, data: &ArrayAndStats, _compress_ctx: CompressorContext, exec_ctx: &mut ExecutionCtx, - ) -> CompressionEstimate { + ) -> SchemeEvaluation { // If the run length is below the threshold, drop it. if data.integer_stats(exec_ctx).average_run_length() < RUN_END_THRESHOLD { - return CompressionEstimate::Verdict(EstimateVerdict::Skip); + return SchemeEvaluation::Skip; } - CompressionEstimate::Deferred(DeferredEstimate::Sample) + SchemeEvaluation::Deferred(DeferredEvaluation::Sample) } fn compress( diff --git a/vortex-btrblocks/src/schemes/integer/sequence.rs b/vortex-btrblocks/src/schemes/integer/sequence.rs index fda9995b510..61f2b7f6c04 100644 --- a/vortex-btrblocks/src/schemes/integer/sequence.rs +++ b/vortex-btrblocks/src/schemes/integer/sequence.rs @@ -11,11 +11,11 @@ use vortex_compressor::builtins::FloatDictScheme; use vortex_compressor::builtins::IntDictScheme; use vortex_compressor::builtins::StringDictScheme; use vortex_compressor::scheme::AncestorExclusion; +use vortex_compressor::scheme::CandidateEstimate; use vortex_compressor::scheme::ChildSelection; -use vortex_compressor::scheme::CompressionEstimate; -use vortex_compressor::scheme::DeferredEstimate; -use vortex_compressor::scheme::EstimateScore; -use vortex_compressor::scheme::EstimateVerdict; +use vortex_compressor::scheme::DeferredEvaluation; +use vortex_compressor::scheme::ResolvedEvaluation; +use vortex_compressor::scheme::SchemeEvaluation; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; @@ -64,22 +64,22 @@ impl Scheme for SequenceScheme { ] } - fn expected_compression_ratio( + fn evaluate( &self, data: &ArrayAndStats, compress_ctx: CompressorContext, exec_ctx: &mut ExecutionCtx, - ) -> CompressionEstimate { + ) -> SchemeEvaluation { // It is pointless checking if a sample is a sequence since it will not correspond to the // entire array. if compress_ctx.is_sample() { - return CompressionEstimate::Verdict(EstimateVerdict::Skip); + return SchemeEvaluation::Skip; } let stats = data.integer_stats(exec_ctx); // `SequenceArray` does not support nulls. if stats.null_count() > 0 { - return CompressionEstimate::Verdict(EstimateVerdict::Skip); + return SchemeEvaluation::Skip; } // If the distinct_values_count was computed, and not all values are unique, then this @@ -88,32 +88,27 @@ impl Scheme for SequenceScheme { .distinct_count() .is_some_and(|count| count as usize != data.array_len()) { - return CompressionEstimate::Verdict(EstimateVerdict::Skip); + return SchemeEvaluation::Skip; } // TODO(connor): `sequence_encode` allocates the encoded array just to confirm feasibility. // A cheaper `is_sequence` probe would let us skip the allocation entirely. - CompressionEstimate::Deferred(DeferredEstimate::Callback(Box::new( - |_compressor, data, best_so_far, _ctx, exec_ctx| { + SchemeEvaluation::Deferred(DeferredEvaluation::Callback(Box::new( + |_compressor, data, _ctx, exec_ctx| { // `SequenceArray` stores exactly two scalars (base and multiplier), so the best // achievable compression ratio is `array_len / 2`. let compressed_size = 2usize; let max_ratio = data.array_len() as f64 / compressed_size as f64; - // If we cannot beat the best so far, then we do not want to even try sequence - // encoding the data. - let threshold = best_so_far.and_then(EstimateScore::finite_ratio); - if threshold.is_some_and(|t| max_ratio <= t) { - return Ok(EstimateVerdict::Skip); - } - // TODO(connor): We should pass this array back to the compressor in the case that // we do want to sequence encode this so that we do not need to recompress. if sequence_encode(data.array_as_primitive(), exec_ctx)?.is_none() { - return Ok(EstimateVerdict::Skip); + return Ok(ResolvedEvaluation::Skip); } // TODO(connor): Should we get the actual ratio here? - Ok(EstimateVerdict::Ratio(max_ratio)) + Ok(ResolvedEvaluation::Candidate( + CandidateEstimate::from_compression_ratio(max_ratio), + )) }, ))) } diff --git a/vortex-btrblocks/src/schemes/integer/sparse.rs b/vortex-btrblocks/src/schemes/integer/sparse.rs index e843d9d7999..dad49d21fde 100644 --- a/vortex-btrblocks/src/schemes/integer/sparse.rs +++ b/vortex-btrblocks/src/schemes/integer/sparse.rs @@ -12,10 +12,10 @@ use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::primitive::PrimitiveArrayExt; use vortex_array::scalar::Scalar; use vortex_compressor::builtins::IntDictScheme; +use vortex_compressor::scheme::CandidateEstimate; use vortex_compressor::scheme::ChildSelection; -use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::DescendantExclusion; -use vortex_compressor::scheme::EstimateVerdict; +use vortex_compressor::scheme::SchemeEvaluation; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_sparse::Sparse; @@ -77,24 +77,26 @@ impl Scheme for SparseScheme { ] } - fn expected_compression_ratio( + fn evaluate( &self, data: &ArrayAndStats, _compress_ctx: CompressorContext, exec_ctx: &mut ExecutionCtx, - ) -> CompressionEstimate { + ) -> SchemeEvaluation { let len = data.array_len() as f64; let stats = data.integer_stats(exec_ctx); let value_count = stats.value_count(); // All-null arrays should be compressed as constant instead anyways. if value_count == 0 { - return CompressionEstimate::Verdict(EstimateVerdict::Skip); + return SchemeEvaluation::Skip; } // If the majority (90%) of values is null, this will compress well. if stats.null_count() as f64 / len > 0.9 { - return CompressionEstimate::Verdict(EstimateVerdict::Ratio(len / value_count as f64)); + return SchemeEvaluation::Candidate(CandidateEstimate::from_compression_ratio( + len / value_count as f64, + )); } let (_, most_frequent_count) = stats @@ -106,18 +108,18 @@ impl Scheme for SparseScheme { // If the most frequent value is the only value, we should compress as constant instead. if most_frequent_count == value_count { - return CompressionEstimate::Verdict(EstimateVerdict::Skip); + return SchemeEvaluation::Skip; } debug_assert!(value_count > most_frequent_count); // See if the most frequent value accounts for >= 90% of the set values. let freq = most_frequent_count as f64 / value_count as f64; if freq < 0.9 { - return CompressionEstimate::Verdict(EstimateVerdict::Skip); + return SchemeEvaluation::Skip; } // We only store the positions of the non-top values. - CompressionEstimate::Verdict(EstimateVerdict::Ratio( + SchemeEvaluation::Candidate(CandidateEstimate::from_compression_ratio( value_count as f64 / (value_count - most_frequent_count) as f64, )) } diff --git a/vortex-btrblocks/src/schemes/integer/zigzag.rs b/vortex-btrblocks/src/schemes/integer/zigzag.rs index 959889dabac..98f5a3d27be 100644 --- a/vortex-btrblocks/src/schemes/integer/zigzag.rs +++ b/vortex-btrblocks/src/schemes/integer/zigzag.rs @@ -14,10 +14,9 @@ use vortex_compressor::builtins::IntDictScheme; use vortex_compressor::builtins::StringDictScheme; use vortex_compressor::scheme::AncestorExclusion; use vortex_compressor::scheme::ChildSelection; -use vortex_compressor::scheme::CompressionEstimate; -use vortex_compressor::scheme::DeferredEstimate; +use vortex_compressor::scheme::DeferredEvaluation; use vortex_compressor::scheme::DescendantExclusion; -use vortex_compressor::scheme::EstimateVerdict; +use vortex_compressor::scheme::SchemeEvaluation; use vortex_error::VortexResult; use vortex_zigzag::ZigZag; use vortex_zigzag::ZigZagArrayExt; @@ -91,25 +90,25 @@ impl Scheme for ZigZagScheme { ] } - fn expected_compression_ratio( + fn evaluate( &self, data: &ArrayAndStats, compress_ctx: CompressorContext, exec_ctx: &mut ExecutionCtx, - ) -> CompressionEstimate { + ) -> SchemeEvaluation { // ZigZag only transforms negative values to positive. Without further compression, // the output is the same size. if compress_ctx.finished_cascading() { - return CompressionEstimate::Verdict(EstimateVerdict::Skip); + return SchemeEvaluation::Skip; } let stats = data.integer_stats(exec_ctx); // ZigZag is only useful when there are negative values. if !stats.erased().min_is_negative() { - return CompressionEstimate::Verdict(EstimateVerdict::Skip); + return SchemeEvaluation::Skip; } - CompressionEstimate::Deferred(DeferredEstimate::Sample) + SchemeEvaluation::Deferred(DeferredEvaluation::Sample) } fn compress( diff --git a/vortex-btrblocks/src/schemes/string/fsst.rs b/vortex-btrblocks/src/schemes/string/fsst.rs index bd5bd010396..f41cf53d025 100644 --- a/vortex-btrblocks/src/schemes/string/fsst.rs +++ b/vortex-btrblocks/src/schemes/string/fsst.rs @@ -11,8 +11,8 @@ use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBinArray; use vortex_array::arrays::primitive::PrimitiveArrayExt; use vortex_array::arrays::varbin::VarBinArrayExt; -use vortex_compressor::scheme::CompressionEstimate; -use vortex_compressor::scheme::DeferredEstimate; +use vortex_compressor::scheme::DeferredEvaluation; +use vortex_compressor::scheme::SchemeEvaluation; use vortex_error::VortexResult; use vortex_fsst::FSST; use vortex_fsst::FSSTArrayExt; @@ -48,13 +48,13 @@ impl Scheme for FSSTScheme { 2 } - fn expected_compression_ratio( + fn evaluate( &self, _data: &ArrayAndStats, _compress_ctx: CompressorContext, _exec_ctx: &mut ExecutionCtx, - ) -> CompressionEstimate { - CompressionEstimate::Deferred(DeferredEstimate::Sample) + ) -> SchemeEvaluation { + SchemeEvaluation::Deferred(DeferredEvaluation::Sample) } fn compress( diff --git a/vortex-btrblocks/src/schemes/string/onpair.rs b/vortex-btrblocks/src/schemes/string/onpair.rs index f0a8966089c..d490c717ca9 100644 --- a/vortex-btrblocks/src/schemes/string/onpair.rs +++ b/vortex-btrblocks/src/schemes/string/onpair.rs @@ -9,8 +9,8 @@ use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::primitive::PrimitiveArrayExt; -use vortex_compressor::scheme::CompressionEstimate; -use vortex_compressor::scheme::DeferredEstimate; +use vortex_compressor::scheme::DeferredEvaluation; +use vortex_compressor::scheme::SchemeEvaluation; use vortex_compressor::scheme::SchemeId; use vortex_error::VortexResult; use vortex_onpair::DEFAULT_DICT12_CONFIG; @@ -56,13 +56,13 @@ impl Scheme for OnPairScheme { 4 } - fn expected_compression_ratio( + fn evaluate( &self, _data: &ArrayAndStats, _compress_ctx: CompressorContext, _exec_ctx: &mut ExecutionCtx, - ) -> CompressionEstimate { - CompressionEstimate::Deferred(DeferredEstimate::Sample) + ) -> SchemeEvaluation { + SchemeEvaluation::Deferred(DeferredEvaluation::Sample) } fn compress( diff --git a/vortex-btrblocks/src/schemes/string/sparse.rs b/vortex-btrblocks/src/schemes/string/sparse.rs index 750dd978030..a6b49e514a7 100644 --- a/vortex-btrblocks/src/schemes/string/sparse.rs +++ b/vortex-btrblocks/src/schemes/string/sparse.rs @@ -9,10 +9,10 @@ use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::primitive::PrimitiveArrayExt; +use vortex_compressor::scheme::CandidateEstimate; use vortex_compressor::scheme::ChildSelection; -use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::DescendantExclusion; -use vortex_compressor::scheme::EstimateVerdict; +use vortex_compressor::scheme::SchemeEvaluation; use vortex_error::VortexResult; use vortex_sparse::Sparse; use vortex_sparse::SparseExt as _; @@ -59,28 +59,30 @@ impl Scheme for NullDominatedSparseScheme { ] } - fn expected_compression_ratio( + fn evaluate( &self, data: &ArrayAndStats, _compress_ctx: CompressorContext, exec_ctx: &mut ExecutionCtx, - ) -> CompressionEstimate { + ) -> SchemeEvaluation { let len = data.array_len() as f64; let stats = data.varbinview_stats(exec_ctx); let value_count = stats.value_count(); // All-null arrays should be compressed as constant instead anyways. if value_count == 0 { - return CompressionEstimate::Verdict(EstimateVerdict::Skip); + return SchemeEvaluation::Skip; } // If the majority (90%) of values is null, this will compress well. if stats.null_count() as f64 / len > 0.9 { - return CompressionEstimate::Verdict(EstimateVerdict::Ratio(len / value_count as f64)); + return SchemeEvaluation::Candidate(CandidateEstimate::from_compression_ratio( + len / value_count as f64, + )); } // Otherwise we don't go this route. - CompressionEstimate::Verdict(EstimateVerdict::Skip) + SchemeEvaluation::Skip } fn compress( diff --git a/vortex-btrblocks/src/schemes/string/zstd.rs b/vortex-btrblocks/src/schemes/string/zstd.rs index 177acfa1543..2408d14ece9 100644 --- a/vortex-btrblocks/src/schemes/string/zstd.rs +++ b/vortex-btrblocks/src/schemes/string/zstd.rs @@ -7,8 +7,8 @@ use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; -use vortex_compressor::scheme::CompressionEstimate; -use vortex_compressor::scheme::DeferredEstimate; +use vortex_compressor::scheme::DeferredEvaluation; +use vortex_compressor::scheme::SchemeEvaluation; use vortex_error::VortexResult; use crate::ArrayAndStats; @@ -29,13 +29,13 @@ impl Scheme for ZstdScheme { canonical.dtype().is_utf8() } - fn expected_compression_ratio( + fn evaluate( &self, _data: &ArrayAndStats, _compress_ctx: CompressorContext, _exec_ctx: &mut ExecutionCtx, - ) -> CompressionEstimate { - CompressionEstimate::Deferred(DeferredEstimate::Sample) + ) -> SchemeEvaluation { + SchemeEvaluation::Deferred(DeferredEvaluation::Sample) } fn compress( diff --git a/vortex-btrblocks/src/schemes/string/zstd_buffers.rs b/vortex-btrblocks/src/schemes/string/zstd_buffers.rs index a8d16cc837f..94d51269cf1 100644 --- a/vortex-btrblocks/src/schemes/string/zstd_buffers.rs +++ b/vortex-btrblocks/src/schemes/string/zstd_buffers.rs @@ -7,8 +7,8 @@ use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; -use vortex_compressor::scheme::CompressionEstimate; -use vortex_compressor::scheme::DeferredEstimate; +use vortex_compressor::scheme::DeferredEvaluation; +use vortex_compressor::scheme::SchemeEvaluation; use vortex_error::VortexResult; use crate::ArrayAndStats; @@ -29,13 +29,13 @@ impl Scheme for ZstdBuffersScheme { canonical.dtype().is_utf8() } - fn expected_compression_ratio( + fn evaluate( &self, _data: &ArrayAndStats, _compress_ctx: CompressorContext, _exec_ctx: &mut ExecutionCtx, - ) -> CompressionEstimate { - CompressionEstimate::Deferred(DeferredEstimate::Sample) + ) -> SchemeEvaluation { + SchemeEvaluation::Deferred(DeferredEvaluation::Sample) } fn compress( diff --git a/vortex-btrblocks/src/schemes/temporal.rs b/vortex-btrblocks/src/schemes/temporal.rs index 6989b6a2aba..9eccce6d195 100644 --- a/vortex-btrblocks/src/schemes/temporal.rs +++ b/vortex-btrblocks/src/schemes/temporal.rs @@ -15,8 +15,7 @@ use vortex_array::arrays::primitive::PrimitiveArrayExt; use vortex_array::dtype::extension::Matcher; use vortex_array::extension::datetime::AnyTemporal; use vortex_array::extension::datetime::TemporalMetadata; -use vortex_compressor::scheme::CompressionEstimate; -use vortex_compressor::scheme::EstimateVerdict; +use vortex_compressor::scheme::SchemeEvaluation; use vortex_datetime_parts::DateTimeParts; use vortex_datetime_parts::TemporalParts; use vortex_datetime_parts::split_temporal; @@ -58,14 +57,14 @@ impl Scheme for TemporalScheme { 3 } - fn expected_compression_ratio( + fn evaluate( &self, _data: &ArrayAndStats, _compress_ctx: CompressorContext, _exec_ctx: &mut ExecutionCtx, - ) -> CompressionEstimate { + ) -> SchemeEvaluation { // Temporal compression (splitting into parts) is almost always beneficial. - CompressionEstimate::Verdict(EstimateVerdict::AlwaysUse) + SchemeEvaluation::AlwaysUse } fn compress( diff --git a/vortex-compressor/src/builtins/dict/binary.rs b/vortex-compressor/src/builtins/dict/binary.rs index 72b5f01141a..50bb954ebc3 100644 --- a/vortex-compressor/src/builtins/dict/binary.rs +++ b/vortex-compressor/src/builtins/dict/binary.rs @@ -22,12 +22,11 @@ use vortex_error::VortexResult; use crate::CascadingCompressor; use crate::builtins::IntDictScheme; use crate::scheme::ChildSelection; -use crate::scheme::CompressionEstimate; use crate::scheme::CompressorContext; -use crate::scheme::DeferredEstimate; +use crate::scheme::DeferredEvaluation; use crate::scheme::DescendantExclusion; -use crate::scheme::EstimateVerdict; use crate::scheme::Scheme; +use crate::scheme::SchemeEvaluation; use crate::scheme::SchemeExt; use crate::stats::ArrayAndStats; use crate::stats::GenerateStatsOptions; @@ -68,16 +67,16 @@ impl Scheme for BinaryDictScheme { }] } - fn expected_compression_ratio( + fn evaluate( &self, data: &ArrayAndStats, _compress_ctx: CompressorContext, exec_ctx: &mut ExecutionCtx, - ) -> CompressionEstimate { + ) -> SchemeEvaluation { let stats = data.varbinview_stats(exec_ctx); if stats.value_count() == 0 { - return CompressionEstimate::Verdict(EstimateVerdict::Skip); + return SchemeEvaluation::Skip; } let estimated_distinct_values_count = stats.estimated_distinct_count().vortex_expect( @@ -86,11 +85,11 @@ impl Scheme for BinaryDictScheme { // If > 50% of the values are distinct, skip dictionary scheme. if estimated_distinct_values_count > stats.value_count() / 2 { - return CompressionEstimate::Verdict(EstimateVerdict::Skip); + return SchemeEvaluation::Skip; } // Let sampling determine the expected ratio. - CompressionEstimate::Deferred(DeferredEstimate::Sample) + SchemeEvaluation::Deferred(DeferredEvaluation::Sample) } fn compress( diff --git a/vortex-compressor/src/builtins/dict/float.rs b/vortex-compressor/src/builtins/dict/float.rs index 074d3ce5e07..a909881adfe 100644 --- a/vortex-compressor/src/builtins/dict/float.rs +++ b/vortex-compressor/src/builtins/dict/float.rs @@ -26,12 +26,11 @@ use vortex_error::VortexResult; use crate::CascadingCompressor; use crate::builtins::IntDictScheme; use crate::scheme::ChildSelection; -use crate::scheme::CompressionEstimate; use crate::scheme::CompressorContext; -use crate::scheme::DeferredEstimate; +use crate::scheme::DeferredEvaluation; use crate::scheme::DescendantExclusion; -use crate::scheme::EstimateVerdict; use crate::scheme::Scheme; +use crate::scheme::SchemeEvaluation; use crate::scheme::SchemeExt; use crate::stats::ArrayAndStats; use crate::stats::FloatErasedStats; @@ -83,16 +82,16 @@ impl Scheme for FloatDictScheme { ] } - fn expected_compression_ratio( + fn evaluate( &self, data: &ArrayAndStats, _compress_ctx: CompressorContext, exec_ctx: &mut ExecutionCtx, - ) -> CompressionEstimate { + ) -> SchemeEvaluation { let stats = data.float_stats(exec_ctx); if stats.value_count() == 0 { - return CompressionEstimate::Verdict(EstimateVerdict::Skip); + return SchemeEvaluation::Skip; } let distinct_values_count = stats.distinct_count().vortex_expect( @@ -101,11 +100,11 @@ impl Scheme for FloatDictScheme { // If > 50% of the values are distinct, skip dictionary scheme. if distinct_values_count > stats.value_count() / 2 { - return CompressionEstimate::Verdict(EstimateVerdict::Skip); + return SchemeEvaluation::Skip; } // Let sampling determine the expected ratio. - CompressionEstimate::Deferred(DeferredEstimate::Sample) + SchemeEvaluation::Deferred(DeferredEvaluation::Sample) } fn compress( diff --git a/vortex-compressor/src/builtins/dict/integer.rs b/vortex-compressor/src/builtins/dict/integer.rs index 4e91bace3ac..b11c7b36f7f 100644 --- a/vortex-compressor/src/builtins/dict/integer.rs +++ b/vortex-compressor/src/builtins/dict/integer.rs @@ -23,10 +23,10 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use crate::CascadingCompressor; -use crate::scheme::CompressionEstimate; +use crate::scheme::CandidateEstimate; use crate::scheme::CompressorContext; -use crate::scheme::EstimateVerdict; use crate::scheme::Scheme; +use crate::scheme::SchemeEvaluation; use crate::scheme::SchemeExt; use crate::stats::ArrayAndStats; use crate::stats::GenerateStatsOptions; @@ -57,17 +57,17 @@ impl Scheme for IntDictScheme { 2 } - fn expected_compression_ratio( + fn evaluate( &self, data: &ArrayAndStats, _compress_ctx: CompressorContext, exec_ctx: &mut ExecutionCtx, - ) -> CompressionEstimate { + ) -> SchemeEvaluation { let bit_width = data.array_as_primitive().ptype().bit_width(); let stats = data.integer_stats(exec_ctx); if stats.value_count() == 0 { - return CompressionEstimate::Verdict(EstimateVerdict::Skip); + return SchemeEvaluation::Skip; } let distinct_values_count = stats.distinct_count().vortex_expect( @@ -76,7 +76,7 @@ impl Scheme for IntDictScheme { // If > 50% of the values are distinct, skip dictionary scheme. if distinct_values_count > stats.value_count() / 2 { - return CompressionEstimate::Verdict(EstimateVerdict::Skip); + return SchemeEvaluation::Skip; } // Ignore nulls encoding for the estimate. We only focus on values. @@ -97,7 +97,7 @@ impl Scheme for IntDictScheme { let before = stats.value_count() as usize * bit_width; - CompressionEstimate::Verdict(EstimateVerdict::Ratio( + SchemeEvaluation::Candidate(CandidateEstimate::from_compression_ratio( before as f64 / (values_size + codes_size) as f64, )) } diff --git a/vortex-compressor/src/builtins/dict/string.rs b/vortex-compressor/src/builtins/dict/string.rs index 64ed469dde9..9daff5e1c6c 100644 --- a/vortex-compressor/src/builtins/dict/string.rs +++ b/vortex-compressor/src/builtins/dict/string.rs @@ -22,12 +22,11 @@ use vortex_error::VortexResult; use crate::CascadingCompressor; use crate::builtins::IntDictScheme; use crate::scheme::ChildSelection; -use crate::scheme::CompressionEstimate; use crate::scheme::CompressorContext; -use crate::scheme::DeferredEstimate; +use crate::scheme::DeferredEvaluation; use crate::scheme::DescendantExclusion; -use crate::scheme::EstimateVerdict; use crate::scheme::Scheme; +use crate::scheme::SchemeEvaluation; use crate::scheme::SchemeExt; use crate::stats::ArrayAndStats; use crate::stats::GenerateStatsOptions; @@ -68,16 +67,16 @@ impl Scheme for StringDictScheme { }] } - fn expected_compression_ratio( + fn evaluate( &self, data: &ArrayAndStats, _compress_ctx: CompressorContext, exec_ctx: &mut ExecutionCtx, - ) -> CompressionEstimate { + ) -> SchemeEvaluation { let stats = data.varbinview_stats(exec_ctx); if stats.value_count() == 0 { - return CompressionEstimate::Verdict(EstimateVerdict::Skip); + return SchemeEvaluation::Skip; } let estimated_distinct_values_count = stats.estimated_distinct_count().vortex_expect( @@ -86,11 +85,11 @@ impl Scheme for StringDictScheme { // If > 50% of the values are distinct, skip dictionary scheme. if estimated_distinct_values_count > stats.value_count() / 2 { - return CompressionEstimate::Verdict(EstimateVerdict::Skip); + return SchemeEvaluation::Skip; } // Let sampling determine the expected ratio. - CompressionEstimate::Deferred(DeferredEstimate::Sample) + SchemeEvaluation::Deferred(DeferredEvaluation::Sample) } fn compress( diff --git a/vortex-compressor/src/candidate.rs b/vortex-compressor/src/candidate.rs index 8806e75b9a6..72680017732 100644 --- a/vortex-compressor/src/candidate.rs +++ b/vortex-compressor/src/candidate.rs @@ -1,41 +1,90 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Selection-time candidate bookkeeping. +//! Selection-time candidate view. use vortex_array::ArrayRef; -use crate::scheme::EstimateScore; +use crate::scheme::CandidateEstimate; use crate::scheme::Scheme; +use crate::scheme::SchemeExt; use crate::scheme::SchemeId; +use crate::stats::ArrayAndStats; -/// The mechanical facts the compressor knows about one (scheme, estimate) option during -/// scheme selection. +/// A borrowed view of one scheme estimate during selection. /// -/// A `Candidate` is the argument to [`CostModel::cost`]: it carries everything a model can -/// price without any new [`Scheme`] API. Candidates are built by the compressor during -/// selection and dropped when selection ends. +/// The compressor constructs a `Candidate` immediately before calling [`CostModel::cost`]. +/// It is valid only for that call: models can inspect it but cannot construct or retain it. +/// Accessors expose the stable, cheap facts available during selection without coupling the +/// public cost-model API to the selector's internal winner bookkeeping. /// /// [`CostModel::cost`]: crate::cost::CostModel::cost -#[derive(Debug)] -pub struct Candidate { - /// The scheme that produced this estimate. - pub scheme: &'static dyn Scheme, +pub struct Candidate<'a> { + /// The scheme that produced the estimate. + scheme: &'static dyn Scheme, + /// Model-independent evidence produced by the scheme. + estimate: CandidateEstimate, + /// The canonical input and its lazily generated statistics. + data: &'a ArrayAndStats, + /// The compressed sample for sampling-based estimates. + sampled: Option<&'a ArrayRef>, + /// The cascade ancestry of the selection site. + cascade: &'a [(SchemeId, usize)], +} + +impl<'a> Candidate<'a> { + /// Creates a candidate view for one cost-model call. + pub(crate) fn new( + scheme: &'static dyn Scheme, + estimate: CandidateEstimate, + data: &'a ArrayAndStats, + sampled: Option<&'a ArrayRef>, + cascade: &'a [(SchemeId, usize)], + ) -> Self { + Self { + scheme, + estimate, + data, + sampled, + cascade, + } + } + + /// Returns the ID of the scheme that produced the estimate. + pub fn scheme_id(&self) -> SchemeId { + self.scheme.id() + } + + /// Returns the model-independent evidence produced by the scheme. + pub fn estimate(&self) -> &CandidateEstimate { + &self.estimate + } - /// The ranked estimate: an estimated or sample-measured compression ratio, or the - /// zero-byte special case. - pub score: EstimateScore, + /// Returns the canonical input array being selected over. + pub fn array(&self) -> &ArrayRef { + self.data.array() + } - /// Uncompressed size in bytes of the array under selection. - pub input_nbytes: u64, + /// Returns the uncompressed size of the input array in bytes. + pub fn input_nbytes(&self) -> u64 { + self.data.array().nbytes() + } - /// Number of values in the array under selection. - pub n_values: u64, + /// Returns the number of values in the input array. + pub fn n_values(&self) -> u64 { + self.data.array().len() as u64 + } - /// The compressed sample array, when this estimate came from sampling. Its encoding tree - /// is the best available prediction of the full-array encoding tree. - pub sampled: Option, + /// Returns the compressed sample when this estimate came from sampling. + /// + /// The sample is borrowed only for the duration of the model call. Its encoding tree is + /// the best available prediction of the full-array encoding tree. + pub fn sampled(&self) -> Option<&ArrayRef> { + self.sampled + } - /// The cascade ancestry `(scheme_id, child_index)` of the selection site. - pub cascade: Vec<(SchemeId, usize)>, + /// Returns the cascade ancestry as `(scheme_id, child_index)` pairs. + pub fn cascade(&self) -> &[(SchemeId, usize)] { + self.cascade + } } diff --git a/vortex-compressor/src/compressor/cascade.rs b/vortex-compressor/src/compressor/cascade.rs index 32ee935add2..b1ab307dad7 100644 --- a/vortex-compressor/src/compressor/cascade.rs +++ b/vortex-compressor/src/compressor/cascade.rs @@ -211,8 +211,10 @@ impl CascadingCompressor { /// The main scheme-selection entry point for a single leaf array. /// /// Filters allowed schemes by [`matches`] and exclusion rules, merges their [`stats_options`] - /// into a single [`GenerateStatsOptions`], and picks the winner by estimated compression - /// ratio. + /// into a single [`GenerateStatsOptions`], and asks the configured [`CostModel`] to rank + /// the candidates. + /// + /// [`CostModel`]: crate::cost::CostModel /// /// If a winner is found and its compressed output is actually smaller, that output is /// returned. Otherwise, the original array is returned unchanged. @@ -270,7 +272,7 @@ impl CascadingCompressor { let actual_ratio = (after_nbytes != 0).then(|| before_nbytes as f64 / after_nbytes as f64); let accepted = after_nbytes < before_nbytes; - trace::record_winner_compress_result(after_nbytes, None, None, actual_ratio, accepted); + trace::record_winner_compress_result(after_nbytes, None, actual_ratio, accepted); return if accepted { Ok(compressed) @@ -310,7 +312,6 @@ impl CascadingCompressor { trace::record_winner_compress_result( after_nbytes, - winner_estimate.trace_ratio(), winner_estimate.trace_cost(), actual_ratio, accepted, diff --git a/vortex-compressor/src/compressor/mod.rs b/vortex-compressor/src/compressor/mod.rs index 63f2bb189db..fbbb838d592 100644 --- a/vortex-compressor/src/compressor/mod.rs +++ b/vortex-compressor/src/compressor/mod.rs @@ -34,7 +34,7 @@ pub(crate) const ROOT_SCHEME_ID: SchemeId = SchemeId { /// The compressor works by: /// 1. Canonicalizing input arrays to a standard representation. /// 2. Pre-filtering schemes by [`Scheme::matches`] and exclusion rules. -/// 3. Evaluating each matching scheme's compression estimate and resolving deferred work. +/// 3. Evaluating each matching scheme's candidate and resolving deferred work. /// 4. Compressing with the best scheme and verifying the result is smaller. /// /// No scheme may appear twice in a cascade chain. The compressor enforces this automatically @@ -51,7 +51,7 @@ pub struct CascadingCompressor { /// list offsets). root_exclusions: Vec, - /// The cost model pricing candidates during scheme selection. + /// The cost model used during scheme selection. cost_model: Arc, } @@ -74,9 +74,10 @@ impl CascadingCompressor { } } - /// Replaces the cost model used to price candidates during scheme selection. + /// Replaces the cost model used during scheme selection. /// - /// The default is [`SizeCost`], which maximizes estimated compression ratio. + /// The default is [`SizeCost`], which assigns lower cost to higher estimated compression + /// ratios. pub fn with_cost_model(mut self, cost_model: Arc) -> Self { self.cost_model = cost_model; self diff --git a/vortex-compressor/src/compressor/sample.rs b/vortex-compressor/src/compressor/sample.rs index a10eebccaaf..4037d8eb81e 100644 --- a/vortex-compressor/src/compressor/sample.rs +++ b/vortex-compressor/src/compressor/sample.rs @@ -15,8 +15,8 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use crate::CascadingCompressor; +use crate::scheme::CandidateEstimate; use crate::scheme::CompressorContext; -use crate::scheme::EstimateScore; use crate::scheme::Scheme; use crate::scheme::SchemeExt; use crate::stats::ArrayAndStats; @@ -141,35 +141,34 @@ fn partition_indices(length: usize, num_partitions: u32) -> Vec<(usize, usize)> .collect() } -/// A sampling-based estimate: the ranked score together with the compressed sample array it -/// was measured on. -pub(crate) struct SampledEstimate { - /// The ranked estimate measured on the sample. - pub(crate) score: EstimateScore, +/// A candidate produced through sampling. +pub(crate) struct SampledCandidate { + /// Model-independent evidence measured from the sample. + pub(crate) estimate: CandidateEstimate, /// The compressed sample array. Its encoding tree is the best available prediction of /// the full-array encoding tree. pub(crate) sampled: ArrayRef, } -/// Estimates compression ratio by compressing a ~1% sample of the data. +/// Produces candidate evidence by compressing a ~1% sample of the data. /// /// Creates a new [`ArrayAndStats`] for the sample so that stats are generated from the sample, not /// the full array. /// -/// Returns the compressed sample alongside its score so the selection loop can keep the -/// sample's encoding tree for the winning candidate. +/// Returns the compressed sample alongside evidence containing its measured ratio so the cost +/// model can inspect both while computing the candidate's cost. /// /// # Errors /// /// Returns an error if sample compression fails. -pub(super) fn estimate_compression_ratio_with_sampling( +pub(super) fn estimate_candidate_with_sampling( compressor: &CascadingCompressor, scheme: &S, array: &ArrayRef, compress_ctx: CompressorContext, exec_ctx: &mut ExecutionCtx, -) -> VortexResult { +) -> VortexResult { let sample_array = if compress_ctx.is_sample() { array.clone() } else { @@ -194,14 +193,15 @@ pub(super) fn estimate_compression_ratio_with_sampling( let after = compressed.nbytes(); let before = sample_data.array().nbytes(); - let score = EstimateScore::from_sample_sizes(before, after); - - if matches!(score, EstimateScore::ZeroBytes) { + let estimate = if after == 0 { trace::zero_byte_sample_result(scheme.id(), before); - } + CandidateEstimate::zero_bytes() + } else { + CandidateEstimate::from_compression_ratio(before as f64 / after as f64) + }; - Ok(SampledEstimate { - score, + Ok(SampledCandidate { + estimate, sampled: compressed, }) } diff --git a/vortex-compressor/src/compressor/select.rs b/vortex-compressor/src/compressor/select.rs index 8ebadf29e3e..1387f9cc8d4 100644 --- a/vortex-compressor/src/compressor/select.rs +++ b/vortex-compressor/src/compressor/select.rs @@ -1,166 +1,157 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Scheme selection: estimating each eligible scheme and choosing the winner. +//! Scheme selection: evaluating each eligible scheme and choosing the winner. use vortex_array::ExecutionCtx; use vortex_error::VortexResult; use super::ROOT_SCHEME_ID; -use super::sample::estimate_compression_ratio_with_sampling; +use super::sample::estimate_candidate_with_sampling; use crate::CascadingCompressor; use crate::candidate::Candidate; use crate::cost::Cost; -use crate::scheme::CompressionEstimate; use crate::scheme::CompressorContext; -use crate::scheme::DeferredEstimate; -use crate::scheme::EstimateScore; -use crate::scheme::EstimateVerdict; +use crate::scheme::DeferredEvaluation; +use crate::scheme::ResolvedEvaluation; use crate::scheme::Scheme; +use crate::scheme::SchemeEvaluation; use crate::scheme::SchemeExt; use crate::stats::ArrayAndStats; use crate::trace; -/// Winner estimate carried from scheme selection into result tracing. +/// The selector state retained for the current winner. +/// +/// Unlike [`Candidate`], this owns no sampled array or cascade history. +struct BestCandidate { + /// The currently selected scheme. + scheme: &'static dyn Scheme, + /// The model-defined cost used for later comparisons and tracing. + cost: Cost, +} + +/// Selection result carried into winner tracing. #[derive(Debug, Clone, Copy, PartialEq)] -pub(super) enum WinnerEstimate { +pub(super) enum SelectionOutcome { /// The scheme must be used immediately. AlwaysUse, - /// The scheme won by a ranked estimate, priced by the compressor's cost model. - Score { - /// The winning candidate's ranked estimate. - score: EstimateScore, - /// The winning candidate's cost under the compressor's cost model. - cost: Cost, - }, + /// The cost assigned to the winning candidate. + Cost(Cost), } -impl WinnerEstimate { - /// Returns the traceable numeric ratio for the winning estimate. - pub(super) fn trace_ratio(self) -> Option { - match self { - Self::AlwaysUse => None, - Self::Score { score, .. } => score.finite_ratio(), - } - } - +impl SelectionOutcome { /// Returns the traceable cost for the winning estimate. pub(super) fn trace_cost(self) -> Option { match self { Self::AlwaysUse => None, - Self::Score { cost, .. } => Some(cost.value()), + Self::Cost(cost) => Some(cost.value()), } } } impl CascadingCompressor { - /// Calls [`expected_compression_ratio`] on each candidate and returns the winning scheme along - /// with its resolved winner estimate, or `None` if no scheme beats the canonical encoding. + /// Calls [`Scheme::evaluate`] and returns the winning scheme with its selection outcome, or + /// `None` if no candidate beats the canonical encoding. /// - /// Each scored candidate is priced by the compressor's [`CostModel`]; the winner is the - /// candidate with the minimum cost, and only candidates pricing strictly below + /// The compressor's [`CostModel`] computes a cost for each candidate; the winner is the + /// candidate with the minimum cost, and only candidates with a cost strictly below /// [`CostModel::canonical_cost`] are eligible. /// - /// [`CostModel`]: crate::cost::CostModel - /// [`CostModel::canonical_cost`]: crate::cost::CostModel::canonical_cost - /// - /// Selection runs in two passes. Pass 1 evaluates every immediate - /// [`CompressionEstimate::Verdict`] and tracks the running best. [`Scheme`]s returning - /// [`CompressionEstimate::Deferred`] are stashed for pass 2 so that we do not make any - /// expensive computations if we don't have to. - /// - /// Pass 2 evaluates the deferred work and, for each [`DeferredEstimate::Callback`], passes the - /// current best [`EstimateScore`] as an early-exit hint so the callback can return - /// [`EstimateVerdict::Skip`] without doing expensive work when it cannot beat the threshold. + /// Selection runs in two passes. Pass 1 evaluates every immediate candidate and tracks the + /// running best cost. Deferred evaluations are stashed for pass 2. Every resolved candidate is + /// handed to the configured cost model; scheme evaluation never observes the current winner. /// /// Ties are broken by registration order within each pass (displacing the running best /// requires a strictly lower cost). /// - /// [`expected_compression_ratio`]: Scheme::expected_compression_ratio + /// [`CostModel`]: crate::cost::CostModel + /// [`CostModel::canonical_cost`]: crate::cost::CostModel::canonical_cost pub(super) fn choose_best_scheme( &self, schemes: &[&'static dyn Scheme], data: &ArrayAndStats, compress_ctx: CompressorContext, exec_ctx: &mut ExecutionCtx, - ) -> VortexResult> { - let mut best: Option<(Cost, Candidate)> = None; - let mut deferred: Vec<(&'static dyn Scheme, DeferredEstimate)> = Vec::new(); + ) -> VortexResult> { + let mut best: Option = None; + let mut deferred: Vec<(&'static dyn Scheme, DeferredEvaluation)> = Vec::new(); let canonical_cost = self .cost_model .canonical_cost(data, data.array().len() as u64); - let input_nbytes = data.array().nbytes(); - let n_values = data.array().len() as u64; - let new_candidate = |scheme, score, sampled| Candidate { - scheme, - score, - input_nbytes, - n_values, - sampled, - cascade: compress_ctx.cascade_history().to_vec(), - }; - - // Pass 1: evaluate every immediate verdict. Stash deferred work for pass 2. + let cascade = compress_ctx.cascade_history(); + + // Pass 1: evaluate every immediate candidate. Stash deferred work for pass 2. { - let _verdict_pass = trace::verdict_pass_span().entered(); + let _immediate_pass = trace::immediate_evaluation_pass_span().entered(); for &scheme in schemes { - match scheme.expected_compression_ratio(data, compress_ctx.clone(), exec_ctx) { - CompressionEstimate::Verdict(EstimateVerdict::Skip) => {} - CompressionEstimate::Verdict(EstimateVerdict::AlwaysUse) => { - return Ok(Some((scheme, WinnerEstimate::AlwaysUse))); + match scheme.evaluate(data, compress_ctx.clone(), exec_ctx) { + SchemeEvaluation::Skip => {} + SchemeEvaluation::AlwaysUse => { + return Ok(Some((scheme, SelectionOutcome::AlwaysUse))); } - CompressionEstimate::Verdict(EstimateVerdict::Ratio(ratio)) => { - let score = EstimateScore::FiniteCompression(ratio); - let candidate = new_candidate(scheme, score, None); - - if let Some(cost) = self.price(&candidate, canonical_cost, best.as_ref()) { - best = Some((cost, candidate)); + SchemeEvaluation::Candidate(estimate) => { + let candidate = Candidate::new(scheme, estimate, data, None, cascade); + + if let Some(cost) = self.cost_if_better( + &candidate, + canonical_cost, + best.as_ref().map(|b| b.cost), + ) { + best = Some(BestCandidate { scheme, cost }); } } - CompressionEstimate::Deferred(deferred_estimate) => { - deferred.push((scheme, deferred_estimate)); + SchemeEvaluation::Deferred(deferred_evaluation) => { + deferred.push((scheme, deferred_evaluation)); } } } } - // Pass 2: run deferred work. Callbacks receive the current best as a threshold so they can - // short-circuit with `Skip` when they cannot beat it. - for (scheme, deferred_estimate) in deferred { + // Pass 2: resolve deferred candidates without exposing current selection state. + for (scheme, deferred_evaluation) in deferred { let _span = trace::scheme_eval_span(scheme.id()).entered(); - let threshold: Option = - best.as_ref().map(|(_, candidate)| candidate.score); - match deferred_estimate { - DeferredEstimate::Sample => { - let estimate = estimate_compression_ratio_with_sampling( + match deferred_evaluation { + DeferredEvaluation::Sample => { + let sampled_candidate = estimate_candidate_with_sampling( self, scheme, data.array(), compress_ctx.clone(), exec_ctx, )?; - let candidate = new_candidate(scheme, estimate.score, Some(estimate.sampled)); - - if let Some(cost) = self.price(&candidate, canonical_cost, best.as_ref()) { - best = Some((cost, candidate)); + let candidate = Candidate::new( + scheme, + sampled_candidate.estimate, + data, + Some(&sampled_candidate.sampled), + cascade, + ); + + if let Some(cost) = self.cost_if_better( + &candidate, + canonical_cost, + best.as_ref().map(|b| b.cost), + ) { + best = Some(BestCandidate { scheme, cost }); } } - DeferredEstimate::Callback(callback) => { - match callback(self, data, threshold, compress_ctx.clone(), exec_ctx)? { - EstimateVerdict::Skip => {} - EstimateVerdict::AlwaysUse => { - return Ok(Some((scheme, WinnerEstimate::AlwaysUse))); + DeferredEvaluation::Callback(callback) => { + match callback(self, data, compress_ctx.clone(), exec_ctx)? { + ResolvedEvaluation::Skip => {} + ResolvedEvaluation::AlwaysUse => { + return Ok(Some((scheme, SelectionOutcome::AlwaysUse))); } - EstimateVerdict::Ratio(ratio) => { - let score = EstimateScore::FiniteCompression(ratio); - let candidate = new_candidate(scheme, score, None); - - if let Some(cost) = - self.price(&candidate, canonical_cost, best.as_ref()) - { - best = Some((cost, candidate)); + ResolvedEvaluation::Candidate(estimate) => { + let candidate = Candidate::new(scheme, estimate, data, None, cascade); + + if let Some(cost) = self.cost_if_better( + &candidate, + canonical_cost, + best.as_ref().map(|b| b.cost), + ) { + best = Some(BestCandidate { scheme, cost }); } } } @@ -168,30 +159,20 @@ impl CascadingCompressor { } } - // The sampled arrays carried by candidates are dropped here; only the winning scheme - // and its estimate survive selection. - Ok(best.map(|(cost, candidate)| { - ( - candidate.scheme, - WinnerEstimate::Score { - score: candidate.score, - cost, - }, - ) - })) + Ok(best.map(|candidate| (candidate.scheme, SelectionOutcome::Cost(candidate.cost)))) } - /// Prices a candidate and returns its cost iff it becomes the new best: the model must - /// price it (`Some`), and the cost must be strictly below both the canonical baseline and - /// the best so far. Strict `<` preserves registration-order tie-breaking. - fn price( + /// Computes a candidate's cost and returns it iff the candidate becomes the new best: the + /// model must return `Some`, and the cost must be strictly below both the canonical baseline + /// and the best so far. Strict `<` preserves evaluation-order tie-breaking. + fn cost_if_better( &self, - candidate: &Candidate, + candidate: &Candidate<'_>, canonical_cost: Cost, - best: Option<&(Cost, Candidate)>, + best_cost: Option, ) -> Option { let cost = self.cost_model.cost(candidate)?; - (cost < canonical_cost && best.is_none_or(|&(best_cost, _)| cost < best_cost)) + (cost < canonical_cost && best_cost.is_none_or(|best_cost| cost < best_cost)) .then_some(cost) } diff --git a/vortex-compressor/src/compressor/tests.rs b/vortex-compressor/src/compressor/tests.rs index d8ba5b15c73..ade3b208a40 100644 --- a/vortex-compressor/src/compressor/tests.rs +++ b/vortex-compressor/src/compressor/tests.rs @@ -1,9 +1,11 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::sync::Arc; use std::sync::LazyLock; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; -use parking_lot::Mutex; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; @@ -20,18 +22,21 @@ use vortex_session::VortexSession; use super::CascadingCompressor; use super::ROOT_SCHEME_ID; -use super::sample::estimate_compression_ratio_with_sampling; -use super::select::WinnerEstimate; +use super::sample::estimate_candidate_with_sampling; +use super::select::SelectionOutcome; use super::structural; use crate::builtins::FloatDictScheme; use crate::builtins::IntDictScheme; use crate::builtins::StringDictScheme; -use crate::scheme::CompressionEstimate; +use crate::candidate::Candidate; +use crate::cost::Cost; +use crate::cost::CostModel; +use crate::scheme::CandidateEstimate; use crate::scheme::CompressorContext; -use crate::scheme::DeferredEstimate; -use crate::scheme::EstimateScore; -use crate::scheme::EstimateVerdict; +use crate::scheme::DeferredEvaluation; +use crate::scheme::ResolvedEvaluation; use crate::scheme::Scheme; +use crate::scheme::SchemeEvaluation; use crate::scheme::SchemeExt; use crate::stats::ArrayAndStats; use crate::stats::GenerateStatsOptions; @@ -63,13 +68,13 @@ impl Scheme for DirectRatioScheme { matches_integer_primitive(canonical) } - fn expected_compression_ratio( + fn evaluate( &self, _data: &ArrayAndStats, _compress_ctx: CompressorContext, _exec_ctx: &mut ExecutionCtx, - ) -> CompressionEstimate { - CompressionEstimate::Verdict(EstimateVerdict::Ratio(2.0)) + ) -> SchemeEvaluation { + SchemeEvaluation::Candidate(CandidateEstimate::from_compression_ratio(2.0)) } fn compress( @@ -95,13 +100,13 @@ impl Scheme for ImmediateAlwaysUseScheme { matches_integer_primitive(canonical) } - fn expected_compression_ratio( + fn evaluate( &self, _data: &ArrayAndStats, _compress_ctx: CompressorContext, _exec_ctx: &mut ExecutionCtx, - ) -> CompressionEstimate { - CompressionEstimate::Verdict(EstimateVerdict::AlwaysUse) + ) -> SchemeEvaluation { + SchemeEvaluation::AlwaysUse } fn compress( @@ -127,14 +132,14 @@ impl Scheme for CallbackAlwaysUseScheme { matches_integer_primitive(canonical) } - fn expected_compression_ratio( + fn evaluate( &self, _data: &ArrayAndStats, _compress_ctx: CompressorContext, _exec_ctx: &mut ExecutionCtx, - ) -> CompressionEstimate { - CompressionEstimate::Deferred(DeferredEstimate::Callback(Box::new( - |_compressor, _data, _ctx, _exec_ctx, _best_so_far| Ok(EstimateVerdict::AlwaysUse), + ) -> SchemeEvaluation { + SchemeEvaluation::Deferred(DeferredEvaluation::Callback(Box::new( + |_compressor, _data, _ctx, _exec_ctx| Ok(ResolvedEvaluation::AlwaysUse), ))) } @@ -161,14 +166,14 @@ impl Scheme for CallbackSkipScheme { matches_integer_primitive(canonical) } - fn expected_compression_ratio( + fn evaluate( &self, _data: &ArrayAndStats, _compress_ctx: CompressorContext, _exec_ctx: &mut ExecutionCtx, - ) -> CompressionEstimate { - CompressionEstimate::Deferred(DeferredEstimate::Callback(Box::new( - |_compressor, _data, _ctx, _exec_ctx, _best_so_far| Ok(EstimateVerdict::Skip), + ) -> SchemeEvaluation { + SchemeEvaluation::Deferred(DeferredEvaluation::Callback(Box::new( + |_compressor, _data, _ctx, _exec_ctx| Ok(ResolvedEvaluation::Skip), ))) } @@ -195,14 +200,56 @@ impl Scheme for CallbackRatioScheme { matches_integer_primitive(canonical) } - fn expected_compression_ratio( + fn evaluate( &self, _data: &ArrayAndStats, _compress_ctx: CompressorContext, _exec_ctx: &mut ExecutionCtx, - ) -> CompressionEstimate { - CompressionEstimate::Deferred(DeferredEstimate::Callback(Box::new( - |_compressor, _data, _ctx, _exec_ctx, _best_so_far| Ok(EstimateVerdict::Ratio(3.0)), + ) -> SchemeEvaluation { + SchemeEvaluation::Deferred(DeferredEvaluation::Callback(Box::new( + |_compressor, _data, _ctx, _exec_ctx| { + Ok(ResolvedEvaluation::Candidate( + CandidateEstimate::from_compression_ratio(3.0), + )) + }, + ))) + } + + fn compress( + &self, + _compressor: &CascadingCompressor, + _data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> VortexResult { + unreachable!("test helper should never be selected for compression") + } +} + +#[derive(Debug)] +struct DeferredLowerRatioScheme; + +impl Scheme for DeferredLowerRatioScheme { + fn scheme_name(&self) -> &'static str { + "test.deferred_lower_ratio" + } + + fn matches(&self, canonical: &Canonical) -> bool { + matches_integer_primitive(canonical) + } + + fn evaluate( + &self, + _data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> SchemeEvaluation { + SchemeEvaluation::Deferred(DeferredEvaluation::Callback(Box::new( + |_compressor, _data, _compress_ctx, _exec_ctx| { + Ok(ResolvedEvaluation::Candidate( + CandidateEstimate::from_compression_ratio(1.5), + )) + }, ))) } @@ -229,13 +276,13 @@ impl Scheme for HugeRatioScheme { matches_integer_primitive(canonical) } - fn expected_compression_ratio( + fn evaluate( &self, _data: &ArrayAndStats, _compress_ctx: CompressorContext, _exec_ctx: &mut ExecutionCtx, - ) -> CompressionEstimate { - CompressionEstimate::Verdict(EstimateVerdict::Ratio(100.0)) + ) -> SchemeEvaluation { + SchemeEvaluation::Candidate(CandidateEstimate::from_compression_ratio(100.0)) } fn compress( @@ -261,13 +308,13 @@ impl Scheme for ZeroBytesSamplingScheme { matches_integer_primitive(canonical) } - fn expected_compression_ratio( + fn evaluate( &self, _data: &ArrayAndStats, _compress_ctx: CompressorContext, _exec_ctx: &mut ExecutionCtx, - ) -> CompressionEstimate { - CompressionEstimate::Deferred(DeferredEstimate::Sample) + ) -> SchemeEvaluation { + SchemeEvaluation::Deferred(DeferredEvaluation::Sample) } fn compress( @@ -281,6 +328,90 @@ impl Scheme for ZeroBytesSamplingScheme { } } +/// Assigns the lower-ratio direct scheme a lower cost than the high-ratio scheme, proving that +/// selector policy is supplied by the configured model rather than hard-coded ratio ordering. +#[derive(Debug)] +struct PreferDirectCost; + +impl CostModel for PreferDirectCost { + fn cost(&self, candidate: &Candidate<'_>) -> Option { + if candidate.scheme_id() == DirectRatioScheme.id() { + Some(Cost::new(0.0)) + } else { + Some(Cost::new(1.0)) + } + } + + fn canonical_cost(&self, _data: &ArrayAndStats, _n_values: u64) -> Cost { + Cost::new(2.0) + } +} + +#[derive(Debug)] +struct PreferDeferredLowerRatioCost; + +impl CostModel for PreferDeferredLowerRatioCost { + fn cost(&self, candidate: &Candidate<'_>) -> Option { + if candidate.scheme_id() == DeferredLowerRatioScheme.id() { + Some(Cost::new(0.0)) + } else { + Some(Cost::new(1.0)) + } + } + + fn canonical_cost(&self, _data: &ArrayAndStats, _n_values: u64) -> Cost { + Cost::new(2.0) + } +} + +#[derive(Debug)] +struct AboveCanonicalCost; + +impl CostModel for AboveCanonicalCost { + fn cost(&self, _candidate: &Candidate<'_>) -> Option { + Some(Cost::new(1.0)) + } + + fn canonical_cost(&self, _data: &ArrayAndStats, _n_values: u64) -> Cost { + Cost::new(0.0) + } +} + +#[derive(Debug)] +struct ObservingCost { + immediate_without_sample: Arc, + sampled_with_sample: Arc, +} + +impl CostModel for ObservingCost { + fn cost(&self, candidate: &Candidate<'_>) -> Option { + assert_eq!(candidate.array().len(), 4); + assert_eq!(candidate.n_values(), 4); + assert_eq!(candidate.input_nbytes(), candidate.array().nbytes()); + assert!(candidate.cascade().is_empty()); + + if candidate.scheme_id() == DirectRatioScheme.id() { + self.immediate_without_sample.store( + candidate.sampled().is_none() + && candidate.estimate().estimated_compression_ratio() == Some(2.0), + Ordering::Relaxed, + ); + } else if candidate.scheme_id() == ZeroBytesSamplingScheme.id() { + self.sampled_with_sample.store( + candidate.sampled().is_some() + && candidate.estimate().estimated_compression_ratio().is_none(), + Ordering::Relaxed, + ); + } + + None + } + + fn canonical_cost(&self, _data: &ArrayAndStats, _n_values: u64) -> Cost { + Cost::new(0.0) + } +} + #[test] fn test_self_exclusion() { let c = compressor(); @@ -342,7 +473,7 @@ fn immediate_always_use_wins_immediately() -> VortexResult<()> { assert!(matches!( winner, - Some((scheme, WinnerEstimate::AlwaysUse)) + Some((scheme, SelectionOutcome::AlwaysUse)) if scheme.id() == ImmediateAlwaysUseScheme.id() )); Ok(()) @@ -360,7 +491,7 @@ fn callback_always_use_wins_immediately() -> VortexResult<()> { assert!(matches!( winner, - Some((scheme, WinnerEstimate::AlwaysUse)) + Some((scheme, SelectionOutcome::AlwaysUse)) if scheme.id() == CallbackAlwaysUseScheme.id() )); Ok(()) @@ -378,14 +509,14 @@ fn callback_skip_is_ignored() -> VortexResult<()> { assert!(matches!( winner, - Some((scheme, WinnerEstimate::Score { score: EstimateScore::FiniteCompression(2.0), .. })) + Some((scheme, SelectionOutcome::Cost(_))) if scheme.id() == DirectRatioScheme.id() )); Ok(()) } #[test] -fn callback_ratio_competes_numerically() -> VortexResult<()> { +fn deferred_candidate_competes_by_cost() -> VortexResult<()> { let compressor = CascadingCompressor::new(vec![&DirectRatioScheme, &CallbackRatioScheme]); let schemes: [&'static dyn Scheme; 2] = [&DirectRatioScheme, &CallbackRatioScheme]; let data = estimate_test_data(); @@ -396,12 +527,94 @@ fn callback_ratio_competes_numerically() -> VortexResult<()> { assert!(matches!( winner, - Some((scheme, WinnerEstimate::Score { score: EstimateScore::FiniteCompression(3.0), .. })) + Some((scheme, SelectionOutcome::Cost(_))) if scheme.id() == CallbackRatioScheme.id() )); Ok(()) } +#[test] +fn custom_cost_model_changes_the_winner() -> VortexResult<()> { + let schemes: [&'static dyn Scheme; 2] = [&DirectRatioScheme, &HugeRatioScheme]; + let data = estimate_test_data(); + let mut exec_ctx = SESSION.create_execution_ctx(); + + let size_winner = CascadingCompressor::new(schemes.to_vec()).choose_best_scheme( + &schemes, + &data, + CompressorContext::new(), + &mut exec_ctx, + )?; + assert!(matches!( + size_winner, + Some((scheme, _)) if scheme.id() == HugeRatioScheme.id() + )); + + let custom_winner = CascadingCompressor::new(schemes.to_vec()) + .with_cost_model(Arc::new(PreferDirectCost)) + .choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; + assert!(matches!( + custom_winner, + Some((scheme, _)) if scheme.id() == DirectRatioScheme.id() + )); + Ok(()) +} + +#[test] +fn custom_model_can_choose_a_deferred_lower_ratio_candidate() -> VortexResult<()> { + let schemes: [&'static dyn Scheme; 2] = [&DirectRatioScheme, &DeferredLowerRatioScheme]; + let compressor = CascadingCompressor::new(schemes.to_vec()) + .with_cost_model(Arc::new(PreferDeferredLowerRatioCost)); + let data = estimate_test_data(); + let mut exec_ctx = SESSION.create_execution_ctx(); + + let winner = + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; + + assert!(matches!( + winner, + Some((scheme, _)) if scheme.id() == DeferredLowerRatioScheme.id() + )); + Ok(()) +} + +#[test] +fn candidate_must_beat_canonical_cost() -> VortexResult<()> { + let schemes: [&'static dyn Scheme; 1] = [&DirectRatioScheme]; + let compressor = + CascadingCompressor::new(schemes.to_vec()).with_cost_model(Arc::new(AboveCanonicalCost)); + let data = estimate_test_data(); + let mut exec_ctx = SESSION.create_execution_ctx(); + + let winner = + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; + + assert!(winner.is_none()); + Ok(()) +} + +#[test] +fn sampled_array_is_exposed_only_for_sampled_candidates() -> VortexResult<()> { + let immediate_without_sample = Arc::new(AtomicBool::new(false)); + let sampled_with_sample = Arc::new(AtomicBool::new(false)); + let model = ObservingCost { + immediate_without_sample: Arc::clone(&immediate_without_sample), + sampled_with_sample: Arc::clone(&sampled_with_sample), + }; + let schemes: [&'static dyn Scheme; 2] = [&DirectRatioScheme, &ZeroBytesSamplingScheme]; + let compressor = CascadingCompressor::new(schemes.to_vec()).with_cost_model(Arc::new(model)); + let data = estimate_test_data(); + let mut exec_ctx = SESSION.create_execution_ctx(); + + let winner = + compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; + + assert!(winner.is_none()); + assert!(immediate_without_sample.load(Ordering::Relaxed)); + assert!(sampled_with_sample.load(Ordering::Relaxed)); + Ok(()) +} + #[test] fn zero_byte_sample_loses_to_finite_ratio() -> VortexResult<()> { let compressor = CascadingCompressor::new(vec![&HugeRatioScheme, &ZeroBytesSamplingScheme]); @@ -414,7 +627,7 @@ fn zero_byte_sample_loses_to_finite_ratio() -> VortexResult<()> { assert!(matches!( winner, - Some((scheme, WinnerEstimate::Score { score: EstimateScore::FiniteCompression(100.0), .. })) + Some((scheme, SelectionOutcome::Cost(_))) if scheme.id() == HugeRatioScheme.id() )); Ok(()) @@ -432,7 +645,7 @@ fn finite_ratio_displaces_zero_byte_sample() -> VortexResult<()> { assert!(matches!( winner, - Some((scheme, WinnerEstimate::Score { score: EstimateScore::FiniteCompression(100.0), .. })) + Some((scheme, SelectionOutcome::Cost(_))) if scheme.id() == HugeRatioScheme.id() )); Ok(()) @@ -452,49 +665,6 @@ fn zero_byte_sample_alone_selects_no_scheme() -> VortexResult<()> { Ok(()) } -// Observer helper used by threshold-related tests. Captures the `best_so_far` value the -// compressor passes to its deferred callback. `OBSERVER_LOCK` serializes tests that share -// `OBSERVED_THRESHOLD` so they do not race. -static OBSERVER_LOCK: Mutex<()> = Mutex::new(()); -static OBSERVED_THRESHOLD: Mutex>> = Mutex::new(None); - -#[derive(Debug)] -struct ThresholdObservingScheme; - -impl Scheme for ThresholdObservingScheme { - fn scheme_name(&self) -> &'static str { - "test.threshold_observing" - } - - fn matches(&self, canonical: &Canonical) -> bool { - matches_integer_primitive(canonical) - } - - fn expected_compression_ratio( - &self, - _data: &ArrayAndStats, - _compress_ctx: CompressorContext, - _exec_ctx: &mut ExecutionCtx, - ) -> CompressionEstimate { - CompressionEstimate::Deferred(DeferredEstimate::Callback(Box::new( - |_compressor, _data, best_so_far, _ctx, _exec_ctx| { - *OBSERVED_THRESHOLD.lock() = Some(best_so_far); - Ok(EstimateVerdict::Skip) - }, - ))) - } - - fn compress( - &self, - _compressor: &CascadingCompressor, - _data: &ArrayAndStats, - _compress_ctx: CompressorContext, - _exec_ctx: &mut ExecutionCtx, - ) -> VortexResult { - unreachable!("test helper should never be selected for compression") - } -} - #[derive(Debug)] struct CallbackMatchingRatioScheme; @@ -507,14 +677,18 @@ impl Scheme for CallbackMatchingRatioScheme { matches_integer_primitive(canonical) } - fn expected_compression_ratio( + fn evaluate( &self, _data: &ArrayAndStats, _compress_ctx: CompressorContext, _exec_ctx: &mut ExecutionCtx, - ) -> CompressionEstimate { - CompressionEstimate::Deferred(DeferredEstimate::Callback(Box::new( - |_compressor, _data, _ctx, _exec_ctx, _best_so_far| Ok(EstimateVerdict::Ratio(2.0)), + ) -> SchemeEvaluation { + SchemeEvaluation::Deferred(DeferredEvaluation::Callback(Box::new( + |_compressor, _data, _ctx, _exec_ctx| { + Ok(ResolvedEvaluation::Candidate( + CandidateEstimate::from_compression_ratio(2.0), + )) + }, ))) } @@ -531,7 +705,7 @@ impl Scheme for CallbackMatchingRatioScheme { #[test] fn callback_always_use_overrides_pass_one_best() -> VortexResult<()> { - // `HugeRatioScheme` returns an immediate `Ratio(100.0)` in pass 1; + // `HugeRatioScheme` returns an immediate candidate in pass 1; // `CallbackAlwaysUseScheme` returns `AlwaysUse` from its deferred callback in pass 2. // The deferred `AlwaysUse` must still win. let compressor = CascadingCompressor::new(vec![&HugeRatioScheme, &CallbackAlwaysUseScheme]); @@ -544,97 +718,16 @@ fn callback_always_use_overrides_pass_one_best() -> VortexResult<()> { assert!(matches!( winner, - Some((scheme, WinnerEstimate::AlwaysUse)) + Some((scheme, SelectionOutcome::AlwaysUse)) if scheme.id() == CallbackAlwaysUseScheme.id() )); Ok(()) } #[test] -fn threshold_reflects_pass_one_best() -> VortexResult<()> { - let _guard = OBSERVER_LOCK.lock(); - *OBSERVED_THRESHOLD.lock() = None; - - let compressor = CascadingCompressor::new(vec![&DirectRatioScheme, &ThresholdObservingScheme]); - let schemes: [&'static dyn Scheme; 2] = [&DirectRatioScheme, &ThresholdObservingScheme]; - let data = estimate_test_data(); - let mut exec_ctx = SESSION.create_execution_ctx(); - - compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; - - let observed = *OBSERVED_THRESHOLD.lock(); - assert!(matches!( - observed, - Some(Some(EstimateScore::FiniteCompression(r))) if r == 2.0 - )); - Ok(()) -} - -#[test] -fn threshold_is_none_when_only_prior_is_zero_bytes() -> VortexResult<()> { - let _guard = OBSERVER_LOCK.lock(); - *OBSERVED_THRESHOLD.lock() = None; - - let compressor = - CascadingCompressor::new(vec![&ZeroBytesSamplingScheme, &ThresholdObservingScheme]); - let schemes: [&'static dyn Scheme; 2] = [&ZeroBytesSamplingScheme, &ThresholdObservingScheme]; - let data = estimate_test_data(); - let mut exec_ctx = SESSION.create_execution_ctx(); - - compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; - - // The observing callback was invoked (outer `Some`) and `best_so_far` was `None` (inner - // `None`) because the zero-byte sample is never stored as the best. - let observed = *OBSERVED_THRESHOLD.lock(); - assert_eq!(observed, Some(None)); - Ok(()) -} - -#[test] -fn threshold_is_none_when_no_prior_scheme() -> VortexResult<()> { - let _guard = OBSERVER_LOCK.lock(); - *OBSERVED_THRESHOLD.lock() = None; - - let compressor = CascadingCompressor::new(vec![&ThresholdObservingScheme]); - let schemes: [&'static dyn Scheme; 1] = [&ThresholdObservingScheme]; - let data = estimate_test_data(); - let mut exec_ctx = SESSION.create_execution_ctx(); - - compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; - - let observed = *OBSERVED_THRESHOLD.lock(); - assert_eq!(observed, Some(None)); - Ok(()) -} - -#[test] -fn threshold_updates_from_earlier_deferred_callback() -> VortexResult<()> { - let _guard = OBSERVER_LOCK.lock(); - *OBSERVED_THRESHOLD.lock() = None; - - // Both schemes are deferred. The first callback registers `Ratio(3.0)`; the second - // callback must observe it as its threshold. - let compressor = - CascadingCompressor::new(vec![&CallbackRatioScheme, &ThresholdObservingScheme]); - let schemes: [&'static dyn Scheme; 2] = [&CallbackRatioScheme, &ThresholdObservingScheme]; - let data = estimate_test_data(); - let mut exec_ctx = SESSION.create_execution_ctx(); - - compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?; - - let observed = *OBSERVED_THRESHOLD.lock(); - assert!(matches!( - observed, - Some(Some(EstimateScore::FiniteCompression(r))) if r == 3.0 - )); - Ok(()) -} - -#[test] -fn ratio_tie_between_immediate_and_deferred_favors_immediate() -> VortexResult<()> { - // Both schemes produce the same `Ratio(2.0)`, one from pass 1 (immediate) and one from - // pass 2 (deferred callback). Pass 1 locks in first, and strict `>` tie-breaking means - // the deferred callback's equal ratio cannot displace it. +fn equal_cost_between_immediate_and_deferred_favors_immediate() -> VortexResult<()> { + // Both schemes produce candidates with the same SizeCost, one from pass 1 and one from + // pass 2. Pass 1 locks in first, and strict cost comparison keeps that candidate. let compressor = CascadingCompressor::new(vec![&CallbackMatchingRatioScheme, &DirectRatioScheme]); let schemes: [&'static dyn Scheme; 2] = [&CallbackMatchingRatioScheme, &DirectRatioScheme]; @@ -646,8 +739,8 @@ fn ratio_tie_between_immediate_and_deferred_favors_immediate() -> VortexResult<( assert!(matches!( winner, - Some((scheme, WinnerEstimate::Score { score: EstimateScore::FiniteCompression(r), .. })) - if scheme.id() == DirectRatioScheme.id() && r == 2.0 + Some((scheme, SelectionOutcome::Cost(_))) + if scheme.id() == DirectRatioScheme.id() )); Ok(()) } @@ -671,7 +764,7 @@ fn all_null_array_compresses_to_constant() -> VortexResult<()> { /// Regression test for . /// -/// `estimate_compression_ratio_with_sampling` must use the *scheme's* stats options +/// `estimate_candidate_with_sampling` must use the *scheme's* stats options /// (which request distinct-value counting) rather than the context's stats options /// (which may not). With the old code this panicked inside `dictionary_encode` because /// distinct values were never computed for the sample. @@ -694,13 +787,18 @@ fn sampling_uses_scheme_stats_options() -> VortexResult<()> { // Before the fix this panicked with: // "this must be present since `DictScheme` declared that we need distinct values" let mut exec_ctx = SESSION.create_execution_ctx(); - let estimate = estimate_compression_ratio_with_sampling( + let candidate = estimate_candidate_with_sampling( &compressor, &FloatDictScheme, &array, ctx, &mut exec_ctx, )?; - assert!(matches!(estimate.score, EstimateScore::FiniteCompression(ratio) if ratio.is_finite())); + assert!( + candidate + .estimate + .estimated_compression_ratio() + .is_some_and(f64::is_finite) + ); Ok(()) } diff --git a/vortex-compressor/src/cost/mod.rs b/vortex-compressor/src/cost/mod.rs index a3d596e8abc..d75b426a3d9 100644 --- a/vortex-compressor/src/cost/mod.rs +++ b/vortex-compressor/src/cost/mod.rs @@ -3,28 +3,33 @@ //! Pluggable cost models for scheme selection. //! -//! A [`CostModel`] is the policy half of scheme selection: [`Scheme`]s produce mechanical -//! signals (estimated or sample-measured compression ratios, collected into a [`Candidate`]), -//! and the model prices each candidate. The compressor picks the candidate with the lowest -//! cost that (strictly) beats [`CostModel::canonical_cost`] — the price of leaving the array -//! in its canonical encoding. +//! A [`CostModel`] is the policy half of scheme selection: [`Scheme`]s produce model-independent +//! candidate evidence, the compressor adds selection context to construct a [`Candidate`], and the +//! model computes its cost. The compressor picks the candidate with the lowest cost that +//! (strictly) beats [`CostModel::canonical_cost`] — the cost of leaving the array in its canonical +//! encoding. //! -//! The default model is [`SizeCost`], which reproduces the compressor's historical -//! ratio-argmax selection bit-exactly. +//! The default model is [`SizeCost`], which preserves the compressor's historical candidate +//! ordering and canonical-acceptance threshold. //! //! # What sits outside the model //! -//! Two selection mechanisms are deliberately *not* routed through the cost model: +//! The initial cost-model boundary covers candidates inside +//! `CascadingCompressor::choose_best_scheme`. The following pre-existing decisions remain +//! outside it: //! -//! - [`EstimateVerdict::AlwaysUse`] short-circuits selection entirely. It expresses semantic -//! normalization (e.g. decimal byte-parts, temporal decomposition), not a priced trade-off. +//! - Constant-array handling occurs before scheme selection. +//! - [`SchemeEvaluation::AlwaysUse`] is a forced-selection path that short-circuits candidate cost +//! comparison. //! - The byte-acceptance gate: after the winning scheme compresses the full array, the result //! is kept only if it is byte-wise smaller than its input. This is an axiom for **all** //! models — compression never grows bytes. A model that prefers canonical (e.g. for speed) -//! expresses that by pricing every candidate at or above `canonical_cost`, so selection +//! expresses that by assigning every candidate a cost at or above `canonical_cost`, so selection //! returns no winner and the array stays canonical; the gate never forces a *bad* encoding, //! only "no encoding". (The gate's `AnyScalarFn` carve-out is likewise semantic //! denormalization, not cost.) +//! - Extension arrays separately compare scheme-based compression with compression of their +//! storage array by byte size. //! //! # Determinism //! @@ -33,10 +38,11 @@ //! deterministic function of its input and configuration. //! //! [`Scheme`]: crate::scheme::Scheme -//! [`EstimateVerdict::AlwaysUse`]: crate::scheme::EstimateVerdict::AlwaysUse +//! [`SchemeEvaluation::AlwaysUse`]: crate::scheme::SchemeEvaluation::AlwaysUse mod size; +use std::cmp::Ordering; use std::fmt::Debug; pub use size::SizeCost; @@ -50,9 +56,9 @@ use crate::stats::ArrayAndStats; /// compared *within* one model on one compressor instance, so cross-model comparability is /// intentionally not provided. /// -/// Values must be finite. `NaN`/infinite prices are not representable orderings; a model -/// that cannot price a candidate rejects it by returning `None` from [`CostModel::cost`]. -#[derive(Debug, Clone, Copy, PartialEq, PartialOrd)] +/// Values must be finite. `NaN` and infinite costs are not representable orderings; a model +/// that cannot compute a candidate's cost rejects it by returning `None` from [`CostModel::cost`]. +#[derive(Debug, Clone, Copy)] pub struct Cost(f64); impl Cost { @@ -60,10 +66,10 @@ impl Cost { /// /// # Panics /// - /// Panics in debug builds if `value` is not finite. + /// Panics if `value` is not finite. pub fn new(value: f64) -> Self { - debug_assert!(value.is_finite(), "Cost must be finite, got {value}"); - Self(value) + assert!(value.is_finite(), "Cost must be finite, got {value}"); + Self(if value == 0.0 { 0.0 } else { value }) } /// Returns the raw cost value. @@ -72,16 +78,37 @@ impl Cost { } } -/// Prices [`Candidate`]s during scheme selection. +impl PartialEq for Cost { + fn eq(&self, other: &Self) -> bool { + self.cmp(other) == Ordering::Equal + } +} + +impl Eq for Cost {} + +impl PartialOrd for Cost { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +impl Ord for Cost { + fn cmp(&self, other: &Self) -> Ordering { + self.0.total_cmp(&other.0) + } +} + +/// Computes costs for [`Candidate`]s during scheme selection. /// /// Implementations must be deterministic pure functions of their inputs and cheap: `cost` is -/// called once per scored candidate per selection site, i.e. `O(schemes × cascade levels)` +/// called once per candidate per selection site, i.e. `O(schemes × cascade levels)` /// times per compressed chunk. /// -/// Selection semantics: the winner is the candidate with the minimum cost, requiring -/// strictly lower cost both to displace the running best (ties break by scheme registration -/// order) and to beat [`canonical_cost`]. Returning `None` from [`cost`] rejects the -/// candidate outright. +/// Selection semantics: the winner is the candidate with the minimum cost, requiring strictly +/// lower cost both to displace the running best and to beat [`canonical_cost`]. Equal-cost ties +/// favor the candidate evaluated first: scheme registration order within each selection pass, +/// and immediate candidates over deferred candidates across passes. Returning `None` from [`cost`] +/// rejects the candidate outright. /// /// See the [module docs](self) for the acceptance axioms that apply to all models. /// @@ -90,9 +117,44 @@ impl Cost { pub trait CostModel: Debug + Send + Sync + 'static { /// Estimated cost of choosing this candidate. Lower is better; `None` rejects the /// candidate. - fn cost(&self, candidate: &Candidate) -> Option; + fn cost(&self, candidate: &Candidate<'_>) -> Option; /// Cost of leaving the array canonical — the baseline every candidate must strictly /// beat to be selected. fn canonical_cost(&self, data: &ArrayAndStats, n_values: u64) -> Cost; } + +#[cfg(test)] +mod tests { + use std::cmp::Ordering; + + use super::Cost; + + #[test] + fn rejects_non_finite_values() { + for value in [f64::NAN, f64::INFINITY, f64::NEG_INFINITY] { + assert!( + std::panic::catch_unwind(|| Cost::new(value)).is_err(), + "non-finite cost {value} was accepted" + ); + } + } + + #[test] + fn signed_zero_has_one_representation() { + let negative = Cost::new(-0.0); + let positive = Cost::new(0.0); + + assert_eq!(negative, positive); + assert_eq!(negative.cmp(&positive), Ordering::Equal); + assert_eq!(negative.value().to_bits(), 0.0f64.to_bits()); + } + + #[test] + fn costs_have_a_total_order() { + let mut costs = [Cost::new(3.0), Cost::new(-2.0), Cost::new(0.0)]; + costs.sort(); + + assert_eq!(costs.map(Cost::value), [-2.0, 0.0, 3.0]); + } +} diff --git a/vortex-compressor/src/cost/size.rs b/vortex-compressor/src/cost/size.rs index 4057b628db3..209dfdc8ab6 100644 --- a/vortex-compressor/src/cost/size.rs +++ b/vortex-compressor/src/cost/size.rs @@ -6,16 +6,15 @@ use crate::candidate::Candidate; use crate::cost::Cost; use crate::cost::CostModel; -use crate::scheme::EstimateScore; use crate::stats::ArrayAndStats; /// The default cost model: maximize estimated compression ratio. /// -/// `SizeCost` reproduces the compressor's historical ratio-argmax selection **bit-exactly**: +/// `SizeCost` preserves the compressor's historical ratio-argmax ordering: /// -/// - A candidate is priced at `Cost(-ratio)` over the exact ratio signal the candidate -/// carries (an analytic estimate or a sample-measured `before / after`), so argmin cost is -/// argmax ratio with identical tie behavior. Pricing off the ratio rather than recomputing +/// - A candidate's cost is `Cost(-ratio)` over the exact ratio signal the candidate carries +/// (an analytic estimate or a sample-measured `before / after`), so argmin cost is argmax +/// ratio with identical tie behavior. Computing cost from the ratio rather than recomputing /// bytes (`input / ratio`) matters: IEEE division could collapse strict ratio inequalities /// into cost ties and silently flip winners via the registration-order tie-break. /// - The canonical baseline is `Cost(-1.0)`, so `cost < canonical_cost` reproduces the @@ -26,13 +25,9 @@ use crate::stats::ArrayAndStats; pub struct SizeCost; impl CostModel for SizeCost { - fn cost(&self, candidate: &Candidate) -> Option { - match candidate.score { - EstimateScore::FiniteCompression(ratio) => { - (ratio.is_finite() && !ratio.is_subnormal()).then(|| Cost::new(-ratio)) - } - EstimateScore::ZeroBytes => None, - } + fn cost(&self, candidate: &Candidate<'_>) -> Option { + let ratio = candidate.estimate().estimated_compression_ratio()?; + (ratio.is_finite() && !ratio.is_subnormal()).then(|| Cost::new(-ratio)) } fn canonical_cost(&self, _data: &ArrayAndStats, _n_values: u64) -> Cost { @@ -54,10 +49,10 @@ mod tests { use super::*; use crate::CascadingCompressor; - use crate::scheme::CompressionEstimate; + use crate::scheme::CandidateEstimate; use crate::scheme::CompressorContext; - use crate::scheme::EstimateVerdict; use crate::scheme::Scheme; + use crate::scheme::SchemeEvaluation; use crate::stats::GenerateStatsOptions; #[derive(Debug)] @@ -72,13 +67,13 @@ mod tests { true } - fn expected_compression_ratio( + fn evaluate( &self, _data: &ArrayAndStats, _compress_ctx: CompressorContext, _exec_ctx: &mut ExecutionCtx, - ) -> CompressionEstimate { - CompressionEstimate::Verdict(EstimateVerdict::Skip) + ) -> SchemeEvaluation { + SchemeEvaluation::Skip } fn compress( @@ -92,20 +87,22 @@ mod tests { } } - fn candidate(score: EstimateScore) -> Candidate { - Candidate { - scheme: &TestScheme, - score, - input_nbytes: 1024, - n_values: 256, - sampled: None, - cascade: Vec::new(), - } + fn test_data() -> ArrayAndStats { + let array = PrimitiveArray::new(buffer![1i32, 2, 3, 4], Validity::NonNullable).into_array(); + ArrayAndStats::new(array, GenerateStatsOptions::default()) + } + + fn cost(estimated_ratio: Option) -> Option { + let data = test_data(); + let estimate = estimated_ratio.map_or_else( + CandidateEstimate::zero_bytes, + CandidateEstimate::from_compression_ratio, + ); + SizeCost.cost(&Candidate::new(&TestScheme, estimate, &data, None, &[])) } fn canonical_cost() -> Cost { - let array = PrimitiveArray::new(buffer![1i32, 2, 3, 4], Validity::NonNullable).into_array(); - let data = ArrayAndStats::new(array, GenerateStatsOptions::default()); + let data = test_data(); SizeCost.canonical_cost(&data, 4) } @@ -133,13 +130,13 @@ mod tests { f64::NAN, ]; - /// `cost < canonical_cost` must reproduce `legacy_is_valid` exactly, with non-priceable - /// ratios rejected as `None`. + /// `cost < canonical_cost` must reproduce `legacy_is_valid` exactly, with invalid ratios + /// producing no cost. #[test] fn validity_matches_ratio_gate() { let canonical = canonical_cost(); for &ratio in RATIOS { - let cost = SizeCost.cost(&candidate(EstimateScore::FiniteCompression(ratio))); + let cost = cost(Some(ratio)); let valid = cost.is_some_and(|cost| cost < canonical); assert_eq!( valid, @@ -162,12 +159,8 @@ mod tests { .collect(); for &best in &valid { for &challenger in &valid { - let best_cost = SizeCost - .cost(&candidate(EstimateScore::FiniteCompression(best))) - .expect("valid ratio must be priceable"); - let challenger_cost = SizeCost - .cost(&candidate(EstimateScore::FiniteCompression(challenger))) - .expect("valid ratio must be priceable"); + let best_cost = cost(Some(best)).expect("valid ratio must have a cost"); + let challenger_cost = cost(Some(challenger)).expect("valid ratio must have a cost"); let displaces = challenger_cost < best_cost; assert_eq!( @@ -189,20 +182,12 @@ mod tests { #[case::infinite(f64::INFINITY)] #[case::neg_infinite(f64::NEG_INFINITY)] #[case::subnormal(f64::MIN_POSITIVE / 2.0)] - fn unpriceable_ratios_are_rejected(#[case] ratio: f64) { - assert!( - SizeCost - .cost(&candidate(EstimateScore::FiniteCompression(ratio))) - .is_none() - ); + fn invalid_ratios_have_no_cost(#[case] ratio: f64) { + assert!(cost(Some(ratio)).is_none()); } #[test] fn zero_bytes_is_rejected() { - assert!( - SizeCost - .cost(&candidate(EstimateScore::ZeroBytes)) - .is_none() - ); + assert!(cost(None).is_none()); } } diff --git a/vortex-compressor/src/lib.rs b/vortex-compressor/src/lib.rs index c3bd91f5f92..3788f7bf9c3 100644 --- a/vortex-compressor/src/lib.rs +++ b/vortex-compressor/src/lib.rs @@ -53,8 +53,8 @@ //! `sample.result` events for zero-byte sample outputs, and both `*.compress_failed` events. //! //! The winning-compression span carries `scheme_chosen`, `input_nbytes`, `compressed_nbytes`, -//! `estimated_ratio` (absent when the scheme returned `AlwaysUse` or sampled to 0 bytes), -//! `achieved_ratio` (absent when the compressed output is 0 bytes), and `accepted`. +//! `estimated_cost` (absent when the scheme returned `AlwaysUse`), `achieved_ratio` (absent when +//! the compressed output is 0 bytes), and `accepted`. //! //! Failure events additionally carry `cascade_path` and `cascade_depth`, so nested compression //! errors can be tied back to the ancestor branch that triggered them. diff --git a/vortex-compressor/src/scheme/estimate.rs b/vortex-compressor/src/scheme/estimate.rs index 937004d70b2..2e66e9a8ed2 100644 --- a/vortex-compressor/src/scheme/estimate.rs +++ b/vortex-compressor/src/scheme/estimate.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Compression ratio estimation types returned by schemes. +//! Scheme evaluation and model-independent candidate evidence. use std::fmt; @@ -12,55 +12,49 @@ use crate::CascadingCompressor; use crate::scheme::CompressorContext; use crate::stats::ArrayAndStats; -/// Closure type for [`DeferredEstimate::Callback`]. +/// Closure type for [`DeferredEvaluation::Callback`]. /// -/// The compressor calls this with the same arguments it would pass to sampling, plus the best -/// [`EstimateScore`] observed so far (if any). The closure must resolve directly to a terminal -/// [`EstimateVerdict`]. -/// -/// The `best_so_far` threshold is an early-exit hint. If your scheme's maximum achievable -/// compression ratio is not strictly greater than -/// `best_so_far.and_then(EstimateScore::finite_ratio)`, you should return -/// [`EstimateVerdict::Skip`]. Returning a ratio equal to the threshold is permitted but will -/// lose to the prior best due to strict `>` tie-breaking in the selector. Use the threshold -/// only as an early-exit hint, never to perform additional work. +/// The compressor invokes this when a scheme needs fallible or expensive work to produce a +/// terminal [`ResolvedEvaluation`]. Candidate comparison is deliberately absent from this API: +/// the callback describes the candidate, and the configured cost model decides whether it wins. #[rustfmt::skip] -pub type EstimateFn = dyn FnOnce( +pub type DeferredEvaluationFn = dyn FnOnce( &CascadingCompressor, &ArrayAndStats, - Option, CompressorContext, &mut ExecutionCtx, - ) -> VortexResult + ) -> VortexResult + Send + Sync; -/// The result of a [`Scheme`](crate::scheme::Scheme)'s compression ratio estimation. -/// -/// This type is returned by -/// [`Scheme::expected_compression_ratio`](crate::scheme::Scheme::expected_compression_ratio) to -/// tell the compressor how promising this scheme is for a given array without performing any -/// expensive work. +/// A scheme's initial evaluation result. /// -/// [`CompressionEstimate::Verdict`] means the scheme already knows the terminal answer. -/// [`CompressionEstimate::Deferred`] means the compressor must do extra work before the scheme can -/// produce a terminal answer. +/// A scheme either declines the input, forces itself for semantic reasons, produces a candidate +/// estimate immediately, or asks the compressor to resolve deferred work. Candidate estimates are +/// model-independent facts; the compressor wraps them in a [`Candidate`](crate::cost::Candidate) +/// and passes them to the configured cost model. #[derive(Debug)] -pub enum CompressionEstimate { - /// The scheme already knows the terminal estimation verdict. - Verdict(EstimateVerdict), +pub enum SchemeEvaluation { + /// Do not consider this scheme for the input. + Skip, - /// The compressor must perform deferred work to resolve the terminal estimation verdict. - Deferred(DeferredEstimate), + /// Select this scheme immediately without cost-model comparison. + AlwaysUse, + + /// Hand this estimate to the configured cost model as a candidate. + Candidate(CandidateEstimate), + + /// The compressor must perform deferred work to resolve a terminal result. + Deferred(DeferredEvaluation), } -/// The terminal answer to a compression estimate request. +/// A terminal result produced by a deferred scheme evaluation. #[derive(Debug)] -pub enum EstimateVerdict { - /// Do not use this scheme for this array. +pub enum ResolvedEvaluation { + /// Do not consider this scheme for the input. Skip, - /// Always use this scheme, as it is definitively the best choice. + /// Select this scheme immediately without cost-model comparison. /// /// Some examples include decimal byte parts and temporal decomposition. /// @@ -71,72 +65,59 @@ pub enum EstimateVerdict { /// [`Scheme::matches`]: crate::scheme::Scheme::matches AlwaysUse, - /// The estimated compression ratio. This must be greater than `1.0` to be considered by the - /// compressor, otherwise it is worse than the canonical encoding. - Ratio(f64), + /// Hand this estimate to the configured cost model as a candidate. + Candidate(CandidateEstimate), } -/// Deferred work that can resolve to a terminal [`EstimateVerdict`]. -pub enum DeferredEstimate { - /// The scheme cannot cheaply estimate its ratio, so the compressor should compress a small - /// sample to determine effectiveness. +/// Deferred work that can produce a terminal [`ResolvedEvaluation`]. +pub enum DeferredEvaluation { + /// Compress a small sample and expose the resulting candidate to the cost model. Sample, - /// A fallible estimation requiring a custom expensive computation. - /// - /// Use this only when the scheme needs to perform trial encoding or other costly checks to - /// determine its compression ratio. The callback returns an [`EstimateVerdict`] directly, so - /// it cannot request more sampling or another deferred callback. + /// Run a scheme-defined fallible or expensive evaluation. /// - /// The compressor evaluates all immediate [`CompressionEstimate::Verdict`] results before - /// invoking any deferred callback, and passes the best [`EstimateScore`] observed so far to - /// the callback. This lets the callback return [`EstimateVerdict::Skip`] without performing - /// expensive work when its maximum achievable ratio cannot beat the current best. See - /// [`EstimateFn`] for the full contract. - Callback(Box), + /// The callback returns a [`ResolvedEvaluation`] directly, so it cannot request more sampling + /// or another deferred callback. + Callback(Box), } -/// Ranked estimate used for comparing non-terminal compression candidates. -#[derive(Debug, Clone, Copy, PartialEq)] -pub enum EstimateScore { - /// A finite compression ratio. Higher means a smaller amount of data, so it is better. - FiniteCompression(f64), - /// Trial compression produced a 0-byte output. - /// - /// This has no finite ratio and is not eligible for scheme selection. - /// - /// TODO(connor): A zero-byte sample usually means the sampler happened to hit an all-null - /// sample. Improve this logic so we can distinguish real zero-byte wins from sampling artifacts. - ZeroBytes, +/// Model-independent evidence produced by a scheme for one candidate. +/// +/// The estimate is intentionally opaque and evolution-safe. Compression ratio is currently the +/// common signal produced by schemes and consumed by [`SizeCost`](crate::cost::SizeCost); it is not +/// an ordering used by the compressor itself. +#[derive(Debug, Clone, PartialEq)] +pub struct CandidateEstimate { + /// Estimated compression ratio, absent when a sampled candidate produced zero bytes. + estimated_compression_ratio: Option, } -impl EstimateScore { - /// Converts measured sample sizes into a ranked estimate. - pub(crate) fn from_sample_sizes(before_nbytes: u64, after_nbytes: u64) -> Self { - if after_nbytes == 0 { - Self::ZeroBytes - } else { - Self::FiniteCompression(before_nbytes as f64 / after_nbytes as f64) +impl CandidateEstimate { + /// Creates candidate evidence from an estimated compression ratio. + pub fn from_compression_ratio(ratio: f64) -> Self { + Self { + estimated_compression_ratio: Some(ratio), } } - /// Returns the finite compression ratio, or [`None`] for the zero-byte special case. - /// - /// Callers comparing a scheme's maximum achievable ratio against a "best so far" threshold - /// should use this to extract a numeric value from an [`EstimateScore`]. - pub fn finite_ratio(self) -> Option { - match self { - Self::FiniteCompression(ratio) => Some(ratio), - Self::ZeroBytes => None, + /// Returns the estimated compression ratio, or `None` when a sample produced zero bytes. + pub fn estimated_compression_ratio(&self) -> Option { + self.estimated_compression_ratio + } + + /// Creates candidate evidence for a sampled output containing zero bytes. + pub(crate) fn zero_bytes() -> Self { + Self { + estimated_compression_ratio: None, } } } -impl fmt::Debug for DeferredEstimate { +impl fmt::Debug for DeferredEvaluation { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { match self { - DeferredEstimate::Sample => write!(f, "Sample"), - DeferredEstimate::Callback(_) => write!(f, "Callback(..)"), + DeferredEvaluation::Sample => write!(f, "Sample"), + DeferredEvaluation::Callback(_) => write!(f, "Callback(..)"), } } } diff --git a/vortex-compressor/src/scheme/mod.rs b/vortex-compressor/src/scheme/mod.rs index 341acc8acbf..cc5649315ce 100644 --- a/vortex-compressor/src/scheme/mod.rs +++ b/vortex-compressor/src/scheme/mod.rs @@ -15,11 +15,11 @@ use std::fmt::Debug; use std::hash::Hash; use std::hash::Hasher; -pub use estimate::CompressionEstimate; -pub use estimate::DeferredEstimate; -pub use estimate::EstimateFn; -pub use estimate::EstimateScore; -pub use estimate::EstimateVerdict; +pub use estimate::CandidateEstimate; +pub use estimate::DeferredEvaluation; +pub use estimate::DeferredEvaluationFn; +pub use estimate::ResolvedEvaluation; +pub use estimate::SchemeEvaluation; pub use exclusion::AncestorExclusion; pub use exclusion::ChildSelection; pub use exclusion::DescendantExclusion; @@ -56,8 +56,9 @@ impl fmt::Display for SchemeId { /// A single compression encoding that the [`CascadingCompressor`] can select from. /// /// The compressor evaluates every registered scheme whose [`matches`] returns `true` for a given -/// array, picks the one with the highest [`expected_compression_ratio`], and calls [`compress`] on -/// the winner. +/// array, asks the configured [`CostModel`] to compute each candidate's cost, and calls [`compress`] +/// on the lowest-cost eligible candidate. The default [`SizeCost`] model preserves the historical +/// highest-compression-ratio selection. /// /// One of the key features of the compressor in this crate is that schemes may "cascade". A /// scheme's [`compress`] can call back into the compressor via @@ -95,10 +96,9 @@ impl fmt::Display for SchemeId { /// /// # Implementing a scheme /// -/// [`expected_compression_ratio`] should return -/// `CompressionEstimate::Deferred(DeferredEstimate::Sample)` when a cheap heuristic is not -/// available, asking the compressor to estimate via sampling. Implementors should return an -/// immediate [`CompressionEstimate::Verdict`] when possible. +/// [`evaluate`] should return `SchemeEvaluation::Deferred(DeferredEvaluation::Sample)` when a +/// candidate cannot be described cheaply, asking the compressor to evaluate it through sampling. +/// Implementors should return an immediate [`SchemeEvaluation::Candidate`] when possible. /// /// Schemes that need statistics that may be expensive to compute should override [`stats_options`] /// to declare what they require. The compressor merges all eligible schemes' options before @@ -111,7 +111,9 @@ impl fmt::Display for SchemeId { /// [`scheme_name`]: Scheme::scheme_name /// [`matches`]: Scheme::matches /// [`compress`]: Scheme::compress -/// [`expected_compression_ratio`]: Scheme::expected_compression_ratio +/// [`evaluate`]: Scheme::evaluate +/// [`CostModel`]: crate::cost::CostModel +/// [`SizeCost`]: crate::cost::SizeCost /// [`stats_options`]: Scheme::stats_options /// [`num_children`]: Scheme::num_children /// [`descendant_exclusions`]: Scheme::descendant_exclusions @@ -152,22 +154,24 @@ pub trait Scheme: Debug + Send + Sync { Vec::new() } - /// Cheaply estimate the compression ratio for this scheme on the given array. + /// Produces model-independent candidate evidence for this scheme on the given array. /// /// This method should be fast and infallible. Any expensive or fallible work should be /// deferred to the compressor by returning - /// `CompressionEstimate::Deferred(DeferredEstimate::Sample)` or - /// `CompressionEstimate::Deferred(DeferredEstimate::Callback(...))`. + /// `SchemeEvaluation::Deferred(DeferredEvaluation::Sample)` or + /// `SchemeEvaluation::Deferred(DeferredEvaluation::Callback(...))`. /// - /// The compressor will ask all schemes what their expected compression ratio is given the array - /// and statistics. The scheme with the highest estimated ratio will then be applied to the - /// entire array. + /// The compressor combines the returned estimate with compressor-owned selection context to + /// construct a [`crate::cost::Candidate`], then asks the configured + /// [`crate::cost::CostModel`] to compute its cost. The default [`crate::cost::SizeCost`] model + /// interprets the candidate's estimated compression ratio; other models may use the remaining + /// candidate evidence differently. /// - /// [`CompressionEstimate::Verdict`] means the scheme already knows the terminal - /// [`crate::scheme::EstimateVerdict`]. `CompressionEstimate::Deferred(DeferredEstimate::Sample)` - /// asks the compressor to sample. `CompressionEstimate::Deferred(DeferredEstimate::Callback(...))` - /// asks the compressor to run custom deferred work. Deferred callbacks must return a - /// [`crate::scheme::EstimateVerdict`] directly, never another deferred request. + /// [`SchemeEvaluation::Candidate`] means the scheme can immediately describe a candidate. + /// `SchemeEvaluation::Deferred(DeferredEvaluation::Sample)` asks the compressor to sample; + /// `SchemeEvaluation::Deferred(DeferredEvaluation::Callback(...))` asks it to run custom + /// deferred work. Deferred callbacks must return a terminal + /// [`ResolvedEvaluation`], never another deferred request. /// /// Note that the compressor will also use this method when compressing samples, so some /// statistics that might hold for the samples may not hold for the entire array (e.g., @@ -178,12 +182,12 @@ pub trait Scheme: Debug + Send + Sync { /// called, so implementations may assume the array has at least one valid element. Outside of /// sample compression, the compressor also encodes constant arrays itself before evaluating /// schemes, so implementations only see constant arrays when `ctx.is_sample()` is `true`. - fn expected_compression_ratio( + fn evaluate( &self, _data: &ArrayAndStats, _compress_ctx: CompressorContext, _exec_ctx: &mut ExecutionCtx, - ) -> CompressionEstimate; + ) -> SchemeEvaluation; /// Compress the array using this scheme. /// diff --git a/vortex-compressor/src/trace.rs b/vortex-compressor/src/trace.rs index da7bc5cd8a9..d8a4464ca2c 100644 --- a/vortex-compressor/src/trace.rs +++ b/vortex-compressor/src/trace.rs @@ -35,7 +35,8 @@ pub(super) fn compress_span( /// Builds a span covering on-demand materialization of a cached stats type. /// /// Child of whatever span is active when a stats accessor first fires. Typically that's -/// [`verdict_pass_span`]; entering this span disambiguates stats cost from the rest of Pass 1. +/// [`immediate_evaluation_pass_span`]; entering this span disambiguates stats cost from the rest of +/// Pass 1. /// `kind` is usually `std::any::type_name::()` so the args identify which group was generated /// (e.g. `IntegerStats`, `FloatStats`). #[inline] @@ -47,16 +48,16 @@ pub(super) fn generate_stats_span(kind: &'static str) -> tracing::Span { ) } -/// Builds a span covering Pass 1 of scheme selection (the cheap-verdict pass). +/// Builds a span covering immediate scheme evaluation in Pass 1. /// /// Stats batches merged across eligible schemes are materialized lazily by the first -/// `expected_compression_ratio` call that touches them. Grouping those calls under one span makes -/// the stats cost (and unexpectedly slow verdicts) visible independently of per-candidate sampling. +/// scheme `evaluate` call that touches them. Grouping those calls under one span makes +/// the stats cost (and unexpectedly slow evaluations) visible independently of sampling. #[inline] -pub(super) fn verdict_pass_span() -> tracing::Span { +pub(super) fn immediate_evaluation_pass_span() -> tracing::Span { tracing::debug_span!( target: TARGET_TRACE, - "verdict_pass", + "immediate_evaluation_pass", ) } @@ -87,7 +88,7 @@ pub(super) fn zero_byte_sample_result(scheme: SchemeId, sampled_before: u64) { /// Builds a span covering the winning scheme's full-array compression. /// /// `scheme_chosen` and `input_nbytes` are known up front. `compressed_nbytes`, -/// `estimated_ratio`, `estimated_cost`, `achieved_ratio`, and `accepted` are filled in by +/// the model-defined `estimated_cost`, `achieved_ratio`, and `accepted` are filled in by /// [`record_winner_compress_result`] once the encode completes. #[inline] pub(super) fn winner_compress_span(scheme: SchemeId, before_nbytes: u64) -> tracing::Span { @@ -97,7 +98,6 @@ pub(super) fn winner_compress_span(scheme: SchemeId, before_nbytes: u64) -> trac scheme_chosen = %scheme, input_nbytes = before_nbytes, compressed_nbytes = tracing::field::Empty, - estimated_ratio = tracing::field::Empty, estimated_cost = tracing::field::Empty, achieved_ratio = tracing::field::Empty, accepted = tracing::field::Empty, @@ -108,16 +108,12 @@ pub(super) fn winner_compress_span(scheme: SchemeId, before_nbytes: u64) -> trac #[inline] pub(super) fn record_winner_compress_result( compressed_nbytes: u64, - estimated_ratio: Option, estimated_cost: Option, achieved_ratio: Option, accepted: bool, ) { let span = tracing::Span::current(); span.record("compressed_nbytes", compressed_nbytes); - if let Some(r) = estimated_ratio { - span.record("estimated_ratio", r); - } if let Some(c) = estimated_cost { span.record("estimated_cost", c); } diff --git a/vortex-tensor/src/encodings/l2_denorm.rs b/vortex-tensor/src/encodings/l2_denorm.rs index d07734bc6ff..045326e5d36 100644 --- a/vortex-tensor/src/encodings/l2_denorm.rs +++ b/vortex-tensor/src/encodings/l2_denorm.rs @@ -7,10 +7,9 @@ use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::extension::ExtensionArrayExt; use vortex_compressor::CascadingCompressor; -use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::CompressorContext; -use vortex_compressor::scheme::EstimateVerdict; use vortex_compressor::scheme::Scheme; +use vortex_compressor::scheme::SchemeEvaluation; use vortex_compressor::stats::ArrayAndStats; use vortex_error::VortexResult; @@ -32,13 +31,13 @@ impl Scheme for L2DenormScheme { ) } - fn expected_compression_ratio( + fn evaluate( &self, _data: &ArrayAndStats, _compress_ctx: CompressorContext, _exec_ctx: &mut ExecutionCtx, - ) -> CompressionEstimate { - CompressionEstimate::Verdict(EstimateVerdict::AlwaysUse) + ) -> SchemeEvaluation { + SchemeEvaluation::AlwaysUse } fn compress( From 314bf6ec85052c1c7cb5e31ae1a801f3febea58f Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Fri, 10 Jul 2026 14:36:29 +0100 Subject: [PATCH 5/5] Organize candidate evaluation and cost model modules Place the model-facing Candidate, Cost, CostModel, and SizeCost under the cost module. Split Cost and CostModel into cost/model.rs, keep the size submodule private, and re-export every public cost-model type through cost::. Rename the public estimate module to evaluation because it now defines scheme evaluation states, deferred resolution, and candidate evidence rather than selector scoring. Update every scheme implementation and internal link to use the new path. Each concept now has one public path, and the source layout reflects the candidate-to-cost-model boundary. Co-Authored-By: Claude Fable 5 Signed-off-by: Connor Tsui --- vortex-compressor/src/compressor/sample.rs | 2 +- vortex-compressor/src/compressor/select.rs | 6 +- vortex-compressor/src/compressor/tests.rs | 8 +- vortex-compressor/src/{ => cost}/candidate.rs | 0 vortex-compressor/src/cost/mod.rs | 121 +----------------- vortex-compressor/src/cost/model.rs | 119 +++++++++++++++++ vortex-compressor/src/cost/size.rs | 2 +- vortex-compressor/src/lib.rs | 1 - .../src/scheme/{estimate.rs => evaluation.rs} | 0 vortex-compressor/src/scheme/mod.rs | 12 +- 10 files changed, 140 insertions(+), 131 deletions(-) rename vortex-compressor/src/{ => cost}/candidate.rs (100%) create mode 100644 vortex-compressor/src/cost/model.rs rename vortex-compressor/src/scheme/{estimate.rs => evaluation.rs} (100%) diff --git a/vortex-compressor/src/compressor/sample.rs b/vortex-compressor/src/compressor/sample.rs index 4037d8eb81e..88d663bb8ac 100644 --- a/vortex-compressor/src/compressor/sample.rs +++ b/vortex-compressor/src/compressor/sample.rs @@ -162,7 +162,7 @@ pub(crate) struct SampledCandidate { /// # Errors /// /// Returns an error if sample compression fails. -pub(super) fn estimate_candidate_with_sampling( +pub(super) fn evaluate_candidate_with_sampling( compressor: &CascadingCompressor, scheme: &S, array: &ArrayRef, diff --git a/vortex-compressor/src/compressor/select.rs b/vortex-compressor/src/compressor/select.rs index 1387f9cc8d4..d9d531c130a 100644 --- a/vortex-compressor/src/compressor/select.rs +++ b/vortex-compressor/src/compressor/select.rs @@ -7,9 +7,9 @@ use vortex_array::ExecutionCtx; use vortex_error::VortexResult; use super::ROOT_SCHEME_ID; -use super::sample::estimate_candidate_with_sampling; +use super::sample::evaluate_candidate_with_sampling; use crate::CascadingCompressor; -use crate::candidate::Candidate; +use crate::cost::Candidate; use crate::cost::Cost; use crate::scheme::CompressorContext; use crate::scheme::DeferredEvaluation; @@ -114,7 +114,7 @@ impl CascadingCompressor { let _span = trace::scheme_eval_span(scheme.id()).entered(); match deferred_evaluation { DeferredEvaluation::Sample => { - let sampled_candidate = estimate_candidate_with_sampling( + let sampled_candidate = evaluate_candidate_with_sampling( self, scheme, data.array(), diff --git a/vortex-compressor/src/compressor/tests.rs b/vortex-compressor/src/compressor/tests.rs index ade3b208a40..b44b02a13ea 100644 --- a/vortex-compressor/src/compressor/tests.rs +++ b/vortex-compressor/src/compressor/tests.rs @@ -22,13 +22,13 @@ use vortex_session::VortexSession; use super::CascadingCompressor; use super::ROOT_SCHEME_ID; -use super::sample::estimate_candidate_with_sampling; +use super::sample::evaluate_candidate_with_sampling; use super::select::SelectionOutcome; use super::structural; use crate::builtins::FloatDictScheme; use crate::builtins::IntDictScheme; use crate::builtins::StringDictScheme; -use crate::candidate::Candidate; +use crate::cost::Candidate; use crate::cost::Cost; use crate::cost::CostModel; use crate::scheme::CandidateEstimate; @@ -764,7 +764,7 @@ fn all_null_array_compresses_to_constant() -> VortexResult<()> { /// Regression test for . /// -/// `estimate_candidate_with_sampling` must use the *scheme's* stats options +/// `evaluate_candidate_with_sampling` must use the *scheme's* stats options /// (which request distinct-value counting) rather than the context's stats options /// (which may not). With the old code this panicked inside `dictionary_encode` because /// distinct values were never computed for the sample. @@ -787,7 +787,7 @@ fn sampling_uses_scheme_stats_options() -> VortexResult<()> { // Before the fix this panicked with: // "this must be present since `DictScheme` declared that we need distinct values" let mut exec_ctx = SESSION.create_execution_ctx(); - let candidate = estimate_candidate_with_sampling( + let candidate = evaluate_candidate_with_sampling( &compressor, &FloatDictScheme, &array, diff --git a/vortex-compressor/src/candidate.rs b/vortex-compressor/src/cost/candidate.rs similarity index 100% rename from vortex-compressor/src/candidate.rs rename to vortex-compressor/src/cost/candidate.rs diff --git a/vortex-compressor/src/cost/mod.rs b/vortex-compressor/src/cost/mod.rs index d75b426a3d9..fda8c2c0c86 100644 --- a/vortex-compressor/src/cost/mod.rs +++ b/vortex-compressor/src/cost/mod.rs @@ -40,121 +40,12 @@ //! [`Scheme`]: crate::scheme::Scheme //! [`SchemeEvaluation::AlwaysUse`]: crate::scheme::SchemeEvaluation::AlwaysUse -mod size; +mod candidate; +pub use candidate::Candidate; -use std::cmp::Ordering; -use std::fmt::Debug; +mod model; +pub use model::Cost; +pub use model::CostModel; +mod size; pub use size::SizeCost; - -pub use crate::candidate::Candidate; -use crate::stats::ArrayAndStats; - -/// An opaque, totally ordered cost. Lower is better. -/// -/// Units are defined by each [`CostModel`], not by the framework: costs are only ever -/// compared *within* one model on one compressor instance, so cross-model comparability is -/// intentionally not provided. -/// -/// Values must be finite. `NaN` and infinite costs are not representable orderings; a model -/// that cannot compute a candidate's cost rejects it by returning `None` from [`CostModel::cost`]. -#[derive(Debug, Clone, Copy)] -pub struct Cost(f64); - -impl Cost { - /// Creates a cost from a finite value. - /// - /// # Panics - /// - /// Panics if `value` is not finite. - pub fn new(value: f64) -> Self { - assert!(value.is_finite(), "Cost must be finite, got {value}"); - Self(if value == 0.0 { 0.0 } else { value }) - } - - /// Returns the raw cost value. - pub fn value(self) -> f64 { - self.0 - } -} - -impl PartialEq for Cost { - fn eq(&self, other: &Self) -> bool { - self.cmp(other) == Ordering::Equal - } -} - -impl Eq for Cost {} - -impl PartialOrd for Cost { - fn partial_cmp(&self, other: &Self) -> Option { - Some(self.cmp(other)) - } -} - -impl Ord for Cost { - fn cmp(&self, other: &Self) -> Ordering { - self.0.total_cmp(&other.0) - } -} - -/// Computes costs for [`Candidate`]s during scheme selection. -/// -/// Implementations must be deterministic pure functions of their inputs and cheap: `cost` is -/// called once per candidate per selection site, i.e. `O(schemes × cascade levels)` -/// times per compressed chunk. -/// -/// Selection semantics: the winner is the candidate with the minimum cost, requiring strictly -/// lower cost both to displace the running best and to beat [`canonical_cost`]. Equal-cost ties -/// favor the candidate evaluated first: scheme registration order within each selection pass, -/// and immediate candidates over deferred candidates across passes. Returning `None` from [`cost`] -/// rejects the candidate outright. -/// -/// See the [module docs](self) for the acceptance axioms that apply to all models. -/// -/// [`canonical_cost`]: CostModel::canonical_cost -/// [`cost`]: CostModel::cost -pub trait CostModel: Debug + Send + Sync + 'static { - /// Estimated cost of choosing this candidate. Lower is better; `None` rejects the - /// candidate. - fn cost(&self, candidate: &Candidate<'_>) -> Option; - - /// Cost of leaving the array canonical — the baseline every candidate must strictly - /// beat to be selected. - fn canonical_cost(&self, data: &ArrayAndStats, n_values: u64) -> Cost; -} - -#[cfg(test)] -mod tests { - use std::cmp::Ordering; - - use super::Cost; - - #[test] - fn rejects_non_finite_values() { - for value in [f64::NAN, f64::INFINITY, f64::NEG_INFINITY] { - assert!( - std::panic::catch_unwind(|| Cost::new(value)).is_err(), - "non-finite cost {value} was accepted" - ); - } - } - - #[test] - fn signed_zero_has_one_representation() { - let negative = Cost::new(-0.0); - let positive = Cost::new(0.0); - - assert_eq!(negative, positive); - assert_eq!(negative.cmp(&positive), Ordering::Equal); - assert_eq!(negative.value().to_bits(), 0.0f64.to_bits()); - } - - #[test] - fn costs_have_a_total_order() { - let mut costs = [Cost::new(3.0), Cost::new(-2.0), Cost::new(0.0)]; - costs.sort(); - - assert_eq!(costs.map(Cost::value), [-2.0, 0.0, 3.0]); - } -} diff --git a/vortex-compressor/src/cost/model.rs b/vortex-compressor/src/cost/model.rs new file mode 100644 index 00000000000..d91d8cd6ce1 --- /dev/null +++ b/vortex-compressor/src/cost/model.rs @@ -0,0 +1,119 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Definition of [`CostModel`] and [`Cost`]. + +use std::cmp::Ordering; +use std::fmt::Debug; + +use crate::cost::candidate::Candidate; +use crate::stats::ArrayAndStats; + +/// Computes costs for [`Candidate`]s during scheme selection. +/// +/// Implementations must be deterministic pure functions of their inputs and cheap: `cost` is +/// called once per candidate per selection site, i.e. `O(schemes × cascade levels)` +/// times per compressed chunk. +/// +/// Selection semantics: the winner is the candidate with the minimum cost, requiring strictly +/// lower cost both to displace the running best and to beat [`canonical_cost`]. Equal-cost ties +/// favor the candidate evaluated first: scheme registration order within each selection pass, +/// and immediate candidates over deferred candidates across passes. Returning `None` from [`cost`] +/// rejects the candidate outright. +/// +/// See the [module docs](crate::cost) for the acceptance axioms that apply to all models. +/// +/// [`canonical_cost`]: CostModel::canonical_cost +/// [`cost`]: CostModel::cost +pub trait CostModel: Debug + Send + Sync + 'static { + /// Estimated cost of choosing this candidate. Lower is better; `None` rejects the + /// candidate. + fn cost(&self, candidate: &Candidate<'_>) -> Option; + + /// Cost of leaving the array canonical — the baseline every candidate must strictly + /// beat to be selected. + fn canonical_cost(&self, data: &ArrayAndStats, n_values: u64) -> Cost; +} + +/// An opaque, totally ordered cost. Lower is better. +/// +/// Units are defined by each [`CostModel`], not by the framework: costs are only ever +/// compared *within* one model on one compressor instance, so cross-model comparability is +/// intentionally not provided. +/// +/// Values must be finite. `NaN` and infinite costs are not representable orderings; a model +/// that cannot compute a candidate's cost rejects it by returning `None` from [`CostModel::cost`]. +#[derive(Debug, Clone, Copy)] +pub struct Cost(f64); + +impl Cost { + /// Creates a cost from a finite value. + /// + /// # Panics + /// + /// Panics if `value` is not finite. + pub fn new(value: f64) -> Self { + assert!(value.is_finite(), "Cost must be finite, got {value}"); + Self(if value == 0.0 { 0.0 } else { value }) + } + + /// Returns the raw cost value. + pub fn value(self) -> f64 { + self.0 + } +} + +impl PartialEq for Cost { + fn eq(&self, other: &Self) -> bool { + self.cmp(other) == Ordering::Equal + } +} + +impl Eq for Cost {} + +impl PartialOrd for Cost { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +impl Ord for Cost { + fn cmp(&self, other: &Self) -> Ordering { + self.0.total_cmp(&other.0) + } +} + +#[cfg(test)] +mod tests { + use std::cmp::Ordering; + + use super::Cost; + + #[test] + fn rejects_non_finite_values() { + for value in [f64::NAN, f64::INFINITY, f64::NEG_INFINITY] { + assert!( + std::panic::catch_unwind(|| Cost::new(value)).is_err(), + "non-finite cost {value} was accepted" + ); + } + } + + #[test] + fn signed_zero_has_one_representation() { + let negative = Cost::new(-0.0); + let positive = Cost::new(0.0); + + assert_eq!(negative, positive); + assert_eq!(negative.cmp(&positive), Ordering::Equal); + assert_eq!(negative.value().to_bits(), 0.0f64.to_bits()); + } + + #[test] + fn costs_have_a_total_order() { + let mut costs = [Cost::new(3.0), Cost::new(-2.0), Cost::new(0.0)]; + costs.sort(); + + assert_eq!(costs.map(Cost::value), [-2.0, 0.0, 3.0]); + } +} diff --git a/vortex-compressor/src/cost/size.rs b/vortex-compressor/src/cost/size.rs index 209dfdc8ab6..34117e2b557 100644 --- a/vortex-compressor/src/cost/size.rs +++ b/vortex-compressor/src/cost/size.rs @@ -3,7 +3,7 @@ //! The default, size-only cost model. -use crate::candidate::Candidate; +use crate::cost::Candidate; use crate::cost::Cost; use crate::cost::CostModel; use crate::stats::ArrayAndStats; diff --git a/vortex-compressor/src/lib.rs b/vortex-compressor/src/lib.rs index 3788f7bf9c3..54a7588c844 100644 --- a/vortex-compressor/src/lib.rs +++ b/vortex-compressor/src/lib.rs @@ -67,7 +67,6 @@ pub mod cost; pub mod scheme; pub mod stats; -mod candidate; mod compressor; pub use compressor::CascadingCompressor; diff --git a/vortex-compressor/src/scheme/estimate.rs b/vortex-compressor/src/scheme/evaluation.rs similarity index 100% rename from vortex-compressor/src/scheme/estimate.rs rename to vortex-compressor/src/scheme/evaluation.rs diff --git a/vortex-compressor/src/scheme/mod.rs b/vortex-compressor/src/scheme/mod.rs index cc5649315ce..fe10843cf20 100644 --- a/vortex-compressor/src/scheme/mod.rs +++ b/vortex-compressor/src/scheme/mod.rs @@ -8,18 +8,18 @@ mod ctx; pub use ctx::CompressorContext; pub use ctx::MAX_CASCADE; -pub(crate) mod estimate; +pub(crate) mod evaluation; mod exclusion; use std::fmt; use std::fmt::Debug; use std::hash::Hash; use std::hash::Hasher; -pub use estimate::CandidateEstimate; -pub use estimate::DeferredEvaluation; -pub use estimate::DeferredEvaluationFn; -pub use estimate::ResolvedEvaluation; -pub use estimate::SchemeEvaluation; +pub use evaluation::CandidateEstimate; +pub use evaluation::DeferredEvaluation; +pub use evaluation::DeferredEvaluationFn; +pub use evaluation::ResolvedEvaluation; +pub use evaluation::SchemeEvaluation; pub use exclusion::AncestorExclusion; pub use exclusion::ChildSelection; pub use exclusion::DescendantExclusion;