From d322b2070cf58a40707b4138b0319cb80ea00aea Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Thu, 23 Jul 2026 15:25:24 +0100 Subject: [PATCH 1/3] Replace all usages of Registry with ArcSwapMap Signed-off-by: Robert Kruszewski --- docs/developer-guide/internals/session.md | 7 +- encodings/zstd/src/zstd_buffers.rs | 7 +- vortex-array/src/aggregate_fn/session.rs | 2 +- vortex-array/src/arrays/patched/vtable/mod.rs | 2 +- vortex-array/src/dtype/serde/flatbuffers.rs | 2 +- vortex-array/src/dtype/serde/proto.rs | 2 +- vortex-array/src/dtype/serde/serde.rs | 2 +- vortex-array/src/dtype/session.rs | 14 +-- vortex-array/src/expr/proto.rs | 2 +- vortex-array/src/lib.rs | 31 +++-- vortex-array/src/optimizer/kernels.rs | 2 +- vortex-array/src/scalar_fn/session.rs | 14 +-- vortex-array/src/serde.rs | 2 +- vortex-array/src/session/mod.rs | 21 ++-- vortex-arrow/src/session.rs | 2 +- vortex-file/src/lib.rs | 6 +- vortex-file/src/writer.rs | 2 +- vortex-layout/src/children.rs | 10 +- vortex-layout/src/flatbuffers.rs | 2 +- vortex-layout/src/lib.rs | 4 +- vortex-layout/src/session.rs | 47 ++++---- vortex-python-cuda/src/lib.rs | 2 +- vortex-session/Cargo.toml | 2 +- .../src/arc_swap_map.rs | 91 +++++++++++--- vortex-session/src/lib.rs | 2 + vortex-session/src/registry.rs | 111 ++++-------------- vortex-session/src/session.rs | 52 ++++---- 27 files changed, 218 insertions(+), 225 deletions(-) rename {vortex-array => vortex-session}/src/arc_swap_map.rs (64%) diff --git a/docs/developer-guide/internals/session.md b/docs/developer-guide/internals/session.md index 71158eaae20..0312d496cfe 100644 --- a/docs/developer-guide/internals/session.md +++ b/docs/developer-guide/internals/session.md @@ -15,10 +15,11 @@ The session is built on two primitives from the `vortex-session` crate: - **`VortexSession`** -- a cloneable, thread-safe map from Rust `TypeId` to a shared (`Arc`-wrapped) value. Any type that is `Clone + Send + Sync + Debug + 'static` can be stored as a session variable. -- **`Registry`** -- a concurrent map from string IDs to values of type `T`, used by each - component to look up registered plugins at runtime. +- **`ArcSwapMap`** -- a concurrent copy-on-write map, used both as the backing store of + `VortexSession` itself and (keyed by interned string `Id`s) as the plugin registries each + component uses to look up registered plugins at runtime. -Because `VortexSession` is backed by an `ArcSwap`, cloning is cheap and all clones share the +Because both primitives are backed by an `ArcSwap`, cloning is cheap and all clones share the same state, with lock-free reads and copy-on-write writes. This makes it safe to hand the session to multiple threads, tasks, or I/O operations without coordination. diff --git a/encodings/zstd/src/zstd_buffers.rs b/encodings/zstd/src/zstd_buffers.rs index 73361a593a1..ba17407708b 100644 --- a/encodings/zstd/src/zstd_buffers.rs +++ b/encodings/zstd/src/zstd_buffers.rs @@ -109,9 +109,10 @@ impl ZstdBuffers { buffer_handles: &[BufferHandle], session: &VortexSession, ) -> VortexResult { - let registry = session.arrays().registry().clone(); - let inner_vtable = registry - .find(&array.data().inner_encoding_id) + let inner_vtable = session + .arrays() + .registry() + .get(&array.data().inner_encoding_id) .ok_or_else(|| { vortex_err!("Unknown inner encoding: {}", array.data().inner_encoding_id) })?; diff --git a/vortex-array/src/aggregate_fn/session.rs b/vortex-array/src/aggregate_fn/session.rs index 72ef78a400c..effe33c6ca4 100644 --- a/vortex-array/src/aggregate_fn/session.rs +++ b/vortex-array/src/aggregate_fn/session.rs @@ -4,6 +4,7 @@ use std::any::Any; use std::sync::Arc; +use vortex_session::ArcSwapMap; use vortex_session::SessionExt; use vortex_session::SessionGuard; use vortex_session::SessionVar; @@ -35,7 +36,6 @@ use crate::aggregate_fn::fns::sum::Sum; use crate::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes; use crate::aggregate_fn::kernels::DynAggregateKernel; use crate::aggregate_fn::kernels::DynGroupedAggregateKernel; -use crate::arc_swap_map::ArcSwapMap; use crate::array::ArrayId; use crate::array::VTable; use crate::arrays::Chunked; diff --git a/vortex-array/src/arrays/patched/vtable/mod.rs b/vortex-array/src/arrays/patched/vtable/mod.rs index f69b7b73afa..0c9fe8d8f91 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_valid_ids(session.arrays().registry().ids()); + let ctx = ArrayContext::empty().with_allowed_ids(session.arrays().registry().keys()); let serialized = array .serialize(&ctx, &session, &SerializeOptions::default()) .unwrap(); diff --git a/vortex-array/src/dtype/serde/flatbuffers.rs b/vortex-array/src/dtype/serde/flatbuffers.rs index a796490f58a..6344973c2b3 100644 --- a/vortex-array/src/dtype/serde/flatbuffers.rs +++ b/vortex-array/src/dtype/serde/flatbuffers.rs @@ -270,7 +270,7 @@ impl TryFrom for DType { vortex_err!("failed to parse extension metadata from flatbuffer") })? .bytes(); - let ext_dtype = if let Some(vtable) = vfdt.session.dtypes().registry().find(&id) { + let ext_dtype = if let Some(vtable) = vfdt.session.dtypes().registry().get(&id) { vtable.deserialize(metadata, storage_dtype)? } else if vfdt.session.allows_unknown() { ForeignExtDType::from_parts(id, metadata.to_vec(), storage_dtype)? diff --git a/vortex-array/src/dtype/serde/proto.rs b/vortex-array/src/dtype/serde/proto.rs index 8389a90200a..485907e2e95 100644 --- a/vortex-array/src/dtype/serde/proto.rs +++ b/vortex-array/src/dtype/serde/proto.rs @@ -116,7 +116,7 @@ impl DType { .ok_or_else(|| vortex_err!("Extension DType missing storage proto"))?, session, )?; - let ext_dtype = if let Some(vtable) = session.dtypes().registry().find(&id) { + let ext_dtype = if let Some(vtable) = session.dtypes().registry().get(&id) { vtable.deserialize(e.metadata(), storage_dtype)? } else if session.allows_unknown() { ForeignExtDType::from_parts(id, e.metadata().to_vec(), storage_dtype)? diff --git a/vortex-array/src/dtype/serde/serde.rs b/vortex-array/src/dtype/serde/serde.rs index b3870b8c03e..8fd52fda054 100644 --- a/vortex-array/src/dtype/serde/serde.rs +++ b/vortex-array/src/dtype/serde/serde.rs @@ -722,7 +722,7 @@ impl<'de> DeserializeSeed<'de> for DTypeSerde<'_, ExtDTypeRef> { storage_dtype.ok_or_else(|| de::Error::missing_field("storage_dtype"))?; let metadata = metadata.ok_or_else(|| de::Error::missing_field("metadata"))?; - if let Some(vtable) = self.session.dtypes().registry().find(&id) { + if let Some(vtable) = self.session.dtypes().registry().get(&id) { vtable.deserialize(&metadata, storage_dtype).map_err(|e| { de::Error::custom(format!( "failed to deserialize extension dtype {}: {}", diff --git a/vortex-array/src/dtype/session.rs b/vortex-array/src/dtype/session.rs index 2e13b923bc3..05a276589c7 100644 --- a/vortex-array/src/dtype/session.rs +++ b/vortex-array/src/dtype/session.rs @@ -6,10 +6,11 @@ use std::any::Any; use std::sync::Arc; +use vortex_session::ArcSwapMap; use vortex_session::SessionExt; use vortex_session::SessionGuard; use vortex_session::SessionVar; -use vortex_session::registry::Registry; +use vortex_session::registry::Id; use crate::dtype::extension::ExtDTypePluginRef; use crate::dtype::extension::ExtVTable; @@ -18,19 +19,16 @@ use crate::extension::datetime::Time; use crate::extension::datetime::Timestamp; use crate::extension::uuid::Uuid; -/// Registry for extension dtypes. -pub type ExtDTypeRegistry = Registry; - /// Session for managing extension dtypes. #[derive(Clone, Debug)] pub struct DTypeSession { - registry: ExtDTypeRegistry, + registry: ArcSwapMap, } impl Default for DTypeSession { fn default() -> Self { let this = Self { - registry: Registry::default(), + registry: ArcSwapMap::default(), }; // Register built-in temporal extension dtypes @@ -57,11 +55,11 @@ impl DTypeSession { /// Register an extension DType with the Vortex session. pub fn register(&self, vtable: V) { self.registry - .register(vtable.id(), Arc::new(vtable) as ExtDTypePluginRef); + .insert(vtable.id(), Arc::new(vtable) as ExtDTypePluginRef); } /// Return the registry of extension dtypes. - pub fn registry(&self) -> &ExtDTypeRegistry { + pub fn registry(&self) -> &ArcSwapMap { &self.registry } } diff --git a/vortex-array/src/expr/proto.rs b/vortex-array/src/expr/proto.rs index 0a8d71f2954..8febdb389a0 100644 --- a/vortex-array/src/expr/proto.rs +++ b/vortex-array/src/expr/proto.rs @@ -47,7 +47,7 @@ impl Expression { .map(|e| Expression::from_proto(e, session)) .collect::>>()?; - let scalar_fn = if let Some(vtable) = session.scalar_fns().registry().find(&expr_id) { + let scalar_fn = if let Some(vtable) = session.scalar_fns().registry().get(&expr_id) { vtable.deserialize(expr.metadata(), session)? } else if session.allows_unknown() { ForeignScalarFnVTable::make_scalar_fn(expr_id, expr.metadata().to_vec(), children.len()) diff --git a/vortex-array/src/lib.rs b/vortex-array/src/lib.rs index 1cbe2fd2895..ea61c18dfd8 100644 --- a/vortex-array/src/lib.rs +++ b/vortex-array/src/lib.rs @@ -4,7 +4,7 @@ //! //! At the heart of Vortex are [arrays](ArrayRef). //! -//! Arrays are typed views of memory buffers that hold [scalars](crate::scalar::Scalar). These +//! Arrays are typed views of memory buffers that hold [scalars](scalar::Scalar). These //! buffers can be held in a number of physical encodings to perform lightweight compression that //! exploits the particular data distribution of the array's values. //! @@ -51,10 +51,10 @@ //! //! # Nulls and Scalars //! -//! [`Validity`](crate::validity::Validity) separates nullness from values. It can be a cheap +//! [`Validity`](validity::Validity) separates nullness from values. It can be a cheap //! constant state (`NonNullable`, `AllValid`, `AllInvalid`) or a boolean array that may itself be -//! encoded. [`Scalar`](crate::scalar::Scalar) is the single-value counterpart: it pairs a -//! [`DType`] with an optional [`ScalarValue`](crate::scalar::ScalarValue). +//! encoded. [`Scalar`](scalar::Scalar) is the single-value counterpart: it pairs a +//! [`DType`] with an optional [`ScalarValue`](scalar::ScalarValue). //! //! # Extending Vortex //! @@ -66,16 +66,16 @@ //! - [`OperationsVTable`] provides scalar access. //! - [`ValidityVTable`] exposes validity only for nullable arrays. //! -//! New logical extension dtypes implement [`ExtVTable`](crate::dtype::extension::ExtVTable) and +//! New logical extension dtypes implement [`ExtVTable`](dtype::extension::ExtVTable) and //! store values in an ordinary Vortex storage dtype. //! -//! [`PrimitiveArray`]: crate::arrays::PrimitiveArray -//! [`DType`]: crate::dtype::DType -//! [`ChunkedArray`]: crate::arrays::ChunkedArray -//! [`ConstantArray`]: crate::arrays::ConstantArray -//! [`FilterArray`]: crate::arrays::FilterArray -//! [`SliceArray`]: crate::arrays::SliceArray -//! [`ScalarFnArray`]: crate::arrays::ScalarFnArray +//! [`PrimitiveArray`]: arrays::PrimitiveArray +//! [`DType`]: dtype::DType +//! [`ChunkedArray`]: arrays::ChunkedArray +//! [`ConstantArray`]: arrays::ConstantArray +//! [`FilterArray`]: arrays::FilterArray +//! [`SliceArray`]: arrays::SliceArray +//! [`ScalarFnArray`]: arrays::ScalarFnArray extern crate self as vortex_array; @@ -92,7 +92,7 @@ pub use smallvec; pub use vortex_array_macros::array_slots; use vortex_session::SessionExt; use vortex_session::VortexSession; -use vortex_session::registry::Context; +use vortex_session::registry::Interner; use crate::aggregate_fn::session::AggregateFnSession; use crate::dtype::session::DTypeSession; @@ -105,7 +105,6 @@ use crate::stats::session::StatsSession; pub mod aggregate_fn; #[doc(hidden)] pub mod aliases; -pub mod arc_swap_map; mod array; pub mod arrays; pub mod buffer; @@ -149,7 +148,7 @@ pub mod flatbuffers { } /// Register vortex-array's built-in session-scoped kernels into the active -/// [`ArrayKernels`](crate::optimizer::kernels::ArrayKernels) registry. +/// [`ArrayKernels`](optimizer::kernels::ArrayKernels) registry. /// /// If the session contains a [`KernelSession`], this registers into its registry. Sessions that use /// [`KernelSession::default`] already receive these built-in kernels. @@ -190,4 +189,4 @@ pub fn legacy_session() -> &'static VortexSession { &LEGACY_SESSION } -pub type ArrayContext = Context; +pub type ArrayContext = Interner; diff --git a/vortex-array/src/optimizer/kernels.rs b/vortex-array/src/optimizer/kernels.rs index e419e927ca1..c68c009f532 100644 --- a/vortex-array/src/optimizer/kernels.rs +++ b/vortex-array/src/optimizer/kernels.rs @@ -32,6 +32,7 @@ use std::sync::Arc; use std::sync::LazyLock; use vortex_error::VortexResult; +use vortex_session::ArcSwapMap; use vortex_session::SessionExt; use vortex_session::SessionGuard; use vortex_session::SessionVar; @@ -42,7 +43,6 @@ use vortex_utils::aliases::hash_map::HashMap; use crate::ArrayRef; use crate::ExecutionCtx; -use crate::arc_swap_map::ArcSwapMap; use crate::array::VTable; use crate::arrays::Struct; use crate::arrays::struct_::compute::rules::struct_cast_reduce_parent; diff --git a/vortex-array/src/scalar_fn/session.rs b/vortex-array/src/scalar_fn/session.rs index e937ac8d316..ee52b390437 100644 --- a/vortex-array/src/scalar_fn/session.rs +++ b/vortex-array/src/scalar_fn/session.rs @@ -4,10 +4,11 @@ use std::any::Any; use std::sync::Arc; +use vortex_session::ArcSwapMap; use vortex_session::SessionExt; use vortex_session::SessionGuard; use vortex_session::SessionVar; -use vortex_session::registry::Registry; +use vortex_session::registry::Id; use crate::scalar_fn::ScalarFnPluginRef; use crate::scalar_fn::ScalarFnVTable; @@ -32,31 +33,28 @@ use crate::scalar_fn::fns::select::Select; use crate::scalar_fn::fns::stat::StatFn; use crate::scalar_fn::fns::variant_get::VariantGet; -/// Registry of scalar function vtables. -pub type ScalarFnRegistry = Registry; - /// Session state for scalar function vtables and rewrite rules. #[derive(Clone, Debug)] pub struct ScalarFnSession { - registry: ScalarFnRegistry, + registry: ArcSwapMap, } impl ScalarFnSession { - pub fn registry(&self) -> &ScalarFnRegistry { + pub fn registry(&self) -> &ArcSwapMap { &self.registry } /// Register a scalar function vtable in the session, replacing any existing vtable with the same ID. pub fn register(&self, vtable: V) { self.registry - .register(vtable.id(), Arc::new(vtable) as ScalarFnPluginRef); + .insert(vtable.id(), Arc::new(vtable) as ScalarFnPluginRef); } } impl Default for ScalarFnSession { fn default() -> Self { let this = Self { - registry: ScalarFnRegistry::default(), + registry: ArcSwapMap::default(), }; // Register built-in expressions. diff --git a/vortex-array/src/serde.rs b/vortex-array/src/serde.rs index f6907d85391..04291b9eef8 100644 --- a/vortex-array/src/serde.rs +++ b/vortex-array/src/serde.rs @@ -325,7 +325,7 @@ impl SerializedArray { let encoding_id = ctx .resolve(encoding_idx) .ok_or_else(|| vortex_err!("Unknown encoding index: {}", encoding_idx))?; - let Some(plugin) = session.arrays().registry().find(&encoding_id) else { + let Some(plugin) = session.arrays().registry().get(&encoding_id) else { if session.allows_unknown() { return self.decode_foreign(encoding_id, dtype, len, ctx); } diff --git a/vortex-array/src/session/mod.rs b/vortex-array/src/session/mod.rs index 406ab31bdc3..205bd57856f 100644 --- a/vortex-array/src/session/mod.rs +++ b/vortex-array/src/session/mod.rs @@ -6,10 +6,11 @@ use std::sync::Arc; use vortex_error::VortexResult; use vortex_error::vortex_bail; +use vortex_session::ArcSwapMap; use vortex_session::SessionExt; use vortex_session::SessionGuard; use vortex_session::SessionVar; -use vortex_session::registry::Registry; +use vortex_session::registry::Id; use crate::ArrayRef; use crate::array::ArrayPlugin; @@ -33,36 +34,34 @@ use crate::arrays::VarBin; use crate::arrays::VarBinView; use crate::arrays::Variant; -pub type ArrayRegistry = Registry; - #[derive(Clone, Debug)] pub struct ArraySession { /// The set of registered array encodings. - registry: ArrayRegistry, + registry: ArcSwapMap, } impl ArraySession { pub fn empty() -> ArraySession { Self { - registry: ArrayRegistry::default(), + registry: ArcSwapMap::default(), } } - pub fn registry(&self) -> &ArrayRegistry { + pub fn registry(&self) -> &ArcSwapMap { &self.registry } /// Register a new array encoding, replacing any existing encoding with the same ID. pub fn register(&self, plugin: P) { self.registry - .register(plugin.id(), Arc::new(plugin) as ArrayPluginRef); + .insert(plugin.id(), Arc::new(plugin) as ArrayPluginRef); } } impl Default for ArraySession { fn default() -> Self { let this = ArraySession { - registry: ArrayRegistry::default(), + registry: ArcSwapMap::default(), }; // Register the canonical encodings. @@ -110,7 +109,7 @@ pub trait ArraySessionExt: SessionExt { /// Serialize an array using a plugin from the registry. fn array_serialize(&self, array: &ArrayRef) -> VortexResult>> { - let Some(plugin) = self.arrays().registry.find(&array.encoding_id()) else { + let Some(plugin) = self.arrays().registry.get(&array.encoding_id()) else { vortex_bail!( "Array {} is not registered for serializations", array.encoding_id() @@ -136,13 +135,13 @@ mod tests { fn array_session_default_registers_encodings() { let session = VortexSession::empty().with::(); - assert!(session.arrays().registry().find(&Bool.id()).is_some()); + assert!(session.arrays().registry().contains_key(&Bool.id())); } #[test] fn empty_array_session_registers_no_encodings() { let session = VortexSession::empty().with_some(ArraySession::empty()); - assert!(session.arrays().registry().find(&Bool.id()).is_none()); + assert!(!session.arrays().registry().contains_key(&Bool.id())); } } diff --git a/vortex-arrow/src/session.rs b/vortex-arrow/src/session.rs index 745e8afeb47..9eb98e71540 100644 --- a/vortex-arrow/src/session.rs +++ b/vortex-arrow/src/session.rs @@ -39,7 +39,6 @@ use tracing::trace; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; -use vortex_array::arc_swap_map::ArcSwapMap; use vortex_array::arrays::FixedSizeListArray; use vortex_array::arrays::ListArray; use vortex_array::arrays::ListViewArray; @@ -56,6 +55,7 @@ use vortex_array::validity::Validity; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; +use vortex_session::ArcSwapMap; use vortex_session::SessionExt; use vortex_session::SessionGuard; use vortex_session::SessionVar; diff --git a/vortex-file/src/lib.rs b/vortex-file/src/lib.rs index 1a1a751568c..708dc685305 100644 --- a/vortex-file/src/lib.rs +++ b/vortex-file/src/lib.rs @@ -214,14 +214,14 @@ mod default_encoding_tests { fn register_default_encodings_registers_external_execute_parent_kernels() { let session = array_session(); - assert!(session.arrays().registry().find(&FSST.id()).is_none()); + assert!(!session.arrays().registry().contains_key(&FSST.id())); assert!(!session.kernels().has_execute_parent(Filter.id(), FSST.id())); register_default_encodings(&session); - assert!(session.arrays().registry().find(&FSST.id()).is_some()); + assert!(session.arrays().registry().contains_key(&FSST.id())); assert!(session.kernels().has_execute_parent(Filter.id(), FSST.id())); - assert!(session.arrays().registry().find(&OnPair.id()).is_some()); + assert!(session.arrays().registry().contains_key(&OnPair.id())); assert!( session .kernels() diff --git a/vortex-file/src/writer.rs b/vortex-file/src/writer.rs index 58d1cb7cd4f..f315f9ff4ef 100644 --- a/vortex-file/src/writer.rs +++ b/vortex-file/src/writer.rs @@ -201,7 +201,7 @@ impl VortexWriteOptions { // 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_valid_ids(self.session.arrays().registry().ids()); + .with_allowed_ids(self.session.arrays().registry().keys()); let dtype = stream.dtype().clone(); let (mut ptr, eof) = SequenceId::root().split(); diff --git a/vortex-layout/src/children.rs b/vortex-layout/src/children.rs index fa9b3ae47f8..72dcee9d07d 100644 --- a/vortex-layout/src/children.rs +++ b/vortex-layout/src/children.rs @@ -14,14 +14,16 @@ use vortex_error::vortex_bail; use vortex_error::vortex_err; use vortex_flatbuffers::FlatBuffer; use vortex_flatbuffers::layout as fbl; +use vortex_session::ArcSwapMap; use vortex_session::VortexSession; +use vortex_session::registry::Id; use vortex_session::registry::ReadContext; use crate::LayoutBuildContext; +use crate::LayoutEncodingRef; use crate::LayoutRef; use crate::layouts::foreign::new_foreign_layout; use crate::segments::SegmentId; -use crate::session::LayoutRegistry; /// Abstract way of accessing the children of a layout. /// @@ -110,7 +112,7 @@ pub(crate) struct ViewedLayoutChildren { flatbuffer_loc: usize, array_read_ctx: ReadContext, layout_read_ctx: ReadContext, - layouts: LayoutRegistry, + layouts: ArcSwapMap, allow_unknown: bool, session: VortexSession, cache: Arc<[OnceCell]>, @@ -127,7 +129,7 @@ impl ViewedLayoutChildren { flatbuffer_loc: usize, array_read_ctx: ReadContext, layout_read_ctx: ReadContext, - layouts: LayoutRegistry, + layouts: ArcSwapMap, allow_unknown: bool, session: VortexSession, ) -> Self { @@ -226,7 +228,7 @@ impl LayoutChildren for ViewedLayoutChildren { .ok_or_else(|| { vortex_err!("Unknown layout encoding index: {}", fb_child.encoding()) })?; - let Some(encoding) = self.layouts.find(&encoding_id) else { + let Some(encoding) = self.layouts.get(&encoding_id) else { if self.allow_unknown { return viewed_children.foreign_layout_from_fb(fb_child, dtype); } diff --git a/vortex-layout/src/flatbuffers.rs b/vortex-layout/src/flatbuffers.rs index 509b41b12da..af149016462 100644 --- a/vortex-layout/src/flatbuffers.rs +++ b/vortex-layout/src/flatbuffers.rs @@ -70,7 +70,7 @@ pub fn layout_from_flatbuffer_with_options( let encoding_id = layout_ctx .resolve(fb_layout.encoding()) .ok_or_else(|| vortex_err!("Invalid encoding ID: {}", fb_layout.encoding()))?; - let encoding = layouts.find(&encoding_id); + let encoding = layouts.get(&encoding_id); if encoding.is_none() && allow_unknown { return foreign_layout_from_fb(fb_layout, dtype, layout_ctx); diff --git a/vortex-layout/src/lib.rs b/vortex-layout/src/lib.rs index 779076bddac..96a34ce75f3 100644 --- a/vortex-layout/src/lib.rs +++ b/vortex-layout/src/lib.rs @@ -23,7 +23,7 @@ pub use layout::*; pub use reader::*; pub use reader_context::*; pub use strategy::*; -use vortex_session::registry::Context; +use vortex_session::registry::Interner; pub use vtable::*; pub mod aliases; mod children; @@ -42,4 +42,4 @@ mod strategy; mod test; mod vtable; -pub type LayoutContext = Context; +pub type LayoutContext = Interner; diff --git a/vortex-layout/src/session.rs b/vortex-layout/src/session.rs index 53ced66e955..059b313844f 100644 --- a/vortex-layout/src/session.rs +++ b/vortex-layout/src/session.rs @@ -3,62 +3,61 @@ use std::any::Any; +use vortex_session::ArcSwapMap; use vortex_session::SessionExt; use vortex_session::SessionGuard; use vortex_session::SessionVar; -use vortex_session::registry::Registry; +use vortex_session::registry::Id; use crate::LayoutEncoding; use crate::LayoutEncodingRef; -use crate::layouts::chunked::Chunked; -use crate::layouts::dict::Dict; -use crate::layouts::flat::Flat; -use crate::layouts::list::List; -use crate::layouts::struct_::Struct; -use crate::layouts::zoned::LegacyStats; -use crate::layouts::zoned::Zoned; - -pub type LayoutRegistry = Registry; +use crate::layouts::chunked::ChunkedLayoutEncoding; +use crate::layouts::dict::DictLayoutEncoding; +use crate::layouts::flat::FlatLayoutEncoding; +use crate::layouts::list::ListLayoutEncoding; +use crate::layouts::struct_::StructLayoutEncoding; +use crate::layouts::zoned::{LegacyStatsLayoutEncoding, ZonedLayoutEncoding}; /// Session state for layout encodings. #[derive(Clone, Debug)] pub struct LayoutSession { - registry: LayoutRegistry, + registry: ArcSwapMap, } impl LayoutSession { /// Register a layout encoding in the session, replacing any existing encoding with the same ID. pub fn register(&self, layout: LayoutEncodingRef) { - self.registry.register(layout.id(), layout); + self.registry.insert(layout.id(), layout); } /// Register layout encodings in the session, replacing any existing encodings with the same IDs. pub fn register_many(&self, layouts: impl IntoIterator) { for layout in layouts { - self.registry.register(layout.id(), layout); + self.registry.insert(layout.id(), layout); } } /// Returns the layout encoding registry. - pub fn registry(&self) -> &LayoutRegistry { + pub fn registry(&self) -> &ArcSwapMap { &self.registry } } impl Default for LayoutSession { fn default() -> Self { - let layouts = LayoutRegistry::default(); + let this = Self { + registry: ArcSwapMap::default(), + }; // Register the built-in layout encodings. - layouts.register(Chunked.id(), &Chunked as &dyn LayoutEncoding); - layouts.register(Flat.id(), &Flat as &dyn LayoutEncoding); - layouts.register(Struct.id(), &Struct as &dyn LayoutEncoding); - layouts.register(Zoned.id(), &Zoned as &dyn LayoutEncoding); - layouts.register(LegacyStats.id(), &LegacyStats as &dyn LayoutEncoding); - layouts.register(Dict.id(), &Dict as &dyn LayoutEncoding); - layouts.register(List.id(), &List as &dyn LayoutEncoding); - - Self { registry: layouts } + this.register(ChunkedLayoutEncoding.as_ref().into()); + this.register(FlatLayoutEncoding.as_ref().into()); + this.register(StructLayoutEncoding.as_ref().into()); + this.register(ZonedLayoutEncoding.as_ref().into()); + this.register(LegacyStatsLayoutEncoding.as_ref().into()); + this.register(DictLayoutEncoding.as_ref().into()); + this.register(ListLayoutEncoding.as_ref().into()); + this } } diff --git a/vortex-python-cuda/src/lib.rs b/vortex-python-cuda/src/lib.rs index 2ac3166b103..dff0a0fd03d 100644 --- a/vortex-python-cuda/src/lib.rs +++ b/vortex-python-cuda/src/lib.rs @@ -297,7 +297,7 @@ fn deserialize_metadata_tree( let plugin = session .arrays() .registry() - .find(&encoding_id) + .get(&encoding_id) .ok_or_else(|| vortex_err!("Unknown array encoding: {}", metadata.encoding_id))?; let decoded = plugin.deserialize( &dtype, diff --git a/vortex-session/Cargo.toml b/vortex-session/Cargo.toml index 026f9b9e528..b243d3e4f63 100644 --- a/vortex-session/Cargo.toml +++ b/vortex-session/Cargo.toml @@ -24,4 +24,4 @@ arc-swap = { workspace = true } lasso = { workspace = true } parking_lot = { workspace = true } vortex-error = { workspace = true } -vortex-utils = { workspace = true, features = ["dashmap"] } +vortex-utils = { workspace = true } diff --git a/vortex-array/src/arc_swap_map.rs b/vortex-session/src/arc_swap_map.rs similarity index 64% rename from vortex-array/src/arc_swap_map.rs rename to vortex-session/src/arc_swap_map.rs index 1788f704feb..f63e62563aa 100644 --- a/vortex-array/src/arc_swap_map.rs +++ b/vortex-session/src/arc_swap_map.rs @@ -6,11 +6,15 @@ use std::borrow::Borrow; use std::fmt::Debug; use std::fmt::Formatter; +use std::hash::BuildHasher; use std::hash::Hash; use std::sync::Arc; use arc_swap::ArcSwap; +use arc_swap::Guard; +use vortex_utils::aliases::hash_map::DefaultHashBuilder; use vortex_utils::aliases::hash_map::HashMap; +use vortex_utils::aliases::hash_set::HashSet; /// A concurrent [`HashMap`] backed by an [`ArcSwap`], offering lock-free reads /// and copy-on-write writes. @@ -20,19 +24,21 @@ use vortex_utils::aliases::hash_map::HashMap; /// reader always observes a consistent snapshot and writers never block readers. /// /// This is the shared building block behind the session-scoped registries (the -/// optimizer-kernel and aggregate-function registries). Because every write -/// clones the entire map, it is intended for maps that are written rarely -/// (typically only while a session is being configured) and read often. +/// plugin registries as well as the optimizer-kernel and aggregate-function +/// registries) and the [`VortexSession`](crate::VortexSession) type-map itself. +/// Because every write clones the entire map, it is intended for maps that are +/// written rarely (typically only while a session is being configured) and read +/// often. /// /// The map is held behind an [`Arc`] so that [`Clone`] shares the same /// underlying cell: a registry mutated through one clone is observed by all /// others. Session variables rely on this so that encodings registered after a /// session is built remain visible to clones of that session. -pub struct ArcSwapMap { - inner: Arc>>, +pub struct ArcSwapMap { + inner: Arc>>, } -impl Default for ArcSwapMap { +impl Default for ArcSwapMap { fn default() -> Self { Self { inner: Arc::new(ArcSwap::from_pointee(HashMap::default())), @@ -40,7 +46,7 @@ impl Default for ArcSwapMap { } } -impl Clone for ArcSwapMap { +impl Clone for ArcSwapMap { fn clone(&self) -> Self { Self { inner: Arc::clone(&self.inner), @@ -48,15 +54,15 @@ impl Clone for ArcSwapMap { } } -impl Debug for ArcSwapMap { +impl Debug for ArcSwapMap { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { self.read(|map| f.debug_tuple("ArcSwapMap").field(map).finish()) } } -impl ArcSwapMap { +impl ArcSwapMap { /// Return the currently published map snapshot. - pub fn snapshot(&self) -> Arc> { + pub fn snapshot(&self) -> Arc> { self.inner.load_full() } @@ -64,18 +70,29 @@ impl ArcSwapMap { /// /// Every lookup inside `f` observes the same snapshot, which matters when a /// single logical read consults more than one key. - pub fn read(&self, f: impl FnOnce(&HashMap) -> R) -> R { + pub fn read(&self, f: impl FnOnce(&HashMap) -> R) -> R { f(&self.inner.load()) } + /// Return a lock-free guard to the current snapshot without cloning the + /// [`Arc`]. + /// + /// Cheaper than [`snapshot`](Self::snapshot) for short-lived reads, but the + /// guard pins an internal arc-swap slot, so it should stay on the stack + /// rather than be stored in a long-lived data structure. + pub(crate) fn load(&self) -> Guard>> { + self.inner.load() + } + /// Replace the map with the result of applying `f` to a private copy. /// /// Writes are copy-on-write via [`ArcSwap::rcu`], so `f` may run more than /// once under contention and must not move out of its captures. - fn modify(&self, f: impl Fn(&mut HashMap)) + fn modify(&self, f: impl Fn(&mut HashMap)) where K: Clone, V: Clone, + S: Clone, { self.inner.rcu(|existing| { let mut map = existing.as_ref().clone(); @@ -85,28 +102,64 @@ impl ArcSwapMap { } } -impl ArcSwapMap { +impl ArcSwapMap { /// Return a clone of the value stored under `key`, if present. pub fn get(&self, key: &Q) -> Option where K: Borrow, Q: Eq + Hash + ?Sized, + V: Clone, { self.inner.load().get(key).cloned() } + /// Returns whether a value is stored under `key`. + pub fn contains_key(&self, key: &Q) -> bool + where + K: Borrow, + Q: Eq + Hash + ?Sized, + { + self.inner.load().contains_key(key) + } + + /// Return a clone of the keys in the current snapshot. + pub fn keys(&self) -> HashSet + where + K: Clone, + { + self.read(|map| map.keys().cloned().collect()) + } + /// Insert `value` under `key`, replacing any existing value. pub fn insert(&self, key: K, value: V) where K: Clone, + V: Clone, + S: Clone, { self.modify(|map| { map.insert(key.clone(), value.clone()); }); } + + /// Insert `value` under `key` only if no value is stored there yet. + /// + /// If a concurrent writer publishes a value under `key` first, that value + /// is kept and `value` is dropped. `value` is constructed by the caller + /// before this call, so no user code runs while the map is being updated. + pub fn insert_if_absent(&self, key: K, value: V) + where + K: Clone, + V: Clone, + S: Clone, + { + self.modify(|map| { + map.entry(key.clone()).or_insert_with(|| value.clone()); + }); + } } -impl ArcSwapMap> { +impl ArcSwapMap, S> { /// Append `values` to the list stored under `key`, creating it if absent. /// /// Each key maps to an immutable `Arc<[T]>`; appending rebuilds that slice @@ -136,9 +189,19 @@ mod tests { fn get_and_insert() { let map = ArcSwapMap::::default(); assert_eq!(map.get(&1), None); + assert!(!map.contains_key(&1)); map.insert(1, 10); map.insert(1, 20); assert_eq!(map.get(&1), Some(20)); + assert!(map.contains_key(&1)); + } + + #[test] + fn insert_if_absent_keeps_first_value() { + let map = ArcSwapMap::::default(); + map.insert_if_absent(1, 10); + map.insert_if_absent(1, 20); + assert_eq!(map.get(&1), Some(10)); } #[test] diff --git a/vortex-session/src/lib.rs b/vortex-session/src/lib.rs index f69dbc8b3ee..a2f6f44ba60 100644 --- a/vortex-session/src/lib.rs +++ b/vortex-session/src/lib.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +mod arc_swap_map; pub mod registry; mod session; @@ -8,6 +9,7 @@ use std::any::Any; use std::fmt::Debug; use std::hash::Hasher; +pub use arc_swap_map::ArcSwapMap; pub use session::SessionGuard; pub use session::SessionMut; pub use session::VortexSession; diff --git a/vortex-session/src/registry.rs b/vortex-session/src/registry.rs index e75ce01d214..330ae01f07a 100644 --- a/vortex-session/src/registry.rs +++ b/vortex-session/src/registry.rs @@ -1,9 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Many session types use a registry of objects that can be looked up by name to construct -//! contexts. This module provides a generic registry type for that purpose. - use std::cmp::Ordering; use std::fmt; use std::fmt::Debug; @@ -20,7 +17,6 @@ use lasso::ThreadedRodeo; 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. @@ -147,56 +143,6 @@ impl Deref for CachedId { } } -/// A registry of items that are keyed by a string identifier. -#[derive(Clone, Debug)] -pub struct Registry(Arc>); - -impl Default for Registry { - fn default() -> Self { - Self(Default::default()) - } -} - -impl Registry { - pub fn empty() -> Self { - Self(Default::default()) - } - - /// List the IDs in the registry. - pub fn ids(&self) -> impl Iterator + '_ { - self.0.iter().map(|i| *i.key()) - } - - /// List the items in the registry. - pub fn items(&self) -> impl Iterator + '_ { - self.0.iter().map(|i| i.value().clone()) - } - - /// Return the items with the given IDs. - pub fn find_many<'a>( - &self, - ids: impl IntoIterator, - ) -> impl Iterator>> { - ids.into_iter().map(|id| self.0.get(id)) - } - - /// Find the item with the given ID. - pub fn find(&self, id: &Id) -> Option { - self.0.get(id).as_deref().cloned() - } - - /// Register a new item, replacing any existing item with the same ID. - pub fn register(&self, id: impl Into, item: impl Into) { - self.0.insert(id.into(), item.into()); - } - - /// Register a new item, replacing any existing item with the same ID, and return self for - pub fn with(self, id: impl Into, item: impl Into) -> Self { - self.register(id, item.into()); - self - } -} - /// A [`ReadContext`] holds a set of interned IDs for use during deserialization, mapping /// u16 indices to IDs. #[derive(Clone, Debug)] @@ -220,60 +166,53 @@ impl ReadContext { } } -/// A [`Context`] holds a set of interned IDs for use during serialization/deserialization, mapping -/// IDs to u16 indices. +/// An [`Interner`] holds a set of interned IDs for use during serialization/deserialization, +/// mapping IDs to u16 indices. /// /// ## Upcoming Changes /// -/// 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. -#[derive(Clone, Debug)] -pub struct Context { +/// 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. +#[derive(Clone, Debug, Default)] +pub struct Interner { // 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>>, + // Optional set of permissible IDs; when present, only these may be interned. + allowed: Option>>, } -impl Default for Context { - fn default() -> Self { - Self { - ids: Arc::new(RwLock::new(Vec::new())), - valid_ids: None, - } - } -} - -impl Context { - /// Create a context with the given initial IDs. +impl Interner { + /// Create an interner with the given initial IDs. pub fn new(ids: Vec) -> Self { Self { ids: Arc::new(RwLock::new(ids)), - valid_ids: None, + allowed: None, } } - /// Create an empty context. + /// Create an empty interner. pub fn empty() -> Self { Self::default() } - /// 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())); + /// Restrict the permissible set of interned IDs to `allowed`. + /// + /// The set is snapshotted at this call: IDs registered elsewhere afterwards are not + /// permitted. + pub fn with_allowed_ids(mut self, allowed: HashSet) -> Self { + self.allowed = Some(Arc::new(allowed)); self } /// Intern an ID, returning its index. pub fn intern(&self, id: &Id) -> Option { - if let Some(valid_ids) = &self.valid_ids - && !valid_ids.contains(id) + if let Some(allowed) = &self.allowed + && !allowed.contains(id) { - // ID is not valid, cannot intern. + // ID not permitted, cannot intern. return None; } @@ -299,8 +238,10 @@ impl Context { #[cfg(test)] mod tests { + use vortex_utils::aliases::hash_set::HashSet; + use super::CachedId; - use super::Context; + use super::Interner; static VALID: CachedId = CachedId::new("vortex.test.valid"); static INVALID: CachedId = CachedId::new("vortex.test.invalid"); @@ -309,7 +250,7 @@ mod tests { fn context_filters_interned_ids() { let valid = *VALID; let invalid = *INVALID; - let context = Context::empty().with_valid_ids([valid]); + let context = Interner::empty().with_allowed_ids(HashSet::from([valid])); assert_eq!(context.intern(&valid), Some(0)); assert_eq!(context.intern(&valid), Some(0)); diff --git a/vortex-session/src/session.rs b/vortex-session/src/session.rs index 5745842f3d7..ac88a971716 100644 --- a/vortex-session/src/session.rs +++ b/vortex-session/src/session.rs @@ -4,7 +4,7 @@ //! The [`VortexSession`] container. //! //! A [`VortexSession`] is a type-map of [`VortexSessionVar`]s keyed by [`TypeId`]. It is backed by -//! an [`ArcSwap`], giving lock-free reads and copy-on-write writes: +//! an [`ArcSwapMap`], giving lock-free reads and copy-on-write writes: //! //! * **Reads** ([`SessionExt::get`], [`SessionExt::get_opt`]) load the current snapshot //! without taking any lock and hand back a [`SessionGuard`] that derefs to the variable. Because a @@ -13,11 +13,11 @@ //! //! * **Writes** ([`VortexSession::with_some`], [`SessionExt::get`] on a missing default) are //! copy-on-write: the map is cloned, the change applied to the private copy, and the new map -//! atomically published via [`ArcSwap::rcu`]. The closure passed to `rcu` only clones the map and -//! inserts an already-constructed value, so no user code (in particular, no `Default::default` -//! implementation) ever runs while a lock is held. This is the key difference from the previous -//! `DashMap`-backed session, where `entry().or_insert_with(f)` ran `f` while holding the shard's -//! write lock and could deadlock if `f` re-entered the session. +//! atomically published. The value is constructed *before* the map is updated, so no user code +//! (in particular, no `Default::default` implementation) ever runs while a lock is held. This is +//! the key difference from the previous `DashMap`-backed session, where +//! `entry().or_insert_with(f)` ran `f` while holding the shard's write lock and could deadlock if +//! `f` re-entered the session. //! //! A modified session is produced by mutating it **in place**: [`VortexSession::with_some`] — and //! the configuration `with_*` helpers built on it — apply their change copy-on-write to the shared @@ -39,12 +39,12 @@ use std::ops::Deref; use std::ops::DerefMut; use std::sync::Arc; -use arc_swap::ArcSwap; use arc_swap::Guard; use vortex_error::VortexExpect; use vortex_error::vortex_panic; use vortex_utils::aliases::hash_map::HashMap; +use crate::ArcSwapMap; use crate::IdHasher; use crate::SessionExt; use crate::SessionVar; @@ -63,8 +63,11 @@ pub trait VortexSessionVar: SessionVar {} impl VortexSessionVar for V {} +/// The hasher for the session type-map; [`TypeId`]s are already hashes. +type IdHashBuilder = BuildHasherDefault; + /// The immutable type-map backing a published [`VortexSession`] snapshot. -type SessionVars = HashMap, BuildHasherDefault>; +type SessionVars = HashMap, IdHashBuilder>; /// A reference to a session variable of type `V`, returned by [`SessionExt::get`] and /// [`SessionExt::get_opt`]. @@ -155,38 +158,30 @@ impl Debug for SessionMut<'_, V> { /// clone (via [`VortexSession::with_some`] or one of the `with_*` helpers) is observed by all /// clones. To build an *independent* session, start from [`VortexSession::empty`]. #[derive(Clone)] -pub struct VortexSession(Arc>); +pub struct VortexSession(ArcSwapMap, IdHashBuilder>); impl Debug for VortexSession { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - f.debug_tuple("VortexSession") - .field(&self.0.load().as_ref()) - .finish() + self.0 + .read(|vars| f.debug_tuple("VortexSession").field(vars).finish()) } } impl VortexSession { /// Create a new [`VortexSession`] with no session state. pub fn empty() -> Self { - Self(Arc::new(ArcSwap::from_pointee(SessionVars::default()))) + Self(ArcSwapMap::default()) } /// Inserts `V::default()` if no variable of type `V` is present yet, copy-on-write. /// - /// The default is constructed *outside* the [`ArcSwap::rcu`] closure, so `V::default()` never - /// runs under a lock and is never run more than once — it may therefore freely re-enter the - /// session (read or even register other variables) without risk of deadlock. The closure only - /// clones the map and inserts the already-built value, so if a concurrent writer published a - /// value (or the closure is retried under contention) the first-published value is kept and the - /// prebuilt default is simply dropped. + /// The default is constructed *before* the map is updated, so `V::default()` never runs under a + /// lock and is never run more than once — it may therefore freely re-enter the session (read or + /// even register other variables) without risk of deadlock. If a concurrent writer published a + /// value first, that value is kept and the prebuilt default is simply dropped. fn insert_default(&self) { let default: Arc = Arc::new(V::default()); - self.0.rcu(|current| { - let mut next = SessionVars::clone(current); - next.entry(TypeId::of::()) - .or_insert_with(|| Arc::clone(&default)); - next - }); + self.0.insert_if_absent(TypeId::of::(), default); } /// Inserts a session variable of type `V`, replacing any existing variable of that type. @@ -196,12 +191,7 @@ impl VortexSession { /// through those (or through a default inserted by [`get`](SessionExt::get)). The mutation is /// applied in place to the shared backing store, so it is visible through every clone. pub fn register(&self, var: V) { - let var: Arc = Arc::new(var); - self.0.rcu(|current| { - let mut next = SessionVars::clone(current); - next.insert(TypeId::of::(), Arc::clone(&var)); - next - }); + self.0.insert(TypeId::of::(), Arc::new(var)); } /// Inserts a new session variable of type `V` with its default value, mutating this session in From de06de127f624906d173d930c47e1cce782284ce Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Fri, 24 Jul 2026 22:43:21 +0100 Subject: [PATCH 2/3] less Signed-off-by: Robert Kruszewski --- vortex-layout/src/session.rs | 29 +++++++++++++++-------------- 1 file changed, 15 insertions(+), 14 deletions(-) diff --git a/vortex-layout/src/session.rs b/vortex-layout/src/session.rs index 059b313844f..304a63730d7 100644 --- a/vortex-layout/src/session.rs +++ b/vortex-layout/src/session.rs @@ -11,12 +11,12 @@ use vortex_session::registry::Id; use crate::LayoutEncoding; use crate::LayoutEncodingRef; -use crate::layouts::chunked::ChunkedLayoutEncoding; -use crate::layouts::dict::DictLayoutEncoding; -use crate::layouts::flat::FlatLayoutEncoding; -use crate::layouts::list::ListLayoutEncoding; -use crate::layouts::struct_::StructLayoutEncoding; -use crate::layouts::zoned::{LegacyStatsLayoutEncoding, ZonedLayoutEncoding}; +use crate::layouts::chunked::Chunked; +use crate::layouts::dict::Dict; +use crate::layouts::flat::Flat; +use crate::layouts::list::List; +use crate::layouts::struct_::Struct; +use crate::layouts::zoned::{LegacyStats, Zoned}; /// Session state for layout encodings. #[derive(Clone, Debug)] @@ -26,7 +26,8 @@ pub struct LayoutSession { impl LayoutSession { /// Register a layout encoding in the session, replacing any existing encoding with the same ID. - pub fn register(&self, layout: LayoutEncodingRef) { + pub fn register(&self, layout_ref: impl Into) { + let layout = layout_ref.into(); self.registry.insert(layout.id(), layout); } @@ -50,13 +51,13 @@ impl Default for LayoutSession { }; // Register the built-in layout encodings. - this.register(ChunkedLayoutEncoding.as_ref().into()); - this.register(FlatLayoutEncoding.as_ref().into()); - this.register(StructLayoutEncoding.as_ref().into()); - this.register(ZonedLayoutEncoding.as_ref().into()); - this.register(LegacyStatsLayoutEncoding.as_ref().into()); - this.register(DictLayoutEncoding.as_ref().into()); - this.register(ListLayoutEncoding.as_ref().into()); + this.register(&Chunked as &dyn LayoutEncoding); + this.register(&Flat as &dyn LayoutEncoding); + this.register(&Struct as &dyn LayoutEncoding); + this.register(&Zoned as &dyn LayoutEncoding); + this.register(&LegacyStats as &dyn LayoutEncoding); + this.register(&Dict as &dyn LayoutEncoding); + this.register(&List as &dyn LayoutEncoding); this } } From 4a7a22577889f5e94b32c98898e8fd9fdf78a394 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Fri, 24 Jul 2026 23:02:11 +0100 Subject: [PATCH 3/3] nits Signed-off-by: Robert Kruszewski --- vortex-layout/src/session.rs | 3 ++- vortex/src/editions/tests.rs | 2 +- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/vortex-layout/src/session.rs b/vortex-layout/src/session.rs index 304a63730d7..d7cfd12654d 100644 --- a/vortex-layout/src/session.rs +++ b/vortex-layout/src/session.rs @@ -16,7 +16,8 @@ use crate::layouts::dict::Dict; use crate::layouts::flat::Flat; use crate::layouts::list::List; use crate::layouts::struct_::Struct; -use crate::layouts::zoned::{LegacyStats, Zoned}; +use crate::layouts::zoned::LegacyStats; +use crate::layouts::zoned::Zoned; /// Session state for layout encodings. #[derive(Clone, Debug)] diff --git a/vortex/src/editions/tests.rs b/vortex/src/editions/tests.rs index 1da24b904ef..fcd8cfcdacc 100644 --- a/vortex/src/editions/tests.rs +++ b/vortex/src/editions/tests.rs @@ -141,7 +141,7 @@ fn core_edition_ids_are_registered_array_encodings() { let registry = session.arrays().registry().clone(); for inclusion in session.editions().encodings_in(&CORE_2026_07_0) { assert!( - registry.find(&inclusion.encoding_id).is_some(), + registry.contains_key(&inclusion.encoding_id), "{} is declared in core but not registered as an array encoding", inclusion.encoding_id );