LazyReaderChildren only allocate once

Signed-off-by: Andrew Duffy <[email protected]>

Author: a10y

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

@@ -7,7 +7,7 @@ use std::sync::Arc;
 
 use futures::future::BoxFuture;
 use futures::stream::FuturesOrdered;
-use futures::{FutureExt, TryStreamExt};
+use futures::{FutureExt, StreamExt, TryStreamExt};
 use vortex_array::arrays::ChunkedArray;
 use vortex_array::{ArrayRef, MaskFuture};
 use vortex_dtype::{DType, FieldMask};
@@ -43,7 +43,12 @@ impl ChunkedReader {
         }
         chunk_offsets[nchildren] = layout.row_count();
 
-        let lazy_children = LazyReaderChildren::new(layout.children.clone(), segment_source);
+        let dtypes = vec![layout.dtype.clone(); nchildren];
+        let names = (0..nchildren)
+            .map(|idx| Arc::from(format!("{name}.[{idx}]")))
+            .collect();
+        let lazy_children =
+            LazyReaderChildren::new(layout.children.clone(), dtypes, names, segment_source);
 
         Self {
             layout,
@@ -55,11 +60,7 @@ impl ChunkedReader {
 
     /// Return the [`LayoutReader`] for the given chunk.
     fn chunk_reader(&self, idx: usize) -> VortexResult<&LayoutReaderRef> {
-        self.lazy_children.get(
-            idx,
-            self.layout.dtype(),
-            &format!("{}.[{}]", self.name, idx).into(),
-        )
+        self.lazy_children.get(idx)
     }
 
     fn chunk_offset(&self, idx: usize) -> u64 {

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

@@ -58,8 +58,24 @@ impl StructReader {
                 .collect()
         });
 
-        let lazy_children =
-            LazyReaderChildren::new(layout.children.clone(), segment_source.clone());
+        let mut dtypes: Vec<DType> = struct_dt.fields().collect();
+        let mut names: Vec<Arc<str>> = struct_dt
+            .names()
+            .iter()
+            .map(|x| x.inner().clone())
+            .collect();
+
+        if layout.dtype.is_nullable() {
+            dtypes.insert(0, DType::Bool(Nullability::NonNullable));
+            names.insert(0, Arc::from("validity"));
+        }
+
+        let lazy_children = LazyReaderChildren::new(
+            layout.children.clone(),
+            dtypes,
+            names,
+            segment_source.clone(),
+        );
 
         // Create an expanded root expression that contains all fields of the struct.
         let expanded_root_expr = replace_root_fields(root(), struct_dt);
@@ -105,24 +121,14 @@ impl StructReader {
             .field_by_index(idx)
             .ok_or_else(|| vortex_err!("Missing field {idx}"))?;
         let name = &self.struct_fields().names()[idx];
-        self.lazy_children.get(
-            child_index,
-            &field_dtype,
-            &format!("{}.{}", self.name, name).into(),
-        )
+        self.lazy_children.get(child_index)
     }
 
     /// Return the reader for the struct validity, if present
     fn validity(&self) -> VortexResult<Option<&LayoutReaderRef>> {
         self.dtype()
             .is_nullable()
-            .then(|| {
-                self.lazy_children.get(
-                    0,
-                    &DType::Bool(Nullability::NonNullable),
-                    &"validity".into(),
-                )
-            })
+            .then(|| self.lazy_children.get(0))
             .transpose()
     }
 

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

@@ -50,8 +50,17 @@ impl ZonedReader {
         name: Arc<str>,
         segment_source: Arc<dyn SegmentSource>,
     ) -> VortexResult<Self> {
-        let lazy_children =
-            LazyReaderChildren::new(layout.children.clone(), segment_source.clone());
+        let dtypes = vec![
+            layout.dtype.clone(),
+            ZoneMap::dtype_for_stats_table(layout.dtype(), layout.present_stats()),
+        ];
+        let names = vec![name.clone(), format!("{}.zones", name).into()];
+        let lazy_children = LazyReaderChildren::new(
+            layout.children.clone(),
+            dtypes,
+            names,
+            segment_source.clone(),
+        );
 
         Ok(Self {
             layout,
@@ -65,7 +74,7 @@ impl ZonedReader {
 
     #[inline]
     fn data_child(&self) -> VortexResult<&LayoutReaderRef> {
-        self.lazy_children.get(0, self.layout.dtype(), &self.name)
+        self.lazy_children.get(0)
     }
 
     /// Get or create the pruning predicate for a given expression.
@@ -97,14 +106,7 @@ impl ZonedReader {
 
                 let zones_eval = self
                     .lazy_children
-                    .get(
-                        1,
-                        &ZoneMap::dtype_for_stats_table(
-                            self.layout.dtype(),
-                            self.layout.present_stats(),
-                        ),
-                        &format!("{}.zones", self.name).into(),
-                    )
+                    .get(1)
                     .vortex_expect("failed to get zone child")
                     .projection_evaluation(
                         &(0..nzones as u64),

vortex-layout/src/reader.rs

@@ -101,36 +101,41 @@ impl ArrayFutureExt for ArrayFuture {
 
 pub struct LazyReaderChildren {
     children: Arc<dyn LayoutChildren>,
+    dtypes: Vec<DType>,
+    names: Vec<Arc<str>>,
     segment_source: Arc<dyn SegmentSource>,
 
     // TODO(ngates): we may want a hash map of some sort here?
     cache: Vec<OnceCell<LayoutReaderRef>>,
 }
 
 impl LazyReaderChildren {
-    pub fn new(children: Arc<dyn LayoutChildren>, segment_source: Arc<dyn SegmentSource>) -> Self {
+    pub fn new(
+        children: Arc<dyn LayoutChildren>,
+        dtypes: Vec<DType>,
+        names: Vec<Arc<str>>,
+        segment_source: Arc<dyn SegmentSource>,
+    ) -> Self {
         let nchildren = children.nchildren();
         let cache = (0..nchildren).map(|_| OnceCell::new()).collect();
         Self {
             children,
+            dtypes,
+            names,
             segment_source,
             cache,
         }
     }
 
-    pub fn get(
-        &self,
-        idx: usize,
-        dtype: &DType,
-        name: &Arc<str>,
-    ) -> VortexResult<&LayoutReaderRef> {
+    pub fn get(&self, idx: usize) -> VortexResult<&LayoutReaderRef> {
         if idx >= self.cache.len() {
             vortex_bail!("Child index out of bounds: {} of {}", idx, self.cache.len());
         }
 
         self.cache[idx].get_or_try_init(|| {
+            let dtype = &self.dtypes[idx];
             let child = self.children.child(idx, dtype)?;
-            child.new_reader(name.clone(), self.segment_source.clone())
+            child.new_reader(Arc::clone(&self.names[idx]), self.segment_source.clone())
         })
     }
 }