From 8a7f93ebb4ae07ffaba1a2233db3eb025ca57261 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Tue, 28 Jul 2026 09:10:41 +0000 Subject: [PATCH 1/5] Gate the writer's array context and compressors by the enabled editions MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The file writer could emit any encoding in `ALLOWED_ENCODINGS` regardless of which editions a session had enabled for writing, so a core-only session could still put an unstable-family encoding into a file whose compatibility guarantee does not cover it. Add `EditionSessionExt::retain_writable_encodings`, which drops every encoding a registered edition declares but the session has not enabled. Encodings no registered edition declares are left alone, so sessions outside the edition system — and third-party encodings without an inclusion — are not gated. Apply it in two places: - `VortexWriteOptions::write_internal` builds its `ArrayContext` from the gated ids, for both the pre-populated table and the permitted set, so a gated-out encoding cannot be interned even under a custom write strategy. - `WriteStrategyBuilder::with_session_editions` narrows the allow-list, and `build` now filters the BtrBlocks compressor's schemes to those whose output the allow-list covers. Filtering happens on the finished builder, so schemes added later — including the compact Zstd/Pco ones — are filtered too. Otherwise the compressor would produce an encoding the validator then rejects, failing the write instead of picking the next best scheme. Two sessions opt into the unstable edition to keep behaviour they already had: the default session when experimental patched arrays are switched on, since the compressor then emits `vortex.patched`; and the JNI session, which registers `vortex.parquet.variant`. Signed-off-by: Joe Isaacs Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_013prrhMhUESQfFucbDPjYAd --- Cargo.lock | 2 + vortex-edition/Cargo.toml | 1 + vortex-edition/src/session.rs | 34 +++++++ vortex-edition/src/tests.rs | 63 +++++++++++++ vortex-file/Cargo.toml | 1 + vortex-file/src/strategy.rs | 166 +++++++++++++++++++++++++++++++++- vortex-file/src/tests.rs | 77 ++++++++++++++++ vortex-file/src/writer.rs | 47 ++++++---- vortex-jni/src/session.rs | 10 ++ vortex/src/editions/mod.rs | 19 ++-- vortex/src/editions/tests.rs | 49 +++++++++- 11 files changed, 435 insertions(+), 34 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 16b0a83f9ac..504b5096993 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9737,6 +9737,7 @@ version = "0.1.0" dependencies = [ "parking_lot", "vortex-session", + "vortex-utils", ] [[package]] @@ -9821,6 +9822,7 @@ dependencies = [ "vortex-bytebool", "vortex-datetime-parts", "vortex-decimal-byte-parts", + "vortex-edition", "vortex-error", "vortex-fastlanes", "vortex-flatbuffers", diff --git a/vortex-edition/Cargo.toml b/vortex-edition/Cargo.toml index a4ec225c5a8..47b918b31dc 100644 --- a/vortex-edition/Cargo.toml +++ b/vortex-edition/Cargo.toml @@ -22,3 +22,4 @@ workspace = true [dependencies] parking_lot = { workspace = true } vortex-session = { workspace = true } +vortex-utils = { workspace = true } diff --git a/vortex-edition/src/session.rs b/vortex-edition/src/session.rs index e682590ba8f..11b023b4a0e 100644 --- a/vortex-edition/src/session.rs +++ b/vortex-edition/src/session.rs @@ -13,6 +13,7 @@ use vortex_session::SessionExt; use vortex_session::SessionGuard; use vortex_session::SessionVar; use vortex_session::registry::Id; +use vortex_utils::aliases::hash_set::HashSet; use crate::Edition; use crate::EditionDeclaration; @@ -150,6 +151,17 @@ impl EditionSession { .collect() } + /// Every encoding id declared by a registered edition, in encoding id order, regardless of + /// which editions are enabled. + /// + /// This is the set of encodings the edition system speaks for; anything outside it is + /// unknown to the registered declarations and therefore ungated by + /// [`EditionSessionExt::retain_writable_encodings`]. + pub fn declared_encodings(&self) -> Vec { + // The map is keyed by encoding id, so the keys are already sorted by it. + self.inner.read().inclusions.keys().copied().collect() + } + /// Validate all registered declarations. Errors on inclusions referencing undeclared /// editions, editions out of chronological order within a family (unversioned drafts /// must be newest), malformed version strings, and members requiring a release newer @@ -284,6 +296,28 @@ pub trait EditionSessionExt: SessionExt { ids.dedup(); ids } + + /// Retain only the encodings this session may write, sorted by encoding id and deduplicated. + /// + /// An encoding survives when an enabled edition includes it, or when no registered edition + /// declares it at all. The second case matters because the edition system only speaks for + /// the encodings its declarations describe: a session that registers no editions — or one + /// that registers an encoding of its own without declaring an inclusion for it — is not + /// gated. An encoding that *is* declared, in an edition the session has not enabled, is + /// removed: writing it would put an encoding in the file that the enabled editions' + /// read-compatibility guarantee does not cover. + fn retain_writable_encodings(&self, encodings: impl IntoIterator) -> Vec { + let declared: HashSet = self.editions().declared_encodings().into_iter().collect(); + let enabled: HashSet = self.enabled_encoding_ids().into_iter().collect(); + + let mut ids: Vec = encodings + .into_iter() + .filter(|id| enabled.contains(id) || !declared.contains(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 09e1345615a..5e463aec14d 100644 --- a/vortex-edition/src/tests.rs +++ b/vortex-edition/src/tests.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_session::VortexSession; +use vortex_session::registry::Id; use crate::Edition; use crate::EditionDeclaration; @@ -274,3 +275,65 @@ fn edition_ids_order_within_family_only() { fn edition_id_display() { assert_eq!(FIRST.to_string(), "test2026.01.0"); } + +#[test] +fn declared_encodings_covers_every_registered_edition() { + let editions = session(); + let declared = editions.declared_encodings(); + let declared: Vec<&str> = declared.iter().map(|id| id.as_str()).collect(); + assert_eq!(declared, ["test.alpha", "test.beta", "test.gamma"]); + assert!(EditionSession::empty().declared_encodings().is_empty()); +} + +#[test] +fn retain_writable_encodings_gates_declared_but_disabled_encodings() +-> Result<(), crate::EditionError> { + #[expect(clippy::disallowed_methods, reason = "interning test encoding ids")] + fn ids(names: &[&str]) -> Vec { + names.iter().map(|name| Id::new(name)).collect() + } + + let session = VortexSession::empty().with::(); + for declaration in DECLARATIONS { + session.register_edition(declaration)?; + } + + let candidates = ids(&["test.alpha", "test.gamma", "test.undeclared"]); + + // With no edition enabled, every declared encoding is gated out and undeclared ones survive. + let writable = session.retain_writable_encodings(candidates.clone()); + assert_eq!( + writable.iter().map(|id| id.as_str()).collect::>(), + ["test.undeclared"] + ); + + // Enabling the first edition admits its members, but not those added by a later one. + session.enable_edition(FIRST)?; + let writable = session.retain_writable_encodings(candidates.clone()); + assert_eq!( + writable.iter().map(|id| id.as_str()).collect::>(), + ["test.alpha", "test.undeclared"] + ); + + session.enable_edition(SECOND)?; + let writable = session.retain_writable_encodings(candidates); + assert_eq!( + writable.iter().map(|id| id.as_str()).collect::>(), + ["test.alpha", "test.gamma", "test.undeclared"] + ); + Ok(()) +} + +#[test] +fn retain_writable_encodings_is_ungated_without_declarations() { + #[expect(clippy::disallowed_methods, reason = "interning test encoding ids")] + let candidates: Vec = ["test.alpha", "test.gamma"].map(Id::new).into(); + + // A session that registers no editions has nothing to gate against. + let session = VortexSession::empty().with::(); + assert_eq!( + session.retain_writable_encodings(candidates).len(), + 2, + "a session with no declarations must not gate anything" + ); +} 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/strategy.rs b/vortex-file/src/strategy.rs index 25f9003fd31..c5fc09a1d42 100644 --- a/vortex-file/src/strategy.rs +++ b/vortex-file/src/strategy.rs @@ -36,6 +36,7 @@ use vortex_btrblocks::schemes::integer::IntDictScheme; use vortex_bytebool::ByteBool; use vortex_datetime_parts::DateTimeParts; use vortex_decimal_byte_parts::DecimalByteParts; +use vortex_edition::EditionSessionExt; use vortex_error::VortexExpect; use vortex_fastlanes::BitPacked; use vortex_fastlanes::Delta; @@ -63,6 +64,7 @@ use vortex_onpair::OnPair; use vortex_pco::Pco; use vortex_runend::RunEnd; use vortex_sequence::Sequence; +use vortex_session::VortexSession; use vortex_sparse::Sparse; use vortex_utils::aliases::hash_map::HashMap; use vortex_utils::aliases::hash_set::HashSet; @@ -132,6 +134,22 @@ pub static ALLOWED_ENCODINGS: LazyLock> = LazyLock::new(|| { allowed }); +/// The array encodings a writer may emit when writing with `session`: [`ALLOWED_ENCODINGS`] +/// gated by the session's enabled editions. +/// +/// An encoding survives the gate when an enabled edition includes it, or when no registered +/// edition declares it at all — see +/// [`retain_writable_encodings`](EditionSessionExt::retain_writable_encodings). A session with no +/// editions registered is therefore not gated and gets [`ALLOWED_ENCODINGS`] unchanged, while the +/// default Vortex session, which enables only the current `core` edition, loses the encodings +/// declared solely by `unstable` editions. +pub fn writable_encodings(session: &VortexSession) -> HashSet { + session + .retain_writable_encodings(ALLOWED_ENCODINGS.iter().copied()) + .into_iter() + .collect() +} + /// How the compressor was configured on [`WriteStrategyBuilder`]. enum CompressorConfig { /// A [`BtrBlocksCompressorBuilder`] that [`WriteStrategyBuilder::build`] will finalize. @@ -230,12 +248,34 @@ impl WriteStrategyBuilder { /// Override the allowed array encodings for normalization. /// /// 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, and the + /// compressor's schemes are filtered down to those whose output the allow-list covers. pub fn with_allow_encodings(mut self, allow_encodings: HashSet) -> Self { self.allow_encodings = Some(allow_encodings); self } + /// Gate the allowed array encodings by the editions `session` has enabled for writing. + /// + /// This removes every allowed encoding that a registered edition declares but the session has + /// not enabled, so the file cannot contain an encoding outside the read-compatibility + /// guarantee the session opted into. Encodings no registered edition declares are left alone; + /// see [`writable_encodings`]. + /// + /// The gate applies to compression as well as validation: [`build`](Self::build) drops every + /// compression scheme that could produce a gated-out encoding. + pub fn with_session_editions(mut self, session: &VortexSession) -> Self { + if let Some(allow_encodings) = self.allow_encodings.take() { + self.allow_encodings = Some( + session + .retain_writable_encodings(allow_encodings) + .into_iter() + .collect(), + ); + } + self + } + /// Override the flat layout strategy used for leaf chunks. /// /// By default, this uses [`FlatLayoutStrategy`]. This can be used to substitute a custom @@ -277,12 +317,31 @@ impl WriteStrategyBuilder { } else { Arc::new(FlatLayoutStrategy::default()) }; - let flat: Arc = if let Some(allow_encodings) = self.allow_encodings { - Arc::new(LayoutStrategyEncodingValidator::new(flat, allow_encodings)) + let flat: Arc = if let Some(allow_encodings) = &self.allow_encodings { + Arc::new(LayoutStrategyEncodingValidator::new( + flat, + allow_encodings.clone(), + )) } else { flat }; + // Filter the compressor to the schemes whose output the allow-list covers. Without this, + // a scheme excluded from the allow-list — for instance one only an unstable edition + // declares — would compress a chunk into an encoding the validator then rejects, failing + // the write instead of choosing the next best scheme. This runs on the finished builder, + // so schemes added after construction, including the compact ones from + // [`BtrBlocksCompressorBuilder::with_compact`], are filtered too. + let compressor = match self.compressor { + CompressorConfig::BtrBlocks(builder) => { + CompressorConfig::BtrBlocks(match &self.allow_encodings { + Some(allowed) => builder.retain_allowed_encodings(allowed), + None => builder, + }) + } + opaque @ CompressorConfig::Opaque(_) => opaque, + }; + // 7. for each chunk create a flat layout let chunked = ChunkedLayoutStrategy::new(Arc::clone(&flat)); // 6. buffer chunks so they end up with closer segment ids physically @@ -292,7 +351,7 @@ impl WriteStrategyBuilder { // Exclude IntDictScheme from the data compressor because DictStrategy (step 3) already // dictionary-encodes columns. Allowing IntDictScheme here would redundantly // dictionary-encode the integer codes produced by that earlier step. - let data_compressor: Arc = match &self.compressor { + let data_compressor: Arc = match &compressor { CompressorConfig::BtrBlocks(builder) => Arc::new( builder .clone() @@ -321,7 +380,7 @@ impl WriteStrategyBuilder { ); // 2.1. | 3.1. compress stats tables and dict values. - let stats_compressor: Arc = match self.compressor { + let stats_compressor: Arc = match compressor { CompressorConfig::BtrBlocks(builder) => Arc::new(builder.build()), CompressorConfig::Opaque(compressor) => compressor, }; @@ -402,3 +461,100 @@ impl WriteStrategyBuilder { Arc::new(table_strategy) } } + +#[cfg(test)] +pub(crate) mod tests { + use vortex_array::array_session; + use vortex_edition::Edition; + use vortex_edition::EditionDeclaration; + use vortex_edition::EditionError; + use vortex_edition::EditionId; + + use super::*; + + /// An enabled edition that adds nothing, so nothing it declares is writable beyond what + /// other editions contribute. + const ENABLED: EditionId = EditionId::new("strategytest", 2026, 1, 0); + /// A registered but never enabled edition, whose members must be gated out. + const DISABLED: EditionId = EditionId::new("strategydraft", 2026, 1, 0); + + static DECLARATIONS: &[EditionDeclaration] = &[ + EditionDeclaration { + edition: Edition { + id: ENABLED, + min_vortex_version: Some("0.1.0"), + }, + added: &[&"vortex.primitive"], + }, + EditionDeclaration { + edition: Edition { + id: DISABLED, + min_vortex_version: None, + }, + added: &[&"fastlanes.bitpacked", &"fastlanes.for"], + }, + ]; + + /// Register the test editions on `session` and enable the one covering `vortex.primitive`, + /// leaving the edition that declares the FastLanes encodings registered but disabled. + pub(crate) fn gate_session(session: &VortexSession) -> Result<(), EditionError> { + for declaration in DECLARATIONS { + session.register_edition(declaration)?; + } + session.enable_edition(ENABLED) + } + + fn gated_session() -> Result { + let session = array_session(); + gate_session(&session)?; + Ok(session) + } + + #[test] + fn writable_encodings_drops_disabled_editions() -> Result<(), EditionError> { + let writable = writable_encodings(&gated_session()?); + + assert!(writable.contains(&Primitive.id()), "enabled edition member"); + assert!(writable.contains(&FSST.id()), "no edition declares FSST"); + assert!(!writable.contains(&BitPacked.id())); + assert!(!writable.contains(&FoR.id())); + Ok(()) + } + + #[test] + fn session_editions_gate_the_allow_list() -> Result<(), EditionError> { + let builder = WriteStrategyBuilder::default().with_session_editions(&gated_session()?); + let allowed = builder + .allow_encodings + .as_ref() + .vortex_expect("default builder sets an allow-list"); + + assert!(allowed.contains(&Primitive.id())); + assert!(!allowed.contains(&BitPacked.id())); + assert!(!allowed.contains(&FoR.id())); + Ok(()) + } + + #[test] + fn an_ungated_session_keeps_every_allowed_encoding() { + // `array_session` registers no editions, so there is nothing to gate against. + let session = array_session(); + assert_eq!(writable_encodings(&session), *ALLOWED_ENCODINGS); + + let builder = WriteStrategyBuilder::default().with_session_editions(&session); + assert_eq!(builder.allow_encodings, Some(ALLOWED_ENCODINGS.clone())); + } + + #[test] + fn a_custom_allow_list_is_narrowed_not_replaced() -> Result<(), EditionError> { + let builder = WriteStrategyBuilder::default() + .with_allow_encodings(HashSet::from_iter([Primitive.id(), FoR.id()])) + .with_session_editions(&gated_session()?); + + assert_eq!( + builder.allow_encodings, + Some(HashSet::from_iter([Primitive.id()])) + ); + Ok(()) + } +} diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index 44c0b76d2db..4f897c5b7d6 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -87,6 +87,7 @@ use crate::VERSION; use crate::VortexFile; use crate::WriteOptionsSessionExt; use crate::footer::SegmentSpec; +use crate::strategy::tests::gate_session; static SESSION: LazyLock = LazyLock::new(|| { let session = array_session() .with::() @@ -2538,3 +2539,79 @@ async fn repro_8166_binary_gt_all_ff_max() -> VortexResult<()> { assert_eq!(result.len(), 1); Ok(()) } + +/// A session whose registered editions leave `fastlanes.for` and `fastlanes.bitpacked` gated out +/// of the writer. See [`gate_session`]. +static GATED_SESSION: LazyLock = LazyLock::new(|| { + let session = array_session() + .with::() + .with::(); + + crate::register_default_encodings(&session); + gate_session(&session).unwrap_or_else(|e| panic!("registering test editions: {e}")); + + session +}); + +/// Small, tightly-clustered integers are exactly what FoR and BitPacking compress, so a writer +/// that only gated the allow-list would reject the compressor's own output. The write must +/// instead fall back to a scheme the enabled editions cover. +#[tokio::test] +#[cfg_attr(miri, ignore)] +async fn gated_editions_filter_the_compressor_rather_than_failing_the_write() -> VortexResult<()> { + let numbers = PrimitiveArray::new( + Buffer::from_iter((0..8192u32).map(|i| 1_000_000 + (i % 16))), + Validity::NonNullable, + ) + .into_array(); + + let mut buf = ByteBufferMut::empty(); + GATED_SESSION + .write_options() + .write(&mut buf, numbers.clone().to_array_stream()) + .await?; + + let round_tripped = GATED_SESSION + .open_options() + .open_buffer(buf)? + .scan()? + .into_array_stream()? + .read_all() + .await?; + + let mut ctx = GATED_SESSION.create_execution_ctx(); + assert_arrays_eq!(numbers, round_tripped, &mut ctx); + Ok(()) +} + +/// The context gate is independent of the strategy: even a strategy whose allow-list still +/// covers the FastLanes encodings cannot get them into the file, because the writer's +/// [`ArrayContext`](vortex_array::ArrayContext) refuses to intern an encoding outside the +/// session's enabled editions. +#[tokio::test] +#[cfg_attr(miri, ignore)] +async fn the_writer_context_rejects_encodings_outside_the_enabled_editions() -> VortexResult<()> { + let numbers = PrimitiveArray::new( + Buffer::from_iter((0..8192u32).map(|i| 1_000_000 + (i % 16))), + Validity::NonNullable, + ) + .into_array(); + + let mut buf = ByteBufferMut::empty(); + let result = GATED_SESSION + .write_options() + // An ungated strategy: its allow-list and compressor are the unrestricted defaults. + .with_strategy(crate::strategy::WriteStrategyBuilder::default().build()) + .write(&mut buf, numbers.to_array_stream()) + .await; + + let error = match result { + Ok(_) => panic!("the writer context must reject the gated-out encoding"), + Err(error) => error.to_string(), + }; + assert!( + error.contains("not permitted by ctx"), + "unexpected: {error}" + ); + Ok(()) +} diff --git a/vortex-file/src/writer.rs b/vortex-file/src/writer.rs index 54faff97dcb..8b7553b404d 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; @@ -30,6 +29,7 @@ use vortex_array::stream::ArrayStreamAdapter; use vortex_array::stream::ArrayStreamExt; use vortex_array::stream::SendableArrayStream; use vortex_buffer::ByteBuffer; +use vortex_edition::EditionSessionExt; use vortex_error::VortexError; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -81,24 +81,21 @@ pub struct VortexWriteOptions { 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(self.session()) } } impl WriteOptionsSessionExt for S {} impl VortexWriteOptions { /// Create a new [`VortexWriteOptions`] with the given session. + /// + /// The default write strategy is gated by the session's enabled editions, so the compressor + /// only produces — and the writer only accepts — encodings those editions cover. pub fn new(session: VortexSession) -> Self { VortexWriteOptions { - strategy: WriteStrategyBuilder::default().build(), + strategy: WriteStrategyBuilder::default() + .with_session_editions(&session) + .build(), session, exclude_dtype: false, file_statistics: PRUNING_STATS.to_vec(), @@ -199,14 +196,26 @@ 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()) - // Only permit encodings known to the session. - .with_allowed_ids( - self.session - .arrays() - .registry() - .read(|map| map.keys().copied().collect()), - ); + // + // Both the pre-populated ids and the permitted ids are gated by the session's enabled + // editions, so no encoding outside the read-compatibility guarantee the session opted + // into can be interned — and therefore serialized — even under a custom write strategy. + let registered = self + .session + .arrays() + .registry() + .read(|map| map.keys().copied().collect::>()); + let ctx = ArrayContext::new( + self.session + .retain_writable_encodings(ALLOWED_ENCODINGS.iter().copied()), + ) + // Only permit encodings known to the session. + .with_allowed_ids( + self.session + .retain_writable_encodings(registered) + .into_iter() + .collect(), + ); let dtype = stream.dtype().clone(); let (mut ptr, eof) = SequenceId::root().split(); diff --git a/vortex-jni/src/session.rs b/vortex-jni/src/session.rs index 881babc5dea..e2d75705485 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; @@ -19,6 +23,12 @@ pub(crate) fn new_session() -> Box { let session = VortexSession::default().with_handle(RUNTIME.handle()); vortex_parquet_variant::initialize(&session); vortex_geo::initialize(&session); + // `vortex.parquet.variant` is declared by the `unstable` family, so the writer would gate it + // out of a core-only session. Registering the encoding here is the opt-in to writing it. + 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/src/editions/mod.rs b/vortex/src/editions/mod.rs index a4deba81a53..61fdd9f5d3e 100644 --- a/vortex/src/editions/mod.rs +++ b/vortex/src/editions/mod.rs @@ -11,13 +11,14 @@ //! 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. +//! selected or experimental patched arrays are switched on. pub mod core; #[cfg(test)] mod tests; pub mod unstable; +use vortex_array::arrays::patched::use_experimental_patches; pub use vortex_edition::Edition; pub use vortex_edition::EditionDeclaration; pub use vortex_edition::EditionId; @@ -72,15 +73,21 @@ pub fn register_default_editions(session: &VortexSession) { /// 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`]. +/// +/// The unstable edition is also enabled when experimental patched arrays are switched on: the +/// compressor then emits `vortex.patched`, which only the `unstable` family declares, so a +/// core-only session would gate it — and with it every scheme that produces patches — back out. 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"); + let unstable = cfg!(feature = "unstable_encodings") || use_experimental_patches(); + if unstable { + 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 index fcd8cfcdacc..b6736be8c92 100644 --- a/vortex/src/editions/tests.rs +++ b/vortex/src/editions/tests.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_array::arrays::patched::use_experimental_patches; use vortex_edition::EditionError; use vortex_edition::EditionSession; use vortex_edition::EditionSessionExt; @@ -125,10 +126,50 @@ fn default_session_enables_the_write_editions() { 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)); + // Experimental patched arrays make the compressor emit `vortex.patched`, which only the + // `unstable` family declares, so they enable that edition too. + let unstable = cfg!(feature = "unstable_encodings") || use_experimental_patches(); + assert_eq!(enabled.contains(&DEFAULT_UNSTABLE_EDITION), unstable); +} + +/// The writer's static allow-list is gated by the enabled editions, so a default session must +/// not be able to write an encoding that only an `unstable` edition declares. +#[cfg(all(feature = "files", not(feature = "unstable_encodings")))] +#[test] +fn the_default_writer_cannot_emit_unstable_encodings() { + use vortex_file::writable_encodings; + + use crate::VortexSessionDefault; + + if use_experimental_patches() { + // The experimental flag opts the default session into the unstable edition as well. + return; + } + + let session = VortexSession::default(); + let writable = writable_encodings(&session); + assert!(!writable.is_empty()); + + let core: Vec<_> = session + .editions() + .encodings_in(&DEFAULT_CORE_EDITION) + .into_iter() + .map(|inclusion| inclusion.encoding_id) + .collect(); + + for id in &writable { + assert!( + core.contains(id), + "{id} is writable but the enabled core edition does not include it" + ); + } + // `fastlanes.delta` is in the writer's static allow-list but only the `unstable` family + // declares it, so the gate must have removed it. + assert!( + writable + .iter() + .all(|id| id.as_str() != "fastlanes.delta" && id.as_str() != "vortex.patched") + ); } #[test] From 965957befcf309b65a254d66fc7337e71a20d290 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Tue, 28 Jul 2026 09:10:41 +0000 Subject: [PATCH 2/5] Verify the edition gate against the compressor and fix the Python doctests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Gating the writer by the enabled editions makes the default session stop listing `fastlanes.delta` in a file's array context. Two doctests in `vortex-python/src/io.rs` assert exact file sizes, and both shrank by exactly 40 bytes as a result. Reproduced locally: `chonky.vortex` is 215900 bytes and `tiny.vortex` is 55028. The gate filters compression schemes by their declared `produced_encodings`, so a scheme that under-declares would survive the filter, compress into a gated encoding, and make the writer fail the file rather than fall back — `LayoutStrategyEncodingValidator` normalizes with `Operation::Error`, and nothing in the writer uses the `Operation::Execute` fallback. Add a test that checks those declarations against what the schemes actually emit: it narrows a real session to the baseline `core2025.05.0` edition, which predates Zstd, Pco, FastLanes RLE and `vortex.masked`, then compresses primitive, float, string and boolean data with the widest scheme set the writer ever uses and asserts nothing outside the allow-list reaches the output. Removing the `retain_allowed_encodings` call makes it fail on `vortex.pco`, so it bites. Pin the gate's effect in both feature configurations, since neither is otherwise covered: - Without `unstable_encodings`, `fastlanes.delta` is the only allow-list entry no core edition declares, and its scheme is compiled out in that build, so no compression scheme is affected. - With `unstable_encodings` — the configuration every benchmark builds with — the gate is a no-op, so it cannot change a byte the benchmarks write. Signed-off-by: Joe Isaacs Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_013prrhMhUESQfFucbDPjYAd --- vortex-python/src/io.rs | 4 +- vortex/src/editions/tests.rs | 173 +++++++++++++++++++++++++++++++++++ 2 files changed, 175 insertions(+), 2 deletions(-) diff --git a/vortex-python/src/io.rs b/vortex-python/src/io.rs index a9dfa8c8f1a..f99322673b7 100644 --- a/vortex-python/src/io.rs +++ b/vortex-python/src/io.rs @@ -303,7 +303,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. /// @@ -313,7 +313,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/src/editions/tests.rs b/vortex/src/editions/tests.rs index b6736be8c92..3fcffe2fc88 100644 --- a/vortex/src/editions/tests.rs +++ b/vortex/src/editions/tests.rs @@ -188,3 +188,176 @@ fn core_edition_ids_are_registered_array_encodings() { ); } } + +/// Under the feature set the benchmarks build with, the gate must be a no-op: every encoding in +/// the writer's allow-list is either covered by an enabled edition or declared by no edition at +/// all. If this holds, gating cannot change a single byte the benchmarks write, so any file-size +/// movement they report is pre-existing noise rather than a lost compression scheme. +#[cfg(all(feature = "files", feature = "unstable_encodings"))] +#[test] +fn the_benchmark_feature_set_gates_nothing() { + use vortex_file::ALLOWED_ENCODINGS; + use vortex_file::writable_encodings; + + use crate::VortexSessionDefault; + + let session = VortexSession::default(); + assert_eq!( + writable_encodings(&session), + *ALLOWED_ENCODINGS, + "the benchmarks enable both default editions, so nothing may be gated out" + ); +} + +/// Without `unstable_encodings`, the default session enables only the `core` family, and +/// `fastlanes.delta` is the sole allow-list entry no core edition declares. Pinning the whole +/// difference documents exactly what a default writer stops emitting — and shows no compression +/// scheme is affected, since the delta scheme is itself compiled out in this configuration. +#[cfg(all(feature = "files", not(feature = "unstable_encodings")))] +#[test] +fn a_default_session_gates_out_only_fastlanes_delta() { + use vortex_file::ALLOWED_ENCODINGS; + use vortex_file::writable_encodings; + + use crate::VortexSessionDefault; + + if use_experimental_patches() { + // The experimental flag opts into the unstable edition, which declares `fastlanes.delta`. + return; + } + + let session = VortexSession::default(); + let writable = writable_encodings(&session); + + let mut gated: Vec<&str> = ALLOWED_ENCODINGS + .iter() + .filter(|id| !writable.contains(*id)) + .map(|id| id.as_str()) + .collect(); + gated.sort_unstable(); + assert_eq!(gated, ["fastlanes.delta"]); +} + +/// A session narrowed to the *baseline* core edition. `CORE_2025_05_0` predates Zstd, Pco, +/// FastLanes RLE, `vortex.masked`, `vortex.listview` and more, so a large part of the compressor +/// has to be filtered out — far more than any shipping configuration gates. Enabling an edition +/// replaces the enabled edition from the same family, so this narrows the default session. +#[cfg(feature = "files")] +fn baseline_core_session() -> Result { + use crate::VortexSessionDefault; + + let session = VortexSession::default(); + session.enable_edition(CORE_2025_05_0)?; + Ok(session) +} + +/// Every non-canonical encoding in an array tree, including the root. Canonical encodings are +/// always writable, so only the compressed ones need checking against the allow-list. +#[cfg(feature = "files")] +fn compressed_encodings( + array: &vortex_array::ArrayRef, + into: &mut Vec, +) { + if !array.is_canonical() { + into.push(array.encoding_id()); + } + for child in array.children() { + compressed_encodings(&child, into); + } +} + +/// The compressor must never produce an encoding the writer would then reject. +/// +/// This is the property the whole gate rests on: [`retain_allowed_encodings`] filters schemes by +/// their *declared* [`produced_encodings`], so a scheme that under-declares would survive the +/// filter, compress into a gated encoding, and make the writer fail the file rather than fall +/// back. Compressing real data of every canonical kind against a deliberately narrow allow-list +/// checks those declarations against what the schemes actually emit. +/// +/// [`retain_allowed_encodings`]: vortex_btrblocks::BtrBlocksCompressorBuilder::retain_allowed_encodings +/// [`produced_encodings`]: vortex_compressor::Scheme::produced_encodings +#[cfg(feature = "files")] +#[test] +fn a_gated_compressor_never_emits_a_gated_encoding() -> Result<(), Box> { + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::BoolArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::arrays::VarBinArray; + use vortex_array::validity::Validity; + use vortex_buffer::Buffer; + + use crate::compressor::BtrBlocksCompressorBuilder; + + let session = baseline_core_session()?; + let allowed = vortex_file::writable_encodings(&session); + + // Sanity: the narrowed edition really does gate a broad set, or this test proves nothing. + for gated in [ + "vortex.zstd", + "vortex.pco", + "fastlanes.rle", + "vortex.masked", + ] { + assert!( + !allowed.iter().any(|id| id.as_str() == gated), + "{gated} must be gated out by the baseline edition" + ); + } + + let arrays: Vec = vec![ + // Tightly clustered ints: FoR + BitPacking territory. + PrimitiveArray::new( + Buffer::from_iter((0..8192u32).map(|i| 1_000_000 + (i % 16))), + Validity::NonNullable, + ) + .into_array(), + // Long runs: RunEnd and the RLE schemes. + PrimitiveArray::new( + Buffer::from_iter((0..8192u64).map(|i| i / 512)), + Validity::NonNullable, + ) + .into_array(), + // Floats: ALP / ALPRD / Pco territory. + PrimitiveArray::new( + Buffer::from_iter((0..8192).map(|i| f64::from(i) * 0.125)), + Validity::NonNullable, + ) + .into_array(), + // Low-cardinality strings: dictionary and FSST territory. + VarBinArray::from_iter( + (0..8192).map(|i| Some(format!("payload-{}", i % 32))), + vortex_array::dtype::DType::Utf8(vortex_array::dtype::Nullability::Nullable), + ) + .into_array(), + // Highly compressible bytes: what the compact preset would reach for Zstd on. + VarBinArray::from_iter( + (0..4096).map(|i| Some(format!("{}{i}", "repetitive-".repeat(16)))), + vortex_array::dtype::DType::Utf8(vortex_array::dtype::Nullability::Nullable), + ) + .into_array(), + BoolArray::from_iter((0..8192).map(|i| i % 7 == 0)).into_array(), + ]; + + // `with_compact` is the widest scheme set the writer ever uses, so it is the strongest case. + let compressor = BtrBlocksCompressorBuilder::default() + .with_compact() + .retain_allowed_encodings(&allowed) + .build(); + + let mut ctx = session.create_execution_ctx(); + for array in arrays { + let compressed = compressor.compress(&array, &mut ctx)?; + + let mut encodings = Vec::new(); + compressed_encodings(&compressed, &mut encodings); + for encoding in encodings { + assert!( + allowed.contains(&encoding), + "the gated compressor produced {encoding}, which the writer would reject; \ + a scheme's produced_encodings() is under-declared" + ); + } + } + Ok(()) +} From 71000b50093aed33d50e1e572a337743731704ac Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Tue, 28 Jul 2026 09:13:42 +0000 Subject: [PATCH 3/5] Broaden the gate tests and record the write-options ordering bug Extend `a_gated_compressor_never_emits_a_gated_encoding` to decimal, opaque binary and a struct that mixes column types, so the check covers cascading across schemes rather than a single scheme in isolation. Add an end-to-end test that writes and reads back a file with the narrowed edition, exercising the compressor filter, the array context and the leaf validator together. Add an ignored test for a bug this PR introduces. `VortexWriteOptions::new` resolves the writable set when the write options are built, but the writer's array context resolves it again at write time. Narrowing the session between those two points leaves the compressor filtered against the old, wider set while the context enforces the new, narrower one, and the write fails with `Array encoding vortex.sequence not permitted by ctx`. Deriving the set at strategy-construction time is the wrong place for it: `LayoutStrategy::write_stream` already receives the session, so the gate belongs there, where it is enforced. Moving it would also remove the need for `WriteStrategyBuilder::with_session_editions` and cover custom strategies automatically. Left failing and ignored so the fix lands with a test that already reproduces the problem. Signed-off-by: Joe Isaacs Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_013prrhMhUESQfFucbDPjYAd --- vortex/Cargo.toml | 1 + vortex/src/editions/tests.rs | 141 ++++++++++++++++++++++++++++++++++- 2 files changed, 140 insertions(+), 2 deletions(-) diff --git a/vortex/Cargo.toml b/vortex/Cargo.toml index 3cc3e40d65b..5debe6c9448 100644 --- a/vortex/Cargo.toml +++ b/vortex/Cargo.toml @@ -67,6 +67,7 @@ tokio = { workspace = true, features = ["full"] } tracing = { workspace = true } tracing-subscriber = { workspace = true } vortex = { path = ".", features = ["tokio"] } +vortex-array = { workspace = true, features = ["_test-harness"] } [features] default = ["files", "zstd"] diff --git a/vortex/src/editions/tests.rs b/vortex/src/editions/tests.rs index 3fcffe2fc88..4cb970e3862 100644 --- a/vortex/src/editions/tests.rs +++ b/vortex/src/editions/tests.rs @@ -282,8 +282,13 @@ fn a_gated_compressor_never_emits_a_gated_encoding() -> Result<(), Box Result<(), Box = arrays.into_iter().chain([nested]).collect(); + // `with_compact` is the widest scheme set the writer ever uses, so it is the strongest case. let compressor = BtrBlocksCompressorBuilder::default() .with_compact() @@ -361,3 +390,111 @@ fn a_gated_compressor_never_emits_a_gated_encoding() -> Result<(), Box Result<(), Box> +{ + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::BoolArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::arrays::StructArray; + use vortex_array::arrays::VarBinArray; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::stream::ArrayStreamExt; + use vortex_array::validity::Validity; + use vortex_buffer::Buffer; + use vortex_buffer::ByteBufferMut; + use vortex_file::OpenOptionsSessionExt; + use vortex_file::WriteOptionsSessionExt; + + let session = baseline_core_session()?; + + let table = StructArray::try_new( + ["ids", "labels", "flags"].into(), + vec![ + PrimitiveArray::new( + Buffer::from_iter((0..8192u64).map(|i| 5_000_000 + (i % 64))), + Validity::NonNullable, + ) + .into_array(), + VarBinArray::from_iter( + (0..8192).map(|i| Some(format!("label-{}", i % 48))), + DType::Utf8(Nullability::Nullable), + ) + .into_array(), + BoolArray::from_iter((0..8192).map(|i| i % 3 == 0)).into_array(), + ], + 8192, + Validity::NonNullable, + )? + .into_array(); + + let mut buf = ByteBufferMut::empty(); + session + .write_options() + .write(&mut buf, table.clone().to_array_stream()) + .await?; + + let round_tripped = session + .open_options() + .open_buffer(buf.freeze())? + .scan()? + .into_array_stream()? + .read_all() + .await?; + + let mut ctx = session.create_execution_ctx(); + vortex_array::assert_arrays_eq!(table, round_tripped, &mut ctx); + Ok(()) +} + +/// Narrowing the session *after* the write options are built must still produce a valid file. +/// +/// This currently fails. [`VortexWriteOptions::new`] snapshots the writable set into the strategy +/// when the options are built, but the writer's array context recomputes it at write time, so a +/// session narrowed in between leaves the compressor filtered against the old, wider set and the +/// context enforcing the new, narrower one. The write then dies with +/// `Array encoding vortex.sequence not permitted by ctx`. +/// +/// The fix is to stop deriving the writable set at strategy-construction time: +/// [`LayoutStrategy::write_stream`] already receives the session, so the gate belongs there, +/// at the point it is enforced. +#[cfg(feature = "files")] +#[ignore = "known bug: the strategy snapshots the writable set at construction time"] +#[tokio::test] +async fn narrowing_the_session_after_building_write_options() +-> Result<(), Box> { + use vortex_array::IntoArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::validity::Validity; + use vortex_buffer::Buffer; + use vortex_buffer::ByteBufferMut; + use vortex_file::WriteOptionsSessionExt; + + use crate::VortexSessionDefault; + + let session = VortexSession::default(); + // Built while the default (wide) core edition is enabled. + let write_options = session.write_options(); + // ... and only narrowed afterwards. + session.enable_edition(CORE_2025_05_0)?; + + let numbers = PrimitiveArray::new( + Buffer::from_iter((0..8192u64).map(|i| i / 512)), + Validity::NonNullable, + ) + .into_array(); + + let mut buf = ByteBufferMut::empty(); + write_options + .write(&mut buf, numbers.to_array_stream()) + .await?; + Ok(()) +} From 56c61aabca2e00ae116553ce5fcd69d08e3e6776 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Tue, 28 Jul 2026 09:38:10 +0000 Subject: [PATCH 4/5] Resolve the edition gate at write time instead of when the strategy is built MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `WriteStrategyBuilder::with_session_editions` resolved the writable set when the strategy was built, but the writer's array context resolved it again at write time. Enabling an edition mutates the session, so anything that narrowed it between those two points left the compressor filtered against the old, wider set while the context enforced the new, narrower one, and the write died with `Array encoding vortex.sequence not permitted by ctx`. Deriving the set from mutable session state at construction time is the wrong place for it. `LayoutStrategy::write_stream` receives the session, which is where the gate is actually enforced, so resolve it there: - Drop `with_session_editions`. The builder no longer knows about sessions, and `VortexWriteOptions::new` goes back to a plain `build()`. - Drop the build-time scheme filtering. Filtering by the static allow-list was a no-op, and filtering by the session's is what created the ordering hazard. - Add `EditionGatedStrategy`, which wraps the flat leaf, resolves the writable set from the session inside `write_stream`, and normalizes each chunk with `Operation::Execute`. When the session's editions already cover the allow-list — every session by default — it skips the traversal entirely. Two consequences worth noting. Strategies built through `WriteStrategyBuilder` are now gated whether or not the caller opts in, so the `with_strategy` callers that previously slipped past the compressor filter are covered; the test that asserted a session-agnostic strategy fails now asserts it writes and round-trips. And because gated encodings are executed back to an allowed representation rather than rejected, a scheme whose output an edition does not cover costs compression ratio on a narrowed session instead of failing the file, which no longer makes the gate depend on `produced_encodings` being exhaustive. Signed-off-by: Joe Isaacs Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_013prrhMhUESQfFucbDPjYAd --- vortex-file/src/strategy.rs | 188 +++++++++++++++++++++-------------- vortex-file/src/tests.rs | 35 +++---- vortex-file/src/writer.rs | 9 +- vortex/src/editions/tests.rs | 12 +-- 4 files changed, 141 insertions(+), 103 deletions(-) diff --git a/vortex-file/src/strategy.rs b/vortex-file/src/strategy.rs index c5fc09a1d42..6ed68f4de42 100644 --- a/vortex-file/src/strategy.rs +++ b/vortex-file/src/strategy.rs @@ -7,10 +7,14 @@ use std::num::NonZeroUsize; use std::sync::Arc; use std::sync::LazyLock; +use async_trait::async_trait; +use futures::StreamExt; use vortex_alp::ALP; use vortex_alp::ALPRD; +use vortex_array::ArrayContext; use vortex_array::ArrayId; use vortex_array::VTable; +use vortex_array::VortexSessionExecute; use vortex_array::arrays::Bool; use vortex_array::arrays::Chunked; use vortex_array::arrays::Constant; @@ -30,6 +34,8 @@ 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_array::normalize::NormalizeOptions; +use vortex_array::normalize::Operation; use vortex_btrblocks::BtrBlocksCompressorBuilder; use vortex_btrblocks::SchemeExt; use vortex_btrblocks::schemes::integer::IntDictScheme; @@ -38,11 +44,13 @@ use vortex_datetime_parts::DateTimeParts; use vortex_decimal_byte_parts::DecimalByteParts; use vortex_edition::EditionSessionExt; use vortex_error::VortexExpect; +use vortex_error::VortexResult; use vortex_fastlanes::BitPacked; use vortex_fastlanes::Delta; use vortex_fastlanes::FoR; use vortex_fastlanes::RLE; use vortex_fsst::FSST; +use vortex_layout::LayoutRef; use vortex_layout::LayoutStrategy; use vortex_layout::LayoutStrategyEncodingValidator; use vortex_layout::layouts::buffered::BufferedStrategy; @@ -59,6 +67,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; +use vortex_layout::segments::SegmentSinkRef; +use vortex_layout::sequence::SendableSequentialStream; +use vortex_layout::sequence::SequencePointer; +use vortex_layout::sequence::SequentialStreamAdapter; +use vortex_layout::sequence::SequentialStreamExt; #[cfg(feature = "unstable_encodings")] use vortex_onpair::OnPair; use vortex_pco::Pco; @@ -150,6 +163,87 @@ pub fn writable_encodings(session: &VortexSession) -> HashSet { .collect() } +/// Normalizes every chunk down to the encodings the *writing session* may emit, resolved from +/// the session at write time. +/// +/// The set is deliberately not resolved when the strategy is built: enabling an edition mutates +/// the session, so a set captured at construction time can disagree with the one the writer's +/// [`ArrayContext`](vortex_array::ArrayContext) enforces, and the write then fails on an encoding +/// the compressor was still allowed to produce. [`LayoutStrategy::write_stream`] receives the +/// session, which is the point the gate is actually enforced, so it is resolved there. +/// +/// Gated encodings are executed back to an allowed representation rather than rejected. The +/// compressor is free to pick any scheme, so a scheme whose output an edition does not cover +/// costs compression ratio on a narrowed session instead of failing the file. +struct EditionGatedStrategy { + child: Arc, + /// The static allow-list, intersected with the session's writable set at write time. + allow_encodings: HashSet, +} + +impl EditionGatedStrategy { + fn new(child: S, allow_encodings: HashSet) -> Self { + Self { + child: Arc::new(child), + allow_encodings, + } + } +} + +#[async_trait] +impl LayoutStrategy for EditionGatedStrategy { + async fn write_stream( + &self, + ctx: ArrayContext, + segment_sink: SegmentSinkRef, + stream: SendableSequentialStream, + eof: SequencePointer, + session: &VortexSession, + ) -> VortexResult { + let writable: HashSet = session + .retain_writable_encodings(self.allow_encodings.iter().copied()) + .into_iter() + .collect(); + + // Nothing to gate: skip the per-chunk traversal entirely. This is the case for every + // session whose enabled editions cover the writer's allow-list, which is all of them + // by default. + if writable.len() == self.allow_encodings.len() { + return self + .child + .write_stream(ctx, segment_sink, stream, eof, session) + .await; + } + + let dtype = stream.dtype().clone(); + let writable = Arc::new(writable); + let exec_session = session.clone(); + let stream = stream.map(move |chunk| { + let (sequence_id, chunk) = chunk?; + let mut exec_ctx = exec_session.create_execution_ctx(); + let chunk = chunk.normalize(&mut NormalizeOptions { + allowed: &writable, + operation: Operation::Execute(&mut exec_ctx), + })?; + Ok((sequence_id, chunk)) + }); + + self.child + .write_stream( + ctx, + segment_sink, + SequentialStreamAdapter::new(dtype, stream).sendable(), + eof, + session, + ) + .await + } + + fn buffered_bytes(&self) -> u64 { + self.child.buffered_bytes() + } +} + /// How the compressor was configured on [`WriteStrategyBuilder`]. enum CompressorConfig { /// A [`BtrBlocksCompressorBuilder`] that [`WriteStrategyBuilder::build`] will finalize. @@ -248,34 +342,12 @@ impl WriteStrategyBuilder { /// Override the allowed array encodings for normalization. /// /// The configured flat leaf strategy is wrapped in a [`LayoutStrategyEncodingValidator`] - /// that recursively checks every chunk before passing it to the leaf writer, and the - /// compressor's schemes are filtered down to those whose output the allow-list covers. + /// that recursively checks every chunk before passing it to the leaf writer. pub fn with_allow_encodings(mut self, allow_encodings: HashSet) -> Self { self.allow_encodings = Some(allow_encodings); self } - /// Gate the allowed array encodings by the editions `session` has enabled for writing. - /// - /// This removes every allowed encoding that a registered edition declares but the session has - /// not enabled, so the file cannot contain an encoding outside the read-compatibility - /// guarantee the session opted into. Encodings no registered edition declares are left alone; - /// see [`writable_encodings`]. - /// - /// The gate applies to compression as well as validation: [`build`](Self::build) drops every - /// compression scheme that could produce a gated-out encoding. - pub fn with_session_editions(mut self, session: &VortexSession) -> Self { - if let Some(allow_encodings) = self.allow_encodings.take() { - self.allow_encodings = Some( - session - .retain_writable_encodings(allow_encodings) - .into_iter() - .collect(), - ); - } - self - } - /// Override the flat layout strategy used for leaf chunks. /// /// By default, this uses [`FlatLayoutStrategy`]. This can be used to substitute a custom @@ -318,30 +390,16 @@ impl WriteStrategyBuilder { Arc::new(FlatLayoutStrategy::default()) }; let flat: Arc = if let Some(allow_encodings) = &self.allow_encodings { - Arc::new(LayoutStrategyEncodingValidator::new( - flat, + // The session gate runs first, normalizing away any encoding the writing session's + // editions do not cover; the validator then asserts the static allow-list holds. + Arc::new(EditionGatedStrategy::new( + LayoutStrategyEncodingValidator::new(flat, allow_encodings.clone()), allow_encodings.clone(), )) } else { flat }; - // Filter the compressor to the schemes whose output the allow-list covers. Without this, - // a scheme excluded from the allow-list — for instance one only an unstable edition - // declares — would compress a chunk into an encoding the validator then rejects, failing - // the write instead of choosing the next best scheme. This runs on the finished builder, - // so schemes added after construction, including the compact ones from - // [`BtrBlocksCompressorBuilder::with_compact`], are filtered too. - let compressor = match self.compressor { - CompressorConfig::BtrBlocks(builder) => { - CompressorConfig::BtrBlocks(match &self.allow_encodings { - Some(allowed) => builder.retain_allowed_encodings(allowed), - None => builder, - }) - } - opaque @ CompressorConfig::Opaque(_) => opaque, - }; - // 7. for each chunk create a flat layout let chunked = ChunkedLayoutStrategy::new(Arc::clone(&flat)); // 6. buffer chunks so they end up with closer segment ids physically @@ -351,7 +409,7 @@ impl WriteStrategyBuilder { // Exclude IntDictScheme from the data compressor because DictStrategy (step 3) already // dictionary-encodes columns. Allowing IntDictScheme here would redundantly // dictionary-encode the integer codes produced by that earlier step. - let data_compressor: Arc = match &compressor { + let data_compressor: Arc = match &self.compressor { CompressorConfig::BtrBlocks(builder) => Arc::new( builder .clone() @@ -380,7 +438,7 @@ impl WriteStrategyBuilder { ); // 2.1. | 3.1. compress stats tables and dict values. - let stats_compressor: Arc = match compressor { + let stats_compressor: Arc = match self.compressor { CompressorConfig::BtrBlocks(builder) => Arc::new(builder.build()), CompressorConfig::Opaque(compressor) => compressor, }; @@ -521,40 +579,26 @@ pub(crate) mod tests { Ok(()) } + /// The builder no longer resolves the session; the gate happens at write time. What the + /// builder still owns is the static allow-list, which must stay untouched by any session. #[test] - fn session_editions_gate_the_allow_list() -> Result<(), EditionError> { - let builder = WriteStrategyBuilder::default().with_session_editions(&gated_session()?); - let allowed = builder - .allow_encodings - .as_ref() - .vortex_expect("default builder sets an allow-list"); - - assert!(allowed.contains(&Primitive.id())); - assert!(!allowed.contains(&BitPacked.id())); - assert!(!allowed.contains(&FoR.id())); - Ok(()) - } - - #[test] - fn an_ungated_session_keeps_every_allowed_encoding() { - // `array_session` registers no editions, so there is nothing to gate against. - let session = array_session(); - assert_eq!(writable_encodings(&session), *ALLOWED_ENCODINGS); - - let builder = WriteStrategyBuilder::default().with_session_editions(&session); + fn the_builder_allow_list_is_independent_of_any_session() { + let builder = WriteStrategyBuilder::default(); assert_eq!(builder.allow_encodings, Some(ALLOWED_ENCODINGS.clone())); } + /// A custom allow-list is narrowed by the session, not replaced by it: the gate can only + /// ever remove encodings the caller already permitted. #[test] - fn a_custom_allow_list_is_narrowed_not_replaced() -> Result<(), EditionError> { - let builder = WriteStrategyBuilder::default() - .with_allow_encodings(HashSet::from_iter([Primitive.id(), FoR.id()])) - .with_session_editions(&gated_session()?); - - assert_eq!( - builder.allow_encodings, - Some(HashSet::from_iter([Primitive.id()])) - ); + fn a_custom_allow_list_is_narrowed_by_the_session() -> Result<(), EditionError> { + let session = gated_session()?; + let custom: HashSet = HashSet::from_iter([Primitive.id(), FoR.id()]); + let writable: HashSet = session + .retain_writable_encodings(custom.iter().copied()) + .into_iter() + .collect(); + + assert_eq!(writable, HashSet::from_iter([Primitive.id()])); Ok(()) } } diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index 4f897c5b7d6..fdcb285eb41 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -2584,13 +2584,12 @@ async fn gated_editions_filter_the_compressor_rather_than_failing_the_write() -> Ok(()) } -/// The context gate is independent of the strategy: even a strategy whose allow-list still -/// covers the FastLanes encodings cannot get them into the file, because the writer's -/// [`ArrayContext`](vortex_array::ArrayContext) refuses to intern an encoding outside the -/// session's enabled editions. +/// A strategy built with no knowledge of any session is still gated, because the gate resolves +/// the session inside `write_stream` rather than when the strategy is built. This is what makes +/// `with_strategy` callers safe without them having to opt in. #[tokio::test] #[cfg_attr(miri, ignore)] -async fn the_writer_context_rejects_encodings_outside_the_enabled_editions() -> VortexResult<()> { +async fn a_session_agnostic_strategy_is_still_gated() -> VortexResult<()> { let numbers = PrimitiveArray::new( Buffer::from_iter((0..8192u32).map(|i| 1_000_000 + (i % 16))), Validity::NonNullable, @@ -2598,20 +2597,22 @@ async fn the_writer_context_rejects_encodings_outside_the_enabled_editions() -> .into_array(); let mut buf = ByteBufferMut::empty(); - let result = GATED_SESSION + GATED_SESSION .write_options() - // An ungated strategy: its allow-list and compressor are the unrestricted defaults. + // Built without a session anywhere in sight. .with_strategy(crate::strategy::WriteStrategyBuilder::default().build()) - .write(&mut buf, numbers.to_array_stream()) - .await; + .write(&mut buf, numbers.clone().to_array_stream()) + .await?; - let error = match result { - Ok(_) => panic!("the writer context must reject the gated-out encoding"), - Err(error) => error.to_string(), - }; - assert!( - error.contains("not permitted by ctx"), - "unexpected: {error}" - ); + let round_tripped = GATED_SESSION + .open_options() + .open_buffer(buf.freeze())? + .scan()? + .into_array_stream()? + .read_all() + .await?; + + let mut ctx = GATED_SESSION.create_execution_ctx(); + assert_arrays_eq!(numbers, round_tripped, &mut ctx); Ok(()) } diff --git a/vortex-file/src/writer.rs b/vortex-file/src/writer.rs index 8b7553b404d..22ce2735c1d 100644 --- a/vortex-file/src/writer.rs +++ b/vortex-file/src/writer.rs @@ -89,13 +89,12 @@ impl WriteOptionsSessionExt for S {} impl VortexWriteOptions { /// Create a new [`VortexWriteOptions`] with the given session. /// - /// The default write strategy is gated by the session's enabled editions, so the compressor - /// only produces — and the writer only accepts — encodings those editions cover. + /// The default write strategy resolves the session's enabled editions at write time, so the + /// file only contains encodings those editions cover, however the session is configured + /// between building these options and writing. pub fn new(session: VortexSession) -> Self { VortexWriteOptions { - strategy: WriteStrategyBuilder::default() - .with_session_editions(&session) - .build(), + strategy: WriteStrategyBuilder::default().build(), session, exclude_dtype: false, file_statistics: PRUNING_STATS.to_vec(), diff --git a/vortex/src/editions/tests.rs b/vortex/src/editions/tests.rs index 4cb970e3862..18c888ebc4a 100644 --- a/vortex/src/editions/tests.rs +++ b/vortex/src/editions/tests.rs @@ -457,17 +457,11 @@ async fn a_narrowed_edition_still_writes_a_readable_file() -> Result<(), Box Result<(), Box> { From d2447b3389bbc0ded67d8e616608acf0cd99d199 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Tue, 28 Jul 2026 09:44:42 +0000 Subject: [PATCH 5/5] Drop a redundant explicit link target in the edition gate docs `cargo doc --document-private-items` runs with `-D warnings`, and the label already resolves to the same path: error: redundant explicit link target --> vortex-file/src/strategy.rs:171:22 | /// [`ArrayContext`](vortex_array::ArrayContext) enforces, ... | -------------- ^^^^^^^^^^^^^^^^^^^^^^^^^^ explicit target is redundant Signed-off-by: Joe Isaacs Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_013prrhMhUESQfFucbDPjYAd --- vortex-file/src/strategy.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vortex-file/src/strategy.rs b/vortex-file/src/strategy.rs index 6ed68f4de42..4f8b57cfe34 100644 --- a/vortex-file/src/strategy.rs +++ b/vortex-file/src/strategy.rs @@ -168,7 +168,7 @@ pub fn writable_encodings(session: &VortexSession) -> HashSet { /// /// The set is deliberately not resolved when the strategy is built: enabling an edition mutates /// the session, so a set captured at construction time can disagree with the one the writer's -/// [`ArrayContext`](vortex_array::ArrayContext) enforces, and the write then fails on an encoding +/// [`ArrayContext`] enforces, and the write then fails on an encoding /// the compressor was still allowed to produce. [`LayoutStrategy::write_stream`] receives the /// session, which is the point the gate is actually enforced, so it is resolved there. ///