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-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-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-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 ed269dba7eb..e75ce01d214 100644 --- a/vortex-session/src/registry.rs +++ b/vortex-session/src/registry.rs @@ -21,6 +21,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> = @@ -227,32 +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 registry used to filter the permissible interned items. - registry: Option>, + // Optional set used to filter the permissible interned IDs. + valid_ids: Option>>, } -impl Default for Context { +impl Default for Context { fn default() -> Self { Self { ids: Arc::new(RwLock::new(Vec::new())), - registry: None, + valid_ids: None, } } } -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)), - registry: None, + valid_ids: None, } } @@ -261,18 +262,18 @@ 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 } /// 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 +296,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]); + } +}