wip

Signed-off-by: Joe Isaacs <joe.isaacs@live.co.uk>

Author: joseph-isaacs

vortex-edition/src/session.rs

@@ -230,8 +230,8 @@ pub trait EditionSessionExt: SessionExt {
 
     /// Returns the editions enabled for writing.
     ///
-    /// Accessing this method installs the enabled-editions session variable if it is absent,
-    /// opting the session into edition-gated writing with an initially empty selection.
+    /// 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::<EnabledEditions>()
     }
@@ -258,12 +258,11 @@ pub trait EditionSessionExt: SessionExt {
 
     /// Resolve the encodings in all enabled editions.
     ///
-    /// Returns `None` when the enabled-editions variable is absent, which is the explicit
-    /// opt-out used by lower-level sessions. Once the variable is present, an empty enabled
-    /// set resolves to `Some(Vec::new())` and therefore permits no edition encodings.
+    /// 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<Id> {
         let Some(enabled) = self.get_opt::<EnabledEditions>() else {
-            vec![]
+            return vec![];
         };
         let editions = self.editions();
         let mut ids: Vec<Id> = enabled

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::<EditionSession>();
+    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::<Vec<_>>(),
+        ["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::<EnabledEditions>();
+    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::<EditionSession>();
+    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();

vortex-file/src/footer/field_sizes.rs

@@ -141,6 +141,7 @@ mod tests {
             .with::<LayoutSession>()
             .with::<RuntimeSession>();
         crate::register_default_encodings(&session);
+        crate::enable_all_registered_array_encodings(&session);
         session
     });
 

vortex-file/src/lib.rs

@@ -192,6 +192,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 _;
@@ -222,10 +253,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()));
-    }
 }

vortex-file/src/open.rs

@@ -461,6 +461,7 @@ mod tests {
             .with::<LayoutSession>()
             .with::<RuntimeSession>();
         crate::register_default_encodings(&session);
+        crate::enable_all_registered_array_encodings(&session);
         session
     }
 
@@ -517,6 +518,7 @@ mod tests {
             .with::<RuntimeSession>();
 
         crate::register_default_encodings(&session);
+        crate::enable_all_registered_array_encodings(&session);
 
         // Create a large file (> 1MB)
         let mut buf = ByteBufferMut::empty();
@@ -578,6 +580,7 @@ mod tests {
             .with::<RuntimeSession>();
 
         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::<Vec<i32>>()).into_array();

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<HashSet<ArrayId>> = 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,18 @@ 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.
+    /// that recursively checks every chunk before passing it to the leaf writer. A configured
+    /// BtrBlocks builder is restricted to schemes whose output encodings are all in this set.
     pub fn with_allow_encodings(mut self, allow_encodings: HashSet<ArrayId>) -> 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 +147,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
     }

vortex-file/src/tests.rs

@@ -88,6 +88,7 @@ static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
         .with::<RuntimeSession>();
 
     crate::register_default_encodings(&session);
+    crate::enable_all_registered_array_encodings(&session);
 
     session
 });