From 3138b6075b9631dbc54483f5c1950b55283cd132 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Wed, 22 Jul 2026 15:52:20 +0100 Subject: [PATCH 1/2] refactor: store valid IDs in serialization context Signed-off-by: Joe Isaacs --- vortex-array/src/arrays/patched/vtable/mod.rs | 2 +- vortex-file/src/writer.rs | 4 +- vortex-session/src/registry.rs | 53 +++++++++++++++---- 3 files changed, 46 insertions(+), 13 deletions(-) diff --git a/vortex-array/src/arrays/patched/vtable/mod.rs b/vortex-array/src/arrays/patched/vtable/mod.rs index 71b1775cbf4..f69b7b73afa 100644 --- a/vortex-array/src/arrays/patched/vtable/mod.rs +++ b/vortex-array/src/arrays/patched/vtable/mod.rs @@ -588,7 +588,7 @@ mod tests { let session = array_session(); session.arrays().register(Patched); - let ctx = ArrayContext::empty().with_registry(session.arrays().registry().clone()); + let ctx = ArrayContext::empty().with_valid_ids(session.arrays().registry().ids()); let serialized = array .serialize(&ctx, &session, &SerializeOptions::default()) .unwrap(); diff --git a/vortex-file/src/writer.rs b/vortex-file/src/writer.rs index 80da9f4fb89..f009790e2f9 100644 --- a/vortex-file/src/writer.rs +++ b/vortex-file/src/writer.rs @@ -164,8 +164,8 @@ impl VortexWriteOptions { // 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()) - // Configure a registry just to ensure only known encodings are interned. - .with_registry(self.session.arrays().registry().clone()); + // Only permit encodings known to the session. + .with_valid_ids(self.session.arrays().registry().ids()); let dtype = stream.dtype().clone(); let (mut ptr, eof) = SequenceId::root().split(); diff --git a/vortex-session/src/registry.rs b/vortex-session/src/registry.rs index ed269dba7eb..83c6ceb038f 100644 --- a/vortex-session/src/registry.rs +++ b/vortex-session/src/registry.rs @@ -10,6 +10,7 @@ use std::fmt::Debug; use std::fmt::Display; use std::fmt::Formatter; use std::hash::Hash; +use std::marker::PhantomData; use std::ops::Deref; use std::sync::Arc; use std::sync::LazyLock; @@ -21,6 +22,7 @@ use parking_lot::RwLock; use vortex_error::VortexExpect; use vortex_utils::aliases::DefaultHashBuilder; use vortex_utils::aliases::dash_map::DashMap; +use vortex_utils::aliases::hash_set::HashSet; /// Global string interner for [`Id`] values. static INTERNER: LazyLock> = @@ -234,15 +236,17 @@ pub struct Context { // enum of Segment { Array, DType, Buffer } then we don't actually need a mutable context // in the LayoutWriter, therefore we don't need a RwLock here and everyone is happier. ids: Arc>>, - // Optional registry used to filter the permissible interned items. - registry: Option>, + // Optional set used to filter the permissible interned IDs. + valid_ids: Option>>, + _marker: PhantomData T>, } impl Default for Context { fn default() -> Self { Self { ids: Arc::new(RwLock::new(Vec::new())), - registry: None, + valid_ids: None, + _marker: PhantomData, } } } @@ -252,7 +256,8 @@ impl Context { pub fn new(ids: Vec) -> Self { Self { ids: Arc::new(RwLock::new(ids)), - registry: None, + valid_ids: None, + _marker: PhantomData, } } @@ -261,18 +266,25 @@ impl Context { Self::default() } - /// Configure a registry to restrict the permissible set of interned items. - pub fn with_registry(mut self, registry: Registry) -> Self { - self.registry = Some(registry); + /// Restrict the permissible set of interned IDs. + pub fn with_valid_ids(mut self, ids: impl IntoIterator) -> Self { + self.valid_ids = Some(Arc::new(ids.into_iter().collect())); self } + /// Restrict the permissible set of interned IDs to the IDs currently in a registry. + /// + /// IDs registered after this call are not added to the context. + pub fn with_registry(self, registry: Registry) -> Self { + self.with_valid_ids(registry.ids()) + } + /// Intern an ID, returning its index. pub fn intern(&self, id: &Id) -> Option { - if let Some(registry) = &self.registry - && registry.find(id).is_none() + if let Some(valid_ids) = &self.valid_ids + && !valid_ids.contains(id) { - // ID not in registry, cannot intern. + // ID is not valid, cannot intern. return None; } @@ -295,3 +307,24 @@ impl Context { self.ids.read().clone() } } + +#[cfg(test)] +mod tests { + use super::CachedId; + use super::Context; + + static VALID: CachedId = CachedId::new("vortex.test.valid"); + static INVALID: CachedId = CachedId::new("vortex.test.invalid"); + + #[test] + fn context_filters_interned_ids() { + let valid = *VALID; + let invalid = *INVALID; + let context = Context::<()>::empty().with_valid_ids([valid]); + + assert_eq!(context.intern(&valid), Some(0)); + assert_eq!(context.intern(&valid), Some(0)); + assert_eq!(context.intern(&invalid), None); + assert_eq!(context.to_ids(), [valid]); + } +} From fdc2ca33ad2ac0593be9225986d5477324c7db6c Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Wed, 22 Jul 2026 16:11:28 +0100 Subject: [PATCH 2/2] refactor: remove context type parameter Signed-off-by: Joe Isaacs --- vortex-array/src/lib.rs | 2 +- vortex-layout/src/lib.rs | 2 +- vortex-session/src/registry.rs | 21 +++++---------------- 3 files changed, 7 insertions(+), 18 deletions(-) diff --git a/vortex-array/src/lib.rs b/vortex-array/src/lib.rs index 9d3a59eb7a0..1b4298d959a 100644 --- a/vortex-array/src/lib.rs +++ b/vortex-array/src/lib.rs @@ -190,4 +190,4 @@ pub fn legacy_session() -> &'static VortexSession { &LEGACY_SESSION } -pub type ArrayContext = Context; +pub type ArrayContext = Context; diff --git a/vortex-layout/src/lib.rs b/vortex-layout/src/lib.rs index 6612f5e0350..a0f2ef8e2d7 100644 --- a/vortex-layout/src/lib.rs +++ b/vortex-layout/src/lib.rs @@ -42,4 +42,4 @@ mod strategy; mod test; pub mod vtable; -pub type LayoutContext = Context; +pub type LayoutContext = Context; diff --git a/vortex-session/src/registry.rs b/vortex-session/src/registry.rs index 83c6ceb038f..e75ce01d214 100644 --- a/vortex-session/src/registry.rs +++ b/vortex-session/src/registry.rs @@ -10,7 +10,6 @@ use std::fmt::Debug; use std::fmt::Display; use std::fmt::Formatter; use std::hash::Hash; -use std::marker::PhantomData; use std::ops::Deref; use std::sync::Arc; use std::sync::LazyLock; @@ -229,35 +228,32 @@ impl ReadContext { /// 1. This object holds an Arc of RwLock internally because we need concurrent access from the /// layout writer code path. We should update SegmentSink to take an Array rather than /// ByteBuffer such that serializing arrays is done sequentially. -/// 2. The name is terrible. `Interner` is better, but I want to minimize breakage for now. +/// 2. The name is terrible. `Interner` is better, but I want to minimize breakage for now. #[derive(Clone, Debug)] -pub struct Context { +pub struct Context { // TODO(ngates): it's a long story, but if we make SegmentSink and SegmentSource take an // enum of Segment { Array, DType, Buffer } then we don't actually need a mutable context // in the LayoutWriter, therefore we don't need a RwLock here and everyone is happier. ids: Arc>>, // Optional set used to filter the permissible interned IDs. valid_ids: Option>>, - _marker: PhantomData T>, } -impl Default for Context { +impl Default for Context { fn default() -> Self { Self { ids: Arc::new(RwLock::new(Vec::new())), valid_ids: None, - _marker: PhantomData, } } } -impl Context { +impl Context { /// Create a context with the given initial IDs. pub fn new(ids: Vec) -> Self { Self { ids: Arc::new(RwLock::new(ids)), valid_ids: None, - _marker: PhantomData, } } @@ -272,13 +268,6 @@ impl Context { self } - /// Restrict the permissible set of interned IDs to the IDs currently in a registry. - /// - /// IDs registered after this call are not added to the context. - pub fn with_registry(self, registry: Registry) -> Self { - self.with_valid_ids(registry.ids()) - } - /// Intern an ID, returning its index. pub fn intern(&self, id: &Id) -> Option { if let Some(valid_ids) = &self.valid_ids @@ -320,7 +309,7 @@ mod tests { fn context_filters_interned_ids() { let valid = *VALID; let invalid = *INVALID; - let context = Context::<()>::empty().with_valid_ids([valid]); + let context = Context::empty().with_valid_ids([valid]); assert_eq!(context.intern(&valid), Some(0)); assert_eq!(context.intern(&valid), Some(0));