Immutable session + registries

bench-vortex/src/lib.rs

@@ -48,15 +48,16 @@ pub use datasets::file;
 pub use engines::df;
 use vortex::VortexSessionDefault;
 pub use vortex::error::vortex_panic;
-use vortex::io::session::RuntimeSessionExt;
+use vortex::io::session::RuntimeSessionMutExt;
 use vortex::session::VortexSession;
+use vortex::session::VortexSessionRef;
 
 // All benchmarks run with mimalloc for consistency.
 #[global_allocator]
 static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
 
-pub static SESSION: LazyLock<VortexSession> =
-    LazyLock::new(|| VortexSession::default().with_tokio());
+pub static SESSION: LazyLock<VortexSessionRef> =
+    LazyLock::new(|| VortexSession::new_with_defaults().with_tokio().freeze());
 
 #[derive(Clone, Copy, Debug, Hash, PartialEq, Eq, Serialize)]
 pub struct Target {

encodings/alp/src/alp/array.rs

@@ -460,11 +460,12 @@ mod tests {
     use vortex_array::vtable::ValidityHelper;
     use vortex_dtype::PTypeDowncast;
     use vortex_session::VortexSession;
+    use vortex_session::VortexSessionRef;
     use vortex_vector::VectorOps;
 
     use super::*;
 
-    static SESSION: LazyLock<VortexSession> = LazyLock::new(VortexSession::empty);
+    static SESSION: LazyLock<VortexSessionRef> = LazyLock::new(|| VortexSession::empty().freeze());
 
     #[rstest]
     #[case(0)]

encodings/fastlanes/src/bitpacking/array/bitpack_decompress.rs

@@ -211,14 +211,15 @@ mod tests {
     use vortex_buffer::buffer;
     use vortex_dtype::Nullability;
     use vortex_session::VortexSession;
+    use vortex_session::VortexSessionRef;
     use vortex_vector::VectorMutOps;
     use vortex_vector::VectorOps;
 
     use super::*;
     use crate::BitPackedVTable;
     use crate::bitpack_compress::bitpack_encode;
 
-    static SESSION: LazyLock<VortexSession> = LazyLock::new(VortexSession::empty);
+    static SESSION: LazyLock<VortexSessionRef> = LazyLock::new(|| VortexSession::empty().freeze());
 
     fn compression_roundtrip(n: usize) {
         let values = PrimitiveArray::from_iter((0..n).map(|i| (i % 2047) as u16));

fuzz/src/lib.rs

@@ -26,12 +26,16 @@ mod native_runtime {
     use vortex::VortexSessionDefault;
     use vortex_io::runtime::BlockingRuntime;
     use vortex_io::runtime::current::CurrentThreadRuntime;
-    use vortex_io::session::RuntimeSessionExt;
+    use vortex_io::session::RuntimeSessionMutExt;
     use vortex_session::VortexSession;
+    use vortex_session::VortexSessionRef;
 
     pub static RUNTIME: LazyLock<CurrentThreadRuntime> = LazyLock::new(CurrentThreadRuntime::new);
-    pub static SESSION: LazyLock<VortexSession> =
-        LazyLock::new(|| VortexSession::default().with_handle(RUNTIME.handle()));
+    pub static SESSION: LazyLock<VortexSessionRef> = LazyLock::new(|| {
+        VortexSession::new_with_defaults()
+            .with_handle(RUNTIME.handle())
+            .freeze()
+    });
 }
 
 #[cfg(not(target_arch = "wasm32"))]
@@ -46,10 +50,10 @@ mod wasm_runtime {
     use vortex::VortexSessionDefault;
     use vortex_io::runtime::wasm::WasmRuntime;
     use vortex_io::session::RuntimeSessionExt;
-    use vortex_session::VortexSession;
+    use vortex_session::VortexSessionRef;
 
-    pub static SESSION: LazyLock<VortexSession> =
-        LazyLock::new(|| VortexSession::default().with_handle(WasmRuntime::handle()));
+    pub static SESSION: LazyLock<VortexSessionRef> =
+        LazyLock::new(|| VortexSession::empty().with_handle(WasmRuntime::handle()));
 }
 
 #[cfg(target_arch = "wasm32")]

java/testfiles/src/main.rs

@@ -6,13 +6,17 @@
 use std::path::Path;
 
 use vortex::array::arrays::StructArray;
-use vortex::array::builders::{ArrayBuilder, DecimalBuilder, VarBinViewBuilder};
+use vortex::array::builders::ArrayBuilder;
+use vortex::array::builders::DecimalBuilder;
+use vortex::array::builders::VarBinViewBuilder;
 use vortex::array::validity::Validity;
-use vortex::dtype::{DType, DecimalDType, Nullability};
+use vortex::dtype::DType;
+use vortex::dtype::DecimalDType;
+use vortex::dtype::Nullability;
 use vortex::file::WriteOptionsSessionExt;
 use vortex::io::runtime::current::CurrentThreadRuntime;
 use vortex::io::runtime::BlockingRuntime;
-use vortex::io::session::RuntimeSessionExt;
+use vortex::io::session::RuntimeSessionMutExt;
 use vortex::session::VortexSession;
 use vortex::VortexSessionDefault;
 
@@ -32,7 +36,9 @@ use vortex::VortexSessionDefault;
 /// | John   | 10000  | VA    |
 fn main() {
     let runtime = CurrentThreadRuntime::new();
-    let session = VortexSession::default().with_handle(runtime.handle());
+    let session = VortexSession::new_with_defaults()
+        .with_handle(runtime.handle())
+        .freeze();
 
     let mut names = VarBinViewBuilder::with_capacity(DType::Utf8(Nullability::NonNullable), 10);
     names.append_value("Alice");

vortex-array/src/array/mod.rs

@@ -22,7 +22,7 @@ use vortex_error::vortex_ensure;
 use vortex_error::vortex_panic;
 use vortex_mask::Mask;
 use vortex_scalar::Scalar;
-use vortex_session::VortexSession;
+use vortex_session::VortexSessionRef;
 use vortex_vector::Vector;
 use vortex_vector::VectorOps;
 
@@ -366,7 +366,7 @@ impl dyn Array + '_ {
     /// Execute the array and return the resulting vector.
     ///
     /// This entry-point function will choose an appropriate CPU-based execution strategy.
-    pub fn execute(&self, session: &VortexSession) -> VortexResult<Vector> {
+    pub fn execute(&self, session: &VortexSessionRef) -> VortexResult<Vector> {
         let mut ctx = ExecutionCtx::new(session.clone());
         self.batch_execute(&mut ctx)
     }

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

@@ -19,6 +19,7 @@ use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
 use vortex_error::vortex_ensure;
 use vortex_session::VortexSession;
+use vortex_session::VortexSessionRef;
 use vortex_vector::Vector;
 
 use crate::Array;
@@ -43,8 +44,8 @@ use crate::vtable::VTable;
 // TODO(ngates): canonicalize doesn't currently take a session, therefore we cannot dispatch
 //  to registered scalar function kernels. We therefore hold our own non-pluggable session here
 //  that contains all the built-in kernels while we migrate over to "execute" instead of canonicalize.
-static SCALAR_FN_SESSION: LazyLock<VortexSession> =
-    LazyLock::new(|| VortexSession::empty().with::<ArraySession>());
+static SCALAR_FN_SESSION: LazyLock<VortexSessionRef> =
+    LazyLock::new(|| VortexSession::empty().with::<ArraySession>().freeze());
 
 vtable!(ScalarFn);
 

vortex-array/src/execution/mod.rs

@@ -7,24 +7,24 @@ mod validity;
 
 pub use batch::*;
 pub use mask::*;
-use vortex_session::VortexSession;
+use vortex_session::VortexSessionRef;
 
 /// Execution context for batch array compute.
 // NOTE(ngates): This context will eventually hold cached resources for execution, such as CSE
 //  nodes, and may well eventually support a type-map interface for arrays to stash arbitrary
 //  execution-related data.
 pub struct ExecutionCtx {
-    session: VortexSession,
+    session: VortexSessionRef,
 }
 
 impl ExecutionCtx {
     /// Create a new execution context with the given session.
-    pub(crate) fn new(session: VortexSession) -> Self {
+    pub(crate) fn new(session: VortexSessionRef) -> Self {
         Self { session }
     }
 
     /// Get the session associated with this execution context.
-    pub fn session(&self) -> &VortexSession {
+    pub fn session(&self) -> &VortexSessionRef {
         &self.session
     }
 }

vortex-array/src/expr/functions/session.rs

@@ -1,7 +1,6 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-use vortex_session::Ref;
 use vortex_session::SessionExt;
 use vortex_session::registry::Registry;
 
@@ -19,7 +18,7 @@ impl FunctionSession {
 }
 
 pub trait ScalarFuncSessionExt: SessionExt {
-    fn functions(&self) -> Ref<'_, FunctionSession> {
+    fn functions(&self) -> &FunctionSession {
         self.get::<FunctionSession>()
     }
 }

vortex-array/src/expr/session/mod.rs

@@ -4,7 +4,6 @@
 mod rewrite;
 
 pub use rewrite::RewriteRuleRegistry;
-use vortex_session::Ref;
 use vortex_session::SessionExt;
 use vortex_session::registry::Registry;
 
@@ -54,12 +53,12 @@ impl ExprSession {
     }
 
     /// Register an expression vtable in the session, replacing any existing vtable with the same ID.
-    pub fn register(&self, expr: ExprVTable) {
+    pub fn register(&mut self, expr: ExprVTable) {
         self.registry.register(expr)
     }
 
     /// Register expression vtables in the session, replacing any existing vtables with the same IDs.
-    pub fn register_many(&self, exprs: impl IntoIterator<Item = ExprVTable>) {
+    pub fn register_many(&mut self, exprs: impl IntoIterator<Item = ExprVTable>) {
         self.registry.register_many(exprs);
     }
 
@@ -151,7 +150,7 @@ impl ExprSession {
 
 impl Default for ExprSession {
     fn default() -> Self {
-        let expressions = ExprRegistry::default();
+        let mut expressions = ExprRegistry::default();
 
         // Register built-in expressions here if needed.
         expressions.register_many([
@@ -187,7 +186,7 @@ impl Default for ExprSession {
 /// Extension trait for accessing expression session data.
 pub trait ExprSessionExt: SessionExt {
     /// Returns the expression vtable registry.
-    fn expressions(&self) -> Ref<'_, ExprSession> {
+    fn expressions(&self) -> &ExprSession {
         self.get::<ExprSession>()
     }
 }

vortex-array/src/session/mod.rs

@@ -1,7 +1,6 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-use vortex_session::Ref;
 use vortex_session::SessionExt;
 use vortex_session::registry::Registry;
 
@@ -52,19 +51,19 @@ impl ArraySession {
     }
 
     /// Register a new array encoding, replacing any existing encoding with the same ID.
-    pub fn register(&self, encoding: ArrayVTable) {
+    pub fn register(&mut self, encoding: ArrayVTable) {
         self.registry.register(encoding)
     }
 
     /// Register many array encodings, replacing any existing encodings with the same ID.
-    pub fn register_many(&self, encodings: impl IntoIterator<Item = ArrayVTable>) {
+    pub fn register_many(&mut self, encodings: impl IntoIterator<Item = ArrayVTable>) {
         self.registry.register_many(encodings);
     }
 }
 
 impl Default for ArraySession {
     fn default() -> Self {
-        let encodings = ArrayRegistry::default();
+        let mut encodings = ArrayRegistry::default();
 
         // Register the canonical encodings.
         encodings.register_many([
@@ -108,7 +107,7 @@ impl Default for ArraySession {
 /// Session data for Vortex arrays.
 pub trait ArraySessionExt: SessionExt {
     /// Returns the array encoding registry.
-    fn arrays(&self) -> Ref<'_, ArraySession> {
+    fn arrays(&self) -> &ArraySession {
         self.get::<ArraySession>()
     }
 }

vortex-cxx/src/lib.rs

@@ -17,8 +17,9 @@ use scalar::*;
 use vortex::VortexSessionDefault;
 use vortex::io::runtime::BlockingRuntime;
 use vortex::io::runtime::current::CurrentThreadRuntime;
-use vortex::io::session::RuntimeSessionExt;
+use vortex::io::session::RuntimeSessionMutExt;
 use vortex::session::VortexSession;
+use vortex::session::VortexSessionRef;
 use write::*;
 
 /// By default, the C++ API uses a current-thread runtime, providing control of the threading
@@ -28,8 +29,11 @@ use write::*;
 //  this runtime.
 pub(crate) static RUNTIME: LazyLock<CurrentThreadRuntime> =
     LazyLock::new(CurrentThreadRuntime::new);
-pub(crate) static SESSION: LazyLock<VortexSession> =
-    LazyLock::new(|| VortexSession::default().with_handle(RUNTIME.handle()));
+pub(crate) static SESSION: LazyLock<VortexSessionRef> = LazyLock::new(|| {
+    VortexSession::new_with_defaults()
+        .with_handle(RUNTIME.handle())
+        .freeze()
+});
 
 #[cxx::bridge(namespace = "vortex::ffi")]
 #[allow(let_underscore_drop)]

vortex-datafusion/examples/vortex_table.rs

@@ -19,13 +19,13 @@ use vortex::array::validity::Validity;
 use vortex::buffer::buffer;
 use vortex::error::vortex_err;
 use vortex::file::WriteOptionsSessionExt;
-use vortex::io::session::RuntimeSessionExt;
+use vortex::io::session::RuntimeSessionMutExt;
 use vortex::session::VortexSession;
 use vortex_datafusion::VortexFormat;
 
 #[tokio::main]
 async fn main() -> anyhow::Result<()> {
-    let session = VortexSession::default().with_tokio();
+    let session = VortexSession::new_with_defaults().with_tokio().freeze();
 
     let temp_dir = tempdir()?;
     let strings = ChunkedArray::from_iter([

vortex-datafusion/src/persistent/cache.rs

@@ -25,14 +25,14 @@ use vortex::file::VortexFile;
 use vortex::layout::segments::SegmentCache;
 use vortex::layout::segments::SegmentId;
 use vortex::metrics::MetricsSessionExt;
-use vortex::session::VortexSession;
+use vortex::session::VortexSessionRef;
 use vortex::utils::aliases::DefaultHashBuilder;
 
 #[derive(Clone)]
 pub(crate) struct VortexFileCache {
     file_cache: Cache<FileKey, VortexFile, DefaultHashBuilder>,
     segment_cache: Cache<SegmentKey, ByteBuffer, DefaultHashBuilder>,
-    session: VortexSession,
+    session: VortexSessionRef,
 }
 
 /// Cache key for a [`VortexFile`].
@@ -59,7 +59,7 @@ struct SegmentKey {
 }
 
 impl VortexFileCache {
-    pub fn new(size_mb: usize, segment_size_mb: usize, session: VortexSession) -> Self {
+    pub fn new(size_mb: usize, segment_size_mb: usize, session: VortexSessionRef) -> Self {
         let file_cache = Cache::builder()
             .max_capacity(size_mb as u64 * (1 << 20))
             .eviction_listener(|k: Arc<FileKey>, _v: VortexFile, cause| {