diff --git a/Cargo.lock b/Cargo.lock index ab869e730b9..5321b8fd299 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10185,6 +10185,7 @@ dependencies = [ "vortex-bytebool", "vortex-datetime-parts", "vortex-decimal-byte-parts", + "vortex-edition", "vortex-error", "vortex-fastlanes", "vortex-flatbuffers", @@ -10479,6 +10480,7 @@ dependencies = [ "vortex-array", "vortex-arrow", "vortex-buffer", + "vortex-edition", "vortex-error", "vortex-file", "vortex-io", diff --git a/docs/specs/editions.md b/docs/specs/editions.md index d76bf954f1d..6dd87cd250d 100644 --- a/docs/specs/editions.md +++ b/docs/specs/editions.md @@ -37,6 +37,10 @@ Every file you write carries the read-forever guarantee. If a file would contain outside the targeted edition, the write fails immediately; edition violations never surface as someone else's read error later. +The enabled editions are stored on the writer's Vortex session. Registering an edition makes +its declaration available to the session; enabling it separately allows the writer to emit its +encodings. Enabling another edition from the same family replaces the earlier selection. + Two knobs exist when the default is not what you want: - **Pin an older edition** when files must stay readable by deployments running older Vortex. @@ -45,9 +49,10 @@ Two knobs exist when the default is not what you want: example spatial encodings) possible later. A writer targets at most one edition per family and may emit any encoding in their union; each encoding belongs to exactly one family. -You can also opt out of editions entirely to write custom or experimental encodings. Doing so -is an explicit choice that gives up the standardization guarantee — only readers that know your -encodings can read those files. +Lower-level sessions without an enabled-editions store opt out of editions entirely and can write +custom or experimental encodings. A raw `with_allow_encodings` writer policy is another explicit +opt-out. Either choice gives up the standardization guarantee — only readers that know those +encodings can read the files. ## How editions change diff --git a/encodings/parquet-variant/Cargo.toml b/encodings/parquet-variant/Cargo.toml index 6cdc2b3fa8b..e90da6c8ebe 100644 --- a/encodings/parquet-variant/Cargo.toml +++ b/encodings/parquet-variant/Cargo.toml @@ -34,6 +34,7 @@ vortex-proto = { workspace = true, features = ["expr"] } vortex-session = { workspace = true } [dev-dependencies] +vortex-edition = { workspace = true } rstest = { workspace = true } tokio = { workspace = true, features = ["full"] } vortex-array = { workspace = true, features = ["_test-harness"] } diff --git a/encodings/parquet-variant/src/vtable.rs b/encodings/parquet-variant/src/vtable.rs index 4fa3ea10690..cb2fc5a014f 100644 --- a/encodings/parquet-variant/src/vtable.rs +++ b/encodings/parquet-variant/src/vtable.rs @@ -309,7 +309,6 @@ mod tests { use vortex_array::Canonical; use vortex_array::EqMode; use vortex_array::IntoArray; - use vortex_array::VTable; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBinViewArray; @@ -327,6 +326,10 @@ mod tests { use vortex_buffer::BitBuffer; use vortex_buffer::ByteBufferMut; use vortex_buffer::buffer; + use vortex_edition::Edition; + use vortex_edition::EditionId; + use vortex_edition::EditionInclusion; + use vortex_edition::EditionSessionExt; use vortex_error::VortexResult; use vortex_error::vortex_err; use vortex_file::OpenOptionsSessionExt; @@ -389,22 +392,35 @@ mod tests { } #[fixture] - fn parquet_variant_file_session() -> VortexSession { + fn parquet_variant_file_session() -> VortexResult { + const TEST_EDITION: EditionId = EditionId::new("test", 2026, 7, 0); + let session = vortex_array::array_session() .with::() .with::(); vortex_file::register_default_encodings(&session); session.arrays().register(ParquetVariant); + let editions = session.editions(); + editions + .declare_edition(Edition { + id: TEST_EDITION, + min_vortex_version: None, + }) + .map_err(|error| vortex_err!("{error}"))?; + for id in session.arrays().registry().ids() { + editions + .declare_inclusion(EditionInclusion::new(&id, TEST_EDITION)) + .map_err(|error| vortex_err!("{error}"))?; + } session + .enable_edition(TEST_EDITION) + .map_err(|error| vortex_err!("{error}"))?; + Ok(session) } #[fixture] fn write_strategy() -> Arc { - let mut allowed = vortex_file::ALLOWED_ENCODINGS.clone(); - allowed.insert(ParquetVariant.id()); - vortex_file::WriteStrategyBuilder::default() - .with_allow_encodings(allowed) - .build() + vortex_file::WriteStrategyBuilder::default().build() } #[test] @@ -452,9 +468,10 @@ mod tests { #[tokio::test] async fn test_file_roundtrip_typed_value_variant_with_statistics( #[from(typed_value_variant_array)] expected: VortexResult, - parquet_variant_file_session: VortexSession, + parquet_variant_file_session: VortexResult, ) -> VortexResult<()> { let expected = expected?; + let parquet_variant_file_session = parquet_variant_file_session?; let mut bytes = ByteBufferMut::empty(); parquet_variant_file_session @@ -479,10 +496,11 @@ mod tests { #[tokio::test] async fn test_file_roundtrip_typed_value_variant_with_zoned_strategy( #[from(typed_value_variant_array)] expected: VortexResult, - parquet_variant_file_session: VortexSession, + parquet_variant_file_session: VortexResult, write_strategy: Arc, ) -> VortexResult<()> { let expected = expected?; + let parquet_variant_file_session = parquet_variant_file_session?; let mut bytes = ByteBufferMut::empty(); parquet_variant_file_session 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-edition/src/lib.rs b/vortex-edition/src/lib.rs index c4e341450d8..a8eaf8e8287 100644 --- a/vortex-edition/src/lib.rs +++ b/vortex-edition/src/lib.rs @@ -4,8 +4,8 @@ //! Definitions of Vortex *editions*: named, frozen sets of encodings that a writer may put in //! a file, carrying a forever read-compatibility guarantee. //! -//! Editions live on the session, like encodings do: [`EditionSession`] is the session -//! variable holding the edition registry, populated at initialization time. Declarations +//! Editions live on the session, like encodings do: [`EditionSession`] holds the registered +//! editions and [`EnabledEditions`] selects which of them a writer may emit. Declarations //! are plain constants — an [`EditionId`] plus an [`Edition`] record, and one //! [`EditionInclusion`] per encoding stating that it is a member of an edition *and every //! later edition of the same family*. Any crate can register declarations into a session, @@ -18,9 +18,8 @@ //! [`test_harness::validate_edition`] validates one edition's constraints — call it once in //! the `#[cfg(test)]` module of each edition definition. //! -//! 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 +//! The first-party edition declarations live in the public `vortex` crate, which registers +//! and enables them on the default session. See the published spec at //! . mod session; @@ -36,6 +35,7 @@ use std::fmt::Formatter; pub use session::EditionSession; pub use session::EditionSessionExt; +pub use session::EnabledEditions; use vortex_session::registry::Id; /// The identifier of an edition, e.g. `core2026.07.0`. @@ -43,7 +43,7 @@ use vortex_session::registry::Id; /// The `family` names an independently versioned, additive group of encodings (`core` is the /// set the default writer emits). The date components record when the edition was frozen and /// order editions chronologically *within* a family; there is no ordering across families. -#[derive(Clone, Copy, Debug, PartialEq, Eq)] +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] pub struct EditionId { /// The edition family, e.g. `core`. pub family: &'static str, diff --git a/vortex-edition/src/session.rs b/vortex-edition/src/session.rs index a2a9b50066a..a50d17b6b53 100644 --- a/vortex-edition/src/session.rs +++ b/vortex-edition/src/session.rs @@ -1,8 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! The [`EditionSession`] session variable: the per-session registry of editions and -//! edition inclusions. +//! Session variables for registered and enabled editions. use std::any::Any; use std::collections::BTreeMap; @@ -41,6 +40,28 @@ struct Inner { inclusions: BTreeMap, } +/// The editions enabled for writing in a session. +/// +/// At most one edition is enabled per family. Enabling a newer or older edition from the +/// same family replaces the previous selection. This is separate from [`EditionSession`]: +/// registration describes what a session knows how to reason about, while enabling is the +/// explicit writer policy. +#[derive(Clone, Debug, Default)] +pub struct EnabledEditions { + inner: Arc>>, +} + +impl EnabledEditions { + /// Return the enabled editions, sorted by family. + pub fn editions(&self) -> Vec { + self.inner.read().values().copied().collect() + } + + fn enable(&self, edition: EditionId) { + self.inner.write().insert(edition.family, edition); + } +} + impl EditionSession { /// Create a session variable with no declarations. pub fn empty() -> Self { @@ -190,12 +211,70 @@ impl SessionVar for EditionSession { } } +impl SessionVar for EnabledEditions { + fn as_any(&self) -> &dyn Any { + self + } + + fn as_any_mut(&mut self) -> &mut dyn Any { + self + } +} + /// Session data for Vortex editions. pub trait EditionSessionExt: SessionExt { /// Returns the edition registry. fn editions(&self) -> SessionGuard<'_, EditionSession> { self.get::() } + + /// Returns the editions enabled for writing. + /// + /// Accessing this method installs the enabled-editions session variable if it is absent, with + /// an initially empty selection. + fn enabled_editions(&self) -> SessionGuard<'_, EnabledEditions> { + self.get::() + } + + /// Register an edition declaration with this session. + fn register_edition(&self, declaration: &EditionDeclaration) -> Result<(), EditionError> { + self.editions().declare(declaration) + } + + /// Enable a registered edition for writing. + /// + /// Enabling an edition replaces the enabled edition from the same family. An edition + /// must be registered first so a typo or unavailable third-party declaration cannot + /// silently produce an empty writable set. + fn enable_edition(&self, edition: EditionId) -> Result<(), EditionError> { + if self.editions().find(&edition).is_none() { + return Err(EditionError::new(format!( + "cannot enable unregistered edition {edition}" + ))); + } + self.enabled_editions().enable(edition); + Ok(()) + } + + /// Resolve the encodings in all enabled editions. + /// + /// When the enabled-editions variable is absent or no editions are enabled, this returns an + /// empty vector and therefore permits no edition encodings. + fn enabled_encoding_ids(&self) -> Vec { + let Some(enabled) = self.get_opt::() else { + return vec![]; + }; + let editions = self.editions(); + let mut ids: Vec = enabled + .editions() + .iter() + .flat_map(|edition| editions.encodings_in(edition)) + .map(|inclusion| inclusion.encoding_id) + .collect(); + ids.sort_unstable(); + ids.dedup(); + ids + } } impl EditionSessionExt for S {} diff --git a/vortex-edition/src/tests.rs b/vortex-edition/src/tests.rs index 3791173fad8..3d6550982db 100644 --- a/vortex-edition/src/tests.rs +++ b/vortex-edition/src/tests.rs @@ -9,6 +9,7 @@ use crate::EditionId; use crate::EditionInclusion; use crate::EditionSession; use crate::EditionSessionExt; +use crate::EnabledEditions; const FIRST: EditionId = EditionId::new("test", 2026, 1, 0); const SECOND: EditionId = EditionId::new("test", 2026, 7, 0); @@ -126,6 +127,68 @@ fn session_exposes_edition_registry() { assert!(session.editions().find(&FIRST).is_some()); } +#[test] +fn registered_and_enabled_editions_are_separate() -> Result<(), crate::EditionError> { + let session = VortexSession::empty().with::(); + for declaration in DECLARATIONS { + session.register_edition(declaration)?; + } + + assert!(session.enabled_encoding_ids().is_empty()); + session.enable_edition(FIRST)?; + assert_eq!(session.enabled_editions().editions(), [FIRST]); + assert_eq!( + session + .enabled_encoding_ids() + .iter() + .map(|id| id.as_str()) + .collect::>(), + ["test.alpha", "test.beta"] + ); + + session.enable_edition(SECOND)?; + assert_eq!(session.enabled_editions().editions(), [SECOND]); + assert_eq!(session.enabled_encoding_ids().len(), 3); + + // Selecting an older edition in the same family replaces the newer one and removes + // encodings that joined after it. + session.enable_edition(FIRST)?; + assert_eq!(session.enabled_editions().editions(), [FIRST]); + let enabled = session.enabled_encoding_ids(); + assert_eq!(enabled.len(), 2); + assert!(enabled.iter().all(|id| id.as_str() != "test.gamma")); + Ok(()) +} + +#[test] +fn enabling_requires_registration() { + let session = VortexSession::empty().with::(); + assert!(session.enable_edition(FIRST).is_err()); + assert!(session.enabled_editions().editions().is_empty()); +} + +#[test] +fn enabled_editions_are_independent_across_families() -> Result<(), crate::EditionError> { + const OTHER: EditionId = EditionId::new("other", 2026, 4, 0); + static OTHER_DECLARATION: EditionDeclaration = EditionDeclaration { + edition: Edition { + id: OTHER, + min_vortex_version: None, + }, + added: &[&"other.delta"], + }; + + let session = VortexSession::empty().with::(); + session.register_edition(&DECLARATIONS[0])?; + session.register_edition(&OTHER_DECLARATION)?; + session.enable_edition(FIRST)?; + session.enable_edition(OTHER)?; + + assert_eq!(session.enabled_editions().editions(), [OTHER, FIRST]); + assert_eq!(session.enabled_encoding_ids().len(), 3); + Ok(()) +} + #[test] fn duplicate_declarations_error() { let editions = session(); 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/footer/field_sizes.rs b/vortex-file/src/footer/field_sizes.rs index 5d7d2ce7c02..91e794f3729 100644 --- a/vortex-file/src/footer/field_sizes.rs +++ b/vortex-file/src/footer/field_sizes.rs @@ -141,6 +141,7 @@ mod tests { .with::() .with::(); crate::register_default_encodings(&session); + crate::enable_all_registered_array_encodings(&session); session }); diff --git a/vortex-file/src/lib.rs b/vortex-file/src/lib.rs index 1a1a751568c..0e1785ab04f 100644 --- a/vortex-file/src/lib.rs +++ b/vortex-file/src/lib.rs @@ -166,8 +166,8 @@ 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 editions enabled on its session. pub fn register_default_encodings(session: &VortexSession) { vortex_bytebool::initialize(session); vortex_fsst::initialize(session); @@ -198,6 +198,37 @@ pub fn register_default_encodings(session: &VortexSession) { vortex_tensor::initialize(session); } +#[cfg(test)] +pub(crate) fn enable_all_registered_array_encodings(session: &VortexSession) { + use vortex_edition::Edition; + use vortex_edition::EditionId; + use vortex_edition::EditionInclusion; + use vortex_edition::EditionSessionExt; + use vortex_error::VortexExpect; + use vortex_error::vortex_err; + + const TEST_EDITION: EditionId = EditionId::new("test", 2026, 7, 0); + + let editions = session.editions(); + editions + .declare_edition(Edition { + id: TEST_EDITION, + min_vortex_version: None, + }) + .map_err(|error| vortex_err!("{error}")) + .vortex_expect("test edition is valid"); + for id in session.arrays().registry().ids() { + editions + .declare_inclusion(EditionInclusion::new(&id, TEST_EDITION)) + .map_err(|error| vortex_err!("{error}")) + .vortex_expect("registered array encoding has one test-edition inclusion"); + } + session + .enable_edition(TEST_EDITION) + .map_err(|error| vortex_err!("{error}")) + .vortex_expect("test edition is registered"); +} + #[cfg(test)] mod default_encoding_tests { use vortex_array::VTable as _; @@ -228,10 +259,4 @@ mod default_encoding_tests { .has_execute_parent(Filter.id(), OnPair.id()) ); } - - #[cfg(not(feature = "unstable_encodings"))] - #[test] - fn default_writer_does_not_allow_onpair() { - assert!(!crate::ALLOWED_ENCODINGS.contains(&OnPair.id())); - } } diff --git a/vortex-file/src/multi/mod.rs b/vortex-file/src/multi/mod.rs index 019c59a4113..506e75b2aa8 100644 --- a/vortex-file/src/multi/mod.rs +++ b/vortex-file/src/multi/mod.rs @@ -341,6 +341,7 @@ mod tests { .with::() .with::(); crate::register_default_encodings(&session); + crate::enable_all_registered_array_encodings(&session); let expected = ByteBuffer::copy_from(b"cached metadata"); let mut output = ByteBufferMut::empty(); @@ -441,6 +442,7 @@ mod tests { .with::() .with::(); crate::register_default_encodings(&session); + crate::enable_all_registered_array_encodings(&session); let mut out_a = ByteBufferMut::empty(); session diff --git a/vortex-file/src/open.rs b/vortex-file/src/open.rs index 79e8e671240..7d1c75262b1 100644 --- a/vortex-file/src/open.rs +++ b/vortex-file/src/open.rs @@ -528,6 +528,7 @@ mod tests { .with::() .with::(); crate::register_default_encodings(&session); + crate::enable_all_registered_array_encodings(&session); session } @@ -586,6 +587,7 @@ mod tests { .with::(); crate::register_default_encodings(&session); + crate::enable_all_registered_array_encodings(&session); // Create a large file (> 1MB) let mut buf = ByteBufferMut::empty(); @@ -646,6 +648,7 @@ mod tests { .with::() .with::(); crate::register_default_encodings(&session); + crate::enable_all_registered_array_encodings(&session); let metadata = ByteBuffer::copy_from(vec![0x5a; INITIAL_READ_SIZE * 2]); let mut output = ByteBufferMut::empty(); @@ -719,6 +722,7 @@ mod tests { .with::() .with::(); crate::register_default_encodings(&session); + crate::enable_all_registered_array_encodings(&session); let value = ByteBuffer::copy_from(b"supplied-footer metadata"); let mut output = ByteBufferMut::empty(); @@ -760,6 +764,7 @@ mod tests { .with::(); crate::register_default_encodings(&session); + crate::enable_all_registered_array_encodings(&session); let mut buf = ByteBufferMut::empty(); let array = Buffer::from((0i32..16_384).collect::>()).into_array(); diff --git a/vortex-file/src/strategy.rs b/vortex-file/src/strategy.rs index 763db758cb7..22a5d0c7291 100644 --- a/vortex-file/src/strategy.rs +++ b/vortex-file/src/strategy.rs @@ -5,43 +5,13 @@ use std::num::NonZeroUsize; use std::sync::Arc; -use std::sync::LazyLock; -use vortex_alp::ALP; -use vortex_alp::ALPRD; use vortex_array::ArrayId; -use vortex_array::VTable; -use vortex_array::arrays::Bool; -use vortex_array::arrays::Chunked; -use vortex_array::arrays::Constant; -use vortex_array::arrays::Decimal; -use vortex_array::arrays::Dict; -use vortex_array::arrays::Extension; -use vortex_array::arrays::FixedSizeList; -use vortex_array::arrays::List; -use vortex_array::arrays::ListView; -use vortex_array::arrays::Masked; -use vortex_array::arrays::Null; -use vortex_array::arrays::Patched; -use vortex_array::arrays::Primitive; -use vortex_array::arrays::Struct; -use vortex_array::arrays::VarBin; -use vortex_array::arrays::VarBinView; -use vortex_array::arrays::Variant; -use vortex_array::arrays::patched::use_experimental_patches; use vortex_array::dtype::FieldPath; use vortex_btrblocks::BtrBlocksCompressorBuilder; use vortex_btrblocks::SchemeExt; use vortex_btrblocks::schemes::integer::IntDictScheme; -use vortex_bytebool::ByteBool; -use vortex_datetime_parts::DateTimeParts; -use vortex_decimal_byte_parts::DecimalByteParts; use vortex_error::VortexExpect; -use vortex_fastlanes::BitPacked; -use vortex_fastlanes::Delta; -use vortex_fastlanes::FoR; -use vortex_fastlanes::RLE; -use vortex_fsst::FSST; use vortex_layout::LayoutStrategy; use vortex_layout::LayoutStrategyEncodingValidator; use vortex_layout::layouts::buffered::BufferedStrategy; @@ -58,80 +28,11 @@ use vortex_layout::layouts::table::TableStrategy; use vortex_layout::layouts::table::use_experimental_list_layout; use vortex_layout::layouts::zoned::writer::ZonedLayoutOptions; use vortex_layout::layouts::zoned::writer::ZonedStrategy; -#[cfg(feature = "unstable_encodings")] -use vortex_onpair::OnPair; -use vortex_pco::Pco; -use vortex_runend::RunEnd; -use vortex_sequence::Sequence; -use vortex_sparse::Sparse; use vortex_utils::aliases::hash_map::HashMap; use vortex_utils::aliases::hash_set::HashSet; -use vortex_zigzag::ZigZag; -#[cfg(feature = "zstd")] -use vortex_zstd::Zstd; -#[cfg(all(feature = "zstd", feature = "unstable_encodings"))] -use vortex_zstd::ZstdBuffers; const ONE_MEG: u64 = 1 << 20; -/// Static registry of all allowed array encodings for file writing. -/// -/// This 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(); - - // Canonical encodings from vortex-array - allowed.insert(Null.id()); - allowed.insert(Bool.id()); - allowed.insert(Primitive.id()); - allowed.insert(Decimal.id()); - allowed.insert(VarBin.id()); - allowed.insert(VarBinView.id()); - allowed.insert(List.id()); - allowed.insert(ListView.id()); - allowed.insert(FixedSizeList.id()); - allowed.insert(Struct.id()); - allowed.insert(Extension.id()); - allowed.insert(Chunked.id()); - allowed.insert(Constant.id()); - allowed.insert(Masked.id()); - allowed.insert(Dict.id()); - allowed.insert(Variant.id()); - - // Compressed encodings from encoding crates - allowed.insert(ALP.id()); - allowed.insert(ALPRD.id()); - allowed.insert(BitPacked.id()); - allowed.insert(ByteBool.id()); - allowed.insert(DateTimeParts.id()); - allowed.insert(DecimalByteParts.id()); - allowed.insert(Delta.id()); - allowed.insert(FoR.id()); - allowed.insert(FSST.id()); - #[cfg(feature = "unstable_encodings")] - allowed.insert(OnPair.id()); - allowed.insert(Pco.id()); - allowed.insert(RLE.id()); - allowed.insert(RunEnd.id()); - allowed.insert(Sequence.id()); - allowed.insert(Sparse.id()); - allowed.insert(ZigZag.id()); - - // Experimental encodings - - if use_experimental_patches() { - allowed.insert(Patched.id()); - } - - #[cfg(feature = "zstd")] - allowed.insert(Zstd.id()); - #[cfg(all(feature = "zstd", feature = "unstable_encodings"))] - allowed.insert(ZstdBuffers.id()); - - allowed -}); - /// How the compressor was configured on [`WriteStrategyBuilder`]. enum CompressorConfig { /// A [`BtrBlocksCompressorBuilder`] that [`WriteStrategyBuilder::build`] will finalize. @@ -174,7 +75,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(), @@ -216,11 +117,17 @@ impl WriteStrategyBuilder { self } - /// Override the allowed array encodings for normalization. + /// Override the allowed array encodings for file writing. /// /// The configured flat leaf strategy is wrapped in a [`LayoutStrategyEncodingValidator`] /// that recursively checks every chunk before passing it to the leaf writer. pub fn with_allow_encodings(mut self, allow_encodings: HashSet) -> Self { + self.compressor = match self.compressor { + CompressorConfig::BtrBlocks(builder) => { + CompressorConfig::BtrBlocks(builder.retain_allowed_encodings(&allow_encodings)) + } + compressor => compressor, + }; self.allow_encodings = Some(allow_encodings); self } @@ -239,6 +146,11 @@ impl WriteStrategyBuilder { /// The builder is finalized during [`build`](Self::build), producing two compressors: one for /// data (with `IntDictScheme` excluded) and one for stats. pub fn with_btrblocks_builder(mut self, builder: BtrBlocksCompressorBuilder) -> Self { + let builder = if let Some(allow_encodings) = &self.allow_encodings { + builder.retain_allowed_encodings(allow_encodings) + } else { + builder + }; self.compressor = CompressorConfig::BtrBlocks(builder); self } diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index a6c573232b5..dbe02b151da 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -93,6 +93,7 @@ static SESSION: LazyLock = LazyLock::new(|| { .with::(); crate::register_default_encodings(&session); + crate::enable_all_registered_array_encodings(&session); session }); diff --git a/vortex-file/src/writer.rs b/vortex-file/src/writer.rs index 58d1cb7cd4f..0f5ff441081 100644 --- a/vortex-file/src/writer.rs +++ b/vortex-file/src/writer.rs @@ -15,7 +15,6 @@ use futures::future::LocalBoxFuture; use futures::future::ready; use futures::pin_mut; use futures::select; -use itertools::Itertools; use vortex_array::ArrayContext; use vortex_array::ArrayRef; use vortex_array::dtype::DType; @@ -23,13 +22,13 @@ use vortex_array::dtype::FieldPath; use vortex_array::expr::stats::Stat; use vortex_array::iter::ArrayIterator; use vortex_array::iter::ArrayIteratorExt; -use vortex_array::session::ArraySessionExt; use vortex_array::stats::PRUNING_STATS; use vortex_array::stream::ArrayStream; 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; @@ -50,7 +49,6 @@ use vortex_session::VortexSession; use vortex_session::registry::ReadContext; use vortex_utils::aliases::hash_map::HashMap; -use crate::ALLOWED_ENCODINGS; use crate::Footer; use crate::MAGIC_BYTES; use crate::WriteStrategyBuilder; @@ -63,8 +61,7 @@ use crate::segments::writer::BufferedSegmentSink; /// Configure a new writer, which can eventually be used to write an [`ArrayStream`] into a sink /// that implements [`VortexWrite`]. /// -/// Unless overridden, the default [write strategy][crate::WriteStrategyBuilder] will be used with no -/// additional configuration. +/// All write strategies are restricted to the encodings in the session's enabled editions. /// /// Construct with [`WriteOptionsSessionExt::write_options`] for normal use so the writer inherits /// the session's runtime, array registry, and memory configuration. @@ -82,14 +79,7 @@ pub trait WriteOptionsSessionExt: SessionExt { /// Create [`VortexWriteOptions`] for writing to a Vortex file. fn write_options(&self) -> VortexWriteOptions { let session = self.session(); - VortexWriteOptions { - strategy: WriteStrategyBuilder::default().build(), - session, - exclude_dtype: false, - file_statistics: PRUNING_STATS.to_vec(), - max_variable_length_statistics_size: 64, - metadata: HashMap::default(), - } + VortexWriteOptions::new(session) } } impl WriteOptionsSessionExt for S {} @@ -97,8 +87,11 @@ impl WriteOptionsSessionExt for S {} impl VortexWriteOptions { /// Create a new [`VortexWriteOptions`] with the given session. pub fn new(session: VortexSession) -> Self { + let strategy = WriteStrategyBuilder::default() + .with_allow_encodings(session.enabled_encoding_ids().into_iter().collect()) + .build(); VortexWriteOptions { - strategy: WriteStrategyBuilder::default().build(), + strategy, session, exclude_dtype: false, file_statistics: PRUNING_STATS.to_vec(), @@ -111,7 +104,7 @@ impl VortexWriteOptions { /// /// The strategy controls repartitioning, statistics layout, compression, and leaf segment /// emission. Use [`WriteStrategyBuilder`] when only a small part of the default strategy needs - /// customization. + /// customization. Replacing the strategy does not change the enabled-edition encoding policy. pub fn with_strategy(mut self, strategy: Arc) -> Self { self.strategy = strategy; self @@ -199,9 +192,10 @@ impl VortexWriteOptions { // 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()) + let enabled_encoding_ids = self.session.enabled_encoding_ids(); + let ctx = ArrayContext::new(enabled_encoding_ids.clone()) // Only permit encodings known to the session. - .with_valid_ids(self.session.arrays().registry().ids()); + .with_valid_ids(enabled_encoding_ids); let dtype = stream.dtype().clone(); let (mut ptr, eof) = SequenceId::root().split(); @@ -601,3 +595,40 @@ impl WriteSummary { .collect()) } } + +#[cfg(test)] +mod tests { + use vortex_array::ArrayContext; + use vortex_array::VTable; + use vortex_array::array_session; + use vortex_array::arrays::Bool; + use vortex_array::arrays::Primitive; + use vortex_edition::Edition; + use vortex_edition::EditionDeclaration; + use vortex_edition::EditionId; + use vortex_edition::EditionSession; + use vortex_edition::EditionSessionExt; + + #[test] + fn array_context_only_permits_enabled_encodings() -> Result<(), vortex_edition::EditionError> { + const EDITION: EditionId = EditionId::new("test", 2026, 7, 0); + static DECLARATION: EditionDeclaration = EditionDeclaration { + edition: Edition { + id: EDITION, + min_vortex_version: None, + }, + added: &[&"vortex.primitive"], + }; + + let session = array_session().with::(); + session.register_edition(&DECLARATION)?; + session.enable_edition(EDITION)?; + + let enabled_encoding_ids = session.enabled_encoding_ids(); + let ctx = + ArrayContext::new(enabled_encoding_ids.clone()).with_valid_ids(enabled_encoding_ids); + assert_eq!(ctx.to_ids(), [Primitive.id()]); + assert!(ctx.intern(&Bool.id()).is_none()); + Ok(()) + } +} diff --git a/vortex-file/tests/common/mod.rs b/vortex-file/tests/common/mod.rs new file mode 100644 index 00000000000..68e78908098 --- /dev/null +++ b/vortex-file/tests/common/mod.rs @@ -0,0 +1,35 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_array::session::ArraySessionExt; +use vortex_edition::Edition; +use vortex_edition::EditionId; +use vortex_edition::EditionInclusion; +use vortex_edition::EditionSessionExt; +use vortex_error::VortexExpect; +use vortex_error::vortex_err; +use vortex_session::VortexSession; + +/// This is a vortex edition used for testing and shouldn't made public. +const TEST_EDITION: EditionId = EditionId::new("test", 2026, 7, 0); + +pub fn enable_all_registered_array_encodings(session: &VortexSession) { + let editions = session.editions(); + editions + .declare_edition(Edition { + id: TEST_EDITION, + min_vortex_version: None, + }) + .map_err(|error| vortex_err!("{error}")) + .vortex_expect("test edition is valid"); + for id in session.arrays().registry().ids() { + editions + .declare_inclusion(EditionInclusion::new(&id, TEST_EDITION)) + .map_err(|error| vortex_err!("{error}")) + .vortex_expect("registered array encoding has one test-edition inclusion"); + } + session + .enable_edition(TEST_EDITION) + .map_err(|error| vortex_err!("{error}")) + .vortex_expect("test edition is registered"); +} diff --git a/vortex-file/tests/enabled_unregistered_ids.rs b/vortex-file/tests/enabled_unregistered_ids.rs new file mode 100644 index 00000000000..3beb3ca80a0 --- /dev/null +++ b/vortex-file/tests/enabled_unregistered_ids.rs @@ -0,0 +1,80 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! A session's enabled editions may include encodings the session has no registration for, +//! e.g. an optional encoding crate that is not compiled in (vortex-jni enables the unstable +//! edition for `vortex.parquet.variant` without registering its sibling encodings). Files +//! written by such a session must stay readable. + +#![expect(clippy::tests_outside_test_module)] + +use futures::TryStreamExt; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::ChunkedArray; +use vortex_array::assert_arrays_eq; +use vortex_array::session::ArraySessionExt; +use vortex_buffer::ByteBuffer; +use vortex_buffer::ByteBufferMut; +use vortex_buffer::buffer; +use vortex_edition::Edition; +use vortex_edition::EditionId; +use vortex_edition::EditionInclusion; +use vortex_edition::EditionSessionExt; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_file::OpenOptionsSessionExt; +use vortex_file::WriteOptionsSessionExt; +use vortex_io::session::RuntimeSession; +use vortex_layout::session::LayoutSession; +use vortex_session::VortexSession; + +const TEST_EDITION: EditionId = EditionId::new("test", 2026, 7, 0); + +fn session_with_unregistered_edition_member() -> VortexResult { + let session = vortex_array::array_session() + .with::() + .with::(); + vortex_file::register_default_encodings(&session); + + let editions = session.editions(); + editions + .declare_edition(Edition { + id: TEST_EDITION, + min_vortex_version: None, + }) + .map_err(|error| vortex_err!("{error}"))?; + for id in session.arrays().registry().ids() { + editions + .declare_inclusion(EditionInclusion::new(&id, TEST_EDITION)) + .map_err(|error| vortex_err!("{error}"))?; + } + editions + .declare_inclusion(EditionInclusion::new(&"test.not_registered", TEST_EDITION)) + .map_err(|error| vortex_err!("{error}"))?; + session + .enable_edition(TEST_EDITION) + .map_err(|error| vortex_err!("{error}"))?; + Ok(session) +} + +#[tokio::test] +async fn roundtrip_with_enabled_but_unregistered_encoding() -> VortexResult<()> { + let session = session_with_unregistered_edition_member()?; + + let array = buffer![1u32, 2, 3, 4].into_array(); + let mut output = ByteBufferMut::empty(); + session + .write_options() + .write(&mut output, array.clone().to_array_stream()) + .await?; + + let file = session + .open_options() + .open_buffer(ByteBuffer::from(output))?; + let chunks: Vec<_> = file.scan()?.into_stream()?.try_collect().await?; + let actual = ChunkedArray::from_iter(chunks).into_array(); + let mut ctx = session.create_execution_ctx(); + assert_arrays_eq!(array, actual, &mut ctx); + Ok(()) +} diff --git a/vortex-file/tests/issue_8819_footer_segment_oob.rs b/vortex-file/tests/issue_8819_footer_segment_oob.rs index 7b09d7a9c35..93e4ab46db3 100644 --- a/vortex-file/tests/issue_8819_footer_segment_oob.rs +++ b/vortex-file/tests/issue_8819_footer_segment_oob.rs @@ -23,12 +23,17 @@ use vortex_io::session::RuntimeSession; use vortex_layout::session::LayoutSession; use vortex_session::VortexSession; +mod common; + +use common::enable_all_registered_array_encodings; + static SESSION: LazyLock = LazyLock::new(|| { let session = vortex_array::array_session() .with::() .with::(); vortex_file::register_default_encodings(&session); + enable_all_registered_array_encodings(&session); session }); diff --git a/vortex-file/tests/test_write_table.rs b/vortex-file/tests/test_write_table.rs index bd664cc24ed..3f69de67394 100644 --- a/vortex-file/tests/test_write_table.rs +++ b/vortex-file/tests/test_write_table.rs @@ -29,12 +29,18 @@ use vortex_layout::layouts::flat::writer::FlatLayoutStrategy; use vortex_layout::layouts::table::TableStrategy; use vortex_layout::session::LayoutSession; use vortex_session::VortexSession; + +mod common; + +use common::enable_all_registered_array_encodings; + static SESSION: LazyLock = LazyLock::new(|| { let session = vortex_array::array_session() .with::() .with::(); vortex_file::register_default_encodings(&session); + enable_all_registered_array_encodings(&session); session }); diff --git a/vortex-jni/src/session.rs b/vortex-jni/src/session.rs index 9adaf544431..8a19c01d149 100644 --- a/vortex-jni/src/session.rs +++ b/vortex-jni/src/session.rs @@ -7,6 +7,10 @@ use jni::EnvUnowned; use jni::objects::JClass; use jni::sys::jlong; use vortex::VortexSessionDefault; +use vortex::editions::DEFAULT_UNSTABLE_EDITION; +use vortex::editions::EditionSessionExt; +use vortex::error::VortexExpect; +use vortex::error::vortex_err; use vortex::io::runtime::BlockingRuntime; use vortex::io::session::RuntimeSessionExt; use vortex::session::VortexSession; @@ -15,9 +19,17 @@ use crate::RUNTIME; /// Constructs a fresh [`VortexSession`] bound to the JNI-shared tokio runtime and returns /// an opaque pointer that Java must pass to [`Java_dev_vortex_jni_NativeSession_free`]. +/// +/// JNI sessions enable the latest unstable edition so that `vortex.parquet.variant`, an +/// unstable-family encoding registered by [`vortex_parquet_variant::initialize`], stays +/// writable for variant-typed columns. pub(crate) fn new_session() -> Box { let session = VortexSession::default().with_handle(RUNTIME.handle()); vortex_parquet_variant::initialize(&session); + session + .enable_edition(DEFAULT_UNSTABLE_EDITION) + .map_err(|e| vortex_err!("{e}")) + .vortex_expect("default unstable edition is registered"); Box::new(session) } diff --git a/vortex-jni/src/writer.rs b/vortex-jni/src/writer.rs index 53649eef4ab..2e459dac818 100644 --- a/vortex-jni/src/writer.rs +++ b/vortex-jni/src/writer.rs @@ -33,15 +33,12 @@ use jni::sys::jobject; use object_store::ObjectStore; use object_store::path::Path as ObjectStorePath; use vortex::array::ArrayRef; -use vortex::array::VTable; use vortex::array::scalar::PValue; use vortex::array::scalar::Scalar; use vortex::array::scalar::ScalarValue; use vortex::array::stats::StatsSet; use vortex::array::stream::ArrayStreamAdapter; use vortex::dtype::DType; -use vortex::dtype::Field as DTypeField; -use vortex::dtype::FieldPath; use vortex::error::VortexError; use vortex::error::VortexResult; use vortex::error::vortex_err; @@ -49,7 +46,6 @@ use vortex::expr::stats::Stat; use vortex::expr::stats::StatsProvider; use vortex::file::CountingVortexWrite; use vortex::file::WriteOptionsSessionExt; -use vortex::file::WriteStrategyBuilder; use vortex::file::WriteSummary; use vortex::file::multi::parse_uri_or_path; use vortex::io::VortexWrite; @@ -61,7 +57,6 @@ use vortex::io::session::RuntimeSessionExt; use vortex::session::VortexSession; use vortex::utils::aliases::hash_map::HashMap; use vortex_arrow::ArrowSessionExt; -use vortex_parquet_variant::ParquetVariant; use crate::RUNTIME; use crate::dtype::import_arrow_schema; @@ -99,45 +94,6 @@ fn resolve_store( } } -fn write_options_for_schema( - session: &VortexSession, - write_schema: &DType, -) -> vortex::file::VortexWriteOptions { - let variant_paths = variant_field_paths(write_schema); - if variant_paths.is_empty() { - return session.write_options(); - } - - let mut allowed = vortex::file::ALLOWED_ENCODINGS.clone(); - allowed.insert(ParquetVariant.id()); - - let strategy = WriteStrategyBuilder::default().with_allow_encodings(allowed); - - session.write_options().with_strategy(strategy.build()) -} - -fn variant_field_paths(dtype: &DType) -> Vec { - let mut paths = Vec::new(); - collect_variant_field_paths(dtype, FieldPath::root(), &mut paths); - paths -} - -fn collect_variant_field_paths(dtype: &DType, path: FieldPath, paths: &mut Vec) { - match dtype { - DType::Variant(_) => paths.push(path), - DType::Struct(fields, _) => { - for (name, field_dtype) in fields.names().iter().zip(fields.fields()) { - collect_variant_field_paths( - &field_dtype, - path.clone().push(DTypeField::from(name.clone())), - paths, - ); - } - } - _ => {} - } -} - /// Native writer holding a write-task handle and a sender that Java pushes batches into. pub struct NativeWriter { handle: Option>>, @@ -413,7 +369,7 @@ pub extern "system" fn Java_dev_vortex_jni_NativeWriter_create( let resolved = resolve_store(&file_path, &properties)?; let (tx, rx) = mpsc::channel(WRITE_CHANNEL_CAPACITY); let stream = ArrayStreamAdapter::new(write_schema.clone(), rx); - let write_options = write_options_for_schema(session, &write_schema); + let write_options = session.write_options(); let (bytes_written, handle) = match resolved { ResolvedStore::Path(path) => { @@ -491,7 +447,7 @@ pub extern "system" fn Java_dev_vortex_jni_NativeWriter_createStream( let writable = Arc::new(env.new_global_ref(&writable)?); let (tx, rx) = mpsc::channel(WRITE_CHANNEL_CAPACITY); let stream = ArrayStreamAdapter::new(write_schema.clone(), rx); - let write_options = write_options_for_schema(session, &write_schema); + let write_options = session.write_options(); let mut write = CountingVortexWrite::new(JavaWrite::new(vm, writable)); let bytes_written = write.counter(); diff --git a/vortex-layout/src/layouts/struct_/reader.rs b/vortex-layout/src/layouts/struct_/reader.rs index 0dae97f0de0..e92b7b0a774 100644 --- a/vortex-layout/src/layouts/struct_/reader.rs +++ b/vortex-layout/src/layouts/struct_/reader.rs @@ -443,6 +443,7 @@ mod tests { use crate::sequence::SequenceId; use crate::sequence::SequentialArrayStreamExt; use crate::test::SESSION; + use crate::test::test_session; #[fixture] fn empty_struct() -> (Arc, LayoutRef) { @@ -456,7 +457,7 @@ mod tests { ); let segments2 = Arc::::clone(&segments); let layout = block_on(|handle| async move { - let session = SESSION.clone().with_handle(handle); + let session = test_session().with_handle(handle); strategy .write_stream( ctx, @@ -493,7 +494,7 @@ mod tests { ); let segments2 = Arc::::clone(&segments); let layout = block_on(|handle| async move { - let session = SESSION.clone().with_handle(handle); + let session = test_session().with_handle(handle); strategy .write_stream( ctx, @@ -533,7 +534,7 @@ mod tests { ); let segments2 = Arc::::clone(&segments); let layout = block_on(|handle| async move { - let session = SESSION.clone().with_handle(handle); + let session = test_session().with_handle(handle); strategy .write_stream( ctx, @@ -578,7 +579,7 @@ mod tests { ); let segments2 = Arc::::clone(&segments); let layout = block_on(|handle| async move { - let session = SESSION.clone().with_handle(handle); + let session = test_session().with_handle(handle); strategy .write_stream( ctx, @@ -764,19 +765,20 @@ mod tests { // Project out the nested struct field. // The projection should preserve the nulls of the `b` struct when we select out the // child column `c`. - let reader = layout - .new_reader("".into(), segments, &SESSION, &Default::default()) - .unwrap(); let expr = select( vec![FieldName::from("c")], get_item("b", get_item("a", root())), ); - - let project = reader - .projection_evaluation(&(0..3), &expr, MaskFuture::new_true(3)) - .unwrap(); - - let result = block_on(move |_| project).unwrap(); + let result = block_on(move |handle| { + let session = test_session().with_handle(handle); + async move { + layout + .new_reader("".into(), segments, &session, &Default::default())? + .projection_evaluation(&(0..3), &expr, MaskFuture::new_true(3))? + .await + } + }) + .unwrap(); // The result is a nullable struct (because root.a.b is nullable) with a non-nullable // field "c" (because the original field was non-nullable). @@ -857,7 +859,7 @@ mod tests { ); let segments2 = Arc::::clone(&segments); let layout = block_on(|handle| async move { - let session = SESSION.clone().with_handle(handle); + let session = test_session().with_handle(handle); strategy .write_stream( ctx, diff --git a/vortex-layout/src/layouts/table.rs b/vortex-layout/src/layouts/table.rs index 0853478c14a..65d851374f9 100644 --- a/vortex-layout/src/layouts/table.rs +++ b/vortex-layout/src/layouts/table.rs @@ -368,14 +368,15 @@ mod tests { use crate::sequence::SequentialArrayStreamExt; use crate::sequence::SequentialStreamAdapter; use crate::sequence::SequentialStreamExt; - use crate::test::SESSION; + use crate::test::test_session; async fn write(strategy: &S, array: ArrayRef) -> VortexResult { let segments = Arc::new(TestSegments::default()); let (ptr, eof) = SequenceId::root().split(); let stream = array.to_array_stream().sequenced(ptr); + let session = test_session().with_tokio(); strategy - .write_stream(ArrayContext::empty(), segments, stream, eof, &SESSION) + .write_stream(ArrayContext::empty(), segments, stream, eof, &session) .await } @@ -681,7 +682,7 @@ mod tests { ); block_on(|handle| async move { - let session = SESSION.clone().with_handle(handle); + let session = test_session().with_handle(handle); strategy .write_stream( ctx, diff --git a/vortex-layout/src/test.rs b/vortex-layout/src/test.rs index 72258a1161e..568097402a1 100644 --- a/vortex-layout/src/test.rs +++ b/vortex-layout/src/test.rs @@ -8,8 +8,10 @@ use vortex_session::VortexSession; use crate::session::LayoutSession; -pub static SESSION: LazyLock = LazyLock::new(|| { +pub fn test_session() -> VortexSession { vortex_array::array_session() .with::() .with::() -}); +} + +pub static SESSION: LazyLock = LazyLock::new(test_session); 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..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, diff --git a/vortex/src/editions/mod.rs b/vortex/src/editions/mod.rs index 1c3323d077f..0a425b27096 100644 --- a/vortex/src/editions/mod.rs +++ b/vortex/src/editions/mod.rs @@ -3,28 +3,28 @@ //! The Vortex edition declarations. //! -//! [`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`]. +//! [`vortex_edition`] provides the types, session variables, and test harness. The actual +//! first-party declarations live here, one module per edition. The default session first +//! registers them with [`crate::editions::register_default_editions`] and then selects its +//! write policy with [`crate::editions::enable_default_editions`]. //! -//! 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. -//! -//! Note this is currently unused but a future PR will make this public and gate the writer behind -//! editions. +//! The default file writer resolves the session's enabled editions at write time. The +//! facade enables the newest frozen `core` edition, [`crate::editions::CORE_2026_07_0`], +//! and additionally enables the latest unstable edition when the `unstable_encodings` +//! feature is selected. pub mod core; +#[cfg(test)] +mod tests; pub mod unstable; -use vortex_edition::EditionDeclaration; +pub use vortex_edition::Edition; +pub use vortex_edition::EditionDeclaration; +pub use vortex_edition::EditionId; +pub use vortex_edition::EditionInclusion; pub use vortex_edition::EditionSession; -use vortex_edition::EditionSessionExt; +pub use vortex_edition::EditionSessionExt; +pub use vortex_edition::EnabledEditions; use vortex_error::VortexExpect; use vortex_error::vortex_err; use vortex_session::VortexSession; @@ -38,7 +38,14 @@ pub use self::unstable::UNSTABLE_2026_02_0; pub use self::unstable::UNSTABLE_2026_04_0; pub use self::unstable::UNSTABLE_2026_06_0; -/// The Vortex editions, each declared together with the encodings it adds. +/// The `core` edition enabled for writing by the default Vortex session. +pub const DEFAULT_CORE_EDITION: EditionId = CORE_2026_07_0; + +/// The `unstable` edition enabled for writing by the default Vortex session when the +/// `unstable_encodings` feature is selected. +pub const DEFAULT_UNSTABLE_EDITION: EditionId = UNSTABLE_2026_06_0; + +/// The first-party Vortex edition declarations. pub static EDITION_DECLARATIONS: &[&EditionDeclaration] = &[ &core::v2025_05::DECLARATION, &core::v2025_06::DECLARATION, @@ -52,11 +59,28 @@ pub static EDITION_DECLARATIONS: &[&EditionDeclaration] = &[ /// Register the Vortex edition declarations with the session's [`EditionSession`]. pub fn register_default_editions(session: &VortexSession) { - let editions = session.editions(); for declaration in EDITION_DECLARATIONS { - editions - .declare(declaration) + session + .register_edition(declaration) .map_err(|e| vortex_err!("{e}")) .vortex_expect("edition declarations are valid"); } } + +/// Enable the default Vortex editions for writing. +/// +/// This selects the newest frozen `core` edition and, when configured, the newest unstable +/// edition. All declarations must have been registered first with +/// [`register_default_editions`]. +pub fn enable_default_editions(session: &VortexSession) { + session + .enable_edition(DEFAULT_CORE_EDITION) + .map_err(|e| vortex_err!("{e}")) + .vortex_expect("default core edition is registered"); + + #[cfg(feature = "unstable_encodings")] + session + .enable_edition(DEFAULT_UNSTABLE_EDITION) + .map_err(|e| vortex_err!("{e}")) + .vortex_expect("default unstable edition is registered"); +} diff --git a/vortex/src/editions/tests.rs b/vortex/src/editions/tests.rs new file mode 100644 index 00000000000..5cbfc25a744 --- /dev/null +++ b/vortex/src/editions/tests.rs @@ -0,0 +1,483 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::Arc; + +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::array_session; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::StructArray; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::field_path; +use vortex_array::stream::ArrayStreamExt; +use vortex_btrblocks::BtrBlocksCompressorBuilder; +use vortex_buffer::ByteBufferMut; +use vortex_edition::Edition; +use vortex_edition::EditionDeclaration; +use vortex_edition::EditionError; +use vortex_edition::EditionId; +use vortex_edition::EditionSession; +use vortex_edition::EditionSessionExt; +use vortex_edition::test_harness::validate_edition; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_file::OpenOptionsSessionExt; +use vortex_file::WriteOptionsSessionExt; +use vortex_file::WriteStrategyBuilder; +use vortex_io::session::RuntimeSession; +use vortex_layout::LayoutStrategy; +use vortex_layout::layouts::compressed::CompressingStrategy; +use vortex_layout::layouts::flat::writer::FlatLayoutStrategy; +use vortex_layout::session::LayoutSession; +use vortex_sequence::Sequence; +use vortex_session::VortexSession; +use vortex_utils::aliases::hash_set::HashSet; + +use super::CORE_2025_05_0; +use super::CORE_2026_07_0; +use super::DEFAULT_CORE_EDITION; +use super::DEFAULT_UNSTABLE_EDITION; +use super::EDITION_DECLARATIONS; +use super::UNSTABLE_2026_06_0; + +fn session() -> Result { + let session = EditionSession::empty(); + for declaration in EDITION_DECLARATIONS { + session.declare(declaration)?; + } + Ok(session) +} + +#[test] +fn every_declared_edition_validates() -> Result<(), EditionError> { + let session = session()?; + for declaration in EDITION_DECLARATIONS { + validate_edition(&session, &declaration.edition.id)?; + } + Ok(()) +} + +/// 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 session = session().unwrap_or_else(|e| panic!("registering editions: {e}")); + let encodings = session.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 encodings_in_editions_unions_families() { + let session = session().unwrap_or_else(|e| panic!("registering editions: {e}")); + let core_only: Vec<_> = session + .encodings_in(&CORE_2026_07_0) + .into_iter() + .map(|inclusion| inclusion.encoding_id) + .collect(); + let mut both = core_only.clone(); + both.extend( + session + .encodings_in(&UNSTABLE_2026_06_0) + .into_iter() + .map(|inclusion| inclusion.encoding_id), + ); + both.sort_unstable(); + both.dedup(); + + assert!(both.len() > core_only.len()); + assert!(both.iter().any(|id| id.as_str() == "fastlanes.delta")); + assert!(both.iter().any(|id| id.as_str() == "vortex.onpair")); + assert!(core_only.iter().all(|id| both.contains(id))); +} + +#[test] +fn earlier_editions_are_subsets() { + let session = session().unwrap_or_else(|e| panic!("registering editions: {e}")); + let first = session.encodings_in(&CORE_2025_05_0); + let latest = session.encodings_in(&CORE_2026_07_0); + assert!(first.iter().all(|inclusion| { + latest + .iter() + .any(|latest| latest.encoding_id == inclusion.encoding_id) + })); + assert!(first.len() < latest.len()); +} + +#[test] +fn default_session_enables_the_write_editions() { + use crate::VortexSessionDefault; + + let session = VortexSession::default(); + let enabled = session.enabled_editions().editions(); + assert!(enabled.contains(&DEFAULT_CORE_EDITION)); + + #[cfg(feature = "unstable_encodings")] + assert!(enabled.contains(&DEFAULT_UNSTABLE_EDITION)); + #[cfg(not(feature = "unstable_encodings"))] + assert!(!enabled.contains(&DEFAULT_UNSTABLE_EDITION)); +} + +#[test] +fn core_edition_ids_are_registered_array_encodings() { + use vortex_array::session::ArraySessionExt; + + use crate::VortexSessionDefault; + + let session = VortexSession::default(); + let registry = session.arrays().registry().clone(); + for inclusion in session.editions().encodings_in(&CORE_2026_07_0) { + assert!( + registry.find(&inclusion.encoding_id).is_some(), + "{} is declared in core but not registered as an array encoding", + inclusion.encoding_id + ); + } +} + +fn baseline_core_session() -> VortexResult { + use crate::VortexSessionDefault; + + let session = VortexSession::default(); + session + .enable_edition(CORE_2025_05_0) + .map_err(|error| vortex_err!("{error}"))?; + Ok(session) +} + +fn sequential_integers() -> PrimitiveArray { + PrimitiveArray::from_iter(0..65_536i32) +} + +const WRITER_TEST_EDITION: EditionId = EditionId::new("writer-test", 2026, 7, 0); + +static WRITER_TEST_DECLARATION: EditionDeclaration = EditionDeclaration { + edition: Edition { + id: WRITER_TEST_EDITION, + min_vortex_version: None, + }, + added: &[ + &"vortex.chunked", + &"vortex.constant", + &"vortex.primitive", + &"vortex.struct", + ], +}; + +fn writer_test_session() -> VortexResult { + let session = array_session() + .with::() + .with::() + .with::(); + vortex_file::register_default_encodings(&session); + session + .register_edition(&WRITER_TEST_DECLARATION) + .map_err(|error| vortex_err!("{error}"))?; + session + .enable_edition(WRITER_TEST_EDITION) + .map_err(|error| vortex_err!("{error}"))?; + Ok(session) +} + +fn forbidden_sequence(len: usize) -> VortexResult { + Ok(Sequence::try_new_typed(0i32, 1i32, Nullability::NonNullable, len)?.into_array()) +} + +fn forbidden_sequence_compressor( + chunk: &ArrayRef, + _ctx: &mut ExecutionCtx, +) -> VortexResult { + if matches!( + chunk.dtype(), + DType::Primitive(PType::I32, Nullability::NonNullable) + ) { + forbidden_sequence(chunk.len()) + } else { + Ok(chunk.clone()) + } +} + +fn custom_compressing_flat_strategy() -> Arc { + Arc::new(CompressingStrategy::new( + FlatLayoutStrategy::default(), + forbidden_sequence_compressor, + )) +} + +async fn assert_round_trip_encodings_are_enabled( + session: &VortexSession, + strategy: Option>, + array: ArrayRef, +) -> VortexResult<()> { + let mut buffer = ByteBufferMut::empty(); + let write_options = match strategy { + Some(strategy) => session.write_options().with_strategy(strategy), + None => session.write_options(), + }; + if let Err(error) = write_options + .write(&mut buffer, array.to_array_stream()) + .await + { + let message = error.to_string(); + if message.contains("not permitted by ctx") + || message.contains("normalize forbids encoding") + { + return Ok(()); + } + return Err(error); + } + + let round_tripped = session + .open_options() + .open_buffer(buffer)? + .scan()? + .into_array_stream()? + .read_all() + .await?; + let actual: HashSet<_> = round_tripped + .depth_first_traversal() + .map(|array| array.encoding_id()) + .collect(); + let allowed: HashSet<_> = session.enabled_encoding_ids().into_iter().collect(); + let mut forbidden: Vec<_> = actual.difference(&allowed).map(|id| id.as_str()).collect(); + forbidden.sort_unstable(); + if !forbidden.is_empty() { + return Err(vortex_err!( + "round-tripped array contains encodings outside {WRITER_TEST_EDITION}: {forbidden:?}" + )); + } + + Ok(()) +} + +#[tokio::test] +async fn default_strategy_round_trip_uses_only_enabled_encodings() -> VortexResult<()> { + let session = writer_test_session()?; + assert_round_trip_encodings_are_enabled(&session, None, sequential_integers().into_array()) + .await +} + +#[tokio::test] +async fn replacement_default_builder_round_trip_uses_only_enabled_encodings() -> VortexResult<()> { + let session = writer_test_session()?; + assert_round_trip_encodings_are_enabled( + &session, + Some(WriteStrategyBuilder::default().build()), + sequential_integers().into_array(), + ) + .await +} + +#[tokio::test] +async fn replacement_btrblocks_builder_round_trip_uses_only_enabled_encodings() -> VortexResult<()> +{ + let session = writer_test_session()?; + let strategy = WriteStrategyBuilder::default() + .with_btrblocks_builder(BtrBlocksCompressorBuilder::default()) + .build(); + assert_round_trip_encodings_are_enabled( + &session, + Some(strategy), + sequential_integers().into_array(), + ) + .await +} + +#[tokio::test] +async fn opaque_compressor_round_trip_uses_only_enabled_encodings() -> VortexResult<()> { + let session = writer_test_session()?; + let strategy = WriteStrategyBuilder::default() + .with_compressor(forbidden_sequence_compressor) + .build(); + assert_round_trip_encodings_are_enabled( + &session, + Some(strategy), + sequential_integers().into_array(), + ) + .await +} + +#[tokio::test] +async fn custom_flat_strategy_round_trip_uses_only_enabled_encodings() -> VortexResult<()> { + let session = writer_test_session()?; + let strategy = WriteStrategyBuilder::default() + .with_flat_strategy(Arc::new(FlatLayoutStrategy::default())) + .build(); + assert_round_trip_encodings_are_enabled( + &session, + Some(strategy), + sequential_integers().into_array(), + ) + .await +} + +#[tokio::test] +async fn custom_field_writer_round_trip_uses_only_enabled_encodings() -> VortexResult<()> { + let session = writer_test_session()?; + let strategy = WriteStrategyBuilder::default() + .with_field_writer(field_path!(values), custom_compressing_flat_strategy()) + .build(); + let array = + StructArray::from_fields(&[("values", sequential_integers().into_array())])?.into_array(); + assert_round_trip_encodings_are_enabled(&session, Some(strategy), array).await +} + +#[tokio::test] +async fn replacement_strategy_round_trip_uses_only_enabled_encodings() -> VortexResult<()> { + let session = writer_test_session()?; + assert_round_trip_encodings_are_enabled( + &session, + Some(custom_compressing_flat_strategy()), + sequential_integers().into_array(), + ) + .await +} + +#[tokio::test] +async fn replacement_flat_strategy_round_trip_uses_only_enabled_encodings() -> VortexResult<()> { + let session = writer_test_session()?; + assert_round_trip_encodings_are_enabled( + &session, + Some(Arc::new(FlatLayoutStrategy::default())), + forbidden_sequence(65_536)?, + ) + .await +} + +#[tokio::test] +async fn probe_compressor_round_trip_uses_only_enabled_encodings() -> VortexResult<()> { + let session = writer_test_session()?; + let strategy = WriteStrategyBuilder::default() + .with_probe_compressor(forbidden_sequence_compressor) + .build(); + assert_round_trip_encodings_are_enabled( + &session, + Some(strategy), + sequential_integers().into_array(), + ) + .await +} + +#[tokio::test] +async fn default_writer_filters_compressor_to_enabled_editions() -> VortexResult<()> { + let session = baseline_core_session()?; + let mut buffer = ByteBufferMut::empty(); + + session + .write_options() + .write( + &mut buffer, + sequential_integers().into_array().to_array_stream(), + ) + .await?; + + Ok(()) +} + +#[tokio::test] +async fn configured_btrblocks_builder_uses_enabled_editions_in_either_order() -> VortexResult<()> { + let session = baseline_core_session()?; + let allowed: HashSet<_> = session.enabled_encoding_ids().into_iter().collect(); + let strategies = [ + WriteStrategyBuilder::default() + .with_btrblocks_builder(BtrBlocksCompressorBuilder::default()) + .with_allow_encodings(allowed.clone()) + .build(), + WriteStrategyBuilder::default() + .with_allow_encodings(allowed) + .with_btrblocks_builder(BtrBlocksCompressorBuilder::default()) + .build(), + ]; + + for strategy in strategies { + let mut buffer = ByteBufferMut::empty(); + session + .write_options() + .with_strategy(strategy) + .write( + &mut buffer, + sequential_integers().into_array().to_array_stream(), + ) + .await?; + } + + Ok(()) +} + +#[tokio::test] +async fn opaque_compressor_cannot_write_outside_enabled_editions() -> VortexResult<()> { + let session = baseline_core_session()?; + let allowed = session.enabled_encoding_ids().into_iter().collect(); + let strategy = WriteStrategyBuilder::default() + .with_compressor(BtrBlocksCompressorBuilder::default().build()) + .with_allow_encodings(allowed) + .build(); + let mut buffer = ByteBufferMut::empty(); + + let result = session + .write_options() + .with_strategy(strategy) + .write( + &mut buffer, + sequential_integers().into_array().to_array_stream(), + ) + .await; + let error = match result { + Ok(_) => { + return Err(vortex_err!( + "the unrestricted opaque compressor wrote an encoding outside core@2025.05" + )); + } + Err(error) => error, + }; + let message = error.to_string(); + assert!( + message.contains("normalize forbids encoding (vortex.sequence)"), + "unexpected error: {message}" + ); + + Ok(()) +} diff --git a/vortex/src/lib.rs b/vortex/src/lib.rs index 1af8484f230..2f5c3e59aac 100644 --- a/vortex/src/lib.rs +++ b/vortex/src/lib.rs @@ -145,7 +145,6 @@ 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; @@ -309,6 +308,7 @@ impl VortexSessionDefault for VortexSession { .with::(); vortex_arrow::initialize(&session); editions::register_default_editions(&session); + editions::enable_default_editions(&session); // `MultiFileSession` holds a `moka` cache whose clock reads `std::time::Instant::now()` // when constructed. `Instant` is unsupported on `wasm32` and panics with "time not