diff --git a/vortex-btrblocks/src/builder.rs b/vortex-btrblocks/src/builder.rs index e51fac9a9d5..749c131ed05 100644 --- a/vortex-btrblocks/src/builder.rs +++ b/vortex-btrblocks/src/builder.rs @@ -3,6 +3,7 @@ //! Builder for configuring `BtrBlocksCompressor` instances. +use vortex_array::ArrayId; use vortex_utils::aliases::hash_set::HashSet; use crate::BtrBlocksCompressor; @@ -208,6 +209,16 @@ impl BtrBlocksCompressorBuilder { self } + /// Retains only schemes whose produced encodings all belong to `allowed`. + /// + /// The file writer uses this to restrict compression to the encodings of its configured + /// editions. + pub fn retain_allowed_encodings(mut self, allowed: &HashSet) -> Self { + self.schemes + .retain(|s| s.produced_encodings().iter().all(|id| allowed.contains(id))); + self + } + /// Builds the configured [`BtrBlocksCompressor`]. pub fn build(self) -> BtrBlocksCompressor { BtrBlocksCompressor(CascadingCompressor::new(self.schemes)) @@ -216,6 +227,9 @@ impl BtrBlocksCompressorBuilder { #[cfg(test)] mod tests { + use vortex_array::VTable; + use vortex_fastlanes::FoR; + use super::*; #[test] @@ -230,6 +244,27 @@ mod tests { assert_eq!(builder.schemes.len(), ALL_SCHEMES.len()); } + #[test] + fn retain_allowed_encodings_filters_schemes() { + let allowed: HashSet = [FoR.id()].into_iter().collect(); + let builder = BtrBlocksCompressorBuilder::default().retain_allowed_encodings(&allowed); + assert_eq!(builder.schemes.len(), 1); + assert_eq!(builder.schemes[0].id(), integer::FoRScheme.id()); + + let none = BtrBlocksCompressorBuilder::default().retain_allowed_encodings(&HashSet::new()); + assert!(none.schemes.is_empty()); + } + + #[test] + fn retaining_all_declared_outputs_keeps_every_scheme() { + let allowed: HashSet = ALL_SCHEMES + .iter() + .flat_map(|scheme| scheme.produced_encodings()) + .collect(); + let builder = BtrBlocksCompressorBuilder::default().retain_allowed_encodings(&allowed); + assert_eq!(builder.schemes.len(), ALL_SCHEMES.len()); + } + #[test] fn cuda_compatible_excludes_alprd() { let builder = BtrBlocksCompressorBuilder::default().only_cuda_compatible(); diff --git a/vortex-btrblocks/src/schemes/binary/zstd.rs b/vortex-btrblocks/src/schemes/binary/zstd.rs index 83c09d24e94..d652e344db2 100644 --- a/vortex-btrblocks/src/schemes/binary/zstd.rs +++ b/vortex-btrblocks/src/schemes/binary/zstd.rs @@ -3,10 +3,12 @@ //! Zstd compression for binary arrays. +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_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::DeferredEstimate; use vortex_error::VortexResult; @@ -29,6 +31,10 @@ impl Scheme for ZstdScheme { canonical.dtype().is_binary() } + fn produced_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 a84672ef860..3f06d65b061 100644 --- a/vortex-btrblocks/src/schemes/binary/zstd_buffers.rs +++ b/vortex-btrblocks/src/schemes/binary/zstd_buffers.rs @@ -3,10 +3,12 @@ //! Zstd buffer-level binary compression preserving array layout for GPU decompression. +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_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::DeferredEstimate; use vortex_error::VortexResult; @@ -29,6 +31,10 @@ impl Scheme for ZstdBuffersScheme { canonical.dtype().is_binary() } + fn produced_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 20c5cd80a3c..1dff2171f60 100644 --- a/vortex-btrblocks/src/schemes/decimal.rs +++ b/vortex-btrblocks/src/schemes/decimal.rs @@ -3,10 +3,12 @@ //! Decimal compression scheme using byte-part decomposition. +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::arrays::DecimalArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::decimal::narrowed_decimal; @@ -38,6 +40,10 @@ impl Scheme for DecimalScheme { matches!(canonical, Canonical::Decimal(_)) } + fn produced_encodings(&self) -> Vec { + vec![DecimalByteParts.id()] + } + /// Children: primitive=0. fn num_children(&self) -> usize { 1 diff --git a/vortex-btrblocks/src/schemes/float/alp.rs b/vortex-btrblocks/src/schemes/float/alp.rs index b8a26abf348..f9fc7066bf4 100644 --- a/vortex-btrblocks/src/schemes/float/alp.rs +++ b/vortex-btrblocks/src/schemes/float/alp.rs @@ -7,10 +7,12 @@ use vortex_alp::ALP; use vortex_alp::ALPArrayExt; use vortex_alp::ALPArraySlotsExt; use vortex_alp::alp_encode; +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::arrays::Patched; use vortex_array::arrays::patched::use_experimental_patches; use vortex_array::arrays::primitive::PrimitiveArrayExt; @@ -40,6 +42,14 @@ impl Scheme for ALPScheme { canonical.dtype().is_float() } + fn produced_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 662a47a1d9c..4b208b8ead1 100644 --- a/vortex-btrblocks/src/schemes/float/alprd.rs +++ b/vortex-btrblocks/src/schemes/float/alprd.rs @@ -6,10 +6,12 @@ use vortex_alp::ALPRDArrayExt; use vortex_alp::ALPRDArrayOwnedExt; use vortex_alp::RDEncoder; +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::arrays::primitive::PrimitiveArrayExt; use vortex_array::dtype::PType; use vortex_compressor::scheme::CompressionEstimate; @@ -37,6 +39,10 @@ impl Scheme for ALPRDScheme { canonical.dtype().is_float() } + fn produced_encodings(&self) -> Vec { + vec![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 dc9a96133d5..416668c2fd0 100644 --- a/vortex-btrblocks/src/schemes/float/pco.rs +++ b/vortex-btrblocks/src/schemes/float/pco.rs @@ -3,10 +3,12 @@ //! Pco (pcodec) float compression. +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_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::DeferredEstimate; use vortex_error::VortexResult; @@ -29,6 +31,10 @@ impl Scheme for PcoScheme { canonical.dtype().is_float() } + fn produced_encodings(&self) -> Vec { + vec![vortex_pco::Pco.id()] + } + fn expected_compression_ratio( &self, _data: &ArrayAndStats, diff --git a/vortex-btrblocks/src/schemes/float/rle.rs b/vortex-btrblocks/src/schemes/float/rle.rs index c2683ed131c..71158b9dc3b 100644 --- a/vortex-btrblocks/src/schemes/float/rle.rs +++ b/vortex-btrblocks/src/schemes/float/rle.rs @@ -3,15 +3,18 @@ //! Run-length float encoding. +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; +use vortex_array::VTable; use vortex_compressor::scheme::AncestorExclusion; use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::DeferredEstimate; use vortex_compressor::scheme::DescendantExclusion; use vortex_compressor::scheme::EstimateVerdict; use vortex_error::VortexResult; +use vortex_fastlanes::RLE; use crate::ArrayAndStats; use crate::CascadingCompressor; @@ -35,6 +38,10 @@ impl Scheme for FloatRLEScheme { canonical.dtype().is_float() } + fn produced_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 42d09e323c5..3d9c25b18e4 100644 --- a/vortex-btrblocks/src/schemes/float/sparse.rs +++ b/vortex-btrblocks/src/schemes/float/sparse.rs @@ -3,10 +3,12 @@ //! Sparse encoding for null-dominated float arrays. +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::arrays::PrimitiveArray; use vortex_array::arrays::primitive::PrimitiveArrayExt; use vortex_compressor::scheme::ChildSelection; @@ -39,6 +41,10 @@ impl Scheme for NullDominatedSparseScheme { canonical.dtype().is_float() } + fn produced_encodings(&self) -> Vec { + vec![Sparse.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 2424402649e..5ac7d0e4078 100644 --- a/vortex-btrblocks/src/schemes/integer/bitpacking.rs +++ b/vortex-btrblocks/src/schemes/integer/bitpacking.rs @@ -3,10 +3,12 @@ //! BitPacking integer encoding. +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::arrays::Patched; use vortex_array::arrays::patched::use_experimental_patches; use vortex_array::arrays::primitive::PrimitiveArrayExt; @@ -38,6 +40,14 @@ impl Scheme for BitPackingScheme { canonical.dtype().is_int() } + fn produced_encodings(&self) -> Vec { + let mut encodings = vec![BitPacked.id()]; + if use_experimental_patches() { + encodings.push(Patched.id()); + } + encodings + } + fn expected_compression_ratio( &self, data: &ArrayAndStats, diff --git a/vortex-btrblocks/src/schemes/integer/delta.rs b/vortex-btrblocks/src/schemes/integer/delta.rs index ed087f5d845..f0288ef9436 100644 --- a/vortex-btrblocks/src/schemes/integer/delta.rs +++ b/vortex-btrblocks/src/schemes/integer/delta.rs @@ -3,10 +3,12 @@ //! FastLanes Delta integer encoding. +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::arrays::PrimitiveArray; use vortex_compressor::builtins::BinaryDictScheme; use vortex_compressor::builtins::FloatDictScheme; @@ -78,6 +80,10 @@ impl Scheme for DeltaScheme { canonical.dtype().is_int() } + fn produced_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 832ed4819fd..2014582ddcf 100644 --- a/vortex-btrblocks/src/schemes/integer/for_.rs +++ b/vortex-btrblocks/src/schemes/integer/for_.rs @@ -3,10 +3,12 @@ //! Frame of Reference integer encoding. +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::arrays::PrimitiveArray; use vortex_compressor::builtins::BinaryDictScheme; use vortex_compressor::builtins::FloatDictScheme; @@ -41,6 +43,10 @@ impl Scheme for FoRScheme { canonical.dtype().is_int() } + fn produced_encodings(&self) -> Vec { + vec![FoR.id()] + } + /// Dict codes always start at 0, so FoR (which subtracts the min) is a no-op. fn ancestor_exclusions(&self) -> Vec { vec![ diff --git a/vortex-btrblocks/src/schemes/integer/pco.rs b/vortex-btrblocks/src/schemes/integer/pco.rs index d7f182f588e..675a112d44f 100644 --- a/vortex-btrblocks/src/schemes/integer/pco.rs +++ b/vortex-btrblocks/src/schemes/integer/pco.rs @@ -3,10 +3,12 @@ //! Pco (pcodec) integer compression. +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_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::DeferredEstimate; use vortex_compressor::scheme::EstimateVerdict; @@ -30,6 +32,10 @@ impl Scheme for PcoScheme { canonical.dtype().is_int() } + fn produced_encodings(&self) -> Vec { + vec![vortex_pco::Pco.id()] + } + fn expected_compression_ratio( &self, data: &ArrayAndStats, diff --git a/vortex-btrblocks/src/schemes/integer/rle.rs b/vortex-btrblocks/src/schemes/integer/rle.rs index 5bc4e5296f1..778ba5690e4 100644 --- a/vortex-btrblocks/src/schemes/integer/rle.rs +++ b/vortex-btrblocks/src/schemes/integer/rle.rs @@ -3,10 +3,12 @@ //! Run-length integer encoding and shared RLE compression helpers. +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::arrays::PrimitiveArray; use vortex_array::arrays::primitive::PrimitiveArrayExt; use vortex_compressor::scheme::AncestorExclusion; @@ -156,6 +158,10 @@ impl Scheme for IntRLEScheme { canonical.dtype().is_int() } + fn produced_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 175f33e7406..6a97f7ec37d 100644 --- a/vortex-btrblocks/src/schemes/integer/runend.rs +++ b/vortex-btrblocks/src/schemes/integer/runend.rs @@ -3,10 +3,12 @@ //! Run-end integer encoding. +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::arrays::PrimitiveArray; use vortex_compressor::builtins::BinaryDictScheme; use vortex_compressor::builtins::FloatDictScheme; @@ -46,6 +48,10 @@ impl Scheme for RunEndScheme { canonical.dtype().is_int() } + fn produced_encodings(&self) -> Vec { + vec![RunEnd.id()] + } + /// Children: values=0, ends=1. fn num_children(&self) -> usize { 2 diff --git a/vortex-btrblocks/src/schemes/integer/sequence.rs b/vortex-btrblocks/src/schemes/integer/sequence.rs index fda9995b510..edcefb99fc2 100644 --- a/vortex-btrblocks/src/schemes/integer/sequence.rs +++ b/vortex-btrblocks/src/schemes/integer/sequence.rs @@ -3,9 +3,11 @@ //! Sequence integer encoding for sequential patterns. +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; +use vortex_array::VTable; use vortex_compressor::builtins::BinaryDictScheme; use vortex_compressor::builtins::FloatDictScheme; use vortex_compressor::builtins::IntDictScheme; @@ -19,6 +21,7 @@ use vortex_compressor::scheme::EstimateVerdict; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; +use vortex_sequence::Sequence; use vortex_sequence::sequence_encode; use crate::ArrayAndStats; @@ -40,6 +43,10 @@ impl Scheme for SequenceScheme { canonical.dtype().is_int() } + fn produced_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 e843d9d7999..429ff5c1a31 100644 --- a/vortex-btrblocks/src/schemes/integer/sparse.rs +++ b/vortex-btrblocks/src/schemes/integer/sparse.rs @@ -3,10 +3,13 @@ //! Sparse integer encoding for single-value-dominated arrays. +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::arrays::Constant; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::primitive::PrimitiveArrayExt; @@ -43,6 +46,10 @@ impl Scheme for SparseScheme { canonical.dtype().is_int() } + fn produced_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 959889dabac..6082fcc541c 100644 --- a/vortex-btrblocks/src/schemes/integer/zigzag.rs +++ b/vortex-btrblocks/src/schemes/integer/zigzag.rs @@ -3,10 +3,12 @@ //! ZigZag integer encoding for signed integers. +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::arrays::PrimitiveArray; use vortex_compressor::builtins::BinaryDictScheme; use vortex_compressor::builtins::FloatDictScheme; @@ -44,6 +46,10 @@ impl Scheme for ZigZagScheme { canonical.dtype().is_int() } + fn produced_encodings(&self) -> Vec { + vec![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 bd5bd010396..b2e91b8e9a8 100644 --- a/vortex-btrblocks/src/schemes/string/fsst.rs +++ b/vortex-btrblocks/src/schemes/string/fsst.rs @@ -3,11 +3,14 @@ //! FSST (Fast Static Symbol Table) string compression. +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::arrays::PrimitiveArray; +use vortex_array::arrays::VarBin; use vortex_array::arrays::VarBinArray; use vortex_array::arrays::primitive::PrimitiveArrayExt; use vortex_array::arrays::varbin::VarBinArrayExt; @@ -43,6 +46,10 @@ impl Scheme for FSSTScheme { canonical.dtype().is_utf8() } + fn produced_encodings(&self) -> Vec { + vec![FSST.id(), VarBin.id()] + } + /// Children: lengths=0, code_offsets=1. fn num_children(&self) -> usize { 2 diff --git a/vortex-btrblocks/src/schemes/string/onpair.rs b/vortex-btrblocks/src/schemes/string/onpair.rs index 8c7b0561502..a771593efc2 100644 --- a/vortex-btrblocks/src/schemes/string/onpair.rs +++ b/vortex-btrblocks/src/schemes/string/onpair.rs @@ -3,10 +3,12 @@ //! OnPair short-string compression (dict-12). +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::arrays::PrimitiveArray; use vortex_array::arrays::primitive::PrimitiveArrayExt; use vortex_compressor::scheme::CompressionEstimate; @@ -47,6 +49,10 @@ impl Scheme for OnPairScheme { canonical.dtype().is_utf8() } + fn produced_encodings(&self) -> Vec { + vec![OnPair.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 750dd978030..8620c366f77 100644 --- a/vortex-btrblocks/src/schemes/string/sparse.rs +++ b/vortex-btrblocks/src/schemes/string/sparse.rs @@ -3,10 +3,12 @@ //! Sparse encoding for null-dominated string arrays. +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::arrays::PrimitiveArray; use vortex_array::arrays::primitive::PrimitiveArrayExt; use vortex_compressor::scheme::ChildSelection; @@ -40,6 +42,10 @@ impl Scheme for NullDominatedSparseScheme { canonical.dtype().is_utf8() } + fn produced_encodings(&self) -> Vec { + vec![Sparse.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 177acfa1543..84e8860d626 100644 --- a/vortex-btrblocks/src/schemes/string/zstd.rs +++ b/vortex-btrblocks/src/schemes/string/zstd.rs @@ -3,10 +3,12 @@ //! Zstd string compression without dictionaries (nvCOMP compatible). +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_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::DeferredEstimate; use vortex_error::VortexResult; @@ -29,6 +31,10 @@ impl Scheme for ZstdScheme { canonical.dtype().is_utf8() } + fn produced_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 a8d16cc837f..cf691c70fcb 100644 --- a/vortex-btrblocks/src/schemes/string/zstd_buffers.rs +++ b/vortex-btrblocks/src/schemes/string/zstd_buffers.rs @@ -3,10 +3,12 @@ //! Zstd buffer-level string compression preserving array layout for GPU decompression. +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_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::DeferredEstimate; use vortex_error::VortexResult; @@ -29,6 +31,10 @@ impl Scheme for ZstdBuffersScheme { canonical.dtype().is_utf8() } + fn produced_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 6989b6a2aba..79748b69450 100644 --- a/vortex-btrblocks/src/schemes/temporal.rs +++ b/vortex-btrblocks/src/schemes/temporal.rs @@ -3,10 +3,12 @@ //! Temporal compression scheme using datetime-part decomposition. +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::arrays::ExtensionArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::TemporalArray; @@ -53,6 +55,10 @@ impl Scheme for TemporalScheme { ) } + fn produced_encodings(&self) -> Vec { + vec![DateTimeParts.id()] + } + /// Children: days=0, seconds=1, subseconds=2. fn num_children(&self) -> usize { 3 diff --git a/vortex-compressor/src/builtins/dict/binary.rs b/vortex-compressor/src/builtins/dict/binary.rs index 72b5f01141a..c407d0251c6 100644 --- a/vortex-compressor/src/builtins/dict/binary.rs +++ b/vortex-compressor/src/builtins/dict/binary.rs @@ -6,10 +6,13 @@ //! Vortex encoders must always produce unsigned integer codes; signed codes are only accepted //! for external compatibility. +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::arrays::Dict; use vortex_array::arrays::DictArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::dict::DictArrayExt; @@ -45,6 +48,10 @@ impl Scheme for BinaryDictScheme { canonical.dtype().is_binary() } + fn produced_encodings(&self) -> Vec { + vec![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 074d3ce5e07..f962c3ff967 100644 --- a/vortex-compressor/src/builtins/dict/float.rs +++ b/vortex-compressor/src/builtins/dict/float.rs @@ -6,11 +6,14 @@ //! Vortex encoders must always produce unsigned integer codes; signed codes are only accepted for //! external compatibility. +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::ArrayView; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; +use vortex_array::VTable; +use vortex_array::arrays::Dict; use vortex_array::arrays::DictArray; use vortex_array::arrays::Primitive; use vortex_array::arrays::PrimitiveArray; @@ -51,6 +54,10 @@ impl Scheme for FloatDictScheme { canonical.dtype().is_float() } + fn produced_encodings(&self) -> Vec { + vec![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 4e91bace3ac..27a17ef94ad 100644 --- a/vortex-compressor/src/builtins/dict/integer.rs +++ b/vortex-compressor/src/builtins/dict/integer.rs @@ -6,11 +6,14 @@ //! Vortex encoders must always produce unsigned integer codes; signed codes are only accepted //! for external compatibility. +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::ArrayView; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; +use vortex_array::VTable; +use vortex_array::arrays::Dict; use vortex_array::arrays::DictArray; use vortex_array::arrays::Primitive; use vortex_array::arrays::PrimitiveArray; @@ -46,6 +49,10 @@ impl Scheme for IntDictScheme { canonical.dtype().is_int() } + fn produced_encodings(&self) -> Vec { + vec![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 64ed469dde9..f5cbcd54d89 100644 --- a/vortex-compressor/src/builtins/dict/string.rs +++ b/vortex-compressor/src/builtins/dict/string.rs @@ -6,10 +6,13 @@ //! Vortex encoders must always produce unsigned integer codes; signed codes are only accepted //! for external compatibility. +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::arrays::Dict; use vortex_array::arrays::DictArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::dict::DictArrayExt; @@ -45,6 +48,10 @@ impl Scheme for StringDictScheme { canonical.dtype().is_utf8() } + fn produced_encodings(&self) -> Vec { + vec![Dict.id()] + } + fn stats_options(&self) -> GenerateStatsOptions { GenerateStatsOptions { count_distinct_values: true, diff --git a/vortex-compressor/src/compressor/tests.rs b/vortex-compressor/src/compressor/tests.rs index 08eafa369de..b23152d7ad3 100644 --- a/vortex-compressor/src/compressor/tests.rs +++ b/vortex-compressor/src/compressor/tests.rs @@ -4,6 +4,7 @@ use std::sync::LazyLock; use parking_lot::Mutex; +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; @@ -63,6 +64,10 @@ impl Scheme for DirectRatioScheme { matches_integer_primitive(canonical) } + fn produced_encodings(&self) -> Vec { + Vec::new() + } + fn expected_compression_ratio( &self, _data: &ArrayAndStats, @@ -95,6 +100,10 @@ impl Scheme for ImmediateAlwaysUseScheme { matches_integer_primitive(canonical) } + fn produced_encodings(&self) -> Vec { + Vec::new() + } + fn expected_compression_ratio( &self, _data: &ArrayAndStats, @@ -127,6 +136,10 @@ impl Scheme for CallbackAlwaysUseScheme { matches_integer_primitive(canonical) } + fn produced_encodings(&self) -> Vec { + Vec::new() + } + fn expected_compression_ratio( &self, _data: &ArrayAndStats, @@ -161,6 +174,10 @@ impl Scheme for CallbackSkipScheme { matches_integer_primitive(canonical) } + fn produced_encodings(&self) -> Vec { + Vec::new() + } + fn expected_compression_ratio( &self, _data: &ArrayAndStats, @@ -195,6 +212,10 @@ impl Scheme for CallbackRatioScheme { matches_integer_primitive(canonical) } + fn produced_encodings(&self) -> Vec { + Vec::new() + } + fn expected_compression_ratio( &self, _data: &ArrayAndStats, @@ -229,6 +250,10 @@ impl Scheme for HugeRatioScheme { matches_integer_primitive(canonical) } + fn produced_encodings(&self) -> Vec { + Vec::new() + } + fn expected_compression_ratio( &self, _data: &ArrayAndStats, @@ -261,6 +286,10 @@ impl Scheme for ZeroBytesSamplingScheme { matches_integer_primitive(canonical) } + fn produced_encodings(&self) -> Vec { + Vec::new() + } + fn expected_compression_ratio( &self, _data: &ArrayAndStats, @@ -470,6 +499,10 @@ impl Scheme for ThresholdObservingScheme { matches_integer_primitive(canonical) } + fn produced_encodings(&self) -> Vec { + Vec::new() + } + fn expected_compression_ratio( &self, _data: &ArrayAndStats, @@ -507,6 +540,10 @@ impl Scheme for CallbackMatchingRatioScheme { matches_integer_primitive(canonical) } + fn produced_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 341acc8acbf..de9e67690d4 100644 --- a/vortex-compressor/src/scheme/mod.rs +++ b/vortex-compressor/src/scheme/mod.rs @@ -23,6 +23,7 @@ pub use estimate::EstimateVerdict; pub use exclusion::AncestorExclusion; pub use exclusion::ChildSelection; pub use exclusion::DescendantExclusion; +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; @@ -123,6 +124,13 @@ pub trait Scheme: Debug + Send + Sync { /// Whether this scheme can compress the given canonical array. fn matches(&self, canonical: &Canonical) -> bool; + /// The array encodings this scheme itself may introduce into its compressed output. + /// + /// Cascaded children are compressed by other schemes, which declare their own encodings, + /// so only encodings constructed directly by [`compress`](Scheme::compress) belong here. + /// Canonical arrays the scheme merely rearranges do not need to be declared. + fn produced_encodings(&self) -> Vec; + /// Returns the stats generation options this scheme requires. The compressor merges all /// eligible schemes' options before generating stats so that a single stats pass satisfies /// every scheme. diff --git a/vortex-tensor/src/encodings/l2_denorm.rs b/vortex-tensor/src/encodings/l2_denorm.rs index d07734bc6ff..8bec6292ee0 100644 --- a/vortex-tensor/src/encodings/l2_denorm.rs +++ b/vortex-tensor/src/encodings/l2_denorm.rs @@ -1,11 +1,13 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::extension::ExtensionArrayExt; +use vortex_array::scalar_fn::ScalarFnVTable; use vortex_compressor::CascadingCompressor; use vortex_compressor::scheme::CompressionEstimate; use vortex_compressor::scheme::CompressorContext; @@ -15,6 +17,7 @@ use vortex_compressor::stats::ArrayAndStats; use vortex_error::VortexResult; use crate::matcher::AnyTensor; +use crate::scalar_fns::l2_denorm::L2Denorm; use crate::scalar_fns::l2_denorm::normalize_as_l2_denorm; #[derive(Debug)] @@ -32,6 +35,10 @@ impl Scheme for L2DenormScheme { ) } + fn produced_encodings(&self) -> Vec { + vec![L2Denorm.id()] + } + fn expected_compression_ratio( &self, _data: &ArrayAndStats,