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-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/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/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/compressor/cascade.rs b/vortex-compressor/src/compressor/cascade.rs index 9dbbc611448..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. @@ -310,7 +312,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..fbbb838d592 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; @@ -30,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 @@ -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 used 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,18 @@ impl CascadingCompressor { Self { schemes, root_exclusions, + cost_model: Arc::new(SizeCost), } } + + /// Replaces the cost model used during scheme selection. + /// + /// 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 + } } // NB: Cascading compression logic is located in `vortex-compressor/src/compressor/cascade.rs`. diff --git a/vortex-compressor/src/compressor/sample.rs b/vortex-compressor/src/compressor/sample.rs index ba1271d4a6d..88d663bb8ac 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,21 +141,34 @@ fn partition_indices(length: usize, num_partitions: u32) -> Vec<(usize, usize)> .collect() } -/// Estimates compression ratio by compressing a ~1% sample of the data. +/// 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, +} + +/// 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 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 evaluate_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 { @@ -180,13 +193,17 @@ 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(score) + Ok(SampledCandidate { + estimate, + sampled: compressed, + }) } #[cfg(test)] diff --git a/vortex-compressor/src/compressor/select.rs b/vortex-compressor/src/compressor/select.rs index 3c73d2d4cdb..d9d531c130a 100644 --- a/vortex-compressor/src/compressor/select.rs +++ b/vortex-compressor/src/compressor/select.rs @@ -1,130 +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::evaluate_candidate_with_sampling; use crate::CascadingCompressor; -use crate::scheme::CompressionEstimate; +use crate::cost::Candidate; +use crate::cost::Cost; 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. - Score(EstimateScore), + /// 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 { +impl SelectionOutcome { + /// Returns the traceable cost for the winning estimate. + pub(super) fn trace_cost(self) -> Option { match self { Self::AlwaysUse => None, - Self::Score(score) => score.finite_ratio(), + Self::Cost(cost) => Some(cost.value()), } } } -/// 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)) -} - 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. /// - /// 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. + /// 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. /// - /// 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. + /// 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<(&'static dyn Scheme, EstimateScore)> = 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(); - // Pass 1: evaluate every immediate verdict. Stash deferred work for pass 2. + let canonical_cost = self + .cost_model + .canonical_cost(data, data.array().len() as u64); + + 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); - - if is_better_score(score, best.as_ref()) { - best = Some((scheme, score)); + 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.map(|(_, score)| score); - match deferred_estimate { - DeferredEstimate::Sample => { - let score = estimate_compression_ratio_with_sampling( + match deferred_evaluation { + DeferredEvaluation::Sample => { + let sampled_candidate = evaluate_candidate_with_sampling( self, scheme, data.array(), compress_ctx.clone(), exec_ctx, )?; - - if is_better_score(score, best.as_ref()) { - best = Some((scheme, score)); + 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); - - if is_better_score(score, best.as_ref()) { - best = Some((scheme, score)); + 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 }); } } } @@ -132,7 +159,21 @@ impl CascadingCompressor { } } - Ok(best.map(|(scheme, score)| (scheme, WinnerEstimate::Score(score)))) + Ok(best.map(|candidate| (candidate.scheme, SelectionOutcome::Cost(candidate.cost)))) + } + + /// 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<'_>, + canonical_cost: Cost, + best_cost: Option, + ) -> Option { + let cost = self.cost_model.cost(candidate)?; + (cost < canonical_cost && best_cost.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 08eafa369de..b44b02a13ea 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::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::scheme::CompressionEstimate; +use crate::cost::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(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(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(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(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(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 +/// `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. @@ -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 score = estimate_compression_ratio_with_sampling( + let candidate = evaluate_candidate_with_sampling( &compressor, &FloatDictScheme, &array, ctx, &mut exec_ctx, )?; - assert!(matches!(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/candidate.rs b/vortex-compressor/src/cost/candidate.rs new file mode 100644 index 00000000000..72680017732 --- /dev/null +++ b/vortex-compressor/src/cost/candidate.rs @@ -0,0 +1,90 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Selection-time candidate view. + +use vortex_array::ArrayRef; + +use crate::scheme::CandidateEstimate; +use crate::scheme::Scheme; +use crate::scheme::SchemeExt; +use crate::scheme::SchemeId; +use crate::stats::ArrayAndStats; + +/// A borrowed view of one scheme estimate during selection. +/// +/// 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 +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 + } + + /// Returns the canonical input array being selected over. + pub fn array(&self) -> &ArrayRef { + self.data.array() + } + + /// Returns the uncompressed size of the input array in bytes. + pub fn input_nbytes(&self) -> u64 { + self.data.array().nbytes() + } + + /// Returns the number of values in the input array. + pub fn n_values(&self) -> u64 { + self.data.array().len() as u64 + } + + /// 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 + } + + /// Returns the cascade ancestry as `(scheme_id, child_index)` pairs. + pub fn cascade(&self) -> &[(SchemeId, usize)] { + self.cascade + } +} diff --git a/vortex-compressor/src/cost/mod.rs b/vortex-compressor/src/cost/mod.rs new file mode 100644 index 00000000000..fda8c2c0c86 --- /dev/null +++ b/vortex-compressor/src/cost/mod.rs @@ -0,0 +1,51 @@ +// 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 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 preserves the compressor's historical candidate +//! ordering and canonical-acceptance threshold. +//! +//! # What sits outside the model +//! +//! The initial cost-model boundary covers candidates inside +//! `CascadingCompressor::choose_best_scheme`. The following pre-existing decisions remain +//! outside it: +//! +//! - 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 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 +//! +//! 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 +//! [`SchemeEvaluation::AlwaysUse`]: crate::scheme::SchemeEvaluation::AlwaysUse + +mod candidate; +pub use candidate::Candidate; + +mod model; +pub use model::Cost; +pub use model::CostModel; + +mod size; +pub use size::SizeCost; 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 new file mode 100644 index 00000000000..34117e2b557 --- /dev/null +++ b/vortex-compressor/src/cost/size.rs @@ -0,0 +1,193 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! The default, size-only cost model. + +use crate::cost::Candidate; +use crate::cost::Cost; +use crate::cost::CostModel; +use crate::stats::ArrayAndStats; + +/// The default cost model: maximize estimated compression ratio. +/// +/// `SizeCost` preserves the compressor's historical ratio-argmax ordering: +/// +/// - 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 +/// 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 { + 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 { + 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::CandidateEstimate; + use crate::scheme::CompressorContext; + use crate::scheme::Scheme; + use crate::scheme::SchemeEvaluation; + 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 evaluate( + &self, + _data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> SchemeEvaluation { + SchemeEvaluation::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 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 data = test_data(); + 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 invalid ratios + /// producing no cost. + #[test] + fn validity_matches_ratio_gate() { + let canonical = canonical_cost(); + for &ratio in RATIOS { + let cost = cost(Some(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 = 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!( + 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 invalid_ratios_have_no_cost(#[case] ratio: f64) { + assert!(cost(Some(ratio)).is_none()); + } + + #[test] + fn zero_bytes_is_rejected() { + assert!(cost(None).is_none()); + } +} diff --git a/vortex-compressor/src/lib.rs b/vortex-compressor/src/lib.rs index 55bb9b188f6..54a7588c844 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. @@ -63,6 +63,7 @@ //! with a short `jq` query. pub mod builtins; +pub mod cost; pub mod scheme; pub mod stats; diff --git a/vortex-compressor/src/scheme/estimate.rs b/vortex-compressor/src/scheme/estimate.rs deleted file mode 100644 index 20b9dc77e9b..00000000000 --- a/vortex-compressor/src/scheme/estimate.rs +++ /dev/null @@ -1,163 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -//! Compression ratio estimation types returned by schemes. - -use std::fmt; - -use vortex_array::ExecutionCtx; -use vortex_error::VortexResult; - -use crate::CascadingCompressor; -use crate::scheme::CompressorContext; -use crate::stats::ArrayAndStats; - -/// Closure type for [`DeferredEstimate::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. -#[rustfmt::skip] -pub type EstimateFn = dyn FnOnce( - &CascadingCompressor, - &ArrayAndStats, - Option, - CompressorContext, - &mut ExecutionCtx, - ) -> 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. -/// -/// [`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. -#[derive(Debug)] -pub enum CompressionEstimate { - /// The scheme already knows the terminal estimation verdict. - Verdict(EstimateVerdict), - - /// The compressor must perform deferred work to resolve the terminal estimation verdict. - Deferred(DeferredEstimate), -} - -/// The terminal answer to a compression estimate request. -#[derive(Debug)] -pub enum EstimateVerdict { - /// Do not use this scheme for this array. - Skip, - - /// Always use this scheme, as it is definitively the best choice. - /// - /// Some examples include decimal byte parts and temporal decomposition. - /// - /// The compressor will select this scheme immediately without evaluating further candidates. - /// Schemes that return `AlwaysUse` must be mutually exclusive per canonical type (enforced by - /// [`Scheme::matches`]), otherwise the winner depends silently on registration order. - /// - /// [`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), -} - -/// 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. - 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. - /// - /// 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), -} - -/// 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, -} - -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) - } - } - - /// 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 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 { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - match self { - DeferredEstimate::Sample => write!(f, "Sample"), - DeferredEstimate::Callback(_) => write!(f, "Callback(..)"), - } - } -} diff --git a/vortex-compressor/src/scheme/evaluation.rs b/vortex-compressor/src/scheme/evaluation.rs new file mode 100644 index 00000000000..2e66e9a8ed2 --- /dev/null +++ b/vortex-compressor/src/scheme/evaluation.rs @@ -0,0 +1,123 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Scheme evaluation and model-independent candidate evidence. + +use std::fmt; + +use vortex_array::ExecutionCtx; +use vortex_error::VortexResult; + +use crate::CascadingCompressor; +use crate::scheme::CompressorContext; +use crate::stats::ArrayAndStats; + +/// Closure type for [`DeferredEvaluation::Callback`]. +/// +/// 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 DeferredEvaluationFn = dyn FnOnce( + &CascadingCompressor, + &ArrayAndStats, + CompressorContext, + &mut ExecutionCtx, + ) -> VortexResult + + Send + + Sync; + +/// A scheme's initial evaluation result. +/// +/// 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 SchemeEvaluation { + /// Do not consider this scheme for the input. + Skip, + + /// 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), +} + +/// A terminal result produced by a deferred scheme evaluation. +#[derive(Debug)] +pub enum ResolvedEvaluation { + /// Do not consider this scheme for the input. + Skip, + + /// Select this scheme immediately without cost-model comparison. + /// + /// Some examples include decimal byte parts and temporal decomposition. + /// + /// The compressor will select this scheme immediately without evaluating further candidates. + /// Schemes that return `AlwaysUse` must be mutually exclusive per canonical type (enforced by + /// [`Scheme::matches`]), otherwise the winner depends silently on registration order. + /// + /// [`Scheme::matches`]: crate::scheme::Scheme::matches + AlwaysUse, + + /// Hand this estimate to the configured cost model as a candidate. + Candidate(CandidateEstimate), +} + +/// 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, + + /// Run a scheme-defined fallible or expensive evaluation. + /// + /// The callback returns a [`ResolvedEvaluation`] directly, so it cannot request more sampling + /// or another deferred callback. + Callback(Box), +} + +/// 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 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 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 DeferredEvaluation { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + 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..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::CompressionEstimate; -pub use estimate::DeferredEstimate; -pub use estimate::EstimateFn; -pub use estimate::EstimateScore; -pub use estimate::EstimateVerdict; +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; @@ -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 ee594621199..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`, `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,7 @@ 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, ) @@ -107,14 +108,14 @@ 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); } if let Some(r) = achieved_ratio { span.record("achieved_ratio", r); 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(