Move array context inside FlatLayout (#3854)

This makes a bit more sense I believe, and is good preparation for
changing the SegmentSource lifetime.

Signed-off-by: Nicholas Gates <[email protected]>

Author: gatesn

vortex-file/src/file.rs

@@ -81,7 +81,7 @@ impl VortexFile {
         self.footer
             .layout()
             // TODO(ngates): we may want to allow the user pass in a name here?
-            .new_reader("".into(), segment_source, self.footer().ctx().clone())
+            .new_reader("".into(), segment_source)
     }
 
     /// Initiate a scan of the file, returning a builder for configuring the scan.

vortex-file/src/footer/mod.rs

@@ -22,18 +22,16 @@ use flatbuffers::root;
 use itertools::Itertools;
 pub(crate) use postscript::*;
 pub use segment::*;
+use vortex_array::ArrayRegistry;
 use vortex_array::stats::StatsSet;
-use vortex_array::{ArrayContext, ArrayRegistry};
 use vortex_dtype::DType;
 use vortex_error::{VortexResult, vortex_bail, vortex_err};
 use vortex_flatbuffers::{FlatBuffer, footer as fb};
-use vortex_layout::{LayoutContext, LayoutRef, LayoutRegistry, layout_from_flatbuffer};
+use vortex_layout::{LayoutRef, LayoutRegistry, layout_from_flatbuffer};
 
 /// Captures the layout information of a Vortex file.
 #[derive(Debug, Clone)]
 pub struct Footer {
-    array_ctx: ArrayContext,
-    layout_ctx: LayoutContext,
     root_layout: LayoutRef,
     segments: Arc<[SegmentSpec]>,
     statistics: Option<FileStatistics>,
@@ -67,7 +65,7 @@ impl Footer {
             .map(|encoding| encoding.id());
         let array_ctx = array_registry.new_context(array_ids)?;
 
-        let root_layout = layout_from_flatbuffer(layout_bytes, &dtype, &layout_ctx)?;
+        let root_layout = layout_from_flatbuffer(layout_bytes, &dtype, &layout_ctx, &array_ctx)?;
 
         let segments: Arc<[SegmentSpec]> = fb_footer
             .segment_specs()
@@ -82,24 +80,12 @@ impl Footer {
         }
 
         Ok(Self {
-            array_ctx,
-            layout_ctx,
             root_layout,
             segments,
             statistics,
         })
     }
 
-    /// Returns the array [`ArrayContext`] of the file.
-    pub fn ctx(&self) -> &ArrayContext {
-        &self.array_ctx
-    }
-
-    /// Returns the [`LayoutContext`] of the file.
-    pub fn layout_ctx(&self) -> &LayoutContext {
-        &self.layout_ctx
-    }
-
     /// Returns the root [`LayoutRef`] of the file.
     pub fn layout(&self) -> &LayoutRef {
         &self.root_layout

vortex-layout/src/children.rs

@@ -6,6 +6,7 @@ use std::sync::Arc;
 
 use flatbuffers::Follow;
 use itertools::Itertools;
+use vortex_array::ArrayContext;
 use vortex_dtype::DType;
 use vortex_error::{VortexResult, vortex_bail, vortex_err, vortex_panic};
 use vortex_flatbuffers::{FlatBuffer, layout as fbl};
@@ -94,7 +95,8 @@ impl LayoutChildren for OwnedLayoutChildren {
 pub(crate) struct ViewedLayoutChildren {
     flatbuffer: FlatBuffer,
     flatbuffer_loc: usize,
-    ctx: LayoutContext,
+    array_ctx: ArrayContext,
+    layout_ctx: LayoutContext,
 }
 
 impl ViewedLayoutChildren {
@@ -106,12 +108,14 @@ impl ViewedLayoutChildren {
     pub(super) unsafe fn new_unchecked(
         flatbuffer: FlatBuffer,
         flatbuffer_loc: usize,
-        ctx: LayoutContext,
+        array_ctx: ArrayContext,
+        layout_ctx: LayoutContext,
     ) -> Self {
         Self {
             flatbuffer,
             flatbuffer_loc,
-            ctx,
+            array_ctx,
+            layout_ctx,
         }
     }
 
@@ -135,10 +139,11 @@ impl LayoutChildren for ViewedLayoutChildren {
         let viewed_children = ViewedLayoutChildren {
             flatbuffer: self.flatbuffer.clone(),
             flatbuffer_loc: fb_child._tab.loc(),
-            ctx: self.ctx.clone(),
+            array_ctx: self.array_ctx.clone(),
+            layout_ctx: self.layout_ctx.clone(),
         };
         let encoding = self
-            .ctx
+            .layout_ctx
             .lookup_encoding(fb_child.encoding())
             .ok_or_else(|| vortex_err!("Encoding not found: {}", fb_child.encoding()))?;
 
@@ -156,6 +161,7 @@ impl LayoutChildren for ViewedLayoutChildren {
                 .map(SegmentId::from)
                 .collect_vec(),
             &viewed_children,
+            self.array_ctx.clone(),
         )
     }
 

vortex-layout/src/encoding.rs

@@ -5,7 +5,7 @@ use std::any::Any;
 use std::fmt::{Debug, Display, Formatter};
 
 use arcref::ArcRef;
-use vortex_array::DeserializeMetadata;
+use vortex_array::{ArrayContext, DeserializeMetadata};
 use vortex_dtype::DType;
 use vortex_error::{VortexExpect, VortexResult, vortex_panic};
 
@@ -27,6 +27,7 @@ pub trait LayoutEncoding: 'static + Send + Sync + Debug + private::Sealed {
         metadata: &[u8],
         segment_ids: Vec<SegmentId>,
         children: &dyn LayoutChildren,
+        ctx: ArrayContext,
     ) -> VortexResult<LayoutRef>;
 }
 
@@ -49,9 +50,18 @@ impl<V: VTable> LayoutEncoding for LayoutEncodingAdapter<V> {
         metadata: &[u8],
         segment_ids: Vec<SegmentId>,
         children: &dyn LayoutChildren,
+        ctx: ArrayContext,
     ) -> VortexResult<LayoutRef> {
         let metadata = <V::Metadata as DeserializeMetadata>::deserialize(metadata)?;
-        let layout = V::build(&self.0, dtype, row_count, &metadata, segment_ids, children)?;
+        let layout = V::build(
+            &self.0,
+            dtype,
+            row_count,
+            &metadata,
+            segment_ids,
+            children,
+            ctx,
+        )?;
 
         // Validate that the builder function returned the expected values.
         if layout.row_count() != row_count {

vortex-layout/src/flatbuffers.rs

@@ -5,6 +5,7 @@ use std::env;
 use std::sync::LazyLock;
 
 use flatbuffers::{FlatBufferBuilder, VerifierOptions, WIPOffset, root_with_opts};
+use vortex_array::ArrayContext;
 use vortex_dtype::DType;
 use vortex_error::{VortexExpect, VortexResult, vortex_err};
 use vortex_flatbuffers::{FlatBuffer, FlatBufferRoot, WriteFlatBuffer, layout};
@@ -34,16 +35,22 @@ static LAYOUT_VERIFIER: LazyLock<VerifierOptions> = LazyLock::new(|| {
 pub fn layout_from_flatbuffer(
     flatbuffer: FlatBuffer,
     dtype: &DType,
-    ctx: &LayoutContext,
+    layout_ctx: &LayoutContext,
+    array_ctx: &ArrayContext,
 ) -> VortexResult<LayoutRef> {
     let fb_layout = root_with_opts::<layout::Layout>(&LAYOUT_VERIFIER, &flatbuffer)?;
-    let encoding = ctx
+    let encoding = layout_ctx
         .lookup_encoding(fb_layout.encoding())
         .ok_or_else(|| vortex_err!("Invalid encoding ID: {}", fb_layout.encoding()))?;
 
     // SAFETY: we validate the flatbuffer above in the `root` call, and extract a loc.
     let viewed_children = unsafe {
-        ViewedLayoutChildren::new_unchecked(flatbuffer.clone(), fb_layout._tab.loc(), ctx.clone())
+        ViewedLayoutChildren::new_unchecked(
+            flatbuffer.clone(),
+            fb_layout._tab.loc(),
+            array_ctx.clone(),
+            layout_ctx.clone(),
+        )
     };
 
     let layout = encoding.build(
@@ -60,6 +67,7 @@ pub fn layout_from_flatbuffer(
             .map(SegmentId::from)
             .collect(),
         &viewed_children,
+        array_ctx.clone(),
     )?;
 
     Ok(layout)

vortex-layout/src/layout.rs

@@ -7,7 +7,7 @@ use std::sync::Arc;
 
 use arcref::ArcRef;
 use itertools::Itertools;
-use vortex_array::{ArrayContext, SerializeMetadata};
+use vortex_array::SerializeMetadata;
 use vortex_dtype::{DType, FieldName};
 use vortex_error::{VortexExpect, VortexResult, vortex_err};
 
@@ -54,7 +54,6 @@ pub trait Layout: 'static + Send + Sync + Debug + private::Sealed {
         &self,
         name: Arc<str>,
         segment_source: Arc<dyn SegmentSource>,
-        ctx: ArrayContext,
     ) -> VortexResult<LayoutReaderRef>;
 }
 
@@ -248,9 +247,8 @@ impl<V: VTable> Layout for LayoutAdapter<V> {
         &self,
         name: Arc<str>,
         segment_source: Arc<dyn SegmentSource>,
-        ctx: ArrayContext,
     ) -> VortexResult<LayoutReaderRef> {
-        V::new_reader(&self.0, name, segment_source, ctx)
+        V::new_reader(&self.0, name, segment_source)
     }
 }
 

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

@@ -64,13 +64,11 @@ impl VTable for ChunkedVTable {
         layout: &Self::Layout,
         name: Arc<str>,
         segment_source: Arc<dyn SegmentSource>,
-        ctx: ArrayContext,
     ) -> VortexResult<LayoutReaderRef> {
         Ok(Arc::new(ChunkedReader::new(
             layout.clone(),
             name,
             segment_source,
-            ctx,
         )))
     }
 
@@ -81,6 +79,7 @@ impl VTable for ChunkedVTable {
         _metadata: &<Self::Metadata as DeserializeMetadata>::Output,
         _segment_ids: Vec<SegmentId>,
         children: &dyn LayoutChildren,
+        _ctx: ArrayContext,
     ) -> VortexResult<Self::Layout> {
         Ok(ChunkedLayout::new(
             row_count,

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

@@ -10,9 +10,9 @@ use futures::future::ready;
 use futures::stream::FuturesOrdered;
 use futures::{FutureExt, TryStreamExt};
 use itertools::Itertools;
+use vortex_array::ArrayRef;
 use vortex_array::arrays::ChunkedArray;
 use vortex_array::stats::Precision;
-use vortex_array::{ArrayContext, ArrayRef};
 use vortex_dtype::{DType, FieldMask};
 use vortex_error::{VortexExpect, VortexResult};
 use vortex_expr::ExprRef;
@@ -39,7 +39,6 @@ impl ChunkedReader {
         layout: ChunkedLayout,
         name: Arc<str>,
         segment_source: Arc<dyn SegmentSource>,
-        ctx: ArrayContext,
     ) -> Self {
         let nchildren = layout.nchildren();
 
@@ -49,7 +48,7 @@ impl ChunkedReader {
         }
         chunk_offsets.push(layout.row_count());
 
-        let lazy_children = LazyReaderChildren::new(layout.children.clone(), segment_source, ctx);
+        let lazy_children = LazyReaderChildren::new(layout.children.clone(), segment_source);
 
         Self {
             layout,
@@ -351,7 +350,7 @@ mod test {
 
     #[fixture]
     /// Create a chunked layout with three chunks of primitive arrays.
-    fn chunked_layout() -> (ArrayContext, Arc<dyn SegmentSource>, LayoutRef) {
+    fn chunked_layout() -> (Arc<dyn SegmentSource>, LayoutRef) {
         let ctx = ArrayContext::empty();
         let segments = TestSegments::default();
         let sequence_writer = SequenceWriter::new(Box::new(segments.clone()));
@@ -374,20 +373,16 @@ mod test {
         )
         .unwrap();
 
-        (ctx, Arc::new(segments), layout)
+        (Arc::new(segments), layout)
     }
 
     #[rstest]
     fn test_chunked_evaluator(
-        #[from(chunked_layout)] (ctx, segments, layout): (
-            ArrayContext,
-            Arc<dyn SegmentSource>,
-            LayoutRef,
-        ),
+        #[from(chunked_layout)] (segments, layout): (Arc<dyn SegmentSource>, LayoutRef),
     ) {
         block_on(async {
             let result = layout
-                .new_reader("".into(), segments, ctx)
+                .new_reader("".into(), segments)
                 .unwrap()
                 .projection_evaluation(&(0..layout.row_count()), &root())
                 .unwrap()

vortex-layout/src/layouts/dict/mod.rs

@@ -74,13 +74,11 @@ impl VTable for DictVTable {
         layout: &Self::Layout,
         name: Arc<str>,
         segment_source: Arc<dyn SegmentSource>,
-        ctx: ArrayContext,
     ) -> VortexResult<LayoutReaderRef> {
         Ok(Arc::new(DictReader::try_new(
             layout.clone(),
             name,
             segment_source,
-            ctx,
         )?))
     }
 
@@ -91,6 +89,7 @@ impl VTable for DictVTable {
         metadata: &<Self::Metadata as DeserializeMetadata>::Output,
         _segment_ids: Vec<SegmentId>,
         children: &dyn LayoutChildren,
+        _ctx: ArrayContext,
     ) -> VortexResult<Self::Layout> {
         let values = children.child(0, dtype)?;
         let codes = children.child(

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

@@ -8,9 +8,9 @@ use std::sync::{Arc, OnceLock};
 use async_trait::async_trait;
 use dashmap::DashMap;
 use futures::{FutureExt, join};
+use vortex_array::ArrayRef;
 use vortex_array::compute::{MinMaxResult, filter, min_max};
 use vortex_array::stats::Precision;
-use vortex_array::{ArrayContext, ArrayRef};
 use vortex_dict::DictArray;
 use vortex_dtype::{DType, FieldMask};
 use vortex_error::{VortexExpect, VortexResult};
@@ -46,17 +46,14 @@ impl DictReader {
         layout: DictLayout,
         name: Arc<str>,
         segment_source: Arc<dyn SegmentSource>,
-        ctx: ArrayContext,
     ) -> VortexResult<Self> {
         let values_len = usize::try_from(layout.values.row_count())?;
-        let values = layout.values.new_reader(
-            format!("{name}.values").into(),
-            segment_source.clone(),
-            ctx.clone(),
-        )?;
+        let values = layout
+            .values
+            .new_reader(format!("{name}.values").into(), segment_source.clone())?;
         let codes = layout
             .codes
-            .new_reader(format!("{name}.codes").into(), segment_source, ctx)?;
+            .new_reader(format!("{name}.codes").into(), segment_source)?;
 
         Ok(Self {
             layout,
@@ -315,7 +312,7 @@ mod tests {
         );
         assert!(layout.encoding_id() == LayoutId::new_ref("vortex.dict"));
         let actual = layout
-            .new_reader("".into(), Arc::from(segments), ctx)
+            .new_reader("".into(), Arc::from(segments))
             .unwrap()
             .projection_evaluation(&(0..layout.row_count()), &expression)
             .unwrap()
@@ -391,7 +388,7 @@ mod tests {
         let expression = not(is_null(root())); // easier to test not_is_null b/c that's the validity array
         assert!(layout.encoding_id() == LayoutId::new_ref("vortex.dict"));
         let actual = layout
-            .new_reader("".into(), Arc::from(segments), ctx)
+            .new_reader("".into(), Arc::from(segments))
             .unwrap()
             .projection_evaluation(&(0..layout.row_count()), &expression)
             .unwrap()