more

Signed-off-by: Adam Gutglick <adam@spiraldb.com>

Author: AdamGS

fuzz/src/lib.rs

@@ -43,22 +43,23 @@ 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::RuntimeSessionBuilderExt;
     use vortex_session::VortexSession;
 
     pub static RUNTIME: LazyLock<CurrentThreadRuntime> = LazyLock::new(CurrentThreadRuntime::new);
     pub static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
         #[allow(unused_mut)]
-        let mut session = VortexSession::default().with_handle(RUNTIME.handle());
+        let mut builder = VortexSession::default_builder().with_handle(RUNTIME.handle());
         #[cfg(all(feature = "cuda", target_os = "linux"))]
         // Even if the CUDA feature is enabled we need to check at
         // runtime whether CUDA is available in the current environment.
+        if vortex_cuda::cuda_available() {
+            builder = builder.with::<vortex_cuda::CudaSession>();
+        }
+        let session = builder.build();
+        #[cfg(all(feature = "cuda", target_os = "linux"))]
         if vortex_cuda::cuda_available() {
             use vortex_cuda::CudaSessionExt;
-            session = session
-                .into_builder()
-                .with::<vortex_cuda::CudaSession>()
-                .build();
             vortex_cuda::initialize_cuda(session.cuda_session());
         }
         session
@@ -78,11 +79,14 @@ mod wasm_runtime {
 
     use vortex::VortexSessionDefault;
     use vortex_io::runtime::wasm::WasmRuntime;
-    use vortex_io::session::RuntimeSessionExt;
+    use vortex_io::session::RuntimeSessionBuilderExt;
     use vortex_session::VortexSession;
 
-    pub static SESSION: LazyLock<VortexSession> =
-        LazyLock::new(|| VortexSession::default().with_handle(WasmRuntime::handle()));
+    pub static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
+        VortexSession::default_builder()
+            .with_handle(WasmRuntime::handle())
+            .build()
+    });
 }
 
 #[cfg(all(target_arch = "wasm32", target_os = "unknown"))]

vortex-array/src/arrays/chunked/vtable/canonical.rs

@@ -320,7 +320,7 @@ mod tests {
     use crate::dtype::PType::I32;
     use crate::memory::DefaultHostAllocator;
     use crate::memory::HostAllocator;
-    use crate::memory::MemorySessionExt;
+    use crate::memory::MemorySessionBuilderExt;
     use crate::memory::WritableHostBuffer;
     use crate::scalar::Scalar;
     use crate::validity::Validity;
@@ -639,10 +639,10 @@ mod tests {
     fn list_canonicalize_uses_memory_session_allocator() {
         let allocations = Arc::new(AtomicUsize::new(0));
         let session = crate::array_session()
-            .build()
             .with_allocator(Arc::new(CountingAllocator {
                 allocations: Arc::clone(&allocations),
-            }));
+            }))
+            .build();
         let mut ctx = ExecutionCtx::new(session);
 
         let l1 = ListArray::try_new(

vortex-array/src/expr/proto.rs

@@ -117,8 +117,8 @@ mod tests {
     fn unknown_expression_id_allow_unknown() {
         let session = VortexSession::builder()
             .with::<ScalarFnSession>()
-            .build()
-            .allow_unknown();
+            .allow_unknown()
+            .build();
 
         let expr_proto = pb::Expr {
             id: "vortex.test.foreign_scalar_fn".to_string(),

vortex-array/src/memory.rs

@@ -18,7 +18,7 @@ use vortex_error::vortex_ensure;
 use vortex_error::vortex_err;
 use vortex_session::SessionExt;
 use vortex_session::SessionVar;
-use vortex_session::VortexSession;
+use vortex_session::VortexSessionBuilder;
 
 /// Mutable host buffer contract used by [`WritableHostBuffer`].
 pub trait HostBufferMut: Send + 'static {
@@ -221,17 +221,23 @@ pub trait MemorySessionExt: SessionExt {
     fn allocator(&self) -> HostAllocatorRef {
         self.memory().allocator()
     }
-
-    /// Returns a new session configured to use `allocator` as its host allocator.
-    fn with_allocator(self, allocator: HostAllocatorRef) -> VortexSession {
-        let mut builder = self.session().into_builder();
-        builder.get_mut::<MemorySession>().set_allocator(allocator);
-        builder.build()
-    }
 }
 
 impl<S: SessionExt> MemorySessionExt for S {}
 
+/// Extension trait for configuring session-scoped memory before a session is built.
+pub trait MemorySessionBuilderExt {
+    /// Configure the host allocator.
+    fn with_allocator(self, allocator: HostAllocatorRef) -> Self;
+}
+
+impl MemorySessionBuilderExt for VortexSessionBuilder {
+    fn with_allocator(mut self, allocator: HostAllocatorRef) -> Self {
+        self.get_mut::<MemorySession>().set_allocator(allocator);
+        self
+    }
+}
+
 /// Default host allocator.
 #[derive(Debug, Default)]
 pub struct DefaultHostAllocator;

vortex-array/src/session/mod.rs

@@ -157,13 +157,14 @@ mod tests {
 
     #[test]
     fn initialize_registers_builtin_kernels_into_empty_array_session() {
-        let session = VortexSession::builder()
-            .with_some(ArraySession::empty())
-            .build();
-
-        assert!(!session.kernels().has_execute_parent(Binary.id(), Bool.id()));
+        let mut builder = VortexSession::builder().with_some(ArraySession::empty());
 
-        let mut builder = session.into_builder();
+        assert!(
+            !builder
+                .get_mut::<ArraySession>()
+                .kernels()
+                .has_execute_parent(Binary.id(), Bool.id())
+        );
         crate::initialize(&mut builder);
         let session = builder.build();
 

vortex-bench/src/lib.rs

@@ -65,15 +65,15 @@ pub use output::BenchmarkOutput;
 pub use output::create_output_writer;
 use vortex::VortexSessionDefault;
 pub use vortex::error::vortex_panic;
-use vortex::io::session::RuntimeSessionExt;
+use vortex::io::session::RuntimeSessionBuilderExt;
 use vortex::session::VortexSession;
 
 // 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());
+    LazyLock::new(|| VortexSession::default_builder().with_tokio().build());
 
 #[derive(Clone, Copy, Debug, Hash, PartialEq, Eq, Serialize, Deserialize)]
 pub struct Target {

vortex-cuda/ffi/src/lib.rs

@@ -14,6 +14,7 @@ use std::ptr;
 use arrow_schema::ffi::FFI_ArrowSchema;
 use vortex::error::VortexResult;
 use vortex::error::vortex_ensure;
+use vortex::error::vortex_err;
 use vortex::session::VortexSession;
 use vortex_cuda::CudaSession;
 use vortex_cuda::arrow::ArrowDeviceArray;
@@ -34,20 +35,14 @@ use vortex_ffi::vx_session_ref;
 const VX_CUDA_OK: c_int = 0;
 const VX_CUDA_ERR: c_int = 1;
 
-/// Return a Vortex session with a [`CudaSession`] session variable.
-///
-/// If `session` already has CUDA support, this returns a clone of it. Otherwise it
-/// returns a new session cloned from `session` with a default [`CudaSession`] attached.
-fn session_with_cuda(session: &VortexSession) -> VortexResult<VortexSession> {
-    if session.get_opt::<CudaSession>().is_some() {
-        return Ok(session.clone());
+fn session_with_cuda(session: &VortexSession) -> VortexResult<&VortexSession> {
+    if session.get_opt::<CudaSession>().is_none() {
+        return Err(vortex_err!(
+            InvalidArgument: "CUDA session state is not configured; create the session with vx_cuda_session_new"
+        ));
     }
 
-    Ok(session
-        .clone()
-        .into_builder()
-        .with_some(CudaSession::try_default()?)
-        .build())
+    Ok(session)
 }
 
 /// Create a CUDA Vortex session.
@@ -65,7 +60,7 @@ pub unsafe extern "C-unwind" fn vx_cuda_session_new(
     try_or(error_out, ptr::null_mut(), || {
         let cuda_session = CudaSession::try_default()?;
         Ok(vx_session_new_with(|session| {
-            session.into_builder().with_some(cuda_session).build()
+            session.with_some(cuda_session)
         }))
     })
 }
@@ -101,7 +96,7 @@ pub unsafe extern "C-unwind" fn vx_cuda_array_export_arrow_device(
 
         let session = session_with_cuda(unsafe { vx_session_ref(session) }?)?;
         let array = unsafe { vx_array_ref(array) }?.clone();
-        let mut ctx = CudaSession::create_execution_ctx(&session)?;
+        let mut ctx = CudaSession::create_execution_ctx(session)?;
         let exported =
             futures::executor::block_on(array.export_device_array_with_schema(&mut ctx))?;
 
@@ -146,7 +141,7 @@ pub unsafe extern "C-unwind" fn vx_cuda_partition_scan_arrow_device_stream(
 
         let session = session_with_cuda(unsafe { vx_session_ref(session) }?)?;
         // Drive the stream on the same runtime the partition's scan spawned its work onto.
-        let device_stream = array_stream.export_device_array_stream(&session, ffi_runtime())?;
+        let device_stream = array_stream.export_device_array_stream(session, ffi_runtime())?;
 
         unsafe { ptr::write(out_stream, device_stream) };
         Ok(VX_CUDA_OK)
@@ -177,6 +172,12 @@ mod tests {
         Box::into_raw(Box::new(session)).cast::<vx_session>()
     }
 
+    fn test_cuda_session() -> VortexResult<VortexSession> {
+        Ok(vortex::array::array_session()
+            .with_some(CudaSession::try_default()?)
+            .build())
+    }
+
     unsafe fn free_test_session(session: *mut vx_session) {
         unsafe { drop(Box::from_raw(session.cast::<VortexSession>())) };
     }
@@ -216,9 +217,9 @@ mod tests {
     }
 
     #[cuda_test]
-    fn test_export_primitive_arrow_device() {
+    fn test_export_primitive_arrow_device() -> VortexResult<()> {
         let mut error = ptr::null_mut();
-        let session = test_session(VortexSession::default());
+        let session = test_session(test_cuda_session()?);
         let array = test_array(PrimitiveArray::from_iter(0u32..5));
         let mut schema = FFI_ArrowSchema::empty();
         let mut device_array = empty_device_array();
@@ -235,7 +236,7 @@ mod tests {
         assert_eq!(status, VX_CUDA_OK);
         assert!(error.is_null());
 
-        let field = Field::try_from(&schema).expect("schema should be a field");
+        let field = Field::try_from(&schema)?;
         assert_eq!(field.name(), "");
         assert_eq!(device_array.array.length, 5);
         assert_eq!(device_array.array.n_buffers, 2);
@@ -249,12 +250,13 @@ mod tests {
             free_test_array(array);
             free_test_session(session);
         }
+        Ok(())
     }
 
     #[cuda_test]
     fn test_export_struct_arrow_device_table() -> VortexResult<()> {
         let mut error = ptr::null_mut();
-        let session = test_session(VortexSession::default());
+        let session = test_session(test_cuda_session()?);
         let array = test_array(StructArray::try_new(
             ["ids", "values"].into(),
             vec![

vortex-cuda/src/hybrid_dispatch/mod.rs

@@ -302,7 +302,7 @@ mod tests {
         use vortex::encodings::fastlanes;
         use vortex::encodings::zstd::ZstdBuffers;
 
-        let mut builder = crate::cuda_session().into_builder();
+        let mut builder = crate::cuda_session_builder();
         fastlanes::initialize(&mut builder);
         builder.get_mut::<ArraySession>().register(ZstdBuffers);
         let session = builder.build();

vortex-cuda/src/lib.rs

@@ -142,7 +142,12 @@ pub fn initialize_cuda(session: &CudaSession) {
 /// # Panics
 ///
 /// Panics if CUDA device 0 cannot be initialized (the same contract as [`CudaSession::default`]).
+#[cfg(any(test, feature = "_test-harness"))]
+pub fn cuda_session_builder() -> vortex::session::VortexSessionBuilder {
+    vortex::array::array_session().with::<CudaSession>()
+}
+
 #[cfg(any(test, feature = "_test-harness"))]
 pub fn cuda_session() -> vortex::session::VortexSession {
-    vortex::array::array_session().with::<CudaSession>().build()
+    cuda_session_builder().build()
 }

vortex-cxx/src/lib.rs

@@ -18,7 +18,7 @@ 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::RuntimeSessionBuilderExt;
 use vortex::session::VortexSession;
 use write::*;
 
@@ -28,8 +28,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<VortexSession> = LazyLock::new(|| {
+    VortexSession::default_builder()
+        .with_handle(RUNTIME.handle())
+        .build()
+});
 
 #[cxx::bridge(namespace = "vortex::ffi")]
 #[allow(let_underscore_drop)]