More Blobs

Signed-off-by: Nicholas Gates <nick@nickgates.com>

Author: gatesn

encodings/bytebool/src/array.rs

@@ -352,6 +352,7 @@ impl From<Vec<Option<bool>>> for ByteBoolData {
 mod tests {
     use vortex_array::ArrayContext;
     use vortex_array::IntoArray;
+    use vortex_array::LEGACY_SESSION;
     use vortex_array::assert_arrays_eq;
     use vortex_array::serde::SerializeOptions;
     use vortex_array::serde::SerializedArray;
@@ -406,11 +407,7 @@ mod tests {
         let serialized = array
             .clone()
             .into_array()
-            .serialize(
-                &ctx,
-                &vortex_session::VortexSession::empty(),
-                &SerializeOptions::default(),
-            )
+            .serialize(&ctx, &LEGACY_SESSION, &SerializeOptions::default())
             .unwrap();
 
         let mut concat = ByteBufferMut::empty();

vortex-cuda/src/layout.rs

@@ -445,7 +445,7 @@ impl LayoutStrategy for CudaFlatLayoutStrategy {
         segment_sink: SegmentSinkRef,
         mut stream: SendableSequentialStream,
         _eof: SequencePointer,
-        _handle: vortex::io::runtime::Handle,
+        session: &VortexSession,
     ) -> VortexResult<LayoutRef> {
         let ctx = ctx.clone();
         let options = self.clone();
@@ -507,10 +507,9 @@ impl LayoutStrategy for CudaFlatLayoutStrategy {
         // Scan for constant array buffers before serialization (while data is still on host).
         let host_buffers = extract_constant_buffers(&chunk);
 
-        let session = VortexSession::empty();
         let buffers = chunk.serialize(
             &ctx,
-            &session,
+            session,
             &SerializeOptions {
                 offset: 0,
                 include_padding: options.include_padding,

vortex-file/src/v2/file_stats_reader.rs

@@ -219,6 +219,7 @@ mod tests {
     use vortex_error::VortexResult;
     use vortex_io::runtime::single::block_on;
     use vortex_io::session::RuntimeSession;
+    use vortex_io::session::RuntimeSessionExt;
     use vortex_layout::LayoutReader;
     use vortex_layout::LayoutStrategy;
     use vortex_layout::layouts::flat::writer::FlatLayoutStrategy;
@@ -253,6 +254,7 @@ mod tests {
     #[test]
     fn pruning_when_filter_out_of_range() -> VortexResult<()> {
         block_on(|handle| async {
+            let session = SESSION.clone().with_handle(handle);
             let ctx = ArrayContext::empty();
             let segments = Arc::new(TestSegments::default());
             let (ptr, eof) = SequenceId::root().split();
@@ -269,7 +271,7 @@ mod tests {
                     Arc::<TestSegments>::clone(&segments),
                     struct_array.into_array().to_array_stream().sequenced(ptr),
                     eof,
-                    handle,
+                    &session,
                 )
                 .await?;
 
@@ -291,6 +293,7 @@ mod tests {
     #[test]
     fn no_pruning_when_filter_in_range() -> VortexResult<()> {
         block_on(|handle| async {
+            let session = SESSION.clone().with_handle(handle);
             let ctx = ArrayContext::empty();
             let segments = Arc::new(TestSegments::default());
             let (ptr, eof) = SequenceId::root().split();
@@ -307,7 +310,7 @@ mod tests {
                     Arc::<TestSegments>::clone(&segments),
                     struct_array.into_array().to_array_stream().sequenced(ptr),
                     eof,
-                    handle,
+                    &session,
                 )
                 .await?;
 

vortex-file/src/writer.rs

@@ -179,15 +179,17 @@ impl VortexWriteOptions {
         // We spawn the layout future so it is driven in the background while we write the
         // buffer stream, so we don't need to poll it until all buffers have been drained.
         let ctx2 = ctx.clone();
-        let layout_fut = self.session.handle().spawn_nested(|h| async move {
+        let session = self.session.clone();
+        let layout_fut = self.session.handle().spawn_nested(move |h| async move {
+            let session = session.with_handle(h);
             let layout = self
                 .strategy
                 .write_stream(
                     ctx2,
                     Arc::<BufferedSegmentSink>::clone(&segments),
                     stream,
                     eof,
-                    h,
+                    &session,
                 )
                 .await?;
             Ok::<_, VortexError>((layout, segments.segment_specs()))

vortex-ipc/src/messages/encoder.rs

@@ -14,6 +14,7 @@ use vortex_error::vortex_err;
 use vortex_flatbuffers::FlatBuffer;
 use vortex_flatbuffers::WriteFlatBufferExt;
 use vortex_flatbuffers::message as fb;
+use vortex_session::VortexSession;
 
 /// An IPC message ready to be passed to the encoder.
 pub enum EncoderMessage<'a> {
@@ -25,17 +26,17 @@ pub enum EncoderMessage<'a> {
 pub struct MessageEncoder {
     /// A reusable buffer of zeros used for padding.
     zeros: Bytes,
+    session: VortexSession,
 }
 
-impl Default for MessageEncoder {
-    fn default() -> Self {
+impl MessageEncoder {
+    pub fn new(session: VortexSession) -> Self {
         Self {
+            session,
             zeros: BytesMut::zeroed(u16::MAX as usize).freeze(),
         }
     }
-}
 
-impl MessageEncoder {
     /// Encode an IPC message for writing to a byte stream.
     ///
     /// The returned buffers should be written contiguously to the stream.

vortex-layout/src/display.rs

@@ -230,6 +230,7 @@ mod tests {
     use vortex_buffer::BitBufferMut;
     use vortex_buffer::buffer;
     use vortex_io::runtime::single::block_on;
+    use vortex_io::session::RuntimeSessionExt;
     use vortex_utils::env::EnvVarGuard;
 
     use crate::IntoLayout;
@@ -242,12 +243,14 @@ mod tests {
     use crate::sequence::SequenceId;
     use crate::sequence::SequentialArrayStreamExt;
     use crate::strategy::LayoutStrategy;
+    use crate::test::SESSION;
 
     /// Test display_tree with inline array_tree metadata (no segment source needed).
     #[test]
     fn test_display_tree_inline_array_tree() {
         let _guard = EnvVarGuard::set("FLAT_LAYOUT_INLINE_ARRAY_NODE", "1");
         block_on(|handle| async move {
+            let session = SESSION.clone().with_handle(handle);
             let ctx = ArrayContext::empty();
             let segments = Arc::new(TestSegments::default());
 
@@ -267,7 +270,7 @@ mod tests {
                     Arc::<TestSegments>::clone(&segments),
                     array1.into_array().to_array_stream().sequenced(ptr1),
                     eof1,
-                    handle.clone(),
+                    &session,
                 )
                 .await
                 .unwrap();
@@ -294,7 +297,7 @@ mod tests {
                         .to_array_stream()
                         .sequenced(ptr2),
                     eof2,
-                    handle.clone(),
+                    &session,
                 )
                 .await
                 .unwrap();
@@ -342,6 +345,7 @@ vortex.struct, dtype: {numbers=i64?, strings=utf8}, children: 2, rows: 5
         // Ensure inline array node is disabled for this test
         let _guard = EnvVarGuard::remove("FLAT_LAYOUT_INLINE_ARRAY_NODE");
         block_on(|handle| async move {
+            let session = SESSION.clone().with_handle(handle);
             let ctx = ArrayContext::empty();
             let segments = Arc::new(TestSegments::default());
 
@@ -354,7 +358,7 @@ vortex.struct, dtype: {numbers=i64?, strings=utf8}, children: 2, rows: 5
                     Arc::<TestSegments>::clone(&segments),
                     array1.into_array().to_array_stream().sequenced(ptr1),
                     eof1,
-                    handle.clone(),
+                    &session,
                 )
                 .await
                 .unwrap();
@@ -368,7 +372,7 @@ vortex.struct, dtype: {numbers=i64?, strings=utf8}, children: 2, rows: 5
                     Arc::<TestSegments>::clone(&segments),
                     array2.into_array().to_array_stream().sequenced(ptr2),
                     eof2,
-                    handle.clone(),
+                    &session,
                 )
                 .await
                 .unwrap();
@@ -407,13 +411,14 @@ vortex.chunked, dtype: i32, children: 2, rows: 10
         // Create a simple primitive array
         let array = PrimitiveArray::new(buffer![1i32, 2, 3, 4, 5], Validity::AllValid);
         let layout = block_on(|handle| async {
+            let session = SESSION.clone().with_handle(handle);
             FlatLayoutStrategy::default()
                 .write_stream(
                     ctx.clone(),
                     Arc::<TestSegments>::clone(&segments),
                     array.into_array().to_array_stream().sequenced(ptr),
                     eof,
-                    handle,
+                    &session,
                 )
                 .await
                 .unwrap()
@@ -449,13 +454,14 @@ vortex.flat, dtype: i32?, segment 0, buffers=[20B], total=20B
         // Create a simple primitive array
         let array = PrimitiveArray::new(buffer![10i64, 20, 30], Validity::NonNullable);
         let layout = block_on(|handle| async {
+            let session = SESSION.clone().with_handle(handle);
             FlatLayoutStrategy::default()
                 .write_stream(
                     ctx,
                     Arc::<TestSegments>::clone(&segments),
                     array.into_array().to_array_stream().sequenced(ptr),
                     eof,
-                    handle,
+                    &session,
                 )
                 .await
                 .unwrap()

vortex-layout/src/layouts/buffered.rs

@@ -12,7 +12,7 @@ use futures::StreamExt as _;
 use futures::pin_mut;
 use vortex_array::ArrayContext;
 use vortex_error::VortexResult;
-use vortex_io::runtime::Handle;
+use vortex_session::VortexSession;
 
 use crate::LayoutRef;
 use crate::LayoutStrategy;
@@ -47,7 +47,7 @@ impl LayoutStrategy for BufferedStrategy {
         segment_sink: SegmentSinkRef,
         stream: SendableSequentialStream,
         eof: SequencePointer,
-        handle: Handle,
+        session: &VortexSession,
     ) -> VortexResult<LayoutRef> {
         let dtype = stream.dtype().clone();
         let buffer_size = self.buffer_size;
@@ -102,7 +102,7 @@ impl LayoutStrategy for BufferedStrategy {
                 segment_sink,
                 SequentialStreamAdapter::new(dtype, buffered_stream).sendable(),
                 eof,
-                handle,
+                session,
             )
             .await
     }

vortex-layout/src/layouts/chunked/reader.rs

@@ -343,6 +343,7 @@ mod test {
     use vortex_array::expr::root;
     use vortex_buffer::buffer;
     use vortex_io::runtime::single::block_on;
+    use vortex_io::session::RuntimeSessionExt;
 
     use crate::LayoutRef;
     use crate::LayoutStrategy;
@@ -364,6 +365,7 @@ mod test {
         let strategy = ChunkedLayoutStrategy::new(FlatLayoutStrategy::default());
         let (mut sequence_id, eof) = SequenceId::root().split();
         let layout = block_on(|handle| {
+            let session = SESSION.clone().with_handle(handle);
             strategy.write_stream(
                 ctx,
                 Arc::<TestSegments>::clone(&segments),
@@ -377,7 +379,7 @@ mod test {
                 )
                 .sendable(),
                 eof,
-                handle,
+                &session,
             )
         })
         .unwrap();

vortex-layout/src/layouts/chunked/writer.rs

@@ -11,7 +11,8 @@ use futures::stream;
 use vortex_array::ArrayContext;
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
-use vortex_io::runtime::Handle;
+use vortex_io::session::RuntimeSessionExt;
+use vortex_session::VortexSession;
 
 use crate::IntoLayout;
 use crate::LayoutRef;
@@ -46,11 +47,12 @@ impl LayoutStrategy for ChunkedLayoutStrategy {
         segment_sink: SegmentSinkRef,
         stream: SendableSequentialStream,
         mut eof: SequencePointer,
-        handle: Handle,
+        session: &VortexSession,
     ) -> VortexResult<LayoutRef> {
         let dtype = stream.dtype().clone();
         let dtype2 = dtype.clone();
         let chunk_strategy = Arc::clone(&self.chunk_strategy);
+        let handle = session.handle();
 
         // We spawn each child to allow parallelism when processing chunks.
         let stream = stream! {
@@ -62,8 +64,10 @@ impl LayoutStrategy for ChunkedLayoutStrategy {
                 let ctx = ctx.clone();
                 let segment_sink = Arc::clone(&segment_sink);
                 let dtype = dtype2.clone();
+                let session = session.clone();
 
                 yield handle.spawn_nested(move |handle| async move {
+                    let session = session.with_handle(handle);
                     chunk_strategy
                         .write_stream(
                             ctx,
@@ -74,7 +78,7 @@ impl LayoutStrategy for ChunkedLayoutStrategy {
                             )
                             .sendable(),
                             chunk_eof,
-                            handle,
+                            &session,
                         )
                         .await
                 })

vortex-layout/src/layouts/collect.rs

@@ -12,7 +12,7 @@ use vortex_array::IntoArray;
 use vortex_array::arrays::ChunkedArray;
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
-use vortex_io::runtime::Handle;
+use vortex_session::VortexSession;
 
 use crate::LayoutRef;
 use crate::LayoutStrategy;
@@ -44,7 +44,7 @@ impl LayoutStrategy for CollectStrategy {
         segment_sink: SegmentSinkRef,
         stream: SendableSequentialStream,
         eof: SequencePointer,
-        handle: Handle,
+        session: &VortexSession,
     ) -> VortexResult<LayoutRef> {
         // Read the whole stream, then write one Chunked stream to the inner thing
         let dtype = stream.dtype().clone();
@@ -68,7 +68,7 @@ impl LayoutStrategy for CollectStrategy {
         let adapted = Box::pin(SequentialStreamAdapter::new(dtype, collected_stream));
 
         self.child
-            .write_stream(ctx, segment_sink, adapted, eof, handle)
+            .write_stream(ctx, segment_sink, adapted, eof, session)
             .await
     }