Layout readers don't need a mutable context (#6808)

When reading the set of array ids is initialised from the file and then is
immutable for whole scan duration. We don't need to share the mutable
version with immutable one.

---------

Signed-off-by: Robert Kruszewski <github@robertk.io>

Author: robert3005

encodings/fastlanes/src/rle/array/mod.rs

@@ -229,6 +229,7 @@ mod tests {
     use vortex_array::validity::Validity;
     use vortex_buffer::Buffer;
     use vortex_buffer::ByteBufferMut;
+    use vortex_session::registry::ReadContext;
 
     use crate::RLEArray;
     use crate::test::SESSION;
@@ -458,7 +459,7 @@ mod tests {
             .decode(
                 &DType::Primitive(PType::U32, Nullability::NonNullable),
                 2048,
-                &ctx,
+                &ReadContext::new(ctx.to_ids()),
                 &SESSION,
             )
             .unwrap();
@@ -498,7 +499,12 @@ mod tests {
 
         let parts = ArrayParts::try_from(concat).unwrap();
         let decoded = parts
-            .decode(sliced.dtype(), sliced.len(), &ctx, &SESSION)
+            .decode(
+                sliced.dtype(),
+                sliced.len(),
+                &ReadContext::new(ctx.to_ids()),
+                &SESSION,
+            )
             .unwrap();
 
         let original_data = sliced.to_primitive();

encodings/pco/src/test.rs

@@ -27,6 +27,7 @@ use vortex_buffer::BufferMut;
 use vortex_error::VortexResult;
 use vortex_mask::Mask;
 use vortex_session::VortexSession;
+use vortex_session::registry::ReadContext;
 
 static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
     let session = VortexSession::empty().with::<ArraySession>();
@@ -166,7 +167,7 @@ fn test_serde() -> VortexResult<()> {
     let decoded = parts.decode(
         &DType::Primitive(PType::I32, Nullability::NonNullable),
         1_000_000,
-        &context,
+        &ReadContext::new(context.to_ids()),
         &SESSION,
     )?;
     let mut ctx = SESSION.create_execution_ctx();

vortex-array/public-api.lock

@@ -15982,7 +15982,7 @@ pub fn vortex_array::serde::ArrayParts::buffer_lengths(&self) -> alloc::vec::Vec
 
 pub fn vortex_array::serde::ArrayParts::child(&self, idx: usize) -> vortex_array::serde::ArrayParts
 
-pub fn vortex_array::serde::ArrayParts::decode(&self, dtype: &vortex_array::dtype::DType, len: usize, ctx: &vortex_array::ArrayContext, session: &vortex_session::VortexSession) -> vortex_error::VortexResult<vortex_array::ArrayRef>
+pub fn vortex_array::serde::ArrayParts::decode(&self, dtype: &vortex_array::dtype::DType, len: usize, ctx: &vortex_session::registry::ReadContext, session: &vortex_session::VortexSession) -> vortex_error::VortexResult<vortex_array::ArrayRef>
 
 pub fn vortex_array::serde::ArrayParts::encoding_id(&self) -> u16
 

vortex-array/src/arrays/decimal/vtable/mod.rs

@@ -239,6 +239,7 @@ impl DecimalVTable {
 mod tests {
     use vortex_buffer::ByteBufferMut;
     use vortex_buffer::buffer;
+    use vortex_session::registry::ReadContext;
 
     use crate::ArrayContext;
     use crate::IntoArray;
@@ -273,7 +274,9 @@ mod tests {
         let concat = concat.freeze();
 
         let parts = ArrayParts::try_from(concat).unwrap();
-        let decoded = parts.decode(&dtype, 5, &ctx, &LEGACY_SESSION).unwrap();
+        let decoded = parts
+            .decode(&dtype, 5, &ReadContext::new(ctx.to_ids()), &LEGACY_SESSION)
+            .unwrap();
         assert!(decoded.is::<DecimalVTable>());
     }
 }

vortex-array/src/arrays/masked/vtable/mod.rs

@@ -217,6 +217,7 @@ mod tests {
     use rstest::rstest;
     use vortex_buffer::ByteBufferMut;
     use vortex_error::VortexError;
+    use vortex_session::registry::ReadContext;
 
     use crate::ArrayContext;
     use crate::Canonical;
@@ -269,7 +270,14 @@ mod tests {
         let concat = concat.freeze();
 
         let parts = ArrayParts::try_from(concat).unwrap();
-        let decoded = parts.decode(&dtype, len, &ctx, &LEGACY_SESSION).unwrap();
+        let decoded = parts
+            .decode(
+                &dtype,
+                len,
+                &ReadContext::new(ctx.to_ids()),
+                &LEGACY_SESSION,
+            )
+            .unwrap();
 
         assert!(decoded.is::<MaskedVTable>());
         assert_eq!(

vortex-array/src/context.rs

@@ -18,14 +18,15 @@ use std::sync::LazyLock;
 pub use array::*;
 pub use canonical::*;
 pub use columnar::*;
-pub use context::*;
 pub use executor::*;
 pub use hash::*;
 pub use mask_future::*;
 pub use metadata::*;
 use vortex_session::VortexSession;
+use vortex_session::registry::Context;
 
 use crate::session::ArraySession;
+use crate::vtable::DynVTable;
 
 pub mod accessor;
 #[doc(hidden)]
@@ -39,7 +40,6 @@ pub mod builtins;
 mod canonical;
 mod columnar;
 pub mod compute;
-mod context;
 pub mod display;
 pub mod dtype;
 mod executor;
@@ -80,3 +80,5 @@ pub mod flatbuffers {
 //  here...
 pub static LEGACY_SESSION: LazyLock<VortexSession> =
     LazyLock::new(|| VortexSession::empty().with::<ArraySession>());
+
+pub type ArrayContext = Context<&'static dyn DynVTable>;

vortex-array/src/lib.rs

@@ -24,6 +24,7 @@ use vortex_flatbuffers::WriteFlatBuffer;
 use vortex_flatbuffers::array as fba;
 use vortex_flatbuffers::array::Compression;
 use vortex_session::VortexSession;
+use vortex_session::registry::ReadContext;
 use vortex_utils::aliases::hash_map::HashMap;
 
 use crate::ArrayContext;
@@ -318,7 +319,7 @@ impl ArrayParts {
         &self,
         dtype: &DType,
         len: usize,
-        ctx: &ArrayContext,
+        ctx: &ReadContext,
         session: &VortexSession,
     ) -> VortexResult<ArrayRef> {
         let encoding_idx = self.flatbuffer().encoding();
@@ -623,7 +624,7 @@ impl ArrayParts {
 
 struct ArrayPartsChildren<'a> {
     parts: &'a ArrayParts,
-    ctx: &'a ArrayContext,
+    ctx: &'a ReadContext,
     session: &'a VortexSession,
 }
 

vortex-array/src/serde.rs

@@ -64,6 +64,7 @@ use vortex::scalar::ScalarTruncation;
 use vortex::scalar::lower_bound;
 use vortex::scalar::upper_bound;
 use vortex::session::VortexSession;
+use vortex::session::registry::ReadContext;
 use vortex::utils::aliases::hash_map::HashMap;
 
 /// A buffer inlined into layout metadata for host-side access.
@@ -94,7 +95,7 @@ pub struct CudaFlatLayout {
     row_count: u64,
     dtype: DType,
     segment_id: SegmentId,
-    ctx: ArrayContext,
+    ctx: ReadContext,
     array_tree: ByteBuffer,
     /// Small buffers kept on host, keyed by global buffer index.
     host_buffers: Arc<HashMap<u32, ByteBuffer>>,
@@ -107,7 +108,7 @@ impl CudaFlatLayout {
     }
 
     #[inline]
-    pub fn array_ctx(&self) -> &ArrayContext {
+    pub fn array_ctx(&self) -> &ReadContext {
         &self.ctx
     }
 
@@ -195,7 +196,7 @@ impl VTable for CudaFlatVTable {
         metadata: &<Self::Metadata as DeserializeMetadata>::Output,
         segment_ids: Vec<SegmentId>,
         _children: &dyn LayoutChildren,
-        ctx: &ArrayContext,
+        ctx: &ReadContext,
     ) -> VortexResult<Self::Layout> {
         if segment_ids.len() != 1 {
             vortex_bail!("CudaFlatLayout must have exactly one segment ID");
@@ -535,7 +536,7 @@ impl LayoutStrategy for CudaFlatLayoutStrategy {
             row_count,
             dtype: stream.dtype().clone(),
             segment_id,
-            ctx: ctx.clone(),
+            ctx: ReadContext::new(ctx.to_ids()),
             array_tree,
             host_buffers: Arc::new(host_buffer_map),
         }

vortex-cuda/src/layout.rs

@@ -11,12 +11,11 @@ use std::sync::Arc;
 
 use flatbuffers::FlatBufferBuilder;
 use flatbuffers::WIPOffset;
-use vortex_array::ArrayContext;
 use vortex_error::VortexResult;
 use vortex_flatbuffers::FlatBufferRoot;
 use vortex_flatbuffers::WriteFlatBuffer;
 use vortex_flatbuffers::footer as fb;
-use vortex_layout::LayoutContext;
+use vortex_session::registry::ReadContext;
 
 use crate::footer::segment::SegmentSpec;
 
@@ -26,9 +25,9 @@ use crate::footer::segment::SegmentSpec;
 /// which describes the structure of the data in the file.
 pub(crate) struct FooterFlatBufferWriter {
     /// The array context containing encodings used in the file.
-    pub(crate) ctx: ArrayContext,
+    pub(crate) ctx: ReadContext,
     /// The layout context containing the layouts used in the file.
-    pub(crate) layout_ctx: LayoutContext,
+    pub(crate) layout_ctx: ReadContext,
     /// Specifications for all segments in the file.
     pub(crate) segment_specs: Arc<[SegmentSpec]>,
 }
@@ -47,7 +46,7 @@ impl WriteFlatBuffer for FooterFlatBufferWriter {
 
         let array_specs = self
             .ctx
-            .to_ids()
+            .ids()
             .iter()
             .map(|e| {
                 let id = fbb.create_string(e.as_ref());
@@ -58,7 +57,7 @@ impl WriteFlatBuffer for FooterFlatBufferWriter {
 
         let layout_specs = self
             .layout_ctx
-            .to_ids()
+            .ids()
             .iter()
             .map(|e| {
                 let id = fbb.create_string(e.as_ref());