diff --git a/vortex-cuda/src/layout.rs b/vortex-cuda/src/layout.rs index 858bb87e0cc..20ab3f3a50b 100644 --- a/vortex-cuda/src/layout.rs +++ b/vortex-cuda/src/layout.rs @@ -180,6 +180,7 @@ impl VTable for CudaFlat { name: Arc, segment_source: Arc, session: &VortexSession, + _ctx: &vortex::layout::LayoutReaderContext, ) -> VortexResult { Ok(Arc::new(CudaFlatReader { layout: layout.clone(), diff --git a/vortex-file/src/file.rs b/vortex-file/src/file.rs index f321d774197..b8e31f1fea2 100644 --- a/vortex-file/src/file.rs +++ b/vortex-file/src/file.rs @@ -90,7 +90,12 @@ 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.session) + .new_reader( + "".into(), + segment_source, + &self.session, + &Default::default(), + ) } /// Create a [`DataSource`](vortex_scan::DataSource) from this file for scanning. diff --git a/vortex-file/src/v2/file_stats_reader.rs b/vortex-file/src/v2/file_stats_reader.rs index f65c758ee7d..bdb104ddabb 100644 --- a/vortex-file/src/v2/file_stats_reader.rs +++ b/vortex-file/src/v2/file_stats_reader.rs @@ -310,7 +310,7 @@ mod tests { ) .await?; - let child = layout.new_reader("".into(), segments, &SESSION)?; + let child = layout.new_reader("".into(), segments, &SESSION, &Default::default())?; let reader = FileStatsLayoutReader::new(child, test_file_stats(0, 100), SESSION.clone()); @@ -349,7 +349,7 @@ mod tests { ) .await?; - let child = layout.new_reader("".into(), segments, &SESSION)?; + let child = layout.new_reader("".into(), segments, &SESSION, &Default::default())?; let reader = FileStatsLayoutReader::new(child, test_file_stats(0, 100), SESSION.clone()); @@ -400,7 +400,7 @@ mod tests { ) .await?; - let child = layout.new_reader("".into(), segments, &SESSION)?; + let child = layout.new_reader("".into(), segments, &SESSION, &Default::default())?; // File-level stats: 1 null in deleted_at. let mut stats = StatsSet::default(); @@ -449,7 +449,7 @@ mod tests { ) .await?; - let child = layout.new_reader("".into(), segments, &SESSION)?; + let child = layout.new_reader("".into(), segments, &SESSION, &Default::default())?; let reader = FileStatsLayoutReader::new(child, test_file_null_count_stats(5), SESSION.clone()); diff --git a/vortex-layout/public-api.lock b/vortex-layout/public-api.lock index 1c6ae70dd6f..f82986026c9 100644 --- a/vortex-layout/public-api.lock +++ b/vortex-layout/public-api.lock @@ -90,7 +90,7 @@ pub fn vortex_layout::layouts::chunked::Chunked::metadata(&Self::Layout) -> Self pub fn vortex_layout::layouts::chunked::Chunked::nchildren(&Self::Layout) -> usize -pub fn vortex_layout::layouts::chunked::Chunked::new_reader(&Self::Layout, alloc::sync::Arc, alloc::sync::Arc, &vortex_session::VortexSession) -> vortex_error::VortexResult +pub fn vortex_layout::layouts::chunked::Chunked::new_reader(&Self::Layout, alloc::sync::Arc, alloc::sync::Arc, &vortex_session::VortexSession, &vortex_layout::LayoutReaderContext) -> vortex_error::VortexResult pub fn vortex_layout::layouts::chunked::Chunked::row_count(&Self::Layout) -> u64 @@ -308,7 +308,7 @@ pub fn vortex_layout::layouts::dict::Dict::metadata(&Self::Layout) -> Self::Meta pub fn vortex_layout::layouts::dict::Dict::nchildren(&Self::Layout) -> usize -pub fn vortex_layout::layouts::dict::Dict::new_reader(&Self::Layout, alloc::sync::Arc, alloc::sync::Arc, &vortex_session::VortexSession) -> vortex_error::VortexResult +pub fn vortex_layout::layouts::dict::Dict::new_reader(&Self::Layout, alloc::sync::Arc, alloc::sync::Arc, &vortex_session::VortexSession, &vortex_layout::LayoutReaderContext) -> vortex_error::VortexResult pub fn vortex_layout::layouts::dict::Dict::row_count(&Self::Layout) -> u64 @@ -474,7 +474,7 @@ pub fn vortex_layout::layouts::flat::Flat::metadata(&Self::Layout) -> Self::Meta pub fn vortex_layout::layouts::flat::Flat::nchildren(&Self::Layout) -> usize -pub fn vortex_layout::layouts::flat::Flat::new_reader(&Self::Layout, alloc::sync::Arc, alloc::sync::Arc, &vortex_session::VortexSession) -> vortex_error::VortexResult +pub fn vortex_layout::layouts::flat::Flat::new_reader(&Self::Layout, alloc::sync::Arc, alloc::sync::Arc, &vortex_session::VortexSession, &vortex_layout::LayoutReaderContext) -> vortex_error::VortexResult pub fn vortex_layout::layouts::flat::Flat::row_count(&Self::Layout) -> u64 @@ -674,7 +674,7 @@ pub fn vortex_layout::layouts::struct_::Struct::metadata(&Self::Layout) -> Self: pub fn vortex_layout::layouts::struct_::Struct::nchildren(&Self::Layout) -> usize -pub fn vortex_layout::layouts::struct_::Struct::new_reader(&Self::Layout, alloc::sync::Arc, alloc::sync::Arc, &vortex_session::VortexSession) -> vortex_error::VortexResult +pub fn vortex_layout::layouts::struct_::Struct::new_reader(&Self::Layout, alloc::sync::Arc, alloc::sync::Arc, &vortex_session::VortexSession, &vortex_layout::LayoutReaderContext) -> vortex_error::VortexResult pub fn vortex_layout::layouts::struct_::Struct::row_count(&Self::Layout) -> u64 @@ -838,7 +838,7 @@ pub fn vortex_layout::layouts::zoned::Zoned::metadata(&Self::Layout) -> Self::Me pub fn vortex_layout::layouts::zoned::Zoned::nchildren(&Self::Layout) -> usize -pub fn vortex_layout::layouts::zoned::Zoned::new_reader(&Self::Layout, alloc::sync::Arc, alloc::sync::Arc, &vortex_session::VortexSession) -> vortex_error::VortexResult +pub fn vortex_layout::layouts::zoned::Zoned::new_reader(&Self::Layout, alloc::sync::Arc, alloc::sync::Arc, &vortex_session::VortexSession, &vortex_layout::LayoutReaderContext) -> vortex_error::VortexResult pub fn vortex_layout::layouts::zoned::Zoned::row_count(&Self::Layout) -> u64 @@ -1456,7 +1456,7 @@ pub fn vortex_layout::vtable::VTable::metadata(&Self::Layout) -> Self::Metadata pub fn vortex_layout::vtable::VTable::nchildren(&Self::Layout) -> usize -pub fn vortex_layout::vtable::VTable::new_reader(&Self::Layout, alloc::sync::Arc, alloc::sync::Arc, &vortex_session::VortexSession) -> vortex_error::VortexResult +pub fn vortex_layout::vtable::VTable::new_reader(&Self::Layout, alloc::sync::Arc, alloc::sync::Arc, &vortex_session::VortexSession, &vortex_layout::LayoutReaderContext) -> vortex_error::VortexResult pub fn vortex_layout::vtable::VTable::row_count(&Self::Layout) -> u64 @@ -1488,7 +1488,7 @@ pub fn vortex_layout::layouts::chunked::Chunked::metadata(&Self::Layout) -> Self pub fn vortex_layout::layouts::chunked::Chunked::nchildren(&Self::Layout) -> usize -pub fn vortex_layout::layouts::chunked::Chunked::new_reader(&Self::Layout, alloc::sync::Arc, alloc::sync::Arc, &vortex_session::VortexSession) -> vortex_error::VortexResult +pub fn vortex_layout::layouts::chunked::Chunked::new_reader(&Self::Layout, alloc::sync::Arc, alloc::sync::Arc, &vortex_session::VortexSession, &vortex_layout::LayoutReaderContext) -> vortex_error::VortexResult pub fn vortex_layout::layouts::chunked::Chunked::row_count(&Self::Layout) -> u64 @@ -1520,7 +1520,7 @@ pub fn vortex_layout::layouts::dict::Dict::metadata(&Self::Layout) -> Self::Meta pub fn vortex_layout::layouts::dict::Dict::nchildren(&Self::Layout) -> usize -pub fn vortex_layout::layouts::dict::Dict::new_reader(&Self::Layout, alloc::sync::Arc, alloc::sync::Arc, &vortex_session::VortexSession) -> vortex_error::VortexResult +pub fn vortex_layout::layouts::dict::Dict::new_reader(&Self::Layout, alloc::sync::Arc, alloc::sync::Arc, &vortex_session::VortexSession, &vortex_layout::LayoutReaderContext) -> vortex_error::VortexResult pub fn vortex_layout::layouts::dict::Dict::row_count(&Self::Layout) -> u64 @@ -1552,7 +1552,7 @@ pub fn vortex_layout::layouts::flat::Flat::metadata(&Self::Layout) -> Self::Meta pub fn vortex_layout::layouts::flat::Flat::nchildren(&Self::Layout) -> usize -pub fn vortex_layout::layouts::flat::Flat::new_reader(&Self::Layout, alloc::sync::Arc, alloc::sync::Arc, &vortex_session::VortexSession) -> vortex_error::VortexResult +pub fn vortex_layout::layouts::flat::Flat::new_reader(&Self::Layout, alloc::sync::Arc, alloc::sync::Arc, &vortex_session::VortexSession, &vortex_layout::LayoutReaderContext) -> vortex_error::VortexResult pub fn vortex_layout::layouts::flat::Flat::row_count(&Self::Layout) -> u64 @@ -1584,7 +1584,7 @@ pub fn vortex_layout::layouts::struct_::Struct::metadata(&Self::Layout) -> Self: pub fn vortex_layout::layouts::struct_::Struct::nchildren(&Self::Layout) -> usize -pub fn vortex_layout::layouts::struct_::Struct::new_reader(&Self::Layout, alloc::sync::Arc, alloc::sync::Arc, &vortex_session::VortexSession) -> vortex_error::VortexResult +pub fn vortex_layout::layouts::struct_::Struct::new_reader(&Self::Layout, alloc::sync::Arc, alloc::sync::Arc, &vortex_session::VortexSession, &vortex_layout::LayoutReaderContext) -> vortex_error::VortexResult pub fn vortex_layout::layouts::struct_::Struct::row_count(&Self::Layout) -> u64 @@ -1616,7 +1616,7 @@ pub fn vortex_layout::layouts::zoned::Zoned::metadata(&Self::Layout) -> Self::Me pub fn vortex_layout::layouts::zoned::Zoned::nchildren(&Self::Layout) -> usize -pub fn vortex_layout::layouts::zoned::Zoned::new_reader(&Self::Layout, alloc::sync::Arc, alloc::sync::Arc, &vortex_session::VortexSession) -> vortex_error::VortexResult +pub fn vortex_layout::layouts::zoned::Zoned::new_reader(&Self::Layout, alloc::sync::Arc, alloc::sync::Arc, &vortex_session::VortexSession, &vortex_layout::LayoutReaderContext) -> vortex_error::VortexResult pub fn vortex_layout::layouts::zoned::Zoned::row_count(&Self::Layout) -> u64 @@ -1682,7 +1682,7 @@ pub fn vortex_layout::LayoutAdapter::metadata(&self) -> alloc::vec::Vec pub fn vortex_layout::LayoutAdapter::nchildren(&self) -> usize -pub fn vortex_layout::LayoutAdapter::new_reader(&self, alloc::sync::Arc, alloc::sync::Arc, &vortex_session::VortexSession) -> vortex_error::VortexResult +pub fn vortex_layout::LayoutAdapter::new_reader(&self, alloc::sync::Arc, alloc::sync::Arc, &vortex_session::VortexSession, &vortex_layout::LayoutReaderContext) -> vortex_error::VortexResult pub fn vortex_layout::LayoutAdapter::row_count(&self) -> u64 @@ -1704,13 +1704,35 @@ pub fn vortex_layout::LayoutEncodingAdapter::build(&self, &vortex_array::dtyp pub fn vortex_layout::LayoutEncodingAdapter::id(&self) -> vortex_layout::LayoutEncodingId +pub struct vortex_layout::LayoutReaderContext + +impl vortex_layout::LayoutReaderContext + +pub fn vortex_layout::LayoutReaderContext::get(&self) -> core::option::Option> + +pub fn vortex_layout::LayoutReaderContext::new() -> Self + +pub fn vortex_layout::LayoutReaderContext::with(&self, alloc::sync::Arc) -> Self + +impl core::clone::Clone for vortex_layout::LayoutReaderContext + +pub fn vortex_layout::LayoutReaderContext::clone(&self) -> vortex_layout::LayoutReaderContext + +impl core::default::Default for vortex_layout::LayoutReaderContext + +pub fn vortex_layout::LayoutReaderContext::default() -> vortex_layout::LayoutReaderContext + +impl core::fmt::Debug for vortex_layout::LayoutReaderContext + +pub fn vortex_layout::LayoutReaderContext::fmt(&self, &mut core::fmt::Formatter<'_>) -> core::fmt::Result + pub struct vortex_layout::LazyReaderChildren impl vortex_layout::LazyReaderChildren pub fn vortex_layout::LazyReaderChildren::get(&self, usize) -> vortex_error::VortexResult<&vortex_layout::LayoutReaderRef> -pub fn vortex_layout::LazyReaderChildren::new(alloc::sync::Arc, alloc::vec::Vec, alloc::vec::Vec>, alloc::sync::Arc, vortex_session::VortexSession) -> Self +pub fn vortex_layout::LazyReaderChildren::new(alloc::sync::Arc, alloc::vec::Vec, alloc::vec::Vec>, alloc::sync::Arc, vortex_session::VortexSession, vortex_layout::LayoutReaderContext) -> Self pub struct vortex_layout::SplitRange @@ -1798,7 +1820,7 @@ pub fn vortex_layout::Layout::metadata(&self) -> alloc::vec::Vec pub fn vortex_layout::Layout::nchildren(&self) -> usize -pub fn vortex_layout::Layout::new_reader(&self, alloc::sync::Arc, alloc::sync::Arc, &vortex_session::VortexSession) -> vortex_error::VortexResult +pub fn vortex_layout::Layout::new_reader(&self, alloc::sync::Arc, alloc::sync::Arc, &vortex_session::VortexSession, &vortex_layout::LayoutReaderContext) -> vortex_error::VortexResult pub fn vortex_layout::Layout::row_count(&self) -> u64 @@ -1824,7 +1846,7 @@ pub fn vortex_layout::LayoutAdapter::metadata(&self) -> alloc::vec::Vec pub fn vortex_layout::LayoutAdapter::nchildren(&self) -> usize -pub fn vortex_layout::LayoutAdapter::new_reader(&self, alloc::sync::Arc, alloc::sync::Arc, &vortex_session::VortexSession) -> vortex_error::VortexResult +pub fn vortex_layout::LayoutAdapter::new_reader(&self, alloc::sync::Arc, alloc::sync::Arc, &vortex_session::VortexSession, &vortex_layout::LayoutReaderContext) -> vortex_error::VortexResult pub fn vortex_layout::LayoutAdapter::row_count(&self) -> u64 @@ -1994,7 +2016,7 @@ pub fn vortex_layout::VTable::metadata(&Self::Layout) -> Self::Metadata pub fn vortex_layout::VTable::nchildren(&Self::Layout) -> usize -pub fn vortex_layout::VTable::new_reader(&Self::Layout, alloc::sync::Arc, alloc::sync::Arc, &vortex_session::VortexSession) -> vortex_error::VortexResult +pub fn vortex_layout::VTable::new_reader(&Self::Layout, alloc::sync::Arc, alloc::sync::Arc, &vortex_session::VortexSession, &vortex_layout::LayoutReaderContext) -> vortex_error::VortexResult pub fn vortex_layout::VTable::row_count(&Self::Layout) -> u64 @@ -2026,7 +2048,7 @@ pub fn vortex_layout::layouts::chunked::Chunked::metadata(&Self::Layout) -> Self pub fn vortex_layout::layouts::chunked::Chunked::nchildren(&Self::Layout) -> usize -pub fn vortex_layout::layouts::chunked::Chunked::new_reader(&Self::Layout, alloc::sync::Arc, alloc::sync::Arc, &vortex_session::VortexSession) -> vortex_error::VortexResult +pub fn vortex_layout::layouts::chunked::Chunked::new_reader(&Self::Layout, alloc::sync::Arc, alloc::sync::Arc, &vortex_session::VortexSession, &vortex_layout::LayoutReaderContext) -> vortex_error::VortexResult pub fn vortex_layout::layouts::chunked::Chunked::row_count(&Self::Layout) -> u64 @@ -2058,7 +2080,7 @@ pub fn vortex_layout::layouts::dict::Dict::metadata(&Self::Layout) -> Self::Meta pub fn vortex_layout::layouts::dict::Dict::nchildren(&Self::Layout) -> usize -pub fn vortex_layout::layouts::dict::Dict::new_reader(&Self::Layout, alloc::sync::Arc, alloc::sync::Arc, &vortex_session::VortexSession) -> vortex_error::VortexResult +pub fn vortex_layout::layouts::dict::Dict::new_reader(&Self::Layout, alloc::sync::Arc, alloc::sync::Arc, &vortex_session::VortexSession, &vortex_layout::LayoutReaderContext) -> vortex_error::VortexResult pub fn vortex_layout::layouts::dict::Dict::row_count(&Self::Layout) -> u64 @@ -2090,7 +2112,7 @@ pub fn vortex_layout::layouts::flat::Flat::metadata(&Self::Layout) -> Self::Meta pub fn vortex_layout::layouts::flat::Flat::nchildren(&Self::Layout) -> usize -pub fn vortex_layout::layouts::flat::Flat::new_reader(&Self::Layout, alloc::sync::Arc, alloc::sync::Arc, &vortex_session::VortexSession) -> vortex_error::VortexResult +pub fn vortex_layout::layouts::flat::Flat::new_reader(&Self::Layout, alloc::sync::Arc, alloc::sync::Arc, &vortex_session::VortexSession, &vortex_layout::LayoutReaderContext) -> vortex_error::VortexResult pub fn vortex_layout::layouts::flat::Flat::row_count(&Self::Layout) -> u64 @@ -2122,7 +2144,7 @@ pub fn vortex_layout::layouts::struct_::Struct::metadata(&Self::Layout) -> Self: pub fn vortex_layout::layouts::struct_::Struct::nchildren(&Self::Layout) -> usize -pub fn vortex_layout::layouts::struct_::Struct::new_reader(&Self::Layout, alloc::sync::Arc, alloc::sync::Arc, &vortex_session::VortexSession) -> vortex_error::VortexResult +pub fn vortex_layout::layouts::struct_::Struct::new_reader(&Self::Layout, alloc::sync::Arc, alloc::sync::Arc, &vortex_session::VortexSession, &vortex_layout::LayoutReaderContext) -> vortex_error::VortexResult pub fn vortex_layout::layouts::struct_::Struct::row_count(&Self::Layout) -> u64 @@ -2154,7 +2176,7 @@ pub fn vortex_layout::layouts::zoned::Zoned::metadata(&Self::Layout) -> Self::Me pub fn vortex_layout::layouts::zoned::Zoned::nchildren(&Self::Layout) -> usize -pub fn vortex_layout::layouts::zoned::Zoned::new_reader(&Self::Layout, alloc::sync::Arc, alloc::sync::Arc, &vortex_session::VortexSession) -> vortex_error::VortexResult +pub fn vortex_layout::layouts::zoned::Zoned::new_reader(&Self::Layout, alloc::sync::Arc, alloc::sync::Arc, &vortex_session::VortexSession, &vortex_layout::LayoutReaderContext) -> vortex_error::VortexResult pub fn vortex_layout::layouts::zoned::Zoned::row_count(&Self::Layout) -> u64 diff --git a/vortex-layout/src/layout.rs b/vortex-layout/src/layout.rs index 11f0afd629c..e4ad80e2d37 100644 --- a/vortex-layout/src/layout.rs +++ b/vortex-layout/src/layout.rs @@ -19,6 +19,7 @@ use vortex_session::registry::Id; use crate::LayoutEncodingId; use crate::LayoutEncodingRef; +use crate::LayoutReaderContext; use crate::LayoutReaderRef; use crate::VTable; use crate::display::DisplayLayoutTree; @@ -63,11 +64,19 @@ pub trait Layout: 'static + Send + Sync + Debug + private::Sealed { /// Get the segment IDs for this layout. fn segment_ids(&self) -> Vec; + /// Construct a new reader for this layout, using the given dependency context. + /// + /// `ctx` is the typed-data registry threaded through reader construction (see + /// [`LayoutReaderContext`]). Top-level callers (file open, tests) typically pass + /// `&LayoutReaderContext::new()`; recursive callers inside layout implementations + /// must propagate the `ctx` they were handed so ancestor-published values reach + /// descendants. fn new_reader( &self, name: Arc, segment_source: Arc, session: &VortexSession, + ctx: &LayoutReaderContext, ) -> VortexResult; } @@ -311,8 +320,9 @@ impl Layout for LayoutAdapter { name: Arc, segment_source: Arc, session: &VortexSession, + ctx: &LayoutReaderContext, ) -> VortexResult { - V::new_reader(&self.0, name, segment_source, session) + V::new_reader(&self.0, name, segment_source, session, ctx) } } diff --git a/vortex-layout/src/layouts/chunked/mod.rs b/vortex-layout/src/layouts/chunked/mod.rs index 82fcb650607..ed6dd3d75ba 100644 --- a/vortex-layout/src/layouts/chunked/mod.rs +++ b/vortex-layout/src/layouts/chunked/mod.rs @@ -74,12 +74,14 @@ impl VTable for Chunked { name: Arc, segment_source: Arc, session: &VortexSession, + ctx: &crate::LayoutReaderContext, ) -> VortexResult { Ok(Arc::new(ChunkedReader::new( layout.clone(), name, segment_source, session, + ctx.clone(), ))) } diff --git a/vortex-layout/src/layouts/chunked/reader.rs b/vortex-layout/src/layouts/chunked/reader.rs index 94ab3272cd7..a71a1429006 100644 --- a/vortex-layout/src/layouts/chunked/reader.rs +++ b/vortex-layout/src/layouts/chunked/reader.rs @@ -47,6 +47,7 @@ impl ChunkedReader { name: Arc, segment_source: Arc, session: &VortexSession, + ctx: crate::LayoutReaderContext, ) -> Self { let nchildren = layout.nchildren(); @@ -66,6 +67,7 @@ impl ChunkedReader { names, segment_source, session.clone(), + ctx, ); Self { @@ -435,7 +437,12 @@ mod test { ) { let layout = nested_chunked_layout(); let reader = layout - .new_reader("".into(), Arc::new(TestSegments::default()), &SESSION) + .new_reader( + "".into(), + Arc::new(TestSegments::default()), + &SESSION, + &Default::default(), + ) .unwrap(); let splits = SplitBy::Layout @@ -451,7 +458,7 @@ mod test { ) { block_on(|_h| async { let result = layout - .new_reader("".into(), segments, &SESSION) + .new_reader("".into(), segments, &SESSION, &Default::default()) .unwrap() .projection_evaluation( &(0..layout.row_count()), diff --git a/vortex-layout/src/layouts/dict/mod.rs b/vortex-layout/src/layouts/dict/mod.rs index 7928b447fa9..5fb4ea080b4 100644 --- a/vortex-layout/src/layouts/dict/mod.rs +++ b/vortex-layout/src/layouts/dict/mod.rs @@ -92,12 +92,14 @@ impl VTable for Dict { name: Arc, segment_source: Arc, session: &VortexSession, + ctx: &crate::LayoutReaderContext, ) -> VortexResult { Ok(Arc::new(DictReader::try_new( layout.clone(), name, segment_source, session.clone(), + ctx.clone(), )?)) } diff --git a/vortex-layout/src/layouts/dict/reader.rs b/vortex-layout/src/layouts/dict/reader.rs index 96f12d53ece..327d6d234ab 100644 --- a/vortex-layout/src/layouts/dict/reader.rs +++ b/vortex-layout/src/layouts/dict/reader.rs @@ -58,17 +58,21 @@ impl DictReader { name: Arc, segment_source: Arc, session: VortexSession, + ctx: crate::LayoutReaderContext, ) -> VortexResult { let values_len = usize::try_from(layout.values.row_count())?; let values = layout.values.new_reader( format!("{name}.values").into(), Arc::clone(&segment_source), &session, + &ctx, + )?; + let codes = layout.codes.new_reader( + format!("{name}.codes").into(), + segment_source, + &session, + &ctx, )?; - let codes = - layout - .codes - .new_reader(format!("{name}.codes").into(), segment_source, &session)?; Ok(Self { layout, @@ -370,7 +374,7 @@ mod tests { ); assert!(layout.encoding_id() == LayoutId::new("vortex.dict")); let actual = layout - .new_reader("".into(), segments, &session) + .new_reader("".into(), segments, &session, &Default::default()) .unwrap() .projection_evaluation( &(0..layout.row_count()), @@ -454,7 +458,7 @@ mod tests { )), ); let mask = layout - .new_reader("".into(), segments, &session) + .new_reader("".into(), segments, &session, &Default::default()) .unwrap() .filter_evaluation(&(0..3), &filter, MaskFuture::new_true(3)) .unwrap() @@ -515,7 +519,7 @@ mod tests { let expression = is_not_null(root()); assert_eq!(layout.encoding_id(), LayoutId::new("vortex.dict")); let actual = layout - .new_reader("".into(), segments, &session) + .new_reader("".into(), segments, &session, &Default::default()) .unwrap() .projection_evaluation( &(0..layout.row_count()), diff --git a/vortex-layout/src/layouts/flat/mod.rs b/vortex-layout/src/layouts/flat/mod.rs index b167e633170..a652f82fe61 100644 --- a/vortex-layout/src/layouts/flat/mod.rs +++ b/vortex-layout/src/layouts/flat/mod.rs @@ -87,6 +87,7 @@ impl VTable for Flat { name: Arc, segment_source: Arc, session: &VortexSession, + _ctx: &crate::LayoutReaderContext, ) -> VortexResult { Ok(Arc::new(FlatReader::new( layout.clone(), diff --git a/vortex-layout/src/layouts/flat/reader.rs b/vortex-layout/src/layouts/flat/reader.rs index 8817fe5c678..81e86b2f721 100644 --- a/vortex-layout/src/layouts/flat/reader.rs +++ b/vortex-layout/src/layouts/flat/reader.rs @@ -276,7 +276,7 @@ mod test { ); let result = layout - .new_reader("".into(), segments, &SESSION)? + .new_reader("".into(), segments, &SESSION, &Default::default())? .projection_evaluation( &(0..layout.row_count()), &root(), @@ -313,7 +313,7 @@ mod test { let expr = gt(root(), lit(3i32)); let result = layout - .new_reader("".into(), segments, &SESSION) + .new_reader("".into(), segments, &SESSION, &Default::default()) .unwrap() .projection_evaluation( &(0..layout.row_count()), @@ -350,7 +350,7 @@ mod test { .unwrap(); let result = layout - .new_reader("".into(), segments, &SESSION) + .new_reader("".into(), segments, &SESSION, &Default::default()) .unwrap() .projection_evaluation(&(2..4), &root(), MaskFuture::new_true(2)) .unwrap() diff --git a/vortex-layout/src/layouts/flat/writer.rs b/vortex-layout/src/layouts/flat/writer.rs index da250414951..83b245d49f1 100644 --- a/vortex-layout/src/layouts/flat/writer.rs +++ b/vortex-layout/src/layouts/flat/writer.rs @@ -260,7 +260,7 @@ mod tests { .unwrap(); let result = layout - .new_reader("".into(), segments, &SESSION) + .new_reader("".into(), segments, &SESSION, &Default::default()) .unwrap() .projection_evaluation( &(0..layout.row_count()), @@ -311,7 +311,7 @@ mod tests { .unwrap(); let result = layout - .new_reader("".into(), segments, &SESSION) + .new_reader("".into(), segments, &SESSION, &Default::default()) .unwrap() .projection_evaluation( &(0..layout.row_count()), @@ -384,7 +384,7 @@ mod tests { // We should be able to read the array we just wrote. let result: ArrayRef = layout - .new_reader("".into(), segments, &SESSION) + .new_reader("".into(), segments, &SESSION, &Default::default()) .unwrap() .projection_evaluation( &(0..layout.row_count()), diff --git a/vortex-layout/src/layouts/foreign/mod.rs b/vortex-layout/src/layouts/foreign/mod.rs index 9abacc8b26e..6e0a351b41b 100644 --- a/vortex-layout/src/layouts/foreign/mod.rs +++ b/vortex-layout/src/layouts/foreign/mod.rs @@ -171,6 +171,7 @@ impl Layout for ForeignLayout { _name: Arc, _segment_source: Arc, _session: &VortexSession, + _ctx: &crate::LayoutReaderContext, ) -> VortexResult { vortex_bail!( "Cannot read unknown layout encoding '{}'", diff --git a/vortex-layout/src/layouts/row_idx/mod.rs b/vortex-layout/src/layouts/row_idx/mod.rs index 18fc62b05cd..a88c8ac5d69 100644 --- a/vortex-layout/src/layouts/row_idx/mod.rs +++ b/vortex-layout/src/layouts/row_idx/mod.rs @@ -370,7 +370,9 @@ mod tests { let expr = eq(root(), lit(3i32)); let result = RowIdxLayoutReader::new( 0, - layout.new_reader("".into(), segments, &SESSION).unwrap(), + layout + .new_reader("".into(), segments, &SESSION, &Default::default()) + .unwrap(), SESSION.clone(), ) .projection_evaluation( @@ -411,7 +413,9 @@ mod tests { let expr = gt(row_idx(), lit(3u64)); let result = RowIdxLayoutReader::new( 0, - layout.new_reader("".into(), segments, &SESSION).unwrap(), + layout + .new_reader("".into(), segments, &SESSION, &Default::default()) + .unwrap(), SESSION.clone(), ) .projection_evaluation( @@ -456,7 +460,9 @@ mod tests { let result = RowIdxLayoutReader::new( 0, - layout.new_reader("".into(), segments, &SESSION).unwrap(), + layout + .new_reader("".into(), segments, &SESSION, &Default::default()) + .unwrap(), SESSION.clone(), ) .projection_evaluation( diff --git a/vortex-layout/src/layouts/struct_/mod.rs b/vortex-layout/src/layouts/struct_/mod.rs index 29843f096d8..39ada8ccaec 100644 --- a/vortex-layout/src/layouts/struct_/mod.rs +++ b/vortex-layout/src/layouts/struct_/mod.rs @@ -114,12 +114,14 @@ impl VTable for Struct { name: Arc, segment_source: Arc, session: &VortexSession, + ctx: &crate::LayoutReaderContext, ) -> VortexResult { Ok(Arc::new(StructReader::try_new( layout.clone(), name, segment_source, session.session(), + ctx.clone(), )?)) } diff --git a/vortex-layout/src/layouts/struct_/reader.rs b/vortex-layout/src/layouts/struct_/reader.rs index 156c7456ed6..edac60ff6f4 100644 --- a/vortex-layout/src/layouts/struct_/reader.rs +++ b/vortex-layout/src/layouts/struct_/reader.rs @@ -68,6 +68,7 @@ impl StructReader { name: Arc, segment_source: Arc, session: VortexSession, + ctx: crate::LayoutReaderContext, ) -> VortexResult { let struct_dt = layout.struct_fields(); @@ -99,6 +100,7 @@ impl StructReader { names, Arc::clone(&segment_source), session.clone(), + ctx, ); // Create an expanded root expression that contains all fields of the struct. @@ -616,7 +618,9 @@ mod tests { fn test_struct_layout_or( #[from(struct_layout)] (segments, layout): (Arc, LayoutRef), ) { - let reader = layout.new_reader("".into(), segments, &SESSION).unwrap(); + let reader = layout + .new_reader("".into(), segments, &SESSION, &Default::default()) + .unwrap(); let filt = or( eq(col("a"), lit(7)), or(eq(col("b"), lit(5)), eq(col("a"), lit(3))), @@ -634,7 +638,9 @@ mod tests { fn test_struct_layout( #[from(struct_layout)] (segments, layout): (Arc, LayoutRef), ) { - let reader = layout.new_reader("".into(), segments, &SESSION).unwrap(); + let reader = layout + .new_reader("".into(), segments, &SESSION, &Default::default()) + .unwrap(); let expr = gt(get_item("a", root()), get_item("b", root())); let result = block_on(|_| { reader @@ -650,7 +656,9 @@ mod tests { fn test_struct_layout_row_mask( #[from(struct_layout)] (segments, layout): (Arc, LayoutRef), ) { - let reader = layout.new_reader("".into(), segments, &SESSION).unwrap(); + let reader = layout + .new_reader("".into(), segments, &SESSION, &Default::default()) + .unwrap(); let expr = gt(get_item("a", root()), get_item("b", root())); let result = block_on(|_| { reader @@ -672,7 +680,9 @@ mod tests { #[from(struct_layout)] (segments, layout): (Arc, LayoutRef), ) { let mut ctx = LEGACY_SESSION.create_execution_ctx(); - let reader = layout.new_reader("".into(), segments, &SESSION).unwrap(); + let reader = layout + .new_reader("".into(), segments, &SESSION, &Default::default()) + .unwrap(); let expr = pack( [("a", get_item("a", root())), ("b", get_item("b", root()))], Nullability::NonNullable, @@ -711,7 +721,9 @@ mod tests { #[from(null_struct_layout)] (segments, layout): (Arc, LayoutRef), ) { // Read the layout source from the top. - let reader = layout.new_reader("".into(), segments, &SESSION).unwrap(); + let reader = layout + .new_reader("".into(), segments, &SESSION, &Default::default()) + .unwrap(); let expr = get_item("a", root()); let project = reader .projection_evaluation(&(0..3), &expr, MaskFuture::new_true(3)) @@ -742,7 +754,9 @@ mod tests { // Project out the nested struct field. // The projection should preserve the nulls of the `b` struct when we select out the // child column `c`. - let reader = layout.new_reader("".into(), segments, &SESSION).unwrap(); + let reader = layout + .new_reader("".into(), segments, &SESSION, &Default::default()) + .unwrap(); let expr = select( vec![FieldName::from("c")], get_item("b", get_item("a", root())), @@ -803,7 +817,9 @@ mod tests { fn test_empty_struct( #[from(empty_struct)] (segments, layout): (Arc, LayoutRef), ) { - let reader = layout.new_reader("".into(), segments, &SESSION).unwrap(); + let reader = layout + .new_reader("".into(), segments, &SESSION, &Default::default()) + .unwrap(); let expr = pack(Vec::<(String, Expression)>::new(), Nullability::Nullable); let project = reader @@ -854,7 +870,9 @@ mod tests { }) .unwrap(); - let reader = layout.new_reader("".into(), segments, &SESSION).unwrap(); + let reader = layout + .new_reader("".into(), segments, &SESSION, &Default::default()) + .unwrap(); // DType mismatch: "age" is u8 but literal is i32 let filt = eq(col("age"), lit(67i32)); diff --git a/vortex-layout/src/layouts/zoned/mod.rs b/vortex-layout/src/layouts/zoned/mod.rs index 697e55e968c..43f9410e831 100644 --- a/vortex-layout/src/layouts/zoned/mod.rs +++ b/vortex-layout/src/layouts/zoned/mod.rs @@ -115,12 +115,14 @@ impl VTable for Zoned { name: Arc, segment_source: Arc, session: &VortexSession, + ctx: &crate::LayoutReaderContext, ) -> VortexResult { Ok(Arc::new(ZonedReader::try_new( layout.clone(), name, segment_source, session.clone(), + ctx.clone(), )?)) } diff --git a/vortex-layout/src/layouts/zoned/reader.rs b/vortex-layout/src/layouts/zoned/reader.rs index a73ed6a547c..6a544c7af1c 100644 --- a/vortex-layout/src/layouts/zoned/reader.rs +++ b/vortex-layout/src/layouts/zoned/reader.rs @@ -42,6 +42,7 @@ impl ZonedReader { name: Arc, segment_source: Arc, session: VortexSession, + ctx: crate::LayoutReaderContext, ) -> VortexResult { let dtypes = vec![ layout.dtype.clone(), @@ -54,6 +55,7 @@ impl ZonedReader { names, Arc::clone(&segment_source), session.clone(), + ctx, )); Ok(Self { @@ -307,7 +309,7 @@ mod test { block_on(|handle| async { let session = session_with_handle(handle); let result = layout - .new_reader("".into(), segments, &session) + .new_reader("".into(), segments, &session, &Default::default()) .unwrap() .projection_evaluation( &(0..layout.row_count()), @@ -330,7 +332,9 @@ mod test { block_on(|handle| async { let row_count = layout.row_count(); let session = session_with_handle(handle); - let reader = layout.new_reader("".into(), segments, &session).unwrap(); + let reader = layout + .new_reader("".into(), segments, &session, &Default::default()) + .unwrap(); // Choose a prune-able expression let expr = gt(root(), lit(7)); @@ -376,7 +380,8 @@ mod test { block_on(|handle| async { let row_count = legacy_layout.row_count(); let session = session_with_handle(handle); - let reader = legacy_layout.new_reader("".into(), segments, &session)?; + let reader = + legacy_layout.new_reader("".into(), segments, &session, &Default::default())?; let result = reader .pruning_evaluation( diff --git a/vortex-layout/src/reader.rs b/vortex-layout/src/reader.rs index b40339e3946..d22afd4a4f4 100644 --- a/vortex-layout/src/reader.rs +++ b/vortex-layout/src/reader.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::any::Any; +use std::any::TypeId; use std::collections::BTreeSet; use std::ops::Range; use std::sync::Arc; @@ -20,6 +21,7 @@ use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_mask::Mask; use vortex_session::VortexSession; +use vortex_utils::aliases::hash_map::HashMap; use crate::children::LayoutChildren; use crate::segments::SegmentSource; @@ -181,6 +183,7 @@ pub struct LazyReaderChildren { names: Vec>, segment_source: Arc, session: VortexSession, + ctx: LayoutReaderContext, // TODO(ngates): we may want a hash map of some sort here? cache: Vec>, } @@ -192,6 +195,7 @@ impl LazyReaderChildren { names: Vec>, segment_source: Arc, session: VortexSession, + ctx: LayoutReaderContext, ) -> Self { let nchildren = children.nchildren(); let cache = (0..nchildren).map(|_| OnceCell::new()).collect(); @@ -201,6 +205,7 @@ impl LazyReaderChildren { names, segment_source, session, + ctx, cache, } } @@ -217,7 +222,62 @@ impl LazyReaderChildren { Arc::clone(&self.names[idx]), Arc::clone(&self.segment_source), &self.session, + &self.ctx, ) }) } } + +/// Per-reader-tree dependency context, threaded through [`crate::VTable::new_reader`]. +/// +/// Holds a typed-data registry keyed by [`TypeId`]. Ancestors publish values via +/// [`Self::with`]; descendants retrieve them via [`Self::get`]. This is a *read-only* +/// channel from the ancestor's perspective — descendants can only consume what their +/// ancestors chose to publish, +/// +/// [`Self::with`] returns a derived context that copies the existing map and adds (or +/// replaces) one entry. The original is unchanged, allowing concurrent reader-tree +/// constructions to each derive their own context without races. When the same `T` is +/// published at multiple ancestors, the innermost (most-recently-published) wins. +/// +/// Contexts are cheap to clone (shared via `Arc`), so they can be captured by lazy +/// children helpers and survive any single stack frame. +#[derive(Clone, Default)] +pub struct LayoutReaderContext { + values: Arc>>, +} + +impl std::fmt::Debug for LayoutReaderContext { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("LayoutReaderContext") + .field("type_ids", &self.values.keys().collect::>()) + .finish() + } +} + +impl LayoutReaderContext { + /// Creates a new, empty context. + pub fn new() -> Self { + Self::default() + } + + /// Returns a new context that publishes `value` under its concrete type. + /// + /// The original context is unchanged. If a value of the same type was already + /// published, the new one replaces it in the returned context. + pub fn with(&self, value: Arc) -> Self { + let mut values = HashMap::clone(&self.values); + values.insert(TypeId::of::(), value); + Self { + values: Arc::new(values), + } + } + + /// Returns the most-recently-published value of type `T`, or `None` if no ancestor + /// published one. + pub fn get(&self) -> Option> { + self.values + .get(&TypeId::of::()) + .and_then(|v| Arc::clone(v).downcast::().ok()) + } +} diff --git a/vortex-layout/src/scan/split_by.rs b/vortex-layout/src/scan/split_by.rs index 8cc6e2e095a..d17957a739b 100644 --- a/vortex-layout/src/scan/split_by.rs +++ b/vortex-layout/src/scan/split_by.rs @@ -98,7 +98,7 @@ mod test { .unwrap(); layout - .new_reader("".into(), segments, &SCAN_SESSION) + .new_reader("".into(), segments, &SCAN_SESSION, &Default::default()) .unwrap() } diff --git a/vortex-layout/src/vtable.rs b/vortex-layout/src/vtable.rs index 9bb83fcf643..c48a9d33d1f 100644 --- a/vortex-layout/src/vtable.rs +++ b/vortex-layout/src/vtable.rs @@ -19,6 +19,7 @@ use crate::LayoutChildType; use crate::LayoutEncoding; use crate::LayoutEncodingRef; use crate::LayoutId; +use crate::LayoutReaderContext; use crate::LayoutReaderRef; use crate::LayoutRef; use crate::children::LayoutChildren; @@ -58,11 +59,20 @@ pub trait VTable: 'static + Sized + Send + Sync + Debug { fn child_type(layout: &Self::Layout, idx: usize) -> LayoutChildType; /// Create a new reader for the layout. + /// + /// **Layouts with children MUST propagate `ctx` to descendants** by passing it + /// through `Layout::new_reader` (or `LazyReaderChildren::new`) when constructing + /// child readers. If `ctx` is dropped at any link in the chain, ancestor-published + /// values won't reach affected descendants — a silent runtime regression for any + /// descendant that looked up an ancestor-published value via `ctx.get::()`. + /// There is no compile-time check that catches this; reviewer discipline + the + /// integration tests in `vortex-layout` are the only safety net. fn new_reader( layout: &Self::Layout, name: Arc, segment_source: Arc, session: &VortexSession, + ctx: &LayoutReaderContext, ) -> VortexResult; /// Construct a new [`Layout`] from the provided parts. diff --git a/vortex-test/compat-gen/src/adapter.rs b/vortex-test/compat-gen/src/adapter.rs index a97399dac79..bed1300b531 100644 --- a/vortex-test/compat-gen/src/adapter.rs +++ b/vortex-test/compat-gen/src/adapter.rs @@ -136,7 +136,12 @@ pub fn read_layout_tree(bytes: ByteBuffer) -> VortexResult<()> { if row_count == 0 { continue; } - let reader = layout.new_reader("".into(), Arc::clone(&segment_source), &session)?; + let reader = layout.new_reader( + "".into(), + Arc::clone(&segment_source), + &session, + &Default::default(), + )?; let len = usize::try_from(row_count).map_err(|e| vortex_err!("row count overflow: {e}"))?; reader diff --git a/vortex-tui/src/browse/app.rs b/vortex-tui/src/browse/app.rs index 0a7f13ab277..97e17bac205 100644 --- a/vortex-tui/src/browse/app.rs +++ b/vortex-tui/src/browse/app.rs @@ -360,7 +360,12 @@ impl AppState { // Load the array. let reader = layout - .new_reader("".into(), self.vxf.segment_source(), &self.session) + .new_reader( + "".into(), + self.vxf.segment_source(), + &self.session, + &Default::default(), + ) .vortex_expect("Failed to create reader"); let array = reader .projection_evaluation( diff --git a/vortex-tui/src/wasm.rs b/vortex-tui/src/wasm.rs index 8b643161fe2..5fd7dd9054b 100644 --- a/vortex-tui/src/wasm.rs +++ b/vortex-tui/src/wasm.rs @@ -78,7 +78,12 @@ async fn load_flat_array( row_count: u64, ) -> vortex::array::ArrayRef { let reader = layout - .new_reader("".into(), segment_source.clone(), session) + .new_reader( + "".into(), + segment_source.clone(), + session, + &Default::default(), + ) .vortex_expect("Failed to create reader"); reader .projection_evaluation( diff --git a/vortex-web/crate/src/wasm.rs b/vortex-web/crate/src/wasm.rs index 05e328e5341..f1bb4df791e 100644 --- a/vortex-web/crate/src/wasm.rs +++ b/vortex-web/crate/src/wasm.rs @@ -293,7 +293,12 @@ impl VortexFileHandle { let segment_source = self.vxf.segment_source(); let reader = layout - .new_reader(node_id.into(), segment_source, &self.session) + .new_reader( + node_id.into(), + segment_source, + &self.session, + &Default::default(), + ) .map_err(|e| JsValue::from_str(&e.to_string()))?; let stream = ScanBuilder::new(self.session.clone(), reader)