From 72e22e0205c5c7fccde671c4b3c46b5c06b2b3c4 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Sat, 7 Nov 2026 00:01:07 +0000 Subject: [PATCH 1/6] Skip the dict probe compression when no dict scheme can win Add `may_compress_to` to `CascadingCompressor` and `CompressorPlugin`. The check runs only the immediate estimates of the schemes that produce the requested encoding, and does not compress. `DictStrategy` uses it to skip the full sample compression when the result cannot be a dictionary. Signed-off-by: Robert Kruszewski --- vortex-btrblocks/src/tests.rs | 46 +++++++++++ vortex-compressor/src/compressor/cascade.rs | 84 +++++++++++++++++++++ vortex-layout/src/layouts/compressed.rs | 33 ++++++++ vortex-layout/src/layouts/dict/writer.rs | 16 +++- 4 files changed, 176 insertions(+), 3 deletions(-) diff --git a/vortex-btrblocks/src/tests.rs b/vortex-btrblocks/src/tests.rs index fea1433dbe7..43dcec43ee1 100644 --- a/vortex-btrblocks/src/tests.rs +++ b/vortex-btrblocks/src/tests.rs @@ -199,3 +199,49 @@ fn test_cuda_binary_zstd_follows_editions( Ok(()) } + +/// The dict pre-check must never say `false` for an array that compresses to a dictionary, and it +/// must say `false` for arrays where no dictionary scheme can win. +#[rstest] +#[case::unique_strings( + VarBinViewArray::from_iter_str((0..4096).map(|i| format!("{i:04} is a unique story"))).into_array(), + false, +)] +#[case::few_strings( + VarBinViewArray::from_iter_str((0..4096).map(|i| ["red", "green", "blue"][i % 3])).into_array(), + true, +)] +#[case::constant_strings( + VarBinViewArray::from_iter_str(iter::repeat_n("same", 4096)).into_array(), + false, +)] +#[case::null_strings( + VarBinViewArray::from_iter(iter::repeat_n(None::<&str>, 4096), DType::Utf8(Nullability::Nullable)) + .into_array(), + false, +)] +#[case::unique_integers(PrimitiveArray::from_iter(0..4096i64).into_array(), false)] +fn test_may_compress_to_dict(#[case] input: ArrayRef, #[case] expected: bool) -> VortexResult<()> { + use vortex_array::VTable; + use vortex_array::arrays::Dict; + + let compressor = BtrBlocksCompressorBuilder::from_session(&SESSION) + .unrestricted() + .build(); + let mut ctx = SESSION.create_execution_ctx(); + + let may_be_dict = compressor.0.may_compress_to(&input, Dict.id(), &mut ctx)?; + let is_dict = compressor.compress(&input, &mut ctx)?.is::(); + + assert_eq!(may_be_dict, expected); + assert!( + may_be_dict || !is_dict, + "the pre-check said no, but the result is a dictionary" + ); + + if expected { + assert!(is_dict); + } + + Ok(()) +} diff --git a/vortex-compressor/src/compressor/cascade.rs b/vortex-compressor/src/compressor/cascade.rs index f86ab6b136b..3645a948c54 100644 --- a/vortex-compressor/src/compressor/cascade.rs +++ b/vortex-compressor/src/compressor/cascade.rs @@ -3,6 +3,7 @@ //! Core cascading compression flow. +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::CanonicalValidity; @@ -35,7 +36,9 @@ use vortex_error::vortex_err; use super::CascadingCompressor; use super::constant; +use crate::scheme::CompressionEstimate; use crate::scheme::CompressorContext; +use crate::scheme::EstimateVerdict; use crate::scheme::Scheme; use crate::scheme::SchemeExt; use crate::scheme::SchemeId; @@ -349,6 +352,87 @@ impl CascadingCompressor { Ok(Selection::Unchanged(data.into_array())) } } + + /// Returns `false` if [`compress`](Self::compress) cannot return an array of `encoding` for + /// `array`. Returns `true` if it can, or if this check cannot tell. + /// + /// The check does not compress. A leaf array compresses to `encoding` only if a scheme that + /// produces `encoding` wins the selection, and a scheme whose estimate is + /// [`EstimateVerdict::Skip`](crate::scheme::EstimateVerdict::Skip) cannot win. Thus the check + /// runs only the estimates of the schemes that produce `encoding`, and only the immediate + /// ones. A deferred estimate can need a sample compression, so it makes the result `true`. + /// + /// # Errors + /// + /// Returns an error if canonicalization or a statistic fails. + pub fn may_compress_to( + &self, + array: &ArrayRef, + encoding: ArrayId, + exec_ctx: &mut ExecutionCtx, + ) -> VortexResult { + let canonical = array.clone().execute::(exec_ctx)?.0; + let canonical = canonical.compact(exec_ctx)?; + + // Only leaf arrays go directly to scheme selection at the top level. + if !matches!( + canonical, + Canonical::Bool(_) + | Canonical::Primitive(_) + | Canonical::Decimal(_) + | Canonical::VarBinView(_) + ) { + return Ok(true); + } + + let compress_ctx = CompressorContext::new(); + + let eligible_schemes: Vec<&'static dyn Scheme> = self + .schemes + .iter() + .copied() + .filter(|s| s.matches(&canonical) && !self.is_excluded(*s, &compress_ctx)) + .collect(); + + let candidates: Vec<&'static dyn Scheme> = eligible_schemes + .iter() + .copied() + .filter(|s| s.produced_encodings().contains(&encoding)) + .collect(); + + if candidates.is_empty() { + return Ok(false); + } + + // The compressor itself makes empty, all-null, and constant arrays, as in + // `choose_and_compress`. + let array: ArrayRef = canonical.into(); + + if array.is_empty() || array.all_invalid(exec_ctx)? { + return Ok(false); + } + + if constant::is_constant_for_compression(&array, exec_ctx)? { + return Ok(false); + } + + // Use the same statistics as the selection, so that each estimate is the same. + let merged_opts = eligible_schemes + .iter() + .fold(GenerateStatsOptions::default(), |acc, s| { + acc.merge(s.stats_options()) + }); + let compress_ctx = compress_ctx.with_merged_stats_options(merged_opts); + + let data = ArrayAndStats::new(array, merged_opts); + + Ok(candidates.iter().any(|scheme| { + !matches!( + scheme.expected_compression_ratio(&data, compress_ctx.clone(), exec_ctx), + CompressionEstimate::Verdict(EstimateVerdict::Skip) + ) + })) + } } /// The bytes `array`'s rows use once canonical, read from the exact statistic the writer caches diff --git a/vortex-layout/src/layouts/compressed.rs b/vortex-layout/src/layouts/compressed.rs index 07d83bf41c7..6dccf784a2e 100644 --- a/vortex-layout/src/layouts/compressed.rs +++ b/vortex-layout/src/layouts/compressed.rs @@ -5,6 +5,7 @@ use std::sync::Arc; use async_trait::async_trait; use futures::StreamExt as _; +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::VortexSessionExecute; @@ -31,12 +32,35 @@ use crate::sequence::SequentialStreamExt; /// API consumers are free to implement this trait to provide new plugin compressors. pub trait CompressorPlugin: Send + Sync + 'static { fn compress_chunk(&self, chunk: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult; + + /// Returns `false` if [`compress_chunk`](Self::compress_chunk) cannot return an array of + /// `encoding` for `chunk`. Returns `true` if it can, or if the plugin cannot tell without + /// compressing. + /// + /// A writer can use this to skip a compression that it does only to look at the encoding. + fn may_compress_to( + &self, + _chunk: &ArrayRef, + _encoding: ArrayId, + _ctx: &mut ExecutionCtx, + ) -> VortexResult { + Ok(true) + } } impl CompressorPlugin for Arc { fn compress_chunk(&self, chunk: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { self.as_ref().compress_chunk(chunk, ctx) } + + fn may_compress_to( + &self, + chunk: &ArrayRef, + encoding: ArrayId, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + self.as_ref().may_compress_to(chunk, encoding, ctx) + } } impl CompressorPlugin for F @@ -52,6 +76,15 @@ impl CompressorPlugin for BtrBlocksCompressor { fn compress_chunk(&self, chunk: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { self.compress(chunk, ctx) } + + fn may_compress_to( + &self, + chunk: &ArrayRef, + encoding: ArrayId, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + self.0.may_compress_to(chunk, encoding, ctx) + } } /// A layout writer that compresses chunks. diff --git a/vortex-layout/src/layouts/dict/writer.rs b/vortex-layout/src/layouts/dict/writer.rs index 0e854a924af..a5fa528260b 100644 --- a/vortex-layout/src/layouts/dict/writer.rs +++ b/vortex-layout/src/layouts/dict/writer.rs @@ -26,6 +26,7 @@ use futures::try_join; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; +use vortex_array::VTable; use vortex_array::VortexSessionExecute; use vortex_array::arrays::ChunkedArray; use vortex_array::arrays::Dict; @@ -173,10 +174,19 @@ impl LayoutStrategy for DictStrategy { None => true, // empty stream Some(sample) => { let mut exec_ctx = session.create_execution_ctx(); - let compressed = self + + // A full compression is slow, so do it only if it can give a dictionary. + if self .probe_compressor - .compress_chunk(&sample, &mut exec_ctx)?; - !compressed.is::() + .may_compress_to(&sample, Dict.id(), &mut exec_ctx)? + { + let compressed = self + .probe_compressor + .compress_chunk(&sample, &mut exec_ctx)?; + !compressed.is::() + } else { + true + } } }; if should_fallback { From ca200fe81a219501f60725bc500e14b706fcba6c Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Sat, 7 Nov 2026 00:11:09 +0000 Subject: [PATCH 2/6] Fix ambiguous Dict id and redundant doc link Signed-off-by: Robert Kruszewski --- vortex-btrblocks/src/tests.rs | 9 +++++---- vortex-compressor/src/compressor/cascade.rs | 6 +++--- 2 files changed, 8 insertions(+), 7 deletions(-) diff --git a/vortex-btrblocks/src/tests.rs b/vortex-btrblocks/src/tests.rs index 43dcec43ee1..9e64e3f4c0a 100644 --- a/vortex-btrblocks/src/tests.rs +++ b/vortex-btrblocks/src/tests.rs @@ -21,6 +21,7 @@ use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::fns::sum::sum; use vortex_array::arrays::DecimalArray; +use vortex_array::arrays::Dict; use vortex_array::arrays::ListArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; @@ -222,15 +223,15 @@ fn test_cuda_binary_zstd_follows_editions( )] #[case::unique_integers(PrimitiveArray::from_iter(0..4096i64).into_array(), false)] fn test_may_compress_to_dict(#[case] input: ArrayRef, #[case] expected: bool) -> VortexResult<()> { - use vortex_array::VTable; - use vortex_array::arrays::Dict; - let compressor = BtrBlocksCompressorBuilder::from_session(&SESSION) .unrestricted() .build(); let mut ctx = SESSION.create_execution_ctx(); - let may_be_dict = compressor.0.may_compress_to(&input, Dict.id(), &mut ctx)?; + let may_be_dict = + compressor + .0 + .may_compress_to(&input, vortex_array::VTable::id(&Dict), &mut ctx)?; let is_dict = compressor.compress(&input, &mut ctx)?.is::(); assert_eq!(may_be_dict, expected); diff --git a/vortex-compressor/src/compressor/cascade.rs b/vortex-compressor/src/compressor/cascade.rs index 3645a948c54..a3a7c380a64 100644 --- a/vortex-compressor/src/compressor/cascade.rs +++ b/vortex-compressor/src/compressor/cascade.rs @@ -358,9 +358,9 @@ impl CascadingCompressor { /// /// The check does not compress. A leaf array compresses to `encoding` only if a scheme that /// produces `encoding` wins the selection, and a scheme whose estimate is - /// [`EstimateVerdict::Skip`](crate::scheme::EstimateVerdict::Skip) cannot win. Thus the check - /// runs only the estimates of the schemes that produce `encoding`, and only the immediate - /// ones. A deferred estimate can need a sample compression, so it makes the result `true`. + /// [`EstimateVerdict::Skip`] cannot win. Thus the check runs only the estimates of the schemes + /// that produce `encoding`, and only the immediate ones. A deferred estimate can need a sample + /// compression, so it makes the result `true`. /// /// # Errors /// From 97433ad2439edffde0133040c71a5835ba50be4a Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Mon, 9 Nov 2026 18:59:51 +0000 Subject: [PATCH 3/6] Ask the compressor plugin which encoding it produces Replace `CompressorPlugin::may_compress_to(chunk, encoding)` with `encoding_for(chunk)`, which returns the encoding that `compress_chunk` gives. The default compresses the chunk. The dict writer compares the answer with `Dict`. `CascadingCompressor` splits scheme selection from the winner's compression, so `compress` and `encoding_for` share one selection. A new `Scheme::root_encoding` lets `encoding_for` skip the full compression when the winner always produces the same root encoding. Signed-off-by: Robert Kruszewski --- vortex-btrblocks/src/schemes/float/alprd.rs | 4 + vortex-btrblocks/src/schemes/float/pco.rs | 4 + vortex-btrblocks/src/schemes/integer/pco.rs | 4 + .../src/schemes/integer/runend.rs | 4 + .../src/schemes/integer/zigzag.rs | 4 + vortex-btrblocks/src/schemes/string/fsst.rs | 4 + vortex-btrblocks/src/schemes/temporal.rs | 4 + vortex-btrblocks/src/tests.rs | 57 +++-- vortex-compressor/src/builtins/dict/binary.rs | 4 + vortex-compressor/src/builtins/dict/float.rs | 4 + .../src/builtins/dict/integer.rs | 4 + vortex-compressor/src/builtins/dict/string.rs | 4 + vortex-compressor/src/compressor/cascade.rs | 221 ++++++++++-------- vortex-compressor/src/scheme/mod.rs | 9 + vortex-file/src/strategy.rs | 5 + vortex-layout/src/layouts/compressed.rs | 35 +-- vortex-layout/src/layouts/dict/writer.rs | 14 +- 17 files changed, 227 insertions(+), 158 deletions(-) diff --git a/vortex-btrblocks/src/schemes/float/alprd.rs b/vortex-btrblocks/src/schemes/float/alprd.rs index 09c0ee85b0d..612435d3a50 100644 --- a/vortex-btrblocks/src/schemes/float/alprd.rs +++ b/vortex-btrblocks/src/schemes/float/alprd.rs @@ -44,6 +44,10 @@ impl Scheme for ALPRDScheme { vec![vortex_alp::ALPRD.id()] } + fn root_encoding(&self) -> Option { + Some(vortex_alp::ALPRD.id()) + } + fn expected_compression_ratio( &self, data: &ArrayAndStats, diff --git a/vortex-btrblocks/src/schemes/float/pco.rs b/vortex-btrblocks/src/schemes/float/pco.rs index 416668c2fd0..0a4790fd376 100644 --- a/vortex-btrblocks/src/schemes/float/pco.rs +++ b/vortex-btrblocks/src/schemes/float/pco.rs @@ -35,6 +35,10 @@ impl Scheme for PcoScheme { vec![vortex_pco::Pco.id()] } + fn root_encoding(&self) -> Option { + Some(vortex_pco::Pco.id()) + } + fn expected_compression_ratio( &self, _data: &ArrayAndStats, diff --git a/vortex-btrblocks/src/schemes/integer/pco.rs b/vortex-btrblocks/src/schemes/integer/pco.rs index 675a112d44f..5fb55c682d8 100644 --- a/vortex-btrblocks/src/schemes/integer/pco.rs +++ b/vortex-btrblocks/src/schemes/integer/pco.rs @@ -36,6 +36,10 @@ impl Scheme for PcoScheme { vec![vortex_pco::Pco.id()] } + fn root_encoding(&self) -> Option { + Some(vortex_pco::Pco.id()) + } + fn expected_compression_ratio( &self, data: &ArrayAndStats, diff --git a/vortex-btrblocks/src/schemes/integer/runend.rs b/vortex-btrblocks/src/schemes/integer/runend.rs index da388131dab..3ed592dce39 100644 --- a/vortex-btrblocks/src/schemes/integer/runend.rs +++ b/vortex-btrblocks/src/schemes/integer/runend.rs @@ -52,6 +52,10 @@ impl Scheme for RunEndScheme { vec![RunEnd.id()] } + fn root_encoding(&self) -> Option { + Some(RunEnd.id()) + } + /// Children: values=0, ends=1. fn num_children(&self) -> usize { 2 diff --git a/vortex-btrblocks/src/schemes/integer/zigzag.rs b/vortex-btrblocks/src/schemes/integer/zigzag.rs index 0e4be01845a..87beb8064ea 100644 --- a/vortex-btrblocks/src/schemes/integer/zigzag.rs +++ b/vortex-btrblocks/src/schemes/integer/zigzag.rs @@ -50,6 +50,10 @@ impl Scheme for ZigZagScheme { vec![ZigZag.id()] } + fn root_encoding(&self) -> Option { + Some(ZigZag.id()) + } + /// Children: encoded=0. fn num_children(&self) -> usize { 1 diff --git a/vortex-btrblocks/src/schemes/string/fsst.rs b/vortex-btrblocks/src/schemes/string/fsst.rs index 97238d97e94..a4bb13f8f1c 100644 --- a/vortex-btrblocks/src/schemes/string/fsst.rs +++ b/vortex-btrblocks/src/schemes/string/fsst.rs @@ -54,6 +54,10 @@ impl Scheme for FSSTScheme { vec![FSST.id(), VarBin.id()] } + fn root_encoding(&self) -> Option { + Some(FSST.id()) + } + /// Children: lengths=0, code_offsets=1. fn num_children(&self) -> usize { 2 diff --git a/vortex-btrblocks/src/schemes/temporal.rs b/vortex-btrblocks/src/schemes/temporal.rs index 79748b69450..3273a3b3763 100644 --- a/vortex-btrblocks/src/schemes/temporal.rs +++ b/vortex-btrblocks/src/schemes/temporal.rs @@ -59,6 +59,10 @@ impl Scheme for TemporalScheme { vec![DateTimeParts.id()] } + fn root_encoding(&self) -> Option { + Some(DateTimeParts.id()) + } + /// Children: days=0, seconds=1, subseconds=2. fn num_children(&self) -> usize { 3 diff --git a/vortex-btrblocks/src/tests.rs b/vortex-btrblocks/src/tests.rs index 9e64e3f4c0a..e75fbe81bbc 100644 --- a/vortex-btrblocks/src/tests.rs +++ b/vortex-btrblocks/src/tests.rs @@ -51,6 +51,8 @@ use vortex_session::VortexSession; use crate::BtrBlocksCompressor; use crate::BtrBlocksCompressorBuilder; use crate::SESSION; +use crate::schemes::integer::IntDictScheme; +use crate::schemes::string::StringDictScheme; #[track_caller] fn assert_roundtrip(compressor: &BtrBlocksCompressor, input: &ArrayRef) -> VortexResult { @@ -201,48 +203,55 @@ fn test_cuda_binary_zstd_follows_editions( Ok(()) } -/// The dict pre-check must never say `false` for an array that compresses to a dictionary, and it -/// must say `false` for arrays where no dictionary scheme can win. +/// `encoding_for` gives the encoding of the array that `compress` returns. #[rstest] #[case::unique_strings( VarBinViewArray::from_iter_str((0..4096).map(|i| format!("{i:04} is a unique story"))).into_array(), - false, )] #[case::few_strings( VarBinViewArray::from_iter_str((0..4096).map(|i| ["red", "green", "blue"][i % 3])).into_array(), - true, -)] -#[case::constant_strings( - VarBinViewArray::from_iter_str(iter::repeat_n("same", 4096)).into_array(), - false, )] +#[case::constant_strings(VarBinViewArray::from_iter_str(iter::repeat_n("same", 4096)).into_array())] #[case::null_strings( VarBinViewArray::from_iter(iter::repeat_n(None::<&str>, 4096), DType::Utf8(Nullability::Nullable)) .into_array(), - false, )] -#[case::unique_integers(PrimitiveArray::from_iter(0..4096i64).into_array(), false)] -fn test_may_compress_to_dict(#[case] input: ArrayRef, #[case] expected: bool) -> VortexResult<()> { +#[case::few_integers(PrimitiveArray::from_iter((0..4096i64).map(|i| (i % 3) * 1_000_000)).into_array())] +#[case::unique_integers(PrimitiveArray::from_iter(0..4096i64).into_array())] +fn test_encoding_for(#[case] input: ArrayRef) -> VortexResult<()> { let compressor = BtrBlocksCompressorBuilder::from_session(&SESSION) .unrestricted() .build(); let mut ctx = SESSION.create_execution_ctx(); - let may_be_dict = - compressor - .0 - .may_compress_to(&input, vortex_array::VTable::id(&Dict), &mut ctx)?; - let is_dict = compressor.compress(&input, &mut ctx)?.is::(); + let encoding = compressor.0.encoding_for(&input, &mut ctx)?; + let compressed = compressor.compress(&input, &mut ctx)?; - assert_eq!(may_be_dict, expected); - assert!( - may_be_dict || !is_dict, - "the pre-check said no, but the result is a dictionary" - ); + assert_eq!(encoding, compressed.encoding_id()); - if expected { - assert!(is_dict); - } + Ok(()) +} + +/// A compressor that has only the dictionary schemes produces a dictionary or nothing. +#[rstest] +#[case::few_strings( + VarBinViewArray::from_iter_str((0..4096).map(|i| ["red", "green", "blue"][i % 3])).into_array(), + true, +)] +#[case::unique_strings( + VarBinViewArray::from_iter_str((0..4096).map(|i| format!("{i:04} is a unique story"))).into_array(), + false, +)] +fn test_dict_only_encoding_for(#[case] input: ArrayRef, #[case] is_dict: bool) -> VortexResult<()> { + let compressor = BtrBlocksCompressorBuilder::empty() + .with_new_scheme(&StringDictScheme) + .with_new_scheme(&IntDictScheme) + .build(); + let mut ctx = SESSION.create_execution_ctx(); + + let encoding = compressor.0.encoding_for(&input, &mut ctx)?; + + assert_eq!(encoding == vortex_array::VTable::id(&Dict), is_dict); Ok(()) } diff --git a/vortex-compressor/src/builtins/dict/binary.rs b/vortex-compressor/src/builtins/dict/binary.rs index c407d0251c6..b15f6cb5a42 100644 --- a/vortex-compressor/src/builtins/dict/binary.rs +++ b/vortex-compressor/src/builtins/dict/binary.rs @@ -52,6 +52,10 @@ impl Scheme for BinaryDictScheme { vec![Dict.id()] } + fn root_encoding(&self) -> Option { + Some(Dict.id()) + } + fn stats_options(&self) -> GenerateStatsOptions { GenerateStatsOptions { count_distinct_values: true, diff --git a/vortex-compressor/src/builtins/dict/float.rs b/vortex-compressor/src/builtins/dict/float.rs index 92dfc01ec5c..a436446ba66 100644 --- a/vortex-compressor/src/builtins/dict/float.rs +++ b/vortex-compressor/src/builtins/dict/float.rs @@ -58,6 +58,10 @@ impl Scheme for FloatDictScheme { vec![Dict.id()] } + fn root_encoding(&self) -> Option { + Some(Dict.id()) + } + fn stats_options(&self) -> GenerateStatsOptions { GenerateStatsOptions { count_distinct_values: true, diff --git a/vortex-compressor/src/builtins/dict/integer.rs b/vortex-compressor/src/builtins/dict/integer.rs index 9e621719b9a..56dccd71ea0 100644 --- a/vortex-compressor/src/builtins/dict/integer.rs +++ b/vortex-compressor/src/builtins/dict/integer.rs @@ -53,6 +53,10 @@ impl Scheme for IntDictScheme { vec![Dict.id()] } + fn root_encoding(&self) -> Option { + Some(Dict.id()) + } + fn stats_options(&self) -> GenerateStatsOptions { GenerateStatsOptions { count_distinct_values: true, diff --git a/vortex-compressor/src/builtins/dict/string.rs b/vortex-compressor/src/builtins/dict/string.rs index f5cbcd54d89..2ea8d8848b8 100644 --- a/vortex-compressor/src/builtins/dict/string.rs +++ b/vortex-compressor/src/builtins/dict/string.rs @@ -52,6 +52,10 @@ impl Scheme for StringDictScheme { vec![Dict.id()] } + fn root_encoding(&self) -> Option { + Some(Dict.id()) + } + fn stats_options(&self) -> GenerateStatsOptions { GenerateStatsOptions { count_distinct_values: true, diff --git a/vortex-compressor/src/compressor/cascade.rs b/vortex-compressor/src/compressor/cascade.rs index a3a7c380a64..465cb596b5b 100644 --- a/vortex-compressor/src/compressor/cascade.rs +++ b/vortex-compressor/src/compressor/cascade.rs @@ -36,9 +36,8 @@ use vortex_error::vortex_err; use super::CascadingCompressor; use super::constant; -use crate::scheme::CompressionEstimate; +use super::select::WinnerEstimate; use crate::scheme::CompressorContext; -use crate::scheme::EstimateVerdict; use crate::scheme::Scheme; use crate::scheme::SchemeExt; use crate::scheme::SchemeId; @@ -63,6 +62,28 @@ impl Selection { } } +/// The outcome of [`CascadingCompressor::choose`]. +enum Choice { + /// The compressor decided the result without a scheme. + Decided(Selection), + /// A scheme won the selection, but it did not compress the array yet. + Scheme(Winner), +} + +/// A scheme that won the selection, with the state it needs to compress the array. +struct Winner { + /// The scheme that won. + scheme: &'static dyn Scheme, + /// The estimate with which the scheme won. + estimate: WinnerEstimate, + /// The array to compress, with its statistics. + data: ArrayAndStats, + /// The context with the merged statistics options of the eligible schemes. + compress_ctx: CompressorContext, + /// The bytes the array's rows use, which the compressed output must beat. + before_nbytes: u64, +} + impl CascadingCompressor { /// Compresses an array using cascading adaptive compression. /// @@ -271,6 +292,70 @@ impl CascadingCompressor { compress_ctx: CompressorContext, exec_ctx: &mut ExecutionCtx, ) -> VortexResult { + match self.choose(canonical, compress_ctx, exec_ctx)? { + Choice::Decided(selection) => Ok(selection), + Choice::Scheme(winner) => self.compress_with_winner(winner, exec_ctx), + } + } + + /// Returns the encoding of the array that [`compress`](Self::compress) returns for `array`. + /// + /// For a leaf array, this runs the same scheme selection as `compress`, with the statistics + /// and the sample compressions that the selection needs. If the winning scheme declares its + /// [`root_encoding`](Scheme::root_encoding), the full compression does not occur. If it does + /// not, or if `array` is not a leaf array, this compresses `array`. + /// + /// Without the full compression, the check that the winner's output is smaller than `array` + /// cannot occur. Thus, if the output of the winner is not smaller, `compress` returns the + /// uncompressed array, but this method returns the encoding of the winner. + /// + /// # Errors + /// + /// Returns an error if canonicalization or compression fails. + pub fn encoding_for( + &self, + array: &ArrayRef, + exec_ctx: &mut ExecutionCtx, + ) -> VortexResult { + let canonical = array.clone().execute::(exec_ctx)?.0; + let compact = canonical.compact(exec_ctx)?; + + let is_leaf = matches!( + compact, + Canonical::Bool(_) + | Canonical::Primitive(_) + | Canonical::Decimal(_) + | Canonical::VarBinView(_) + ); + + if !is_leaf { + return Ok(self + .compress_canonical(compact, CompressorContext::new(), exec_ctx)? + .encoding_id()); + } + + match self.choose(compact, CompressorContext::new(), exec_ctx)? { + Choice::Decided(selection) => Ok(selection.into_array().encoding_id()), + Choice::Scheme(winner) => match winner.scheme.root_encoding() { + Some(encoding) => Ok(encoding), + None => Ok(self + .compress_with_winner(winner, exec_ctx)? + .into_array() + .encoding_id()), + }, + } + } + + /// Chooses how to compress a single leaf array, but does not run the winning scheme. + /// + /// The compressor decides empty, all-null and constant arrays itself, and arrays for which no + /// scheme beats the canonical encoding. + fn choose( + &self, + canonical: Canonical, + compress_ctx: CompressorContext, + exec_ctx: &mut ExecutionCtx, + ) -> VortexResult { let eligible_schemes: Vec<&'static dyn Scheme> = self .schemes .iter() @@ -281,13 +366,13 @@ impl CascadingCompressor { let array: ArrayRef = canonical.into(); if array.is_empty() { - return Ok(Selection::Unchanged(array)); + return Ok(Choice::Decided(Selection::Unchanged(array))); } if array.all_invalid(exec_ctx)? { - return Ok(Selection::Compressed( + return Ok(Choice::Decided(Selection::Compressed( ConstantArray::new(Scalar::null(array.dtype().clone()), array.len()).into_array(), - )); + ))); } // Constant detection is built into the compressor: a constant leaf always short-circuits @@ -296,7 +381,8 @@ impl CascadingCompressor { // // Compare valid values directly so even nullable constants avoid distinct-value counting. if !compress_ctx.is_sample() && constant::is_constant_for_compression(&array, exec_ctx)? { - return constant::compress_as_constant(array, exec_ctx).map(Selection::Compressed); + return constant::compress_as_constant(array, exec_ctx) + .map(|constant| Choice::Decided(Selection::Compressed(constant))); } // A scheme competes against the bytes the rows reference, not dead bytes in the buffers. @@ -312,26 +398,50 @@ impl CascadingCompressor { let data = ArrayAndStats::new(array, merged_opts); if eligible_schemes.is_empty() { - return Ok(Selection::Unchanged(data.into_array())); + return Ok(Choice::Decided(Selection::Unchanged(data.into_array()))); } - let Some((winner, winner_estimate)) = + let Some((scheme, estimate)) = self.choose_best_scheme(&eligible_schemes, &data, compress_ctx.clone(), exec_ctx)? else { - return Ok(Selection::Unchanged(data.into_array())); + return Ok(Choice::Decided(Selection::Unchanged(data.into_array()))); }; - // Run the winning scheme's `compress`. On failure, emit an ERROR event carrying the - // scheme name and cascade history before propagating. + Ok(Choice::Scheme(Winner { + scheme, + estimate, + data, + compress_ctx, + before_nbytes, + })) + } + + /// Runs the winning scheme's `compress`, and keeps its output only if it is smaller than the + /// input. + fn compress_with_winner( + &self, + winner: Winner, + exec_ctx: &mut ExecutionCtx, + ) -> VortexResult { + let Winner { + scheme, + estimate, + data, + compress_ctx, + before_nbytes, + } = winner; + + // On failure, emit an ERROR event carrying the scheme name and cascade history before + // propagating. let error_ctx = trace::enabled_error_context(&compress_ctx); - let _winner_span = trace::winner_compress_span(winner.id(), before_nbytes).entered(); - let compressed = winner + let _winner_span = trace::winner_compress_span(scheme.id(), before_nbytes).entered(); + let compressed = scheme .compress(self, &data, compress_ctx, exec_ctx) .inspect_err(|err| { // NB: this is the only way we can tell which scheme panicked / bailed on their // data, especially for third-party schemes where the error site may not carry any // compressor context. - trace::scheme_compress_failed(winner.id(), before_nbytes, error_ctx.as_ref(), err); + trace::scheme_compress_failed(scheme.id(), before_nbytes, error_ctx.as_ref(), err); })?; let after_nbytes = compressed.nbytes(); @@ -341,7 +451,7 @@ impl CascadingCompressor { trace::record_winner_compress_result( after_nbytes, - winner_estimate.trace_ratio(), + estimate.trace_ratio(), actual_ratio, accepted, ); @@ -352,87 +462,6 @@ impl CascadingCompressor { Ok(Selection::Unchanged(data.into_array())) } } - - /// Returns `false` if [`compress`](Self::compress) cannot return an array of `encoding` for - /// `array`. Returns `true` if it can, or if this check cannot tell. - /// - /// The check does not compress. A leaf array compresses to `encoding` only if a scheme that - /// produces `encoding` wins the selection, and a scheme whose estimate is - /// [`EstimateVerdict::Skip`] cannot win. Thus the check runs only the estimates of the schemes - /// that produce `encoding`, and only the immediate ones. A deferred estimate can need a sample - /// compression, so it makes the result `true`. - /// - /// # Errors - /// - /// Returns an error if canonicalization or a statistic fails. - pub fn may_compress_to( - &self, - array: &ArrayRef, - encoding: ArrayId, - exec_ctx: &mut ExecutionCtx, - ) -> VortexResult { - let canonical = array.clone().execute::(exec_ctx)?.0; - let canonical = canonical.compact(exec_ctx)?; - - // Only leaf arrays go directly to scheme selection at the top level. - if !matches!( - canonical, - Canonical::Bool(_) - | Canonical::Primitive(_) - | Canonical::Decimal(_) - | Canonical::VarBinView(_) - ) { - return Ok(true); - } - - let compress_ctx = CompressorContext::new(); - - let eligible_schemes: Vec<&'static dyn Scheme> = self - .schemes - .iter() - .copied() - .filter(|s| s.matches(&canonical) && !self.is_excluded(*s, &compress_ctx)) - .collect(); - - let candidates: Vec<&'static dyn Scheme> = eligible_schemes - .iter() - .copied() - .filter(|s| s.produced_encodings().contains(&encoding)) - .collect(); - - if candidates.is_empty() { - return Ok(false); - } - - // The compressor itself makes empty, all-null, and constant arrays, as in - // `choose_and_compress`. - let array: ArrayRef = canonical.into(); - - if array.is_empty() || array.all_invalid(exec_ctx)? { - return Ok(false); - } - - if constant::is_constant_for_compression(&array, exec_ctx)? { - return Ok(false); - } - - // Use the same statistics as the selection, so that each estimate is the same. - let merged_opts = eligible_schemes - .iter() - .fold(GenerateStatsOptions::default(), |acc, s| { - acc.merge(s.stats_options()) - }); - let compress_ctx = compress_ctx.with_merged_stats_options(merged_opts); - - let data = ArrayAndStats::new(array, merged_opts); - - Ok(candidates.iter().any(|scheme| { - !matches!( - scheme.expected_compression_ratio(&data, compress_ctx.clone(), exec_ctx), - CompressionEstimate::Verdict(EstimateVerdict::Skip) - ) - })) - } } /// The bytes `array`'s rows use once canonical, read from the exact statistic the writer caches diff --git a/vortex-compressor/src/scheme/mod.rs b/vortex-compressor/src/scheme/mod.rs index d56389b2ce3..a3a44b193ba 100644 --- a/vortex-compressor/src/scheme/mod.rs +++ b/vortex-compressor/src/scheme/mod.rs @@ -135,6 +135,15 @@ pub trait Scheme: AsDynScheme + Debug + Send + Sync { /// formats declares the wire IDs the scheme writes, which may differ from its in-memory ID. fn produced_encodings(&self) -> Vec; + /// The in-memory encoding ID of the array that [`compress`](Scheme::compress) returns, if it + /// is the same for all inputs. + /// + /// The compressor uses this to tell which encoding it produces without a full compression. + /// The default is `None`, which makes the compressor compress to find the encoding. + fn root_encoding(&self) -> Option { + None + } + /// Returns the variant of this scheme to use given which serialized IDs are `allowed`. /// /// The default returns `self`. A variant must share this scheme's [`SchemeId`]. Every ID the diff --git a/vortex-file/src/strategy.rs b/vortex-file/src/strategy.rs index 8cf7bc5bc78..e0a84ff079b 100644 --- a/vortex-file/src/strategy.rs +++ b/vortex-file/src/strategy.rs @@ -147,6 +147,11 @@ impl WriteStrategyBuilder { } /// Override the compressor used to probe whether a column is dict-eligible. + /// + /// A column gets a dictionary layout if the probe compressor produces a dictionary for its + /// first rows. By default, a dictionary must beat all other schemes. To use a dictionary + /// whenever it beats the uncompressed data, give a compressor that has only dictionary + /// schemes. pub fn with_probe_compressor(mut self, compressor: C) -> Self { self.probe_compressor = Some(Arc::new(compressor)); self diff --git a/vortex-layout/src/layouts/compressed.rs b/vortex-layout/src/layouts/compressed.rs index 6dccf784a2e..76c5fc586f3 100644 --- a/vortex-layout/src/layouts/compressed.rs +++ b/vortex-layout/src/layouts/compressed.rs @@ -33,18 +33,13 @@ use crate::sequence::SequentialStreamExt; pub trait CompressorPlugin: Send + Sync + 'static { fn compress_chunk(&self, chunk: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult; - /// Returns `false` if [`compress_chunk`](Self::compress_chunk) cannot return an array of - /// `encoding` for `chunk`. Returns `true` if it can, or if the plugin cannot tell without - /// compressing. + /// Returns the encoding of the array that [`compress_chunk`](Self::compress_chunk) returns + /// for `chunk`. /// - /// A writer can use this to skip a compression that it does only to look at the encoding. - fn may_compress_to( - &self, - _chunk: &ArrayRef, - _encoding: ArrayId, - _ctx: &mut ExecutionCtx, - ) -> VortexResult { - Ok(true) + /// A writer can use this to choose a layout from the encoding the plugin produces. The default + /// compresses `chunk`. A plugin can override it to give the answer without a full compression. + fn encoding_for(&self, chunk: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + Ok(self.compress_chunk(chunk, ctx)?.encoding_id()) } } @@ -53,13 +48,8 @@ impl CompressorPlugin for Arc { self.as_ref().compress_chunk(chunk, ctx) } - fn may_compress_to( - &self, - chunk: &ArrayRef, - encoding: ArrayId, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - self.as_ref().may_compress_to(chunk, encoding, ctx) + fn encoding_for(&self, chunk: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + self.as_ref().encoding_for(chunk, ctx) } } @@ -77,13 +67,8 @@ impl CompressorPlugin for BtrBlocksCompressor { self.compress(chunk, ctx) } - fn may_compress_to( - &self, - chunk: &ArrayRef, - encoding: ArrayId, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - self.0.may_compress_to(chunk, encoding, ctx) + fn encoding_for(&self, chunk: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + self.0.encoding_for(chunk, ctx) } } diff --git a/vortex-layout/src/layouts/dict/writer.rs b/vortex-layout/src/layouts/dict/writer.rs index a5fa528260b..62ff6871561 100644 --- a/vortex-layout/src/layouts/dict/writer.rs +++ b/vortex-layout/src/layouts/dict/writer.rs @@ -174,19 +174,7 @@ impl LayoutStrategy for DictStrategy { None => true, // empty stream Some(sample) => { let mut exec_ctx = session.create_execution_ctx(); - - // A full compression is slow, so do it only if it can give a dictionary. - if self - .probe_compressor - .may_compress_to(&sample, Dict.id(), &mut exec_ctx)? - { - let compressed = self - .probe_compressor - .compress_chunk(&sample, &mut exec_ctx)?; - !compressed.is::() - } else { - true - } + self.probe_compressor.encoding_for(&sample, &mut exec_ctx)? != Dict.id() } }; if should_fallback { From d8d1a483a44ed5d86991a15f631b780247a46a8a Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Tue, 10 Nov 2026 17:47:35 +0000 Subject: [PATCH 4/6] Return the set of possible encodings from the compressor Replace `Scheme::root_encoding` with `returned_encodings`, which lists the encodings of all arrays that `compress` can return. Every scheme in the crate now declares them. `encoding_for` becomes `encodings_for` and returns a `HashSet`. When the winning scheme declares its encodings, the set holds them and the encoding of the uncompressed input, without a full compression. The dict writer checks if the set contains `Dict`, so every probe skips the full compression of the sample. A debug assertion checks that each scheme returns an encoding that it declares. Signed-off-by: Robert Kruszewski --- vortex-btrblocks/src/schemes/binary/varbin.rs | 4 ++ vortex-btrblocks/src/schemes/binary/zstd.rs | 4 ++ .../src/schemes/binary/zstd_buffers.rs | 4 ++ vortex-btrblocks/src/schemes/decimal.rs | 9 +++ vortex-btrblocks/src/schemes/float/alp.rs | 8 +++ vortex-btrblocks/src/schemes/float/alprd.rs | 4 +- vortex-btrblocks/src/schemes/float/pco.rs | 4 +- vortex-btrblocks/src/schemes/float/rle.rs | 4 ++ vortex-btrblocks/src/schemes/float/sparse.rs | 5 ++ .../src/schemes/integer/bitpacking.rs | 9 +++ vortex-btrblocks/src/schemes/integer/delta.rs | 4 ++ vortex-btrblocks/src/schemes/integer/for_.rs | 5 ++ vortex-btrblocks/src/schemes/integer/pco.rs | 4 +- vortex-btrblocks/src/schemes/integer/rle.rs | 4 ++ .../src/schemes/integer/runend.rs | 4 +- .../src/schemes/integer/sequence.rs | 4 ++ .../src/schemes/integer/sparse.rs | 4 ++ .../src/schemes/integer/zigzag.rs | 4 +- vortex-btrblocks/src/schemes/string/fsst.rs | 4 +- vortex-btrblocks/src/schemes/string/onpair.rs | 5 ++ vortex-btrblocks/src/schemes/string/sparse.rs | 5 ++ vortex-btrblocks/src/schemes/string/zstd.rs | 4 ++ .../src/schemes/string/zstd_buffers.rs | 4 ++ vortex-btrblocks/src/schemes/temporal.rs | 4 +- vortex-btrblocks/src/tests.rs | 19 +++--- vortex-compressor/src/builtins/dict/binary.rs | 4 +- vortex-compressor/src/builtins/dict/float.rs | 4 +- .../src/builtins/dict/integer.rs | 4 +- vortex-compressor/src/builtins/dict/string.rs | 4 +- vortex-compressor/src/compressor/cascade.rs | 59 ++++++++++++------- vortex-compressor/src/scheme/mod.rs | 13 ++-- vortex-layout/src/layouts/compressed.rs | 36 +++++++---- vortex-layout/src/layouts/dict/writer.rs | 5 +- 33 files changed, 195 insertions(+), 67 deletions(-) diff --git a/vortex-btrblocks/src/schemes/binary/varbin.rs b/vortex-btrblocks/src/schemes/binary/varbin.rs index c29d2c7ef6d..c48d036837e 100644 --- a/vortex-btrblocks/src/schemes/binary/varbin.rs +++ b/vortex-btrblocks/src/schemes/binary/varbin.rs @@ -48,6 +48,10 @@ impl Scheme for VarBinScheme { vec![VarBin.id()] } + fn returned_encodings(&self) -> Vec { + vec![VarBin.id()] + } + fn num_children(&self) -> usize { 1 } diff --git a/vortex-btrblocks/src/schemes/binary/zstd.rs b/vortex-btrblocks/src/schemes/binary/zstd.rs index d652e344db2..e58200551e6 100644 --- a/vortex-btrblocks/src/schemes/binary/zstd.rs +++ b/vortex-btrblocks/src/schemes/binary/zstd.rs @@ -35,6 +35,10 @@ impl Scheme for ZstdScheme { vec![vortex_zstd::Zstd.id()] } + fn returned_encodings(&self) -> Vec { + vec![vortex_zstd::Zstd.id()] + } + fn expected_compression_ratio( &self, _data: &ArrayAndStats, diff --git a/vortex-btrblocks/src/schemes/binary/zstd_buffers.rs b/vortex-btrblocks/src/schemes/binary/zstd_buffers.rs index 3f06d65b061..5425e3bf9e2 100644 --- a/vortex-btrblocks/src/schemes/binary/zstd_buffers.rs +++ b/vortex-btrblocks/src/schemes/binary/zstd_buffers.rs @@ -35,6 +35,10 @@ impl Scheme for ZstdBuffersScheme { vec![vortex_zstd::ZstdBuffers.id()] } + fn returned_encodings(&self) -> Vec { + vec![vortex_zstd::ZstdBuffers.id()] + } + fn expected_compression_ratio( &self, _data: &ArrayAndStats, diff --git a/vortex-btrblocks/src/schemes/decimal.rs b/vortex-btrblocks/src/schemes/decimal.rs index e59f2a3a535..0b20edb7d48 100644 --- a/vortex-btrblocks/src/schemes/decimal.rs +++ b/vortex-btrblocks/src/schemes/decimal.rs @@ -8,6 +8,8 @@ use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; +use vortex_array::VTable; +use vortex_array::arrays::Decimal; use vortex_array::arrays::DecimalArray; use vortex_array::arrays::decimal::narrowed_decimal; use vortex_array::dtype::DecimalType; @@ -96,6 +98,13 @@ impl Scheme for DecimalScheme { } } + fn returned_encodings(&self) -> Vec { + match self.mode { + DecimalSchemeMode::V1 => vec![Decimal.id(), DecimalByteParts.id()], + DecimalSchemeMode::V2 => vec![DecimalByteParts.id()], + } + } + fn refine(&self, allowed: &dyn Fn(&ArrayId) -> bool) -> &dyn Scheme { if allowed(&decimal_byte_parts_v2_id()) { &DECIMAL_V2 diff --git a/vortex-btrblocks/src/schemes/float/alp.rs b/vortex-btrblocks/src/schemes/float/alp.rs index f9fc7066bf4..86852e27245 100644 --- a/vortex-btrblocks/src/schemes/float/alp.rs +++ b/vortex-btrblocks/src/schemes/float/alp.rs @@ -50,6 +50,14 @@ impl Scheme for ALPScheme { encodings } + fn returned_encodings(&self) -> Vec { + let mut encodings = vec![ALP.id()]; + if use_experimental_patches() { + encodings.push(Patched.id()); + } + encodings + } + /// Children: encoded_ints=0. fn num_children(&self) -> usize { 1 diff --git a/vortex-btrblocks/src/schemes/float/alprd.rs b/vortex-btrblocks/src/schemes/float/alprd.rs index 612435d3a50..c368e950512 100644 --- a/vortex-btrblocks/src/schemes/float/alprd.rs +++ b/vortex-btrblocks/src/schemes/float/alprd.rs @@ -44,8 +44,8 @@ impl Scheme for ALPRDScheme { vec![vortex_alp::ALPRD.id()] } - fn root_encoding(&self) -> Option { - Some(vortex_alp::ALPRD.id()) + fn returned_encodings(&self) -> Vec { + vec![vortex_alp::ALPRD.id()] } fn expected_compression_ratio( diff --git a/vortex-btrblocks/src/schemes/float/pco.rs b/vortex-btrblocks/src/schemes/float/pco.rs index 0a4790fd376..f1e69da3182 100644 --- a/vortex-btrblocks/src/schemes/float/pco.rs +++ b/vortex-btrblocks/src/schemes/float/pco.rs @@ -35,8 +35,8 @@ impl Scheme for PcoScheme { vec![vortex_pco::Pco.id()] } - fn root_encoding(&self) -> Option { - Some(vortex_pco::Pco.id()) + fn returned_encodings(&self) -> Vec { + vec![vortex_pco::Pco.id()] } fn expected_compression_ratio( diff --git a/vortex-btrblocks/src/schemes/float/rle.rs b/vortex-btrblocks/src/schemes/float/rle.rs index 71158b9dc3b..313f7123b6c 100644 --- a/vortex-btrblocks/src/schemes/float/rle.rs +++ b/vortex-btrblocks/src/schemes/float/rle.rs @@ -42,6 +42,10 @@ impl Scheme for FloatRLEScheme { vec![RLE.id()] } + fn returned_encodings(&self) -> Vec { + vec![RLE.id()] + } + /// Children: values=0, indices=1, offsets=2. fn num_children(&self) -> usize { 3 diff --git a/vortex-btrblocks/src/schemes/float/sparse.rs b/vortex-btrblocks/src/schemes/float/sparse.rs index b23ca87a321..acea68bb0c8 100644 --- a/vortex-btrblocks/src/schemes/float/sparse.rs +++ b/vortex-btrblocks/src/schemes/float/sparse.rs @@ -9,6 +9,7 @@ use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VTable; +use vortex_array::arrays::Constant; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::primitive::PrimitiveArrayExt; use vortex_compressor::scheme::ChildSelection; @@ -45,6 +46,10 @@ impl Scheme for NullDominatedSparseScheme { vec![Sparse.id()] } + fn returned_encodings(&self) -> Vec { + vec![Sparse.id(), Constant.id()] + } + /// Children: indices=0. fn num_children(&self) -> usize { 1 diff --git a/vortex-btrblocks/src/schemes/integer/bitpacking.rs b/vortex-btrblocks/src/schemes/integer/bitpacking.rs index 941c354da7c..39a89bae360 100644 --- a/vortex-btrblocks/src/schemes/integer/bitpacking.rs +++ b/vortex-btrblocks/src/schemes/integer/bitpacking.rs @@ -11,6 +11,7 @@ use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VTable; use vortex_array::arrays::Patched; +use vortex_array::arrays::Primitive; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::patched::use_experimental_patches; use vortex_array::arrays::primitive::PrimitiveArrayExt; @@ -104,6 +105,14 @@ impl Scheme for BitPackingScheme { encodings } + fn returned_encodings(&self) -> Vec { + let mut encodings = vec![Primitive.id(), BitPacked.id()]; + if use_experimental_patches() { + encodings.push(Patched.id()); + } + encodings + } + fn refine(&self, allowed: &dyn Fn(&ArrayId) -> bool) -> &dyn Scheme { if allowed(&bitpacked_v2_id()) { &BITPACKING_V2 diff --git a/vortex-btrblocks/src/schemes/integer/delta.rs b/vortex-btrblocks/src/schemes/integer/delta.rs index 17e430e1967..868390c5fee 100644 --- a/vortex-btrblocks/src/schemes/integer/delta.rs +++ b/vortex-btrblocks/src/schemes/integer/delta.rs @@ -101,6 +101,10 @@ impl Scheme for DeltaScheme { vec![Delta.id()] } + fn returned_encodings(&self) -> Vec { + vec![Delta.id()] + } + fn num_children(&self) -> usize { 2 } diff --git a/vortex-btrblocks/src/schemes/integer/for_.rs b/vortex-btrblocks/src/schemes/integer/for_.rs index c885bd97509..50329f5a82f 100644 --- a/vortex-btrblocks/src/schemes/integer/for_.rs +++ b/vortex-btrblocks/src/schemes/integer/for_.rs @@ -8,6 +8,7 @@ use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; +use vortex_array::VTable; use vortex_array::arrays::PrimitiveArray; use vortex_compressor::builtins::BinaryDictScheme; use vortex_compressor::builtins::FloatDictScheme; @@ -119,6 +120,10 @@ impl Scheme for FoRScheme { encodings } + fn returned_encodings(&self) -> Vec { + vec![FoR.id()] + } + fn refine(&self, allowed: &dyn Fn(&ArrayId) -> bool) -> &dyn Scheme { match (allowed(&for_v2_id()), allowed(&bitpacked_v2_id())) { (false, false) => &FOR_V1, diff --git a/vortex-btrblocks/src/schemes/integer/pco.rs b/vortex-btrblocks/src/schemes/integer/pco.rs index 5fb55c682d8..245c3e56714 100644 --- a/vortex-btrblocks/src/schemes/integer/pco.rs +++ b/vortex-btrblocks/src/schemes/integer/pco.rs @@ -36,8 +36,8 @@ impl Scheme for PcoScheme { vec![vortex_pco::Pco.id()] } - fn root_encoding(&self) -> Option { - Some(vortex_pco::Pco.id()) + fn returned_encodings(&self) -> Vec { + vec![vortex_pco::Pco.id()] } fn expected_compression_ratio( diff --git a/vortex-btrblocks/src/schemes/integer/rle.rs b/vortex-btrblocks/src/schemes/integer/rle.rs index 86a40d0b36a..4451885523e 100644 --- a/vortex-btrblocks/src/schemes/integer/rle.rs +++ b/vortex-btrblocks/src/schemes/integer/rle.rs @@ -110,6 +110,10 @@ impl Scheme for IntRLEScheme { vec![RLE.id()] } + fn returned_encodings(&self) -> Vec { + vec![RLE.id()] + } + /// Children: values=0, indices=1, offsets=2. fn num_children(&self) -> usize { 3 diff --git a/vortex-btrblocks/src/schemes/integer/runend.rs b/vortex-btrblocks/src/schemes/integer/runend.rs index 3ed592dce39..91e14d23705 100644 --- a/vortex-btrblocks/src/schemes/integer/runend.rs +++ b/vortex-btrblocks/src/schemes/integer/runend.rs @@ -52,8 +52,8 @@ impl Scheme for RunEndScheme { vec![RunEnd.id()] } - fn root_encoding(&self) -> Option { - Some(RunEnd.id()) + fn returned_encodings(&self) -> Vec { + vec![RunEnd.id()] } /// Children: values=0, ends=1. diff --git a/vortex-btrblocks/src/schemes/integer/sequence.rs b/vortex-btrblocks/src/schemes/integer/sequence.rs index edcefb99fc2..cf68cc502f4 100644 --- a/vortex-btrblocks/src/schemes/integer/sequence.rs +++ b/vortex-btrblocks/src/schemes/integer/sequence.rs @@ -47,6 +47,10 @@ impl Scheme for SequenceScheme { vec![Sequence.id()] } + fn returned_encodings(&self) -> Vec { + vec![Sequence.id()] + } + /// Sequence encoding on dictionary codes just adds a layer of indirection without compressing /// the data. Dict codes are compact integers that benefit from BitPacking or FoR, not from /// sequence detection. diff --git a/vortex-btrblocks/src/schemes/integer/sparse.rs b/vortex-btrblocks/src/schemes/integer/sparse.rs index 438f1084cda..31300f44d9b 100644 --- a/vortex-btrblocks/src/schemes/integer/sparse.rs +++ b/vortex-btrblocks/src/schemes/integer/sparse.rs @@ -50,6 +50,10 @@ impl Scheme for SparseScheme { vec![Sparse.id(), Constant.id()] } + fn returned_encodings(&self) -> Vec { + vec![Sparse.id(), Constant.id()] + } + fn stats_options(&self) -> GenerateStatsOptions { GenerateStatsOptions { count_distinct_values: true, diff --git a/vortex-btrblocks/src/schemes/integer/zigzag.rs b/vortex-btrblocks/src/schemes/integer/zigzag.rs index 87beb8064ea..1d8607b6dde 100644 --- a/vortex-btrblocks/src/schemes/integer/zigzag.rs +++ b/vortex-btrblocks/src/schemes/integer/zigzag.rs @@ -50,8 +50,8 @@ impl Scheme for ZigZagScheme { vec![ZigZag.id()] } - fn root_encoding(&self) -> Option { - Some(ZigZag.id()) + fn returned_encodings(&self) -> Vec { + vec![ZigZag.id()] } /// Children: encoded=0. diff --git a/vortex-btrblocks/src/schemes/string/fsst.rs b/vortex-btrblocks/src/schemes/string/fsst.rs index a4bb13f8f1c..1c234f58acd 100644 --- a/vortex-btrblocks/src/schemes/string/fsst.rs +++ b/vortex-btrblocks/src/schemes/string/fsst.rs @@ -54,8 +54,8 @@ impl Scheme for FSSTScheme { vec![FSST.id(), VarBin.id()] } - fn root_encoding(&self) -> Option { - Some(FSST.id()) + fn returned_encodings(&self) -> Vec { + vec![FSST.id()] } /// Children: lengths=0, code_offsets=1. diff --git a/vortex-btrblocks/src/schemes/string/onpair.rs b/vortex-btrblocks/src/schemes/string/onpair.rs index 256bc3d2d9e..8adeb0e158e 100644 --- a/vortex-btrblocks/src/schemes/string/onpair.rs +++ b/vortex-btrblocks/src/schemes/string/onpair.rs @@ -9,6 +9,7 @@ use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VTable; +use vortex_array::arrays::Constant; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::primitive::PrimitiveArrayExt; use vortex_compressor::scheme::CompressionEstimate; @@ -85,6 +86,10 @@ impl Scheme for OnPairScheme { vec![OnPair.id()] } + fn returned_encodings(&self) -> Vec { + vec![OnPair.id(), Constant.id()] + } + /// 4 primitive slot children flow through the cascading compressor: /// `dict_offsets` (u32 → typically `FoR`/`BitPacked`), `codes` (u16 → /// usually `FastLanes::BitPacked` after scheme selection), diff --git a/vortex-btrblocks/src/schemes/string/sparse.rs b/vortex-btrblocks/src/schemes/string/sparse.rs index 604018f98ba..dbda8c56440 100644 --- a/vortex-btrblocks/src/schemes/string/sparse.rs +++ b/vortex-btrblocks/src/schemes/string/sparse.rs @@ -9,6 +9,7 @@ use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VTable; +use vortex_array::arrays::Constant; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::primitive::PrimitiveArrayExt; use vortex_compressor::scheme::ChildSelection; @@ -46,6 +47,10 @@ impl Scheme for NullDominatedSparseScheme { vec![Sparse.id()] } + fn returned_encodings(&self) -> Vec { + vec![Sparse.id(), Constant.id()] + } + /// Children: indices=0. fn num_children(&self) -> usize { 1 diff --git a/vortex-btrblocks/src/schemes/string/zstd.rs b/vortex-btrblocks/src/schemes/string/zstd.rs index 84e8860d626..b441bac2d85 100644 --- a/vortex-btrblocks/src/schemes/string/zstd.rs +++ b/vortex-btrblocks/src/schemes/string/zstd.rs @@ -35,6 +35,10 @@ impl Scheme for ZstdScheme { vec![vortex_zstd::Zstd.id()] } + fn returned_encodings(&self) -> Vec { + vec![vortex_zstd::Zstd.id()] + } + fn expected_compression_ratio( &self, _data: &ArrayAndStats, diff --git a/vortex-btrblocks/src/schemes/string/zstd_buffers.rs b/vortex-btrblocks/src/schemes/string/zstd_buffers.rs index cf691c70fcb..f9b5d3d1481 100644 --- a/vortex-btrblocks/src/schemes/string/zstd_buffers.rs +++ b/vortex-btrblocks/src/schemes/string/zstd_buffers.rs @@ -35,6 +35,10 @@ impl Scheme for ZstdBuffersScheme { vec![vortex_zstd::ZstdBuffers.id()] } + fn returned_encodings(&self) -> Vec { + vec![vortex_zstd::ZstdBuffers.id()] + } + fn expected_compression_ratio( &self, _data: &ArrayAndStats, diff --git a/vortex-btrblocks/src/schemes/temporal.rs b/vortex-btrblocks/src/schemes/temporal.rs index 3273a3b3763..98cf6dd7eaa 100644 --- a/vortex-btrblocks/src/schemes/temporal.rs +++ b/vortex-btrblocks/src/schemes/temporal.rs @@ -59,8 +59,8 @@ impl Scheme for TemporalScheme { vec![DateTimeParts.id()] } - fn root_encoding(&self) -> Option { - Some(DateTimeParts.id()) + fn returned_encodings(&self) -> Vec { + vec![DateTimeParts.id()] } /// Children: days=0, seconds=1, subseconds=2. diff --git a/vortex-btrblocks/src/tests.rs b/vortex-btrblocks/src/tests.rs index e75fbe81bbc..e7d45100d9c 100644 --- a/vortex-btrblocks/src/tests.rs +++ b/vortex-btrblocks/src/tests.rs @@ -203,7 +203,8 @@ fn test_cuda_binary_zstd_follows_editions( Ok(()) } -/// `encoding_for` gives the encoding of the array that `compress` returns. +/// `encodings_for` includes the encoding of the array that `compress` returns, and includes +/// `Dict` only when that array is a dictionary. #[rstest] #[case::unique_strings( VarBinViewArray::from_iter_str((0..4096).map(|i| format!("{i:04} is a unique story"))).into_array(), @@ -218,16 +219,20 @@ fn test_cuda_binary_zstd_follows_editions( )] #[case::few_integers(PrimitiveArray::from_iter((0..4096i64).map(|i| (i % 3) * 1_000_000)).into_array())] #[case::unique_integers(PrimitiveArray::from_iter(0..4096i64).into_array())] -fn test_encoding_for(#[case] input: ArrayRef) -> VortexResult<()> { +fn test_encodings_for(#[case] input: ArrayRef) -> VortexResult<()> { let compressor = BtrBlocksCompressorBuilder::from_session(&SESSION) .unrestricted() .build(); let mut ctx = SESSION.create_execution_ctx(); - let encoding = compressor.0.encoding_for(&input, &mut ctx)?; + let encodings = compressor.0.encodings_for(&input, &mut ctx)?; let compressed = compressor.compress(&input, &mut ctx)?; - assert_eq!(encoding, compressed.encoding_id()); + assert!(encodings.contains(&compressed.encoding_id())); + assert_eq!( + encodings.contains(&vortex_array::VTable::id(&Dict)), + compressed.is::() + ); Ok(()) } @@ -242,16 +247,16 @@ fn test_encoding_for(#[case] input: ArrayRef) -> VortexResult<()> { VarBinViewArray::from_iter_str((0..4096).map(|i| format!("{i:04} is a unique story"))).into_array(), false, )] -fn test_dict_only_encoding_for(#[case] input: ArrayRef, #[case] is_dict: bool) -> VortexResult<()> { +fn test_dict_only_encodings_for(#[case] input: ArrayRef, #[case] is_dict: bool) -> VortexResult<()> { let compressor = BtrBlocksCompressorBuilder::empty() .with_new_scheme(&StringDictScheme) .with_new_scheme(&IntDictScheme) .build(); let mut ctx = SESSION.create_execution_ctx(); - let encoding = compressor.0.encoding_for(&input, &mut ctx)?; + let encodings = compressor.0.encodings_for(&input, &mut ctx)?; - assert_eq!(encoding == vortex_array::VTable::id(&Dict), is_dict); + assert_eq!(encodings.contains(&vortex_array::VTable::id(&Dict)), is_dict); Ok(()) } diff --git a/vortex-compressor/src/builtins/dict/binary.rs b/vortex-compressor/src/builtins/dict/binary.rs index b15f6cb5a42..52ff339063b 100644 --- a/vortex-compressor/src/builtins/dict/binary.rs +++ b/vortex-compressor/src/builtins/dict/binary.rs @@ -52,8 +52,8 @@ impl Scheme for BinaryDictScheme { vec![Dict.id()] } - fn root_encoding(&self) -> Option { - Some(Dict.id()) + fn returned_encodings(&self) -> Vec { + vec![Dict.id()] } fn stats_options(&self) -> GenerateStatsOptions { diff --git a/vortex-compressor/src/builtins/dict/float.rs b/vortex-compressor/src/builtins/dict/float.rs index a436446ba66..f20d4ca6659 100644 --- a/vortex-compressor/src/builtins/dict/float.rs +++ b/vortex-compressor/src/builtins/dict/float.rs @@ -58,8 +58,8 @@ impl Scheme for FloatDictScheme { vec![Dict.id()] } - fn root_encoding(&self) -> Option { - Some(Dict.id()) + fn returned_encodings(&self) -> Vec { + vec![Dict.id()] } fn stats_options(&self) -> GenerateStatsOptions { diff --git a/vortex-compressor/src/builtins/dict/integer.rs b/vortex-compressor/src/builtins/dict/integer.rs index 56dccd71ea0..75b0b823a3c 100644 --- a/vortex-compressor/src/builtins/dict/integer.rs +++ b/vortex-compressor/src/builtins/dict/integer.rs @@ -53,8 +53,8 @@ impl Scheme for IntDictScheme { vec![Dict.id()] } - fn root_encoding(&self) -> Option { - Some(Dict.id()) + fn returned_encodings(&self) -> Vec { + vec![Dict.id()] } fn stats_options(&self) -> GenerateStatsOptions { diff --git a/vortex-compressor/src/builtins/dict/string.rs b/vortex-compressor/src/builtins/dict/string.rs index 2ea8d8848b8..73289e28232 100644 --- a/vortex-compressor/src/builtins/dict/string.rs +++ b/vortex-compressor/src/builtins/dict/string.rs @@ -52,8 +52,8 @@ impl Scheme for StringDictScheme { vec![Dict.id()] } - fn root_encoding(&self) -> Option { - Some(Dict.id()) + fn returned_encodings(&self) -> Vec { + vec![Dict.id()] } fn stats_options(&self) -> GenerateStatsOptions { diff --git a/vortex-compressor/src/compressor/cascade.rs b/vortex-compressor/src/compressor/cascade.rs index 465cb596b5b..e3d51477e15 100644 --- a/vortex-compressor/src/compressor/cascade.rs +++ b/vortex-compressor/src/compressor/cascade.rs @@ -33,6 +33,7 @@ use vortex_array::arrays::variant::VariantArraySlotsExt; use vortex_array::scalar::Scalar; use vortex_error::VortexResult; use vortex_error::vortex_err; +use vortex_utils::aliases::hash_set::HashSet; use super::CascadingCompressor; use super::constant; @@ -298,25 +299,25 @@ impl CascadingCompressor { } } - /// Returns the encoding of the array that [`compress`](Self::compress) returns for `array`. + /// Returns the encodings that the array [`compress`](Self::compress) returns for `array` can + /// have. /// /// For a leaf array, this runs the same scheme selection as `compress`, with the statistics /// and the sample compressions that the selection needs. If the winning scheme declares its - /// [`root_encoding`](Scheme::root_encoding), the full compression does not occur. If it does - /// not, or if `array` is not a leaf array, this compresses `array`. - /// - /// Without the full compression, the check that the winner's output is smaller than `array` - /// cannot occur. Thus, if the output of the winner is not smaller, `compress` returns the - /// uncompressed array, but this method returns the encoding of the winner. + /// [`returned_encodings`](Scheme::returned_encodings), the full compression does not occur. + /// The result is then the encodings of the winner and the encoding of the uncompressed array, + /// because `compress` returns the uncompressed array if the output of the winner is not + /// smaller. If the winner does not declare its encodings, or if `array` is not a leaf array, + /// this compresses `array` and returns the one encoding of the result. /// /// # Errors /// /// Returns an error if canonicalization or compression fails. - pub fn encoding_for( + pub fn encodings_for( &self, array: &ArrayRef, exec_ctx: &mut ExecutionCtx, - ) -> VortexResult { + ) -> VortexResult> { let canonical = array.clone().execute::(exec_ctx)?.0; let compact = canonical.compact(exec_ctx)?; @@ -329,21 +330,27 @@ impl CascadingCompressor { ); if !is_leaf { - return Ok(self - .compress_canonical(compact, CompressorContext::new(), exec_ctx)? - .encoding_id()); + let compressed = self.compress_canonical(compact, CompressorContext::new(), exec_ctx)?; + return Ok([compressed.encoding_id()].into_iter().collect()); } - match self.choose(compact, CompressorContext::new(), exec_ctx)? { - Choice::Decided(selection) => Ok(selection.into_array().encoding_id()), - Choice::Scheme(winner) => match winner.scheme.root_encoding() { - Some(encoding) => Ok(encoding), - None => Ok(self - .compress_with_winner(winner, exec_ctx)? - .into_array() - .encoding_id()), - }, + let winner = match self.choose(compact, CompressorContext::new(), exec_ctx)? { + Choice::Decided(selection) => { + return Ok([selection.into_array().encoding_id()].into_iter().collect()); + } + Choice::Scheme(winner) => winner, + }; + + let declared = winner.scheme.returned_encodings(); + if declared.is_empty() { + let compressed = self.compress_with_winner(winner, exec_ctx)?.into_array(); + return Ok([compressed.encoding_id()].into_iter().collect()); } + + let mut encodings: HashSet = declared.into_iter().collect(); + encodings.insert(winner.data.array().encoding_id()); + + Ok(encodings) } /// Chooses how to compress a single leaf array, but does not run the winning scheme. @@ -444,6 +451,16 @@ impl CascadingCompressor { trace::scheme_compress_failed(scheme.id(), before_nbytes, error_ctx.as_ref(), err); })?; + debug_assert!( + { + let declared = scheme.returned_encodings(); + declared.is_empty() || declared.contains(&compressed.encoding_id()) + }, + "scheme {} returned encoding {}, which is not in its returned_encodings", + scheme.id(), + compressed.encoding_id(), + ); + let after_nbytes = compressed.nbytes(); let actual_ratio = (after_nbytes != 0).then(|| before_nbytes as f64 / after_nbytes as f64); diff --git a/vortex-compressor/src/scheme/mod.rs b/vortex-compressor/src/scheme/mod.rs index a3a44b193ba..a9bfb7e421d 100644 --- a/vortex-compressor/src/scheme/mod.rs +++ b/vortex-compressor/src/scheme/mod.rs @@ -135,13 +135,14 @@ pub trait Scheme: AsDynScheme + Debug + Send + Sync { /// formats declares the wire IDs the scheme writes, which may differ from its in-memory ID. fn produced_encodings(&self) -> Vec; - /// The in-memory encoding ID of the array that [`compress`](Scheme::compress) returns, if it - /// is the same for all inputs. + /// The in-memory encoding IDs of all arrays that [`compress`](Scheme::compress) can return. /// - /// The compressor uses this to tell which encoding it produces without a full compression. - /// The default is `None`, which makes the compressor compress to find the encoding. - fn root_encoding(&self) -> Option { - None + /// The list must include every return path, such as a canonical input returned unchanged or + /// a constant array returned early. The compressor uses this to tell which encodings it can + /// produce without a full compression. The default is an empty list, which means the scheme + /// does not declare its encodings, and the compressor compresses to find the encoding. + fn returned_encodings(&self) -> Vec { + Vec::new() } /// Returns the variant of this scheme to use given which serialized IDs are `allowed`. diff --git a/vortex-layout/src/layouts/compressed.rs b/vortex-layout/src/layouts/compressed.rs index 76c5fc586f3..96031662086 100644 --- a/vortex-layout/src/layouts/compressed.rs +++ b/vortex-layout/src/layouts/compressed.rs @@ -16,6 +16,7 @@ use vortex_btrblocks::BtrBlocksCompressor; use vortex_error::VortexResult; use vortex_io::session::RuntimeSessionExt; use vortex_session::VortexSession; +use vortex_utils::aliases::hash_set::HashSet; use vortex_utils::parallelism::get_available_parallelism; use crate::LayoutRef; @@ -33,13 +34,20 @@ use crate::sequence::SequentialStreamExt; pub trait CompressorPlugin: Send + Sync + 'static { fn compress_chunk(&self, chunk: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult; - /// Returns the encoding of the array that [`compress_chunk`](Self::compress_chunk) returns - /// for `chunk`. + /// Returns the encodings that the array [`compress_chunk`](Self::compress_chunk) returns for + /// `chunk` can have. /// - /// A writer can use this to choose a layout from the encoding the plugin produces. The default - /// compresses `chunk`. A plugin can override it to give the answer without a full compression. - fn encoding_for(&self, chunk: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { - Ok(self.compress_chunk(chunk, ctx)?.encoding_id()) + /// A writer can use this to choose a layout from the encodings the plugin can produce. The + /// default compresses `chunk` and returns the one encoding of the result. A plugin can + /// override it to give the answer without a full compression. + fn encodings_for( + &self, + chunk: &ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + Ok([self.compress_chunk(chunk, ctx)?.encoding_id()] + .into_iter() + .collect()) } } @@ -48,8 +56,12 @@ impl CompressorPlugin for Arc { self.as_ref().compress_chunk(chunk, ctx) } - fn encoding_for(&self, chunk: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { - self.as_ref().encoding_for(chunk, ctx) + fn encodings_for( + &self, + chunk: &ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + self.as_ref().encodings_for(chunk, ctx) } } @@ -67,8 +79,12 @@ impl CompressorPlugin for BtrBlocksCompressor { self.compress(chunk, ctx) } - fn encoding_for(&self, chunk: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { - self.0.encoding_for(chunk, ctx) + fn encodings_for( + &self, + chunk: &ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + self.0.encodings_for(chunk, ctx) } } diff --git a/vortex-layout/src/layouts/dict/writer.rs b/vortex-layout/src/layouts/dict/writer.rs index 62ff6871561..f61da19a4aa 100644 --- a/vortex-layout/src/layouts/dict/writer.rs +++ b/vortex-layout/src/layouts/dict/writer.rs @@ -174,7 +174,10 @@ impl LayoutStrategy for DictStrategy { None => true, // empty stream Some(sample) => { let mut exec_ctx = session.create_execution_ctx(); - self.probe_compressor.encoding_for(&sample, &mut exec_ctx)? != Dict.id() + !self + .probe_compressor + .encodings_for(&sample, &mut exec_ctx)? + .contains(&Dict.id()) } }; if should_fallback { From cae8a2641f0e64b80fbcebae485073850ccd2bd8 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Tue, 10 Nov 2026 17:51:23 +0000 Subject: [PATCH 5/6] Rename to top_level_encodings and test the fast path Rename `Scheme::returned_encodings` to `top_level_encodings`. Add a test that `encodings_for` does not compress when the winner declares its top-level encodings, and that it compresses when the winner declares nothing. Remove `test_dict_only_encodings_for`, which repeated cases of `test_encodings_for`. Signed-off-by: Robert Kruszewski --- vortex-btrblocks/src/schemes/binary/varbin.rs | 2 +- vortex-btrblocks/src/schemes/binary/zstd.rs | 2 +- .../src/schemes/binary/zstd_buffers.rs | 2 +- vortex-btrblocks/src/schemes/decimal.rs | 2 +- vortex-btrblocks/src/schemes/float/alp.rs | 2 +- vortex-btrblocks/src/schemes/float/alprd.rs | 2 +- vortex-btrblocks/src/schemes/float/pco.rs | 2 +- vortex-btrblocks/src/schemes/float/rle.rs | 2 +- vortex-btrblocks/src/schemes/float/sparse.rs | 2 +- .../src/schemes/integer/bitpacking.rs | 2 +- vortex-btrblocks/src/schemes/integer/delta.rs | 2 +- vortex-btrblocks/src/schemes/integer/for_.rs | 2 +- vortex-btrblocks/src/schemes/integer/pco.rs | 2 +- vortex-btrblocks/src/schemes/integer/rle.rs | 2 +- .../src/schemes/integer/runend.rs | 2 +- .../src/schemes/integer/sequence.rs | 2 +- .../src/schemes/integer/sparse.rs | 2 +- .../src/schemes/integer/zigzag.rs | 2 +- vortex-btrblocks/src/schemes/string/fsst.rs | 2 +- vortex-btrblocks/src/schemes/string/onpair.rs | 2 +- vortex-btrblocks/src/schemes/string/sparse.rs | 2 +- vortex-btrblocks/src/schemes/string/zstd.rs | 2 +- .../src/schemes/string/zstd_buffers.rs | 2 +- vortex-btrblocks/src/schemes/temporal.rs | 2 +- vortex-btrblocks/src/tests.rs | 31 +------ vortex-compressor/src/builtins/dict/binary.rs | 2 +- vortex-compressor/src/builtins/dict/float.rs | 2 +- .../src/builtins/dict/integer.rs | 2 +- vortex-compressor/src/builtins/dict/string.rs | 2 +- vortex-compressor/src/compressor/cascade.rs | 8 +- vortex-compressor/src/compressor/tests.rs | 91 +++++++++++++++++++ vortex-compressor/src/scheme/mod.rs | 2 +- 32 files changed, 127 insertions(+), 61 deletions(-) diff --git a/vortex-btrblocks/src/schemes/binary/varbin.rs b/vortex-btrblocks/src/schemes/binary/varbin.rs index c48d036837e..051817cb02b 100644 --- a/vortex-btrblocks/src/schemes/binary/varbin.rs +++ b/vortex-btrblocks/src/schemes/binary/varbin.rs @@ -48,7 +48,7 @@ impl Scheme for VarBinScheme { vec![VarBin.id()] } - fn returned_encodings(&self) -> Vec { + fn top_level_encodings(&self) -> Vec { vec![VarBin.id()] } diff --git a/vortex-btrblocks/src/schemes/binary/zstd.rs b/vortex-btrblocks/src/schemes/binary/zstd.rs index e58200551e6..5085d7f9fc6 100644 --- a/vortex-btrblocks/src/schemes/binary/zstd.rs +++ b/vortex-btrblocks/src/schemes/binary/zstd.rs @@ -35,7 +35,7 @@ impl Scheme for ZstdScheme { vec![vortex_zstd::Zstd.id()] } - fn returned_encodings(&self) -> Vec { + fn top_level_encodings(&self) -> Vec { vec![vortex_zstd::Zstd.id()] } diff --git a/vortex-btrblocks/src/schemes/binary/zstd_buffers.rs b/vortex-btrblocks/src/schemes/binary/zstd_buffers.rs index 5425e3bf9e2..0043df7c385 100644 --- a/vortex-btrblocks/src/schemes/binary/zstd_buffers.rs +++ b/vortex-btrblocks/src/schemes/binary/zstd_buffers.rs @@ -35,7 +35,7 @@ impl Scheme for ZstdBuffersScheme { vec![vortex_zstd::ZstdBuffers.id()] } - fn returned_encodings(&self) -> Vec { + fn top_level_encodings(&self) -> Vec { vec![vortex_zstd::ZstdBuffers.id()] } diff --git a/vortex-btrblocks/src/schemes/decimal.rs b/vortex-btrblocks/src/schemes/decimal.rs index 0b20edb7d48..d110834e903 100644 --- a/vortex-btrblocks/src/schemes/decimal.rs +++ b/vortex-btrblocks/src/schemes/decimal.rs @@ -98,7 +98,7 @@ impl Scheme for DecimalScheme { } } - fn returned_encodings(&self) -> Vec { + fn top_level_encodings(&self) -> Vec { match self.mode { DecimalSchemeMode::V1 => vec![Decimal.id(), DecimalByteParts.id()], DecimalSchemeMode::V2 => vec![DecimalByteParts.id()], diff --git a/vortex-btrblocks/src/schemes/float/alp.rs b/vortex-btrblocks/src/schemes/float/alp.rs index 86852e27245..d3809e34a24 100644 --- a/vortex-btrblocks/src/schemes/float/alp.rs +++ b/vortex-btrblocks/src/schemes/float/alp.rs @@ -50,7 +50,7 @@ impl Scheme for ALPScheme { encodings } - fn returned_encodings(&self) -> Vec { + fn top_level_encodings(&self) -> Vec { let mut encodings = vec![ALP.id()]; if use_experimental_patches() { encodings.push(Patched.id()); diff --git a/vortex-btrblocks/src/schemes/float/alprd.rs b/vortex-btrblocks/src/schemes/float/alprd.rs index c368e950512..de2ab85ece6 100644 --- a/vortex-btrblocks/src/schemes/float/alprd.rs +++ b/vortex-btrblocks/src/schemes/float/alprd.rs @@ -44,7 +44,7 @@ impl Scheme for ALPRDScheme { vec![vortex_alp::ALPRD.id()] } - fn returned_encodings(&self) -> Vec { + fn top_level_encodings(&self) -> Vec { vec![vortex_alp::ALPRD.id()] } diff --git a/vortex-btrblocks/src/schemes/float/pco.rs b/vortex-btrblocks/src/schemes/float/pco.rs index f1e69da3182..f994c16dc31 100644 --- a/vortex-btrblocks/src/schemes/float/pco.rs +++ b/vortex-btrblocks/src/schemes/float/pco.rs @@ -35,7 +35,7 @@ impl Scheme for PcoScheme { vec![vortex_pco::Pco.id()] } - fn returned_encodings(&self) -> Vec { + fn top_level_encodings(&self) -> Vec { vec![vortex_pco::Pco.id()] } diff --git a/vortex-btrblocks/src/schemes/float/rle.rs b/vortex-btrblocks/src/schemes/float/rle.rs index 313f7123b6c..9923df1cac7 100644 --- a/vortex-btrblocks/src/schemes/float/rle.rs +++ b/vortex-btrblocks/src/schemes/float/rle.rs @@ -42,7 +42,7 @@ impl Scheme for FloatRLEScheme { vec![RLE.id()] } - fn returned_encodings(&self) -> Vec { + fn top_level_encodings(&self) -> Vec { vec![RLE.id()] } diff --git a/vortex-btrblocks/src/schemes/float/sparse.rs b/vortex-btrblocks/src/schemes/float/sparse.rs index acea68bb0c8..7091c6c57c7 100644 --- a/vortex-btrblocks/src/schemes/float/sparse.rs +++ b/vortex-btrblocks/src/schemes/float/sparse.rs @@ -46,7 +46,7 @@ impl Scheme for NullDominatedSparseScheme { vec![Sparse.id()] } - fn returned_encodings(&self) -> Vec { + fn top_level_encodings(&self) -> Vec { vec![Sparse.id(), Constant.id()] } diff --git a/vortex-btrblocks/src/schemes/integer/bitpacking.rs b/vortex-btrblocks/src/schemes/integer/bitpacking.rs index 39a89bae360..0f76d8d37d2 100644 --- a/vortex-btrblocks/src/schemes/integer/bitpacking.rs +++ b/vortex-btrblocks/src/schemes/integer/bitpacking.rs @@ -105,7 +105,7 @@ impl Scheme for BitPackingScheme { encodings } - fn returned_encodings(&self) -> Vec { + fn top_level_encodings(&self) -> Vec { let mut encodings = vec![Primitive.id(), BitPacked.id()]; if use_experimental_patches() { encodings.push(Patched.id()); diff --git a/vortex-btrblocks/src/schemes/integer/delta.rs b/vortex-btrblocks/src/schemes/integer/delta.rs index 868390c5fee..edf09b5615c 100644 --- a/vortex-btrblocks/src/schemes/integer/delta.rs +++ b/vortex-btrblocks/src/schemes/integer/delta.rs @@ -101,7 +101,7 @@ impl Scheme for DeltaScheme { vec![Delta.id()] } - fn returned_encodings(&self) -> Vec { + fn top_level_encodings(&self) -> Vec { vec![Delta.id()] } diff --git a/vortex-btrblocks/src/schemes/integer/for_.rs b/vortex-btrblocks/src/schemes/integer/for_.rs index 50329f5a82f..5267573dc68 100644 --- a/vortex-btrblocks/src/schemes/integer/for_.rs +++ b/vortex-btrblocks/src/schemes/integer/for_.rs @@ -120,7 +120,7 @@ impl Scheme for FoRScheme { encodings } - fn returned_encodings(&self) -> Vec { + fn top_level_encodings(&self) -> Vec { vec![FoR.id()] } diff --git a/vortex-btrblocks/src/schemes/integer/pco.rs b/vortex-btrblocks/src/schemes/integer/pco.rs index 245c3e56714..e4faa3a5cae 100644 --- a/vortex-btrblocks/src/schemes/integer/pco.rs +++ b/vortex-btrblocks/src/schemes/integer/pco.rs @@ -36,7 +36,7 @@ impl Scheme for PcoScheme { vec![vortex_pco::Pco.id()] } - fn returned_encodings(&self) -> Vec { + fn top_level_encodings(&self) -> Vec { vec![vortex_pco::Pco.id()] } diff --git a/vortex-btrblocks/src/schemes/integer/rle.rs b/vortex-btrblocks/src/schemes/integer/rle.rs index 4451885523e..b55adb16dc5 100644 --- a/vortex-btrblocks/src/schemes/integer/rle.rs +++ b/vortex-btrblocks/src/schemes/integer/rle.rs @@ -110,7 +110,7 @@ impl Scheme for IntRLEScheme { vec![RLE.id()] } - fn returned_encodings(&self) -> Vec { + fn top_level_encodings(&self) -> Vec { vec![RLE.id()] } diff --git a/vortex-btrblocks/src/schemes/integer/runend.rs b/vortex-btrblocks/src/schemes/integer/runend.rs index 91e14d23705..a6e0c24d0ca 100644 --- a/vortex-btrblocks/src/schemes/integer/runend.rs +++ b/vortex-btrblocks/src/schemes/integer/runend.rs @@ -52,7 +52,7 @@ impl Scheme for RunEndScheme { vec![RunEnd.id()] } - fn returned_encodings(&self) -> Vec { + fn top_level_encodings(&self) -> Vec { vec![RunEnd.id()] } diff --git a/vortex-btrblocks/src/schemes/integer/sequence.rs b/vortex-btrblocks/src/schemes/integer/sequence.rs index cf68cc502f4..4fa20e764a7 100644 --- a/vortex-btrblocks/src/schemes/integer/sequence.rs +++ b/vortex-btrblocks/src/schemes/integer/sequence.rs @@ -47,7 +47,7 @@ impl Scheme for SequenceScheme { vec![Sequence.id()] } - fn returned_encodings(&self) -> Vec { + fn top_level_encodings(&self) -> Vec { vec![Sequence.id()] } diff --git a/vortex-btrblocks/src/schemes/integer/sparse.rs b/vortex-btrblocks/src/schemes/integer/sparse.rs index 31300f44d9b..4ebd4be18d3 100644 --- a/vortex-btrblocks/src/schemes/integer/sparse.rs +++ b/vortex-btrblocks/src/schemes/integer/sparse.rs @@ -50,7 +50,7 @@ impl Scheme for SparseScheme { vec![Sparse.id(), Constant.id()] } - fn returned_encodings(&self) -> Vec { + fn top_level_encodings(&self) -> Vec { vec![Sparse.id(), Constant.id()] } diff --git a/vortex-btrblocks/src/schemes/integer/zigzag.rs b/vortex-btrblocks/src/schemes/integer/zigzag.rs index 1d8607b6dde..df809105334 100644 --- a/vortex-btrblocks/src/schemes/integer/zigzag.rs +++ b/vortex-btrblocks/src/schemes/integer/zigzag.rs @@ -50,7 +50,7 @@ impl Scheme for ZigZagScheme { vec![ZigZag.id()] } - fn returned_encodings(&self) -> Vec { + fn top_level_encodings(&self) -> Vec { vec![ZigZag.id()] } diff --git a/vortex-btrblocks/src/schemes/string/fsst.rs b/vortex-btrblocks/src/schemes/string/fsst.rs index 1c234f58acd..36e9f021da5 100644 --- a/vortex-btrblocks/src/schemes/string/fsst.rs +++ b/vortex-btrblocks/src/schemes/string/fsst.rs @@ -54,7 +54,7 @@ impl Scheme for FSSTScheme { vec![FSST.id(), VarBin.id()] } - fn returned_encodings(&self) -> Vec { + fn top_level_encodings(&self) -> Vec { vec![FSST.id()] } diff --git a/vortex-btrblocks/src/schemes/string/onpair.rs b/vortex-btrblocks/src/schemes/string/onpair.rs index 8adeb0e158e..ce1617b87af 100644 --- a/vortex-btrblocks/src/schemes/string/onpair.rs +++ b/vortex-btrblocks/src/schemes/string/onpair.rs @@ -86,7 +86,7 @@ impl Scheme for OnPairScheme { vec![OnPair.id()] } - fn returned_encodings(&self) -> Vec { + fn top_level_encodings(&self) -> Vec { vec![OnPair.id(), Constant.id()] } diff --git a/vortex-btrblocks/src/schemes/string/sparse.rs b/vortex-btrblocks/src/schemes/string/sparse.rs index dbda8c56440..534a780117e 100644 --- a/vortex-btrblocks/src/schemes/string/sparse.rs +++ b/vortex-btrblocks/src/schemes/string/sparse.rs @@ -47,7 +47,7 @@ impl Scheme for NullDominatedSparseScheme { vec![Sparse.id()] } - fn returned_encodings(&self) -> Vec { + fn top_level_encodings(&self) -> Vec { vec![Sparse.id(), Constant.id()] } diff --git a/vortex-btrblocks/src/schemes/string/zstd.rs b/vortex-btrblocks/src/schemes/string/zstd.rs index b441bac2d85..0aa99dcd474 100644 --- a/vortex-btrblocks/src/schemes/string/zstd.rs +++ b/vortex-btrblocks/src/schemes/string/zstd.rs @@ -35,7 +35,7 @@ impl Scheme for ZstdScheme { vec![vortex_zstd::Zstd.id()] } - fn returned_encodings(&self) -> Vec { + fn top_level_encodings(&self) -> Vec { vec![vortex_zstd::Zstd.id()] } diff --git a/vortex-btrblocks/src/schemes/string/zstd_buffers.rs b/vortex-btrblocks/src/schemes/string/zstd_buffers.rs index f9b5d3d1481..3074abdebd0 100644 --- a/vortex-btrblocks/src/schemes/string/zstd_buffers.rs +++ b/vortex-btrblocks/src/schemes/string/zstd_buffers.rs @@ -35,7 +35,7 @@ impl Scheme for ZstdBuffersScheme { vec![vortex_zstd::ZstdBuffers.id()] } - fn returned_encodings(&self) -> Vec { + fn top_level_encodings(&self) -> Vec { vec![vortex_zstd::ZstdBuffers.id()] } diff --git a/vortex-btrblocks/src/schemes/temporal.rs b/vortex-btrblocks/src/schemes/temporal.rs index 98cf6dd7eaa..20779f46c5e 100644 --- a/vortex-btrblocks/src/schemes/temporal.rs +++ b/vortex-btrblocks/src/schemes/temporal.rs @@ -59,7 +59,7 @@ impl Scheme for TemporalScheme { vec![DateTimeParts.id()] } - fn returned_encodings(&self) -> Vec { + fn top_level_encodings(&self) -> Vec { vec![DateTimeParts.id()] } diff --git a/vortex-btrblocks/src/tests.rs b/vortex-btrblocks/src/tests.rs index e7d45100d9c..2dc4d93400c 100644 --- a/vortex-btrblocks/src/tests.rs +++ b/vortex-btrblocks/src/tests.rs @@ -51,8 +51,6 @@ use vortex_session::VortexSession; use crate::BtrBlocksCompressor; use crate::BtrBlocksCompressorBuilder; use crate::SESSION; -use crate::schemes::integer::IntDictScheme; -use crate::schemes::string::StringDictScheme; #[track_caller] fn assert_roundtrip(compressor: &BtrBlocksCompressor, input: &ArrayRef) -> VortexResult { @@ -203,8 +201,9 @@ fn test_cuda_binary_zstd_follows_editions( Ok(()) } -/// `encodings_for` includes the encoding of the array that `compress` returns, and includes -/// `Dict` only when that array is a dictionary. +/// With the real schemes, `encodings_for` agrees with `compress`: the set includes the encoding +/// of the compressed array, and includes `Dict` only when that array is a dictionary. The dict +/// layout writer relies on the second property. #[rstest] #[case::unique_strings( VarBinViewArray::from_iter_str((0..4096).map(|i| format!("{i:04} is a unique story"))).into_array(), @@ -236,27 +235,3 @@ fn test_encodings_for(#[case] input: ArrayRef) -> VortexResult<()> { Ok(()) } - -/// A compressor that has only the dictionary schemes produces a dictionary or nothing. -#[rstest] -#[case::few_strings( - VarBinViewArray::from_iter_str((0..4096).map(|i| ["red", "green", "blue"][i % 3])).into_array(), - true, -)] -#[case::unique_strings( - VarBinViewArray::from_iter_str((0..4096).map(|i| format!("{i:04} is a unique story"))).into_array(), - false, -)] -fn test_dict_only_encodings_for(#[case] input: ArrayRef, #[case] is_dict: bool) -> VortexResult<()> { - let compressor = BtrBlocksCompressorBuilder::empty() - .with_new_scheme(&StringDictScheme) - .with_new_scheme(&IntDictScheme) - .build(); - let mut ctx = SESSION.create_execution_ctx(); - - let encodings = compressor.0.encodings_for(&input, &mut ctx)?; - - assert_eq!(encodings.contains(&vortex_array::VTable::id(&Dict)), is_dict); - - Ok(()) -} diff --git a/vortex-compressor/src/builtins/dict/binary.rs b/vortex-compressor/src/builtins/dict/binary.rs index 52ff339063b..2a33477c733 100644 --- a/vortex-compressor/src/builtins/dict/binary.rs +++ b/vortex-compressor/src/builtins/dict/binary.rs @@ -52,7 +52,7 @@ impl Scheme for BinaryDictScheme { vec![Dict.id()] } - fn returned_encodings(&self) -> Vec { + fn top_level_encodings(&self) -> Vec { vec![Dict.id()] } diff --git a/vortex-compressor/src/builtins/dict/float.rs b/vortex-compressor/src/builtins/dict/float.rs index f20d4ca6659..10427cb5383 100644 --- a/vortex-compressor/src/builtins/dict/float.rs +++ b/vortex-compressor/src/builtins/dict/float.rs @@ -58,7 +58,7 @@ impl Scheme for FloatDictScheme { vec![Dict.id()] } - fn returned_encodings(&self) -> Vec { + fn top_level_encodings(&self) -> Vec { vec![Dict.id()] } diff --git a/vortex-compressor/src/builtins/dict/integer.rs b/vortex-compressor/src/builtins/dict/integer.rs index 75b0b823a3c..6a2865e203d 100644 --- a/vortex-compressor/src/builtins/dict/integer.rs +++ b/vortex-compressor/src/builtins/dict/integer.rs @@ -53,7 +53,7 @@ impl Scheme for IntDictScheme { vec![Dict.id()] } - fn returned_encodings(&self) -> Vec { + fn top_level_encodings(&self) -> Vec { vec![Dict.id()] } diff --git a/vortex-compressor/src/builtins/dict/string.rs b/vortex-compressor/src/builtins/dict/string.rs index 73289e28232..42c9d44fe50 100644 --- a/vortex-compressor/src/builtins/dict/string.rs +++ b/vortex-compressor/src/builtins/dict/string.rs @@ -52,7 +52,7 @@ impl Scheme for StringDictScheme { vec![Dict.id()] } - fn returned_encodings(&self) -> Vec { + fn top_level_encodings(&self) -> Vec { vec![Dict.id()] } diff --git a/vortex-compressor/src/compressor/cascade.rs b/vortex-compressor/src/compressor/cascade.rs index e3d51477e15..768149c52a7 100644 --- a/vortex-compressor/src/compressor/cascade.rs +++ b/vortex-compressor/src/compressor/cascade.rs @@ -304,7 +304,7 @@ impl CascadingCompressor { /// /// For a leaf array, this runs the same scheme selection as `compress`, with the statistics /// and the sample compressions that the selection needs. If the winning scheme declares its - /// [`returned_encodings`](Scheme::returned_encodings), the full compression does not occur. + /// [`top_level_encodings`](Scheme::top_level_encodings), the full compression does not occur. /// The result is then the encodings of the winner and the encoding of the uncompressed array, /// because `compress` returns the uncompressed array if the output of the winner is not /// smaller. If the winner does not declare its encodings, or if `array` is not a leaf array, @@ -341,7 +341,7 @@ impl CascadingCompressor { Choice::Scheme(winner) => winner, }; - let declared = winner.scheme.returned_encodings(); + let declared = winner.scheme.top_level_encodings(); if declared.is_empty() { let compressed = self.compress_with_winner(winner, exec_ctx)?.into_array(); return Ok([compressed.encoding_id()].into_iter().collect()); @@ -453,10 +453,10 @@ impl CascadingCompressor { debug_assert!( { - let declared = scheme.returned_encodings(); + let declared = scheme.top_level_encodings(); declared.is_empty() || declared.contains(&compressed.encoding_id()) }, - "scheme {} returned encoding {}, which is not in its returned_encodings", + "scheme {} returned encoding {}, which is not in its top_level_encodings", scheme.id(), compressed.encoding_id(), ); diff --git a/vortex-compressor/src/compressor/tests.rs b/vortex-compressor/src/compressor/tests.rs index 29e32b433c7..099dd3402fc 100644 --- a/vortex-compressor/src/compressor/tests.rs +++ b/vortex-compressor/src/compressor/tests.rs @@ -5,21 +5,25 @@ use std::sync::Arc; use std::sync::LazyLock; use parking_lot::Mutex; +use rstest::rstest; use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; +use vortex_array::VTable; use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::uncompressed_size_in_bytes; use vortex_array::arrays::BoolArray; use vortex_array::arrays::ChunkedArray; use vortex_array::arrays::Constant; +use vortex_array::arrays::ConstantArray; use vortex_array::arrays::DictArray; use vortex_array::arrays::ExtensionArray; use vortex_array::arrays::FixedSizeListArray; use vortex_array::arrays::Map; use vortex_array::arrays::NullArray; +use vortex_array::arrays::Primitive; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::Slice; use vortex_array::arrays::SliceArray; @@ -44,6 +48,7 @@ use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_session::VortexSession; +use vortex_utils::aliases::hash_set::HashSet; use super::CascadingCompressor; use super::ROOT_SCHEME_ID; @@ -385,6 +390,92 @@ fn test_no_exclusion_without_history() { assert!(!c.is_excluded(&IntDictScheme, &ctx)); } +/// A scheme that always wins and gives a constant array. +/// +/// When `declares` is true, it declares its top-level encoding, and `compress` panics. +#[derive(Debug)] +struct TopLevelEncodingScheme { + declares: bool, +} + +static DECLARED_SCHEME: TopLevelEncodingScheme = TopLevelEncodingScheme { declares: true }; +static UNDECLARED_SCHEME: TopLevelEncodingScheme = TopLevelEncodingScheme { declares: false }; + +impl Scheme for TopLevelEncodingScheme { + fn scheme_name(&self) -> &'static str { + if self.declares { + "test.top_level_declared" + } else { + "test.top_level_undeclared" + } + } + + fn matches(&self, canonical: &Canonical) -> bool { + matches_integer_primitive(canonical) + } + + fn produced_encodings(&self) -> Vec { + Vec::new() + } + + fn top_level_encodings(&self) -> Vec { + if self.declares { + vec![Constant.id()] + } else { + Vec::new() + } + } + + fn expected_compression_ratio( + &self, + _data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> CompressionEstimate { + CompressionEstimate::Verdict(EstimateVerdict::AlwaysUse) + } + + fn compress( + &self, + _compressor: &CascadingCompressor, + data: &ArrayAndStats, + _compress_ctx: CompressorContext, + _exec_ctx: &mut ExecutionCtx, + ) -> VortexResult { + assert!( + !self.declares, + "a scheme that declares its top-level encodings must not be compressed" + ); + + Ok(ConstantArray::new( + Scalar::primitive(0i32, Nullability::NonNullable), + data.array().len(), + ) + .into_array()) + } +} + +/// `encodings_for` does not run the winner if the winner declares its top-level encodings. It +/// then also gives the encoding of the uncompressed array, which `compress` can fall back to. If +/// the winner declares nothing, `encodings_for` compresses and gives the encoding of the result. +#[rstest] +#[case::declared(&DECLARED_SCHEME, vec![Constant.id(), Primitive.id()])] +#[case::undeclared(&UNDECLARED_SCHEME, vec![Constant.id()])] +fn encodings_for_compresses_only_undeclared_winner( + #[case] scheme: &'static TopLevelEncodingScheme, + #[case] expected: Vec, +) -> VortexResult<()> { + let schemes: Vec<&'static dyn Scheme> = vec![scheme]; + let compressor = CascadingCompressor::new(schemes); + let array = PrimitiveArray::new(buffer![1i32, 2, 3, 4], Validity::NonNullable).into_array(); + let mut exec_ctx = SESSION.create_execution_ctx(); + + let encodings = compressor.encodings_for(&array, &mut exec_ctx)?; + + assert_eq!(encodings, expected.into_iter().collect::>()); + Ok(()) +} + #[test] fn immediate_always_use_wins_immediately() -> VortexResult<()> { let compressor = CascadingCompressor::new(vec![&DirectRatioScheme, &ImmediateAlwaysUseScheme]); diff --git a/vortex-compressor/src/scheme/mod.rs b/vortex-compressor/src/scheme/mod.rs index a9bfb7e421d..fd994b3f16b 100644 --- a/vortex-compressor/src/scheme/mod.rs +++ b/vortex-compressor/src/scheme/mod.rs @@ -141,7 +141,7 @@ pub trait Scheme: AsDynScheme + Debug + Send + Sync { /// a constant array returned early. The compressor uses this to tell which encodings it can /// produce without a full compression. The default is an empty list, which means the scheme /// does not declare its encodings, and the compressor compresses to find the encoding. - fn returned_encodings(&self) -> Vec { + fn top_level_encodings(&self) -> Vec { Vec::new() } From 1530593dccd0d404db60110e1c86b5490b9231f4 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Tue, 10 Nov 2026 17:57:00 +0000 Subject: [PATCH 6/6] Require schemes to declare their top-level encodings Remove the default of `Scheme::top_level_encodings`. `encodings_for` now never runs the winning scheme, and the debug assertion checks each compressed result against the declared list. Declare the encodings of the fixed-size-list scheme that came from develop. Signed-off-by: Robert Kruszewski --- .../src/schemes/fixed_size_list.rs | 5 + vortex-compressor/src/compressor/cascade.rs | 27 ++---- vortex-compressor/src/compressor/tests.rs | 95 ++++++++++--------- vortex-compressor/src/scheme/mod.rs | 7 +- 4 files changed, 67 insertions(+), 67 deletions(-) diff --git a/vortex-btrblocks/src/schemes/fixed_size_list.rs b/vortex-btrblocks/src/schemes/fixed_size_list.rs index 318e761b373..bf656501fc8 100644 --- a/vortex-btrblocks/src/schemes/fixed_size_list.rs +++ b/vortex-btrblocks/src/schemes/fixed_size_list.rs @@ -9,6 +9,7 @@ use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VTable; +use vortex_array::arrays::FixedSizeList; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::primitive::PrimitiveArrayExt; use vortex_array::scalar::Scalar; @@ -53,6 +54,10 @@ impl Scheme for FixedSizeListSparseScheme { vec![Sparse.id()] } + fn top_level_encodings(&self) -> Vec { + vec![Sparse.id(), FixedSizeList.id()] + } + /// Children: values=0, indices=1. fn num_children(&self) -> usize { 2 diff --git a/vortex-compressor/src/compressor/cascade.rs b/vortex-compressor/src/compressor/cascade.rs index 329bd503fdd..6bc4e299c97 100644 --- a/vortex-compressor/src/compressor/cascade.rs +++ b/vortex-compressor/src/compressor/cascade.rs @@ -329,12 +329,11 @@ impl CascadingCompressor { /// have. /// /// For a leaf array, this runs the same scheme selection as `compress`, with the statistics - /// and the sample compressions that the selection needs. If the winning scheme declares its - /// [`top_level_encodings`](Scheme::top_level_encodings), the full compression does not occur. - /// The result is then the encodings of the winner and the encoding of the uncompressed array, - /// because `compress` returns the uncompressed array if the output of the winner is not - /// smaller. If the winner does not declare its encodings, or if `array` is not a leaf array, - /// this compresses `array` and returns the one encoding of the result. + /// and the sample compressions that the selection needs. The winning scheme does not run. The + /// result is the [`top_level_encodings`](Scheme::top_level_encodings) of the winner and the + /// encoding of the uncompressed array, because `compress` returns the uncompressed array if + /// the output of the winner is not smaller. If `array` is not a leaf array, this compresses + /// `array` and returns the one encoding of the result. /// /// # Errors /// @@ -367,13 +366,8 @@ impl CascadingCompressor { Choice::Scheme(winner) => winner, }; - let declared = winner.scheme.top_level_encodings(); - if declared.is_empty() { - let compressed = self.compress_with_winner(winner, exec_ctx)?.into_array(); - return Ok([compressed.encoding_id()].into_iter().collect()); - } - - let mut encodings: HashSet = declared.into_iter().collect(); + let mut encodings: HashSet = + winner.scheme.top_level_encodings().into_iter().collect(); encodings.insert(winner.data.array().encoding_id()); Ok(encodings) @@ -478,10 +472,9 @@ impl CascadingCompressor { })?; debug_assert!( - { - let declared = scheme.top_level_encodings(); - declared.is_empty() || declared.contains(&compressed.encoding_id()) - }, + scheme + .top_level_encodings() + .contains(&compressed.encoding_id()), "scheme {} returned encoding {}, which is not in its top_level_encodings", scheme.id(), compressed.encoding_id(), diff --git a/vortex-compressor/src/compressor/tests.rs b/vortex-compressor/src/compressor/tests.rs index 099dd3402fc..f6a772db7cb 100644 --- a/vortex-compressor/src/compressor/tests.rs +++ b/vortex-compressor/src/compressor/tests.rs @@ -5,7 +5,6 @@ use std::sync::Arc; use std::sync::LazyLock; use parking_lot::Mutex; -use rstest::rstest; use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; @@ -17,11 +16,11 @@ use vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::uncompressed_si use vortex_array::arrays::BoolArray; use vortex_array::arrays::ChunkedArray; use vortex_array::arrays::Constant; -use vortex_array::arrays::ConstantArray; use vortex_array::arrays::DictArray; use vortex_array::arrays::ExtensionArray; use vortex_array::arrays::FixedSizeListArray; use vortex_array::arrays::Map; +use vortex_array::arrays::Null; use vortex_array::arrays::NullArray; use vortex_array::arrays::Primitive; use vortex_array::arrays::PrimitiveArray; @@ -99,6 +98,10 @@ impl Scheme for DirectRatioScheme { Vec::new() } + fn top_level_encodings(&self) -> Vec { + Vec::new() + } + fn expected_compression_ratio( &self, _data: &ArrayAndStats, @@ -135,6 +138,10 @@ impl Scheme for ImmediateAlwaysUseScheme { Vec::new() } + fn top_level_encodings(&self) -> Vec { + Vec::new() + } + fn expected_compression_ratio( &self, _data: &ArrayAndStats, @@ -171,6 +178,10 @@ impl Scheme for CallbackAlwaysUseScheme { Vec::new() } + fn top_level_encodings(&self) -> Vec { + Vec::new() + } + fn expected_compression_ratio( &self, _data: &ArrayAndStats, @@ -209,6 +220,10 @@ impl Scheme for CallbackSkipScheme { Vec::new() } + fn top_level_encodings(&self) -> Vec { + Vec::new() + } + fn expected_compression_ratio( &self, _data: &ArrayAndStats, @@ -247,6 +262,10 @@ impl Scheme for CallbackRatioScheme { Vec::new() } + fn top_level_encodings(&self) -> Vec { + Vec::new() + } + fn expected_compression_ratio( &self, _data: &ArrayAndStats, @@ -285,6 +304,10 @@ impl Scheme for HugeRatioScheme { Vec::new() } + fn top_level_encodings(&self) -> Vec { + Vec::new() + } + fn expected_compression_ratio( &self, _data: &ArrayAndStats, @@ -321,6 +344,10 @@ impl Scheme for ZeroBytesSamplingScheme { Vec::new() } + fn top_level_encodings(&self) -> Vec { + vec![Null.id()] + } + fn expected_compression_ratio( &self, _data: &ArrayAndStats, @@ -390,24 +417,13 @@ fn test_no_exclusion_without_history() { assert!(!c.is_excluded(&IntDictScheme, &ctx)); } -/// A scheme that always wins and gives a constant array. -/// -/// When `declares` is true, it declares its top-level encoding, and `compress` panics. +/// A scheme that always wins and declares a constant array. Its `compress` panics. #[derive(Debug)] -struct TopLevelEncodingScheme { - declares: bool, -} - -static DECLARED_SCHEME: TopLevelEncodingScheme = TopLevelEncodingScheme { declares: true }; -static UNDECLARED_SCHEME: TopLevelEncodingScheme = TopLevelEncodingScheme { declares: false }; +struct TopLevelEncodingScheme; impl Scheme for TopLevelEncodingScheme { fn scheme_name(&self) -> &'static str { - if self.declares { - "test.top_level_declared" - } else { - "test.top_level_undeclared" - } + "test.top_level_encoding" } fn matches(&self, canonical: &Canonical) -> bool { @@ -419,11 +435,7 @@ impl Scheme for TopLevelEncodingScheme { } fn top_level_encodings(&self) -> Vec { - if self.declares { - vec![Constant.id()] - } else { - Vec::new() - } + vec![Constant.id()] } fn expected_compression_ratio( @@ -438,41 +450,26 @@ impl Scheme for TopLevelEncodingScheme { fn compress( &self, _compressor: &CascadingCompressor, - data: &ArrayAndStats, + _data: &ArrayAndStats, _compress_ctx: CompressorContext, _exec_ctx: &mut ExecutionCtx, ) -> VortexResult { - assert!( - !self.declares, - "a scheme that declares its top-level encodings must not be compressed" - ); - - Ok(ConstantArray::new( - Scalar::primitive(0i32, Nullability::NonNullable), - data.array().len(), - ) - .into_array()) + unreachable!("`encodings_for` must not compress") } } -/// `encodings_for` does not run the winner if the winner declares its top-level encodings. It -/// then also gives the encoding of the uncompressed array, which `compress` can fall back to. If -/// the winner declares nothing, `encodings_for` compresses and gives the encoding of the result. -#[rstest] -#[case::declared(&DECLARED_SCHEME, vec![Constant.id(), Primitive.id()])] -#[case::undeclared(&UNDECLARED_SCHEME, vec![Constant.id()])] -fn encodings_for_compresses_only_undeclared_winner( - #[case] scheme: &'static TopLevelEncodingScheme, - #[case] expected: Vec, -) -> VortexResult<()> { - let schemes: Vec<&'static dyn Scheme> = vec![scheme]; - let compressor = CascadingCompressor::new(schemes); +/// `encodings_for` does not run the winner. It gives the top-level encodings of the winner and +/// the encoding of the uncompressed array, which `compress` can fall back to. +#[test] +fn encodings_for_does_not_compress() -> VortexResult<()> { + let compressor = CascadingCompressor::new(vec![&TopLevelEncodingScheme]); let array = PrimitiveArray::new(buffer![1i32, 2, 3, 4], Validity::NonNullable).into_array(); let mut exec_ctx = SESSION.create_execution_ctx(); let encodings = compressor.encodings_for(&array, &mut exec_ctx)?; - assert_eq!(encodings, expected.into_iter().collect::>()); + let expected: HashSet = [Constant.id(), Primitive.id()].into_iter().collect(); + assert_eq!(encodings, expected); Ok(()) } @@ -602,6 +599,10 @@ impl Scheme for ThresholdObservingScheme { Vec::new() } + fn top_level_encodings(&self) -> Vec { + Vec::new() + } + fn expected_compression_ratio( &self, _data: &ArrayAndStats, @@ -643,6 +644,10 @@ impl Scheme for CallbackMatchingRatioScheme { Vec::new() } + fn top_level_encodings(&self) -> Vec { + Vec::new() + } + fn expected_compression_ratio( &self, _data: &ArrayAndStats, diff --git a/vortex-compressor/src/scheme/mod.rs b/vortex-compressor/src/scheme/mod.rs index fd994b3f16b..754e51420de 100644 --- a/vortex-compressor/src/scheme/mod.rs +++ b/vortex-compressor/src/scheme/mod.rs @@ -139,11 +139,8 @@ pub trait Scheme: AsDynScheme + Debug + Send + Sync { /// /// The list must include every return path, such as a canonical input returned unchanged or /// a constant array returned early. The compressor uses this to tell which encodings it can - /// produce without a full compression. The default is an empty list, which means the scheme - /// does not declare its encodings, and the compressor compresses to find the encoding. - fn top_level_encodings(&self) -> Vec { - Vec::new() - } + /// produce without a full compression. + fn top_level_encodings(&self) -> Vec; /// Returns the variant of this scheme to use given which serialized IDs are `allowed`. ///