diff --git a/Cargo.lock b/Cargo.lock index 0747e266768..4ab0f9a0f36 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9908,6 +9908,7 @@ dependencies = [ "tracing", "vortex-array", "vortex-buffer", + "vortex-edition", "vortex-error", "vortex-mask", "vortex-session", @@ -10185,6 +10186,7 @@ dependencies = [ "vortex-bytebool", "vortex-datetime-parts", "vortex-decimal-byte-parts", + "vortex-edition", "vortex-error", "vortex-fastlanes", "vortex-flatbuffers", @@ -10401,6 +10403,7 @@ dependencies = [ "vortex-arrow", "vortex-btrblocks", "vortex-buffer", + "vortex-edition", "vortex-error", "vortex-flatbuffers", "vortex-io", diff --git a/vortex-btrblocks/src/builder.rs b/vortex-btrblocks/src/builder.rs index e51fac9a9d5..a3af7d827c0 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; @@ -90,12 +91,14 @@ pub const ALL_SCHEMES: &[&dyn Scheme] = &[ #[derive(Debug, Clone)] pub struct BtrBlocksCompressorBuilder { schemes: Vec<&'static dyn Scheme>, + gate_by_enabled_editions: bool, } impl Default for BtrBlocksCompressorBuilder { fn default() -> Self { Self { schemes: ALL_SCHEMES.to_vec(), + gate_by_enabled_editions: false, } } } @@ -107,6 +110,7 @@ impl BtrBlocksCompressorBuilder { pub fn empty() -> Self { Self { schemes: Vec::new(), + gate_by_enabled_editions: false, } } @@ -208,14 +212,40 @@ impl BtrBlocksCompressorBuilder { self } + /// Retains only schemes whose declared produced encodings (see + /// [`Scheme::produced_encodings`]) all belong to `allowed`. The file writer uses this + /// to restrict compression to an explicitly configured allow set. + pub fn retain_allowed_encodings(mut self, allowed: &HashSet) -> Self { + self.schemes + .retain(|s| s.produced_encodings().iter().all(|id| allowed.contains(id))); + self + } + + /// Gate scheme selection by the session's enabled editions at compression time. + /// + /// See [`CascadingCompressor::with_enabled_editions_gating`]. The file writer enables + /// this by default so compression only chooses encodings from the session's enabled + /// editions. + pub fn gate_by_enabled_editions(mut self) -> Self { + self.gate_by_enabled_editions = true; + self + } + /// Builds the configured [`BtrBlocksCompressor`]. pub fn build(self) -> BtrBlocksCompressor { - BtrBlocksCompressor(CascadingCompressor::new(self.schemes)) + let mut compressor = CascadingCompressor::new(self.schemes); + if self.gate_by_enabled_editions { + compressor = compressor.with_enabled_editions_gating(); + } + BtrBlocksCompressor(compressor) } } #[cfg(test)] mod tests { + use vortex_array::VTable; + use vortex_fastlanes::FoR; + use super::*; #[test] @@ -230,6 +260,17 @@ 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 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/Cargo.toml b/vortex-compressor/Cargo.toml index 24a2d8e40d5..cfe02fbdb39 100644 --- a/vortex-compressor/Cargo.toml +++ b/vortex-compressor/Cargo.toml @@ -22,6 +22,7 @@ rustc-hash = { workspace = true } tracing = { workspace = true, features = ["std", "attributes"] } vortex-array = { workspace = true } vortex-buffer = { workspace = true } +vortex-edition = { workspace = true } vortex-error = { workspace = true } vortex-mask = { workspace = true } vortex-utils = { workspace = true } 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/cascade.rs b/vortex-compressor/src/compressor/cascade.rs index 9dbbc611448..4343bcd0b7a 100644 --- a/vortex-compressor/src/compressor/cascade.rs +++ b/vortex-compressor/src/compressor/cascade.rs @@ -228,11 +228,16 @@ impl CascadingCompressor { compress_ctx: CompressorContext, exec_ctx: &mut ExecutionCtx, ) -> VortexResult { + let allowed_encodings = self.enabled_editions_allow_set(exec_ctx); let eligible_schemes: Vec<&'static dyn Scheme> = self .schemes .iter() .copied() - .filter(|s| s.matches(&canonical) && !self.is_excluded(*s, &compress_ctx)) + .filter(|s| { + s.matches(&canonical) + && !self.is_excluded(*s, &compress_ctx) + && Self::produces_allowed_encodings(*s, allowed_encodings.as_ref()) + }) .collect(); let array: ArrayRef = canonical.into(); diff --git a/vortex-compressor/src/compressor/mod.rs b/vortex-compressor/src/compressor/mod.rs index a661970950c..b4a46be2d65 100644 --- a/vortex-compressor/src/compressor/mod.rs +++ b/vortex-compressor/src/compressor/mod.rs @@ -9,6 +9,14 @@ mod sample; mod select; mod structural; +use vortex_array::ArrayId; +use vortex_array::ExecutionCtx; +use vortex_array::VTable; +use vortex_array::arrays::Patched; +use vortex_array::arrays::patched::use_experimental_patches; +use vortex_edition::EditionSessionExt; +use vortex_utils::aliases::hash_set::HashSet; + use crate::builtins::IntDictScheme; use crate::scheme::ChildSelection; use crate::scheme::DescendantExclusion; @@ -46,6 +54,9 @@ pub struct CascadingCompressor { /// Descendant exclusion rules for the compressor's own cascading (e.g. excluding Dict from /// list offsets). root_exclusions: Vec, + + /// Whether scheme selection is gated by the session's enabled editions. + gate_by_enabled_editions: bool, } impl CascadingCompressor { @@ -63,6 +74,57 @@ impl CascadingCompressor { Self { schemes, root_exclusions, + gate_by_enabled_editions: false, + } + } + + /// Gate scheme selection by the session's enabled editions. + /// + /// When the session compressing an array has editions enabled (see + /// [`vortex_edition::EditionSession::enabled_encodings`]), only schemes whose declared + /// [`Scheme::produced_encodings`] all belong to those editions are eligible. Sessions + /// with no enabled editions are not gated. + /// + /// The file writer enables this so compression only chooses encodings it may write. + pub fn with_enabled_editions_gating(mut self) -> Self { + self.gate_by_enabled_editions = true; + self + } + + /// Resolve the allowed-encoding set from the session's enabled editions, if gating + /// applies to this compressor and session. + pub(crate) fn enabled_editions_allow_set( + &self, + ctx: &ExecutionCtx, + ) -> Option> { + if !self.gate_by_enabled_editions { + return None; + } + let mut allowed: HashSet = ctx + .session() + .editions() + .enabled_encodings()? + .into_iter() + .collect(); + // Experimental patches are a runtime opt-out of the editions guarantee: schemes only + // produce `Patched` arrays when the experimental environment variable is set. + if use_experimental_patches() { + allowed.insert(Patched.id()); + } + Some(allowed) + } + + /// Whether a scheme's declared produced encodings all belong to `allowed`. + pub(crate) fn produces_allowed_encodings( + scheme: &dyn Scheme, + allowed: Option<&HashSet>, + ) -> bool { + match allowed { + None => true, + Some(allowed) => scheme + .produced_encodings() + .iter() + .all(|id| allowed.contains(id)), } } } diff --git a/vortex-compressor/src/compressor/tests.rs b/vortex-compressor/src/compressor/tests.rs index 08eafa369de..74ae4107558 100644 --- a/vortex-compressor/src/compressor/tests.rs +++ b/vortex-compressor/src/compressor/tests.rs @@ -4,19 +4,27 @@ use std::sync::LazyLock; use parking_lot::Mutex; +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; +use vortex_array::array_session; use vortex_array::arrays::BoolArray; use vortex_array::arrays::Constant; use vortex_array::arrays::NullArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::validity::Validity; use vortex_buffer::buffer; +use vortex_edition::AsEncodingId; +use vortex_edition::Edition; +use vortex_edition::EditionDeclaration; +use vortex_edition::EditionId; +use vortex_edition::EditionSessionExt; use vortex_error::VortexResult; use vortex_session::VortexSession; +use vortex_session::registry::CachedId; use super::CascadingCompressor; use super::ROOT_SCHEME_ID; @@ -36,7 +44,7 @@ use crate::scheme::SchemeExt; use crate::stats::ArrayAndStats; use crate::stats::GenerateStatsOptions; -static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); +static SESSION: LazyLock = LazyLock::new(array_session); fn compressor() -> CascadingCompressor { CascadingCompressor::new(vec![&IntDictScheme, &FloatDictScheme, &StringDictScheme]) @@ -63,6 +71,10 @@ impl Scheme for DirectRatioScheme { matches_integer_primitive(canonical) } + fn produced_encodings(&self) -> Vec { + vec![] + } + fn expected_compression_ratio( &self, _data: &ArrayAndStats, @@ -95,6 +107,10 @@ impl Scheme for ImmediateAlwaysUseScheme { matches_integer_primitive(canonical) } + fn produced_encodings(&self) -> Vec { + vec![] + } + fn expected_compression_ratio( &self, _data: &ArrayAndStats, @@ -127,6 +143,10 @@ impl Scheme for CallbackAlwaysUseScheme { matches_integer_primitive(canonical) } + fn produced_encodings(&self) -> Vec { + vec![] + } + fn expected_compression_ratio( &self, _data: &ArrayAndStats, @@ -161,6 +181,10 @@ impl Scheme for CallbackSkipScheme { matches_integer_primitive(canonical) } + fn produced_encodings(&self) -> Vec { + vec![] + } + fn expected_compression_ratio( &self, _data: &ArrayAndStats, @@ -195,6 +219,10 @@ impl Scheme for CallbackRatioScheme { matches_integer_primitive(canonical) } + fn produced_encodings(&self) -> Vec { + vec![] + } + fn expected_compression_ratio( &self, _data: &ArrayAndStats, @@ -229,6 +257,10 @@ impl Scheme for HugeRatioScheme { matches_integer_primitive(canonical) } + fn produced_encodings(&self) -> Vec { + vec![] + } + fn expected_compression_ratio( &self, _data: &ArrayAndStats, @@ -261,6 +293,10 @@ impl Scheme for ZeroBytesSamplingScheme { matches_integer_primitive(canonical) } + fn produced_encodings(&self) -> Vec { + vec![] + } + fn expected_compression_ratio( &self, _data: &ArrayAndStats, @@ -452,6 +488,101 @@ fn zero_byte_sample_alone_selects_no_scheme() -> VortexResult<()> { Ok(()) } +// Observer scheme used by edition-gating tests: records whether it was selected for +// compression. `GATED_LOCK` serializes tests that share `GATED_SELECTED` so they do not +// race. +static GATED_LOCK: Mutex<()> = Mutex::new(()); +static GATED_SELECTED: Mutex = Mutex::new(false); + +/// The encoding id the edition-gating test scheme claims to produce. +const GATED_OUTPUT_ENCODING: &str = "test.gated_output"; + +#[derive(Debug)] +struct EditionGatedScheme; + +impl Scheme for EditionGatedScheme { + fn scheme_name(&self) -> &'static str { + "test.edition_gated" + } + + fn matches(&self, canonical: &Canonical) -> bool { + matches_integer_primitive(canonical) + } + + fn produced_encodings(&self) -> Vec { + static ID: CachedId = CachedId::new(GATED_OUTPUT_ENCODING); + vec![*ID] + } + + 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 { + *GATED_SELECTED.lock() = true; + Ok(data.array().clone()) + } +} + +/// Declare a test edition adding the given encodings on the session, and enable it. +fn enable_test_edition(session: &VortexSession, added: &'static [&'static dyn AsEncodingId]) { + let edition_id = EditionId::new("test", 2026, 1, 0); + let editions = session.editions(); + editions + .declare(&EditionDeclaration { + edition: Edition { + id: edition_id, + min_vortex_version: None, + }, + added, + }) + .unwrap_or_else(|e| panic!("declaring test edition: {e}")); + editions + .enable(edition_id) + .unwrap_or_else(|e| panic!("enabling test edition: {e}")); +} + +/// Run the edition-gated compressor over the test array and report whether the scheme was +/// selected for compression. +fn gated_scheme_selected(session: &VortexSession) -> VortexResult { + let _guard = GATED_LOCK.lock(); + *GATED_SELECTED.lock() = false; + let compressor = + CascadingCompressor::new(vec![&EditionGatedScheme]).with_enabled_editions_gating(); + let array = PrimitiveArray::new(buffer![1i32, 2, 3, 4], Validity::NonNullable).into_array(); + compressor.compress(&array, &mut session.create_execution_ctx())?; + Ok(*GATED_SELECTED.lock()) +} + +#[test] +fn enabled_editions_gate_scheme_selection() -> VortexResult<()> { + // A session with no enabled editions is not gated. + let session = array_session(); + assert!(gated_scheme_selected(&session)?); + + // Enabled editions covering the scheme's produced encoding keep it eligible. + let session = array_session(); + enable_test_edition(&session, &[&GATED_OUTPUT_ENCODING]); + assert!(gated_scheme_selected(&session)?); + + // Enabled editions that do not cover the produced encoding exclude the scheme. + let session = array_session(); + enable_test_edition(&session, &[&"test.other_encoding"]); + assert!(!gated_scheme_selected(&session)?); + Ok(()) +} + // Observer helper used by threshold-related tests. Captures the `best_so_far` value the // compressor passes to its deferred callback. `OBSERVER_LOCK` serializes tests that share // `OBSERVED_THRESHOLD` so they do not race. @@ -470,6 +601,10 @@ impl Scheme for ThresholdObservingScheme { matches_integer_primitive(canonical) } + fn produced_encodings(&self) -> Vec { + vec![] + } + fn expected_compression_ratio( &self, _data: &ArrayAndStats, @@ -507,6 +642,10 @@ impl Scheme for CallbackMatchingRatioScheme { matches_integer_primitive(canonical) } + fn produced_encodings(&self) -> Vec { + vec![] + } + 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..82df6fd119d 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,18 @@ 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; a scheme + /// producing only canonical arrays returns an empty list. + /// + /// Writers that restrict their output to a set of allowed encodings — e.g. the Vortex + /// file writer, which only emits encodings from the session's enabled editions — skip + /// schemes whose produced encodings are not all allowed. + 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-edition/src/lib.rs b/vortex-edition/src/lib.rs index c4e341450d8..3ad53c1e253 100644 --- a/vortex-edition/src/lib.rs +++ b/vortex-edition/src/lib.rs @@ -20,8 +20,8 @@ //! //! This crate defines only the types and the session variable; the first-party edition //! declarations live in the public `vortex` crate (`vortex::editions`), which seeds them -//! into the default session. See the published spec at -//! . +//! into the default session and enables the default write editions. See the published spec +//! at . mod session; pub mod test_harness; diff --git a/vortex-edition/src/session.rs b/vortex-edition/src/session.rs index a2a9b50066a..77af1f82829 100644 --- a/vortex-edition/src/session.rs +++ b/vortex-edition/src/session.rs @@ -39,6 +39,8 @@ struct Inner { editions: BTreeMap, /// Keyed by interned encoding id; ordered by the id's string form. inclusions: BTreeMap, + /// The editions writers of this session target, in the order they were enabled. + enabled: Vec, } impl EditionSession { @@ -107,6 +109,60 @@ impl EditionSession { .next_back() } + /// Enable an edition as a write target for this session. Writers gated by editions only + /// emit encodings belonging to the union of the enabled editions. + /// + /// The edition must already be declared, and at most one edition per family can be + /// enabled — a writer targets one edition per family and may emit any encoding in their + /// union. + pub fn enable(&self, edition: EditionId) -> Result<(), EditionError> { + let mut inner = self.inner.write(); + if !inner.editions.contains_key(&edition.to_string()) { + return Err(EditionError::new(format!( + "cannot enable undeclared edition {edition}" + ))); + } + if let Some(existing) = inner + .enabled + .iter() + .find(|enabled| enabled.family == edition.family) + { + return Err(EditionError::new(format!( + "edition {existing} of family {} is already enabled", + edition.family + ))); + } + inner.enabled.push(edition); + Ok(()) + } + + /// The editions writers of this session target, in the order they were enabled. + /// + /// Empty when no editions are enabled, in which case edition-gated writers fall back to + /// their own default allow set. + pub fn enabled(&self) -> Vec { + self.inner.read().enabled.clone() + } + + /// The current set of encodings a writer may emit: the union of the enabled editions' + /// encoding sets, sorted by encoding id and deduplicated. + /// + /// Returns `None` when no editions are enabled, leaving the write gate to the caller. + pub fn enabled_encodings(&self) -> Option> { + let enabled = self.enabled(); + if enabled.is_empty() { + return None; + } + let mut ids: Vec = enabled + .iter() + .flat_map(|edition| self.encodings_in(edition)) + .map(|inclusion| inclusion.encoding_id) + .collect(); + ids.sort_unstable(); + ids.dedup(); + Some(ids) + } + /// Compute the full encoding set of an edition: every declared inclusion of the /// edition's family whose `since` is at or before it, sorted by encoding id. pub fn encodings_in(&self, edition: &EditionId) -> Vec { diff --git a/vortex-edition/src/tests.rs b/vortex-edition/src/tests.rs index 3791173fad8..55c8d89d72f 100644 --- a/vortex-edition/src/tests.rs +++ b/vortex-edition/src/tests.rs @@ -195,6 +195,28 @@ fn validate_rejects_inconsistent_declarations() -> Result<(), crate::EditionErro Ok(()) } +#[test] +fn enabled_editions_gate_the_write_set() -> Result<(), crate::EditionError> { + let editions = session(); + + // Nothing enabled: the write set is undecided and gating falls back to the caller. + assert!(editions.enabled().is_empty()); + assert!(editions.enabled_encodings().is_none()); + + // Enabling an undeclared edition fails. + assert!(editions.enable(EditionId::new("test", 2027, 1, 0)).is_err()); + + editions.enable(FIRST)?; + assert_eq!(editions.enabled(), [FIRST]); + let enabled = editions.enabled_encodings().expect("editions enabled"); + let ids: Vec<&str> = enabled.iter().map(|id| id.as_str()).collect(); + assert_eq!(ids, ["test.alpha", "test.beta"]); + + // Only one edition per family can be enabled. + assert!(editions.enable(SECOND).is_err()); + Ok(()) +} + #[test] fn edition_ids_order_within_family_only() { assert!(FIRST.is_at_or_before(&SECOND)); diff --git a/vortex-file/Cargo.toml b/vortex-file/Cargo.toml index 347bd0fc69e..5992ff06434 100644 --- a/vortex-file/Cargo.toml +++ b/vortex-file/Cargo.toml @@ -39,6 +39,7 @@ vortex-bytebool = { workspace = true } vortex-datetime-parts = { workspace = true } vortex-decimal-byte-parts = { workspace = true } +vortex-edition = { workspace = true } vortex-error = { workspace = true } vortex-fastlanes = { workspace = true } vortex-flatbuffers = { workspace = true, features = ["file"] } diff --git a/vortex-file/src/lib.rs b/vortex-file/src/lib.rs index 4ece7d9cdab..a790d775cfd 100644 --- a/vortex-file/src/lib.rs +++ b/vortex-file/src/lib.rs @@ -160,8 +160,10 @@ mod forever_constant { /// Register the default encodings use in Vortex files with the provided session. /// -/// NOTE: this function will be changed in the future to encapsulate logic for using different -/// Vortex "Editions" that may support different sets of encodings. +/// Registration covers reading: a session can decode every encoding registered here. The +/// writer is gated separately by the session's enabled editions (see +/// `vortex_edition::EditionSession::enabled_encodings`), falling back to +/// [`ALLOWED_ENCODINGS`] for sessions with no enabled editions. pub fn register_default_encodings(session: &VortexSession) { vortex_bytebool::initialize(session); vortex_fsst::initialize(session); diff --git a/vortex-file/src/strategy.rs b/vortex-file/src/strategy.rs index 3cedbc1d68c..e6099d73a25 100644 --- a/vortex-file/src/strategy.rs +++ b/vortex-file/src/strategy.rs @@ -73,10 +73,15 @@ use vortex_zstd::ZstdBuffers; const ONE_MEG: u64 = 1 << 20; -/// Static registry of all allowed array encodings for file writing. +/// Static fallback registry of allowed array encodings for file writing. /// -/// This includes all canonical encodings from vortex-array plus all compressed -/// encodings from the various encoding crates. +/// This is the allow set used when the writing session has **no enabled editions** (see +/// [`EditionSession::enabled_encodings`](vortex_edition::EditionSession::enabled_encodings)). +/// Sessions created through the `vortex` facade enable the default write editions, and the +/// writer then gates on the session's enabled editions instead of this set. +/// +/// It includes all canonical encodings from vortex-array plus all compressed encodings from +/// the various encoding crates. pub static ALLOWED_ENCODINGS: LazyLock> = LazyLock::new(|| { let mut allowed = HashSet::new(); @@ -173,7 +178,7 @@ impl Default for WriteStrategyBuilder { compressor: CompressorConfig::BtrBlocks(BtrBlocksCompressorBuilder::default()), row_block_size: 8192, field_writers: HashMap::new(), - allow_encodings: Some(ALLOWED_ENCODINGS.clone()), + allow_encodings: None, flat_strategy: None, probe_compressor: None, use_list_layout: use_experimental_list_layout(), @@ -215,10 +220,15 @@ impl WriteStrategyBuilder { self } - /// Override the allowed array encodings for normalization. + /// Override the allowed array encodings for file writing. /// /// The flat leaf writer uses this set when deciding whether an existing encoded array can be - /// written as-is or must be normalized before serialization. + /// written as-is or must fail the write, and the default BtrBlocks compressor is filtered + /// down to schemes that only produce encodings in this set. + /// + /// When not overridden, the writer gates on the session's enabled editions (see + /// [`EditionSession::enabled_encodings`](vortex_edition::EditionSession::enabled_encodings)), + /// falling back to [`ALLOWED_ENCODINGS`] for sessions with no enabled editions. pub fn with_allow_encodings(mut self, allow_encodings: HashSet) -> Self { self.allow_encodings = Some(allow_encodings); self @@ -260,12 +270,30 @@ impl WriteStrategyBuilder { /// Builds the canonical [`LayoutStrategy`] implementation, with the configured overrides /// applied. pub fn build(self) -> Arc { + let allow_encodings = self.allow_encodings; + let flat: Arc = if let Some(flat) = self.flat_strategy { flat - } else if let Some(allow_encodings) = self.allow_encodings { + } else if let Some(allow_encodings) = allow_encodings.clone() { Arc::new(FlatLayoutStrategy::default().with_allow_encodings(allow_encodings)) } else { - Arc::new(FlatLayoutStrategy::default()) + // The default gate: the session's enabled editions decide what may be written, + // with the static registry as fallback for sessions that enable no editions. + Arc::new( + FlatLayoutStrategy::default() + .with_allow_encodings(ALLOWED_ENCODINGS.clone()) + .with_enabled_editions_gating(), + ) + }; + + // Restrict BtrBlocks compression to schemes whose output can be written: every + // scheme must declare its produced encodings, and they must all be allowed. With an + // explicit allow set the schemes are filtered here; otherwise the compressor gates + // itself on the session's enabled editions at compression time. This is what keeps + // compression within the enabled editions. + let gate_schemes = |builder: BtrBlocksCompressorBuilder| match &allow_encodings { + Some(allowed) => builder.retain_allowed_encodings(allowed), + None => builder.gate_by_enabled_editions(), }; // 7. for each chunk create a flat layout @@ -279,10 +307,7 @@ impl WriteStrategyBuilder { // dictionary-encode the integer codes produced by that earlier step. let data_compressor: Arc = match &self.compressor { CompressorConfig::BtrBlocks(builder) => Arc::new( - builder - .clone() - .exclude_schemes([IntDictScheme.id()]) - .build(), + gate_schemes(builder.clone().exclude_schemes([IntDictScheme.id()])).build(), ), CompressorConfig::Opaque(compressor) => Arc::clone(compressor), }; @@ -307,7 +332,7 @@ impl WriteStrategyBuilder { // 2.1. | 3.1. compress stats tables and dict values. let stats_compressor: Arc = match self.compressor { - CompressorConfig::BtrBlocks(builder) => Arc::new(builder.build()), + CompressorConfig::BtrBlocks(builder) => Arc::new(gate_schemes(builder).build()), CompressorConfig::Opaque(compressor) => compressor, }; let compress_then_flat = CompressingStrategy::new(flat, Arc::clone(&stats_compressor)); diff --git a/vortex-file/src/writer.rs b/vortex-file/src/writer.rs index 80da9f4fb89..af7b08c4079 100644 --- a/vortex-file/src/writer.rs +++ b/vortex-file/src/writer.rs @@ -17,7 +17,11 @@ use futures::pin_mut; use futures::select; use itertools::Itertools; use vortex_array::ArrayContext; +use vortex_array::ArrayId; use vortex_array::ArrayRef; +use vortex_array::VTable; +use vortex_array::arrays::Patched; +use vortex_array::arrays::patched::use_experimental_patches; use vortex_array::dtype::DType; use vortex_array::dtype::FieldPath; use vortex_array::expr::stats::Stat; @@ -30,6 +34,7 @@ use vortex_array::stream::ArrayStreamAdapter; use vortex_array::stream::ArrayStreamExt; use vortex_array::stream::SendableArrayStream; use vortex_buffer::ByteBuffer; +use vortex_edition::EditionSessionExt; use vortex_error::VortexError; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -158,12 +163,26 @@ impl VortexWriteOptions { mut write: W, stream: SendableArrayStream, ) -> VortexResult { - // NOTE(os): Setup an array context that already has all known encodings pre-populated. - // This is preferred for now over having an empty context here, because only the - // serialised array order is deterministic. The serialisation of arrays are done - // parallel and with an empty context they can register their encodings to the context - // in different order, changing the written bytes from run to run. - let ctx = ArrayContext::new(ALLOWED_ENCODINGS.iter().cloned().sorted().collect()) + // NOTE(os): Setup an array context that already has all writable encodings + // pre-populated. This is preferred for now over having an empty context here, because + // only the serialised array order is deterministic. The serialisation of arrays are + // done parallel and with an empty context they can register their encodings to the + // context in different order, changing the written bytes from run to run. + // + // The writable encodings come from the session's enabled editions; sessions with no + // enabled editions fall back to the static registry. + let writable: Vec = match self.session.editions().enabled_encodings() { + Some(mut ids) => { + // Experimental patches are a runtime opt-out of the editions guarantee. + if use_experimental_patches() { + ids.push(Patched.id()); + ids.sort_unstable(); + } + ids + } + None => ALLOWED_ENCODINGS.iter().cloned().sorted().collect(), + }; + let ctx = ArrayContext::new(writable) // Configure a registry just to ensure only known encodings are interned. .with_registry(self.session.arrays().registry().clone()); let dtype = stream.dtype().clone(); diff --git a/vortex-layout/Cargo.toml b/vortex-layout/Cargo.toml index f772b9ab639..04508134ad5 100644 --- a/vortex-layout/Cargo.toml +++ b/vortex-layout/Cargo.toml @@ -43,6 +43,7 @@ vortex-array = { workspace = true } vortex-arrow = { workspace = true } vortex-btrblocks = { workspace = true } vortex-buffer = { workspace = true } +vortex-edition = { workspace = true } vortex-error = { workspace = true } vortex-flatbuffers = { workspace = true, features = ["layout"] } vortex-io = { workspace = true } diff --git a/vortex-layout/src/layouts/flat/writer.rs b/vortex-layout/src/layouts/flat/writer.rs index f6b6d5acf02..158ac2a115e 100644 --- a/vortex-layout/src/layouts/flat/writer.rs +++ b/vortex-layout/src/layouts/flat/writer.rs @@ -5,6 +5,9 @@ use async_trait::async_trait; use futures::StreamExt; use vortex_array::ArrayContext; use vortex_array::ArrayId; +use vortex_array::VTable; +use vortex_array::arrays::Patched; +use vortex_array::arrays::patched::use_experimental_patches; use vortex_array::dtype::DType; use vortex_array::expr::stats::Precision; use vortex_array::expr::stats::Stat; @@ -19,6 +22,7 @@ use vortex_array::serde::SerializeOptions; use vortex_array::stats::StatsSetRef; use vortex_buffer::BufferString; use vortex_buffer::ByteBuffer; +use vortex_edition::EditionSessionExt; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -46,6 +50,10 @@ pub struct FlatLayoutStrategy { /// Optional set of allowed array encodings for normalization. /// If None, then all are allowed. pub allowed_encodings: Option>, + /// Whether the allowed encodings are resolved from the session's enabled editions at + /// write time. When the session has editions enabled, they take precedence over + /// [`allowed_encodings`](Self::allowed_encodings), which remains the fallback. + pub gate_by_enabled_editions: bool, } impl Default for FlatLayoutStrategy { @@ -54,6 +62,7 @@ impl Default for FlatLayoutStrategy { include_padding: true, max_variable_length_statistics_size: 64, allowed_encodings: None, + gate_by_enabled_editions: false, } } } @@ -76,6 +85,34 @@ impl FlatLayoutStrategy { self.allowed_encodings = Some(allow_encodings); self } + + /// Resolve the allowed encodings from the session's enabled editions at write time. + /// + /// When the writing session has editions enabled (see + /// [`vortex_edition::EditionSession::enabled_encodings`]), any encoding outside them + /// fails the write; sessions with no enabled editions fall back to + /// [`with_allow_encodings`](Self::with_allow_encodings). + pub fn with_enabled_editions_gating(mut self) -> Self { + self.gate_by_enabled_editions = true; + self + } + + /// The allowed-encoding set effective for a write on the given session. + fn effective_allowed_encodings(&self, session: &VortexSession) -> Option> { + if self.gate_by_enabled_editions + && let Some(enabled) = session.editions().enabled_encodings() + { + let mut allowed: HashSet = enabled.into_iter().collect(); + // Experimental patches are a runtime opt-out of the editions guarantee: + // compression only produces `Patched` arrays when the experimental environment + // variable is set, and only then is the encoding writable. + if use_experimental_patches() { + allowed.insert(Patched.id()); + } + return Some(allowed); + } + self.allowed_encodings.clone() + } } fn truncate_scalar_stat Option<(Scalar, bool)>>( @@ -158,9 +195,9 @@ impl LayoutStrategy for FlatLayoutStrategy { _ => {} } - let chunk = if let Some(allowed) = &self.allowed_encodings { + let chunk = if let Some(allowed) = self.effective_allowed_encodings(session) { chunk.normalize(&mut NormalizeOptions { - allowed, + allowed: &allowed, operation: Operation::Error, })? } else { diff --git a/vortex-python/src/io.rs b/vortex-python/src/io.rs index 182ab7fbe74..c411b4035e9 100644 --- a/vortex-python/src/io.rs +++ b/vortex-python/src/io.rs @@ -258,7 +258,7 @@ impl PyVortexWriteOptions { /// >>> vx.io.VortexWriteOptions.default().write(sprl, "chonky.vortex") /// >>> import os /// >>> os.path.getsize('chonky.vortex') - /// 215940 + /// 215900 /// /// Wow, Vortex manages to use about two bytes per integer! So advanced. So tiny. /// @@ -268,7 +268,7 @@ impl PyVortexWriteOptions { /// /// >>> vx.io.VortexWriteOptions.compact().write(sprl, "tiny.vortex") /// >>> os.path.getsize('tiny.vortex') - /// 55068 + /// 55028 /// /// Random numbers are not (usually) composed of random bytes! #[staticmethod] diff --git a/vortex-tensor/src/encodings/l2_denorm.rs b/vortex-tensor/src/encodings/l2_denorm.rs index d07734bc6ff..beebf322533 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,11 @@ impl Scheme for L2DenormScheme { ) } + fn produced_encodings(&self) -> Vec { + // A `ScalarFnArray`'s encoding id is its scalar function's id. + vec![ScalarFnVTable::id(&L2Denorm)] + } + fn expected_compression_ratio( &self, _data: &ArrayAndStats, diff --git a/vortex/src/editions/mod.rs b/vortex/src/editions/mod.rs index 1c3323d077f..44413aa23f0 100644 --- a/vortex/src/editions/mod.rs +++ b/vortex/src/editions/mod.rs @@ -6,23 +6,26 @@ //! [`vortex_edition`] provides the types, the [`crate::editions::EditionSession`] session //! variable, and the test harness; the actual declarations live here, one module per edition //! (`editions::::`), and are seeded into the default session by -//! [`crate::editions::register_default_editions`]. +//! [`crate::editions::register_default_editions`], which also enables +//! [`crate::editions::DEFAULT_WRITE_EDITIONS`] as the session's write target. //! //! Each edition module declares the edition together with the encodings that join the //! family at it; members of earlier editions are inherited and never restated. Correctness -//! is enforced by unit tests: every edition module calls -//! [`vortex_edition::test_harness::validate_edition`] once from its `#[cfg(test)]` module, -//! and the computed set of a frozen edition is pinned by a golden test, so any change to -//! these declarations that alters a frozen set fails CI. New encodings are staged into the -//! newest draft edition. +//! is enforced by unit tests: every declared edition is checked against +//! [`vortex_edition::test_harness::validate_edition`], and the computed set of the newest +//! frozen edition is pinned by a golden test, so any change to these declarations that +//! alters a frozen set fails CI. New encodings are staged into the newest draft edition. //! -//! Note this is currently unused but a future PR will make this public and gate the writer behind -//! editions. +//! The enabled editions gate the file writer: it only emits encodings from the session's +//! enabled editions (see [`crate::editions::EditionSession::enabled_encodings`]). pub mod core; +#[cfg(test)] +mod tests; pub mod unstable; use vortex_edition::EditionDeclaration; +use vortex_edition::EditionId; pub use vortex_edition::EditionSession; use vortex_edition::EditionSessionExt; use vortex_error::VortexExpect; @@ -50,7 +53,19 @@ pub static EDITION_DECLARATIONS: &[&EditionDeclaration] = &[ &unstable::v2026_06::DECLARATION, ]; -/// Register the Vortex edition declarations with the session's [`EditionSession`]. +/// The editions the default session's file writer targets. +/// +/// By default the writer only emits encodings from the newest frozen `core` edition, +/// [`CORE_2026_07_0`]. The `unstable_encodings` feature additionally opts the writer into +/// the newest draft of the `unstable` family, which carries no compatibility guarantee. +pub static DEFAULT_WRITE_EDITIONS: &[EditionId] = if cfg!(feature = "unstable_encodings") { + &[CORE_2026_07_0, UNSTABLE_2026_06_0] +} else { + &[CORE_2026_07_0] +}; + +/// Register the Vortex edition declarations with the session's [`EditionSession`] and +/// enable [`DEFAULT_WRITE_EDITIONS`] as the session's write target. pub fn register_default_editions(session: &VortexSession) { let editions = session.editions(); for declaration in EDITION_DECLARATIONS { @@ -59,4 +74,10 @@ pub fn register_default_editions(session: &VortexSession) { .map_err(|e| vortex_err!("{e}")) .vortex_expect("edition declarations are valid"); } + for edition in DEFAULT_WRITE_EDITIONS { + editions + .enable(*edition) + .map_err(|e| vortex_err!("{e}")) + .vortex_expect("default write editions are declared"); + } } diff --git a/vortex/src/editions/tests.rs b/vortex/src/editions/tests.rs new file mode 100644 index 00000000000..6277f26cae0 --- /dev/null +++ b/vortex/src/editions/tests.rs @@ -0,0 +1,124 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_btrblocks::ALL_SCHEMES; +use vortex_btrblocks::SchemeExt; +use vortex_edition::EditionError; +use vortex_edition::EditionSession; +use vortex_edition::EditionSessionExt; +use vortex_edition::test_harness::validate_edition; +use vortex_session::VortexSession; +use vortex_session::registry::Id; + +use super::CORE_2025_05_0; +use super::CORE_2026_07_0; +use super::DEFAULT_WRITE_EDITIONS; +use super::EDITION_DECLARATIONS; +use super::register_default_editions; + +fn edition_session() -> EditionSession { + let session = VortexSession::empty(); + register_default_editions(&session); + session.editions().clone() +} + +#[test] +fn every_declared_edition_validates() -> Result<(), EditionError> { + let editions = edition_session(); + for declaration in EDITION_DECLARATIONS { + validate_edition(&editions, &declaration.edition.id)?; + } + Ok(()) +} + +#[test] +fn default_write_editions_are_enabled() { + let editions = edition_session(); + assert_eq!(editions.enabled(), DEFAULT_WRITE_EDITIONS); + assert!(editions.enabled_encodings().is_some()); +} + +/// The full encoding set of the newest frozen `core` edition. This set is frozen: the only +/// way it may change is by declaring a *new* edition, so a failure here means a frozen +/// declaration was edited. +#[test] +fn core_2026_07_encoding_set_is_pinned() { + let editions = edition_session(); + let encodings = editions.encodings_in(&CORE_2026_07_0); + let ids: Vec<&str> = encodings + .iter() + .map(|inclusion| inclusion.encoding_id.as_str()) + .collect(); + assert_eq!( + ids, + [ + "fastlanes.bitpacked", + "fastlanes.for", + "fastlanes.rle", + "vortex.alp", + "vortex.alprd", + "vortex.bool", + "vortex.bytebool", + "vortex.chunked", + "vortex.constant", + "vortex.datetimeparts", + "vortex.decimal", + "vortex.decimal_byte_parts", + "vortex.dict", + "vortex.ext", + "vortex.fixed_size_list", + "vortex.fsst", + "vortex.list", + "vortex.listview", + "vortex.masked", + "vortex.null", + "vortex.pco", + "vortex.primitive", + "vortex.runend", + "vortex.sequence", + "vortex.sparse", + "vortex.struct", + "vortex.varbin", + "vortex.varbinview", + "vortex.variant", + "vortex.zigzag", + "vortex.zstd", + ] + ); +} + +#[test] +fn earlier_editions_are_subsets() { + let editions = edition_session(); + let first: Vec = editions + .encodings_in(&CORE_2025_05_0) + .into_iter() + .map(|inclusion| inclusion.encoding_id) + .collect(); + let latest: Vec = editions + .encodings_in(&CORE_2026_07_0) + .into_iter() + .map(|inclusion| inclusion.encoding_id) + .collect(); + assert!(first.iter().all(|id| latest.contains(id))); + assert!(first.len() < latest.len()); +} + +/// Every default compression scheme's produced encodings must be covered by the default +/// write editions — otherwise the default file writer would silently drop the scheme. +#[test] +fn default_schemes_stay_within_default_write_editions() { + let editions = edition_session(); + let enabled = editions + .enabled_encodings() + .expect("default editions enabled"); + for scheme in ALL_SCHEMES { + for encoding in scheme.produced_encodings() { + assert!( + enabled.contains(&encoding), + "scheme {} produces {encoding}, which is outside the default write editions", + scheme.id() + ); + } + } +} diff --git a/vortex/src/lib.rs b/vortex/src/lib.rs index 1af8484f230..36fe2cd2a59 100644 --- a/vortex/src/lib.rs +++ b/vortex/src/lib.rs @@ -145,10 +145,10 @@ pub mod compressor { pub use vortex_btrblocks::SchemeId; } -/// Logical Vortex data types. /// Vortex editions: named, frozen sets of encodings with a read-compatibility guarantee. pub mod editions; +/// Logical Vortex data types. pub mod dtype { pub use vortex_array::dtype::*; }