diff --git a/vortex-btrblocks/src/builder.rs b/vortex-btrblocks/src/builder.rs index c6faaba93d8..341569f408c 100644 --- a/vortex-btrblocks/src/builder.rs +++ b/vortex-btrblocks/src/builder.rs @@ -200,10 +200,10 @@ impl BtrBlocksCompressorBuilder { self } - /// Retains only schemes whose produced encodings all belong to `allowed`. + /// Retains only schemes whose produced serialized IDs all belong to `allowed`. /// - /// The file writer uses this to restrict compression to the encodings of its configured - /// editions. + /// `allowed` holds serialized IDs. The file writer passes the array IDs its enabled editions + /// permit. pub fn retain_allowed_encodings(mut self, allowed: &HashSet) -> Self { self.schemes .retain(|s| s.produced_encodings().iter().all(|id| allowed.contains(id))); diff --git a/vortex-compressor/src/scheme/mod.rs b/vortex-compressor/src/scheme/mod.rs index de9e67690d4..0ba1c90202a 100644 --- a/vortex-compressor/src/scheme/mod.rs +++ b/vortex-compressor/src/scheme/mod.rs @@ -124,11 +124,15 @@ pub trait Scheme: Debug + Send + Sync { /// Whether this scheme can compress the given canonical array. fn matches(&self, canonical: &Canonical) -> bool; - /// The array encodings this scheme itself may introduce into its compressed output. + /// The serialized IDs this scheme itself may write into its compressed output. /// - /// Cascaded children are compressed by other schemes, which declare their own encodings, - /// so only encodings constructed directly by [`compress`](Scheme::compress) belong here. - /// Canonical arrays the scheme merely rearranges do not need to be declared. + /// Every declared ID must be permitted for the scheme to be used. Cascaded children are + /// compressed by other schemes, which declare their own IDs, so only arrays constructed + /// directly by [`compress`](Scheme::compress) belong here. Canonical arrays the scheme + /// merely rearranges do not need to be declared. + /// + /// For most encodings this is the in-memory encoding ID. An encoding with several wire + /// formats declares the wire IDs the scheme writes, which may differ from its in-memory ID. fn produced_encodings(&self) -> Vec; /// Returns the stats generation options this scheme requires. The compressor merges all diff --git a/vortex-file/src/writer.rs b/vortex-file/src/writer.rs index 8167d31ede2..6c0d01dfa05 100644 --- a/vortex-file/src/writer.rs +++ b/vortex-file/src/writer.rs @@ -239,8 +239,7 @@ impl VortexWriteOptions { let enforce_editions = !self.disable_editions; // The array context is built here, rather than when the options were constructed, so that // encodings registered on the session in between are still eligible for the file. - let (array_ctx, allowed_array_encodings) = - new_array_context(&self.session, enforce_editions); + let array_ctx = new_array_context(&self.session, enforce_editions); let ctx = LayoutWriterContext::new(array_ctx) .with_buffered_bytes_tracker(self.buffered_bytes.clone()); let ctx = if enforce_editions { @@ -248,12 +247,14 @@ impl VortexWriteOptions { } else { ctx }; + let allowed_serialized_ids: HashSet = + ctx.array_ctx().to_ids().into_iter().collect(); let strategy = match self.strategy { Some(strategy) => strategy, None => WriteStrategyBuilder::default() .with_btrblocks_builder( BtrBlocksCompressorBuilder::default() - .retain_allowed_encodings(&allowed_array_encodings), + .retain_allowed_encodings(&allowed_serialized_ids), ) .build(), }; @@ -384,36 +385,30 @@ impl VortexWriteOptions { } } -fn new_array_context( - session: &VortexSession, - enforce_editions: bool, -) -> (ArrayContext, HashSet) { +fn new_array_context(session: &VortexSession, enforce_editions: bool) -> ArrayContext { // NOTE(os): Set up an array context with all eligible serialized IDs pre-populated. // This is preferred for now over having an empty context here, because only the // serialised array order is deterministic. The serialisation of arrays are done // parallel and with an empty context they can register their encodings to the context // in different order, changing the written bytes from run to run. - let arrays = session.arrays(); - let serialized_ids = if enforce_editions { + // + // The seeded IDs are also what the writer may emit: callers read them back with + // `ArrayContext::to_ids` to restrict compression to the same set. + let serialized_ids: Vec = if enforce_editions { session.enabled_component_ids(ComponentKind::Array) } else { - arrays + session + .arrays() .registry() .read(|registry| registry.keys().copied().collect()) }; - let allowed_array_encodings = serialized_ids - .iter() - .filter_map(|serialized_id| arrays.registry().get(serialized_id)) - .map(|plugin| plugin.id()) - .collect(); let array_ctx = ArrayContext::new(serialized_ids.iter().copied().sorted().collect()); - let array_ctx = if enforce_editions { + if enforce_editions { // Only permit serialized IDs in the enabled editions. array_ctx.with_allowed_ids(serialized_ids.into_iter().collect()) } else { array_ctx - }; - (array_ctx, allowed_array_encodings) + } } /// The ids of `kind` the enabled editions permit. @@ -787,29 +782,22 @@ mod tests { session.register_edition(&DECLARATION)?; session.enable_edition(EDITION)?; - let (ctx, allowed_array_encodings) = new_array_context(&session, true); + let ctx = new_array_context(&session, true); assert_eq!(ctx.to_ids(), [Primitive.id()]); assert!(ctx.intern(&Bool.id()).is_none()); - assert_eq!(allowed_array_encodings, HashSet::from([Primitive.id()])); Ok(()) } #[test] fn disabling_editions_allows_all_registered_array_ids() { let session = array_session(); - let (registered_ids, registered_encodings) = session.arrays().registry().read(|registry| { - ( - registry.keys().copied().sorted().collect::>(), - registry - .values() - .map(|plugin| plugin.id()) - .collect::>(), - ) - }); + let registered_ids = session + .arrays() + .registry() + .read(|registry| registry.keys().copied().sorted().collect::>()); - let (ctx, allowed_array_encodings) = new_array_context(&session, false); + let ctx = new_array_context(&session, false); assert_eq!(ctx.to_ids(), registered_ids); - assert_eq!(allowed_array_encodings, registered_encodings); assert!(ctx.intern(&Bool.id()).is_some()); }