diff --git a/Cargo.lock b/Cargo.lock index 51188781c9a..444c7987d79 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9740,6 +9740,7 @@ version = "0.1.0" dependencies = [ "parking_lot", "vortex-session", + "vortex-utils", ] [[package]] @@ -9824,6 +9825,7 @@ dependencies = [ "vortex-bytebool", "vortex-datetime-parts", "vortex-decimal-byte-parts", + "vortex-edition", "vortex-error", "vortex-fastlanes", "vortex-flatbuffers", @@ -10133,6 +10135,7 @@ dependencies = [ "vortex-array", "vortex-arrow", "vortex-buffer", + "vortex-edition", "vortex-error", "vortex-file", "vortex-io", diff --git a/encodings/parquet-variant/Cargo.toml b/encodings/parquet-variant/Cargo.toml index 6cdc2b3fa8b..c306b8e73ce 100644 --- a/encodings/parquet-variant/Cargo.toml +++ b/encodings/parquet-variant/Cargo.toml @@ -37,6 +37,7 @@ vortex-session = { workspace = true } rstest = { workspace = true } tokio = { workspace = true, features = ["full"] } vortex-array = { workspace = true, features = ["_test-harness"] } +vortex-edition = { workspace = true } vortex-file = { workspace = true, features = ["tokio"] } vortex-io = { workspace = true, features = ["tokio"] } vortex-layout = { workspace = true } diff --git a/encodings/parquet-variant/src/vtable.rs b/encodings/parquet-variant/src/vtable.rs index 8f0a7a6c17b..410b4269f15 100644 --- a/encodings/parquet-variant/src/vtable.rs +++ b/encodings/parquet-variant/src/vtable.rs @@ -326,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; @@ -387,14 +391,47 @@ mod tests { ParquetVariant::from_arrow_variant(&ArrowVariantArray::try_new(&arrow_storage)?) } + /// The edition these tests enable. + const TEST_EDITION: EditionId = EditionId::new("parquetvarianttest", 2026, 1, 0); + + /// Declare and enable an edition covering every encoding `session` has registered. + /// + /// The writer admits only encodings an enabled edition includes, and the first-party + /// declarations live in the `vortex` facade, which this crate cannot depend on. A session + /// assembled here therefore has to declare its own edition to write anything. + fn enable_all_registered_encodings(session: &VortexSession) -> VortexResult<()> { + let editions = session.editions(); + editions + .declare_edition(Edition { + id: TEST_EDITION, + min_vortex_version: Some("0.1.0"), + }) + .map_err(|e| vortex_err!("declaring the test edition: {e}"))?; + + let registered = session + .arrays() + .registry() + .read(|map| map.keys().copied().collect::>()); + for id in registered { + editions + .declare_inclusion(EditionInclusion::new(&id, TEST_EDITION)) + .map_err(|e| vortex_err!("declaring an inclusion for {id}: {e}"))?; + } + + session + .enable_edition(TEST_EDITION) + .map_err(|e| vortex_err!("enabling the test edition: {e}")) + } + #[fixture] - fn parquet_variant_file_session() -> VortexSession { + fn parquet_variant_file_session() -> VortexResult { let session = vortex_array::array_session() .with::() .with::(); vortex_file::register_default_encodings(&session); session.arrays().register(ParquetVariant); - session + enable_all_registered_encodings(&session)?; + Ok(session) } #[fixture] @@ -451,9 +488,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 @@ -478,10 +516,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-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-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 184f036b497..e9b5d18d47d 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::tests::enable_all_registered_encodings(&session); session }); diff --git a/vortex-file/src/multi/mod.rs b/vortex-file/src/multi/mod.rs index e94f9069ad1..e60dafc0946 100644 --- a/vortex-file/src/multi/mod.rs +++ b/vortex-file/src/multi/mod.rs @@ -359,6 +359,7 @@ mod tests { .with::() .with::(); crate::register_default_encodings(&session); + crate::tests::enable_all_registered_encodings(&session); let expected = ByteBuffer::copy_from(b"cached metadata"); let mut output = ByteBufferMut::empty(); @@ -459,6 +460,7 @@ mod tests { .with::() .with::(); crate::register_default_encodings(&session); + crate::tests::enable_all_registered_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..b706f5e89ac 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::tests::enable_all_registered_encodings(&session); session } @@ -586,6 +587,7 @@ mod tests { .with::(); crate::register_default_encodings(&session); + crate::tests::enable_all_registered_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::tests::enable_all_registered_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::tests::enable_all_registered_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::tests::enable_all_registered_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 25f9003fd31..a66b1ab7237 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,18 +34,23 @@ 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; 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_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; @@ -58,11 +67,17 @@ 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; 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 +147,105 @@ 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 only when an enabled edition includes it. A session that +/// enables no editions therefore writes nothing; the default Vortex session enables the current +/// `core` edition and loses the encodings declared solely by `unstable` editions. +pub fn writable_encodings(session: &VortexSession) -> HashSet { + let enabled: HashSet = session.enabled_encoding_ids().into_iter().collect(); + ALLOWED_ENCODINGS + .iter() + .copied() + .filter(|id| enabled.contains(id)) + .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`] 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 enabled: HashSet = session.enabled_encoding_ids().into_iter().collect(); + let writable: HashSet = self + .allow_encodings + .iter() + .copied() + .filter(|id| enabled.contains(id)) + .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. @@ -277,8 +391,13 @@ 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 { + // 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 }; @@ -402,3 +521,97 @@ 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"), + }, + // The structural encodings are what a file is built from, so an edition that did + // not cover them could not write at all, gated compression or otherwise. + added: &[ + &"vortex.primitive", + &"vortex.struct", + &"vortex.chunked", + &"vortex.constant", + ], + }, + 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) + } + + /// Only what an enabled edition includes survives. An encoding no edition mentions at all + /// is not writable either: the gate admits, it does not merely exclude. + #[test] + fn writable_encodings_keeps_only_enabled_edition_members() -> 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()), "declared, not enabled"); + assert!(!writable.contains(&FoR.id()), "declared, not enabled"); + 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 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_by_the_session() -> Result<(), EditionError> { + let session = gated_session()?; + let custom: HashSet = HashSet::from_iter([Primitive.id(), FoR.id()]); + let enabled: HashSet = session.enabled_encoding_ids().into_iter().collect(); + let writable: HashSet = custom + .iter() + .copied() + .filter(|id| enabled.contains(id)) + .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 f7088056e1b..a4e583fa791 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -56,6 +56,7 @@ use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::ScalarFnVTableExt; use vortex_array::scalar_fn::fns::pack::Pack; use vortex_array::scalar_fn::fns::pack::PackOptions; +use vortex_array::session::ArraySessionExt; use vortex_array::stats::PRUNING_STATS; use vortex_array::stream::ArrayStreamAdapter; use vortex_array::stream::ArrayStreamExt; @@ -67,6 +68,10 @@ use vortex_buffer::Buffer; 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::VortexExpect; use vortex_error::VortexResult; use vortex_flatbuffers::footer as fb; @@ -87,12 +92,47 @@ use crate::VERSION; use crate::VortexFile; use crate::WriteOptionsSessionExt; use crate::footer::SegmentSpec; +use crate::strategy::tests::gate_session; +/// The edition this crate's test session enables. +const TEST_EDITION: EditionId = EditionId::new("filetest", 2026, 1, 0); + +/// Declare and enable an edition covering every encoding `session` has registered. +/// +/// The writer admits only encodings an enabled edition includes, and the first-party +/// declarations live in the `vortex` facade, which this crate cannot depend on. A session +/// assembled here therefore has to declare its own edition to write anything. +pub(crate) fn enable_all_registered_encodings(session: &VortexSession) { + session + .editions() + .declare_edition(Edition { + id: TEST_EDITION, + min_vortex_version: Some("0.1.0"), + }) + .expect("test edition is undeclared"); + + let registered = session + .arrays() + .registry() + .read(|map| map.keys().copied().collect::>()); + for id in registered { + session + .editions() + .declare_inclusion(EditionInclusion::new(&id, TEST_EDITION)) + .expect("each encoding is registered once"); + } + + session + .enable_edition(TEST_EDITION) + .expect("test edition was just declared"); +} + static SESSION: LazyLock = LazyLock::new(|| { let session = array_session() .with::() .with::(); crate::register_default_encodings(&session); + enable_all_registered_encodings(&session); session }); @@ -2537,3 +2577,82 @@ 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 the +/// compressor reaches for encodings this session's editions gate out. A writer that only +/// narrowed its array context would then reject the compressor's own output; the gate must +/// instead normalize the chunk back to an encoding the enabled editions cover, so the write +/// costs compression ratio rather than failing. +#[tokio::test] +#[cfg_attr(miri, ignore)] +async fn gated_editions_normalize_the_chunk_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(()) +} + +/// 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 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, + ) + .into_array(); + + let mut buf = ByteBufferMut::empty(); + GATED_SESSION + .write_options() + // Built without a session anywhere in sight. + .with_strategy(crate::strategy::WriteStrategyBuilder::default().build()) + .write(&mut buf, numbers.clone().to_array_stream()) + .await?; + + 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 8828a98d0fc..6fb95c65b94 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; @@ -81,21 +79,17 @@ 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 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().build(), @@ -208,14 +202,14 @@ 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()) + // + // The context is built from 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 enabled_encoding_ids = self.session.enabled_encoding_ids(); + let ctx = ArrayContext::new(enabled_encoding_ids.clone()) // Only permit encodings known to the session. - .with_allowed_ids( - self.session - .arrays() - .registry() - .read(|map| map.keys().copied().collect()), - ); + .with_allowed_ids(enabled_encoding_ids.into_iter().collect()); let dtype = stream.dtype().clone(); let (mut ptr, eof) = SequenceId::root().split(); diff --git a/vortex-file/tests/common/mod.rs b/vortex-file/tests/common/mod.rs new file mode 100644 index 00000000000..7e4f4fe8bdf --- /dev/null +++ b/vortex-file/tests/common/mod.rs @@ -0,0 +1,48 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Shared setup for the `vortex-file` integration tests. +#![expect( + clippy::expect_used, + reason = "test setup runs inside a LazyLock, where there is no error to propagate" +)] + +use vortex_array::session::ArraySessionExt; +use vortex_edition::Edition; +use vortex_edition::EditionId; +use vortex_edition::EditionInclusion; +use vortex_edition::EditionSessionExt; +use vortex_session::VortexSession; + +/// The edition these tests enable. +const TEST_EDITION: EditionId = EditionId::new("filetest", 2026, 1, 0); + +/// Declare and enable an edition covering every encoding `session` has registered. +/// +/// The writer admits only encodings an enabled edition includes, and the first-party +/// declarations live in the `vortex` facade, which `vortex-file` cannot depend on. A session +/// assembled here therefore has to declare its own edition to write anything. +pub fn enable_all_registered_encodings(session: &VortexSession) { + session + .editions() + .declare_edition(Edition { + id: TEST_EDITION, + min_vortex_version: Some("0.1.0"), + }) + .expect("test edition is undeclared"); + + let registered = session + .arrays() + .registry() + .read(|map| map.keys().copied().collect::>()); + for id in registered { + session + .editions() + .declare_inclusion(EditionInclusion::new(&id, TEST_EDITION)) + .expect("each encoding is registered once"); + } + + session + .enable_edition(TEST_EDITION) + .expect("test edition was just declared"); +} diff --git a/vortex-file/tests/issue_8819_footer_segment_oob.rs b/vortex-file/tests/issue_8819_footer_segment_oob.rs index 7b09d7a9c35..2e02dd56546 100644 --- a/vortex-file/tests/issue_8819_footer_segment_oob.rs +++ b/vortex-file/tests/issue_8819_footer_segment_oob.rs @@ -23,6 +23,8 @@ use vortex_io::session::RuntimeSession; use vortex_layout::session::LayoutSession; use vortex_session::VortexSession; +mod common; + static SESSION: LazyLock = LazyLock::new(|| { let session = vortex_array::array_session() .with::() @@ -30,6 +32,8 @@ static SESSION: LazyLock = LazyLock::new(|| { vortex_file::register_default_encodings(&session); + common::enable_all_registered_encodings(&session); + session }); diff --git a/vortex-file/tests/test_write_table.rs b/vortex-file/tests/test_write_table.rs index bd664cc24ed..c41fc06991d 100644 --- a/vortex-file/tests/test_write_table.rs +++ b/vortex-file/tests/test_write_table.rs @@ -29,12 +29,15 @@ use vortex_layout::layouts::flat::writer::FlatLayoutStrategy; use vortex_layout::layouts::table::TableStrategy; use vortex_layout::session::LayoutSession; use vortex_session::VortexSession; + +mod common; static SESSION: LazyLock = LazyLock::new(|| { let session = vortex_array::array_session() .with::() .with::(); vortex_file::register_default_encodings(&session); + common::enable_all_registered_encodings(&session); session }); 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/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/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..09697a9f0c9 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] @@ -147,3 +188,310 @@ 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); + } +} + +/// A compressor filtered to an allow-list must never produce an encoding outside it. +/// +/// [`retain_allowed_encodings`] filters schemes by their *declared* [`produced_encodings`], so a +/// scheme that under-declares would survive the filter and compress into an encoding the caller +/// excluded. Compressing real data of every canonical kind against a deliberately narrow +/// allow-list checks those declarations against what the schemes actually emit. +/// +/// The writer does not rely on this filter — it normalizes gated encodings away at write time +/// instead, so an under-declaring scheme costs compression ratio rather than the file — but the +/// same narrow allow-list a gated session produces is the strongest input for the check. +/// +/// [`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::DecimalArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::arrays::StructArray; + use vortex_array::arrays::VarBinArray; + use vortex_array::dtype::DType; + use vortex_array::dtype::DecimalDType; + use vortex_array::dtype::Nullability; + 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))), + DType::Utf8(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)))), + DType::Utf8(Nullability::Nullable), + ) + .into_array(), + BoolArray::from_iter((0..8192).map(|i| i % 7 == 0)).into_array(), + // Decimals exercise `DecimalScheme`, whose output cascades into a primitive child. + DecimalArray::new( + Buffer::from_iter((0..8192i128).map(|i| 100_000 + (i % 97))), + DecimalDType::new(38, 2), + Validity::NonNullable, + ) + .into_array(), + // Opaque bytes exercise the binary dictionary and Zstd schemes. + VarBinArray::from_iter( + (0..4096).map(|i| Some(vec![(i % 251) as u8; 24])), + DType::Binary(Nullability::Nullable), + ) + .into_array(), + ]; + + // A struct mixing the above forces the compressor to cascade across schemes, which is where + // an under-declared child encoding would otherwise slip through. + let nested = StructArray::try_new( + ["ints", "strings", "bools"].into(), + vec![arrays[0].clone(), arrays[3].clone(), arrays[5].clone()], + 8192, + Validity::NonNullable, + )? + .into_array(); + let arrays: Vec = 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() + .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(()) +} + +/// The whole chain end to end: with a narrowed edition the compressor is filtered, the writer's +/// array context refuses the gated ids, and the validator normalizes each chunk — a file must +/// still come out, and read back unchanged. This is the case the benchmarks never cover, since +/// they all build with `unstable_encodings` and enable every default edition. +#[cfg(feature = "files")] +#[tokio::test] +async fn a_narrowed_edition_still_writes_a_readable_file() -> 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. +/// +/// The writable set is resolved inside `write_stream`, so there is no window in which the +/// strategy and the writer's array context can disagree about it. Resolving it when the write +/// options were built used to fail here with +/// `Array encoding vortex.sequence not permitted by ctx`. +#[cfg(feature = "files")] +#[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(()) +}