chore: ZonedLayoutReader children are lazy (#5007)

Avoids constructing pruning data layout if there's no predicate to run

Signed-off-by: Robert Kruszewski <[email protected]>

---------

Signed-off-by: Robert Kruszewski <[email protected]>

Author: robert3005

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

@@ -14,7 +14,7 @@ use vortex_array::{ArrayContext, DeserializeMetadata, SerializeMetadata};
 use vortex_dtype::{DType, TryFromBytes};
 use vortex_error::{VortexExpect, VortexResult, vortex_bail, vortex_panic};
 
-use crate::children::LayoutChildren;
+use crate::children::{LayoutChildren, OwnedLayoutChildren};
 use crate::layouts::zoned::reader::ZonedReader;
 use crate::layouts::zoned::zone_map::ZoneMap;
 use crate::segments::{SegmentId, SegmentSource};
@@ -38,11 +38,11 @@ impl VTable for ZonedVTable {
     }
 
     fn row_count(layout: &Self::Layout) -> u64 {
-        layout.data.row_count()
+        layout.children.child_row_count(0)
     }
 
     fn dtype(layout: &Self::Layout) -> &DType {
-        layout.data.dtype()
+        &layout.dtype
     }
 
     fn metadata(layout: &Self::Layout) -> Self::Metadata {
@@ -62,8 +62,11 @@ impl VTable for ZonedVTable {
 
     fn child(layout: &Self::Layout, idx: usize) -> VortexResult<LayoutRef> {
         match idx {
-            0 => Ok(layout.data.clone()),
-            1 => Ok(layout.zones.clone()),
+            0 => layout.children.child(0, layout.dtype()),
+            1 => layout.children.child(
+                1,
+                &ZoneMap::dtype_for_stats_table(layout.dtype(), &layout.present_stats),
+            ),
             _ => vortex_bail!("Invalid child index: {}", idx),
         }
     }
@@ -92,22 +95,17 @@ impl VTable for ZonedVTable {
         _encoding: &Self::Encoding,
         dtype: &DType,
         _row_count: u64,
-        metadata: &<Self::Metadata as DeserializeMetadata>::Output,
+        metadata: &ZonedMetadata,
         _segment_ids: Vec<SegmentId>,
         children: &dyn LayoutChildren,
         _ctx: ArrayContext,
     ) -> VortexResult<Self::Layout> {
-        let data = children.child(0, dtype)?;
-
-        let zones_dtype = ZoneMap::dtype_for_stats_table(data.dtype(), &metadata.present_stats);
-        let zones = children.child(1, &zones_dtype)?;
-
-        Ok(ZonedLayout::new(
-            data,
-            zones,
-            metadata.zone_len as usize,
-            metadata.present_stats.clone(),
-        ))
+        Ok(ZonedLayout {
+            dtype: dtype.clone(),
+            children: children.to_arc(),
+            zone_len: metadata.zone_len as usize,
+            present_stats: metadata.present_stats.clone(),
+        })
     }
 }
 
@@ -116,8 +114,8 @@ pub struct ZonedLayoutEncoding;
 
 #[derive(Clone, Debug)]
 pub struct ZonedLayout {
-    data: LayoutRef,
-    zones: LayoutRef,
+    dtype: DType,
+    children: Arc<dyn LayoutChildren>,
     zone_len: usize,
     present_stats: Arc<[Stat]>,
 }
@@ -137,15 +135,15 @@ impl ZonedLayout {
             vortex_panic!("Invalid zone map layout: zones dtype does not match expected dtype");
         }
         Self {
-            data,
-            zones,
+            dtype: data.dtype().clone(),
+            children: OwnedLayoutChildren::layout_children(vec![data, zones]),
             zone_len,
             present_stats,
         }
     }
 
     pub fn nzones(&self) -> usize {
-        usize::try_from(self.zones.row_count()).vortex_expect("Invalid number of zones")
+        usize::try_from(self.children.child_row_count(1)).vortex_expect("Invalid number of zones")
     }
 
     /// Returns an array of stats that exist in the layout's data, must be sorted.
@@ -186,7 +184,6 @@ impl SerializeMetadata for ZonedMetadata {
 
 #[cfg(test)]
 mod tests {
-
     use rstest::rstest;
 
     use super::*;

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

@@ -3,7 +3,7 @@
 
 use std::collections::BTreeSet;
 use std::ops::{BitAnd, Range};
-use std::sync::{Arc, OnceLock};
+use std::sync::{Arc, LazyLock, OnceLock};
 
 use arrow_buffer::BooleanBufferBuilder;
 use futures::future::{BoxFuture, Shared};
@@ -20,10 +20,10 @@ use vortex_expr::{ExprRef, root};
 use vortex_mask::Mask;
 use vortex_utils::aliases::dash_map::DashMap;
 
-use crate::LayoutReader;
 use crate::layouts::zoned::ZonedLayout;
 use crate::layouts::zoned::zone_map::ZoneMap;
 use crate::segments::SegmentSource;
+use crate::{LayoutReader, LayoutReaderRef, LazyReaderChildren};
 
 type SharedZoneMap = Shared<BoxFuture<'static, SharedVortexResult<ZoneMap>>>;
 type SharedPruningResult = Shared<BoxFuture<'static, SharedVortexResult<Arc<PruningResult>>>>;
@@ -32,21 +32,17 @@ type PredicateCache = Arc<OnceLock<Option<ExprRef>>>;
 pub struct ZonedReader {
     layout: ZonedLayout,
     name: Arc<str>,
-
-    /// Data layout reader
-    data_child: Arc<dyn LayoutReader>,
-    /// Zone map layout reader.
-    zones_child: Arc<dyn LayoutReader>,
+    lazy_children: LazyReaderChildren,
 
     /// A cache of expr -> optional pruning result (applying the pruning expr to the zone map)
-    pruning_result: DashMap<ExprRef, Option<SharedPruningResult>>,
+    pruning_result: LazyLock<DashMap<ExprRef, Option<SharedPruningResult>>>,
 
     /// Shared zone map
     zone_map: OnceLock<SharedZoneMap>,
 
     /// A cache of expr -> optional pruning predicate.
     /// This also uses the present_stats from the `ZonedLayout`
-    pruning_predicates: Arc<DashMap<ExprRef, PredicateCache>>,
+    pruning_predicates: LazyLock<Arc<DashMap<ExprRef, PredicateCache>>>,
 }
 
 impl ZonedReader {
@@ -55,24 +51,24 @@ impl ZonedReader {
         name: Arc<str>,
         segment_source: Arc<dyn SegmentSource>,
     ) -> VortexResult<Self> {
-        let data_child = layout
-            .data
-            .new_reader(name.clone(), segment_source.clone())?;
-        let zones_child = layout
-            .zones
-            .new_reader(format!("{name}.zones").into(), segment_source)?;
+        let lazy_children =
+            LazyReaderChildren::new(layout.children.clone(), segment_source.clone());
 
         Ok(Self {
             layout,
             name,
-            data_child,
-            zones_child,
+            lazy_children,
             pruning_result: Default::default(),
             zone_map: Default::default(),
             pruning_predicates: Default::default(),
         })
     }
 
+    #[inline]
+    fn data_child(&self) -> VortexResult<&LayoutReaderRef> {
+        self.lazy_children.get(0, self.layout.dtype(), &self.name)
+    }
+
     /// Get or create the pruning predicate for a given expression.
     fn pruning_predicate(&self, expr: ExprRef) -> Option<ExprRef> {
         self.pruning_predicates
@@ -101,7 +97,16 @@ impl ZonedReader {
                 let present_stats = self.layout.present_stats.clone();
 
                 let zones_eval = self
-                    .zones_child
+                    .lazy_children
+                    .get(
+                        1,
+                        &ZoneMap::dtype_for_stats_table(
+                            self.layout.dtype(),
+                            self.layout.present_stats(),
+                        ),
+                        &format!("{}.zones", self.name).into(),
+                    )
+                    .vortex_expect("failed to get zone child")
                     .projection_evaluation(
                         &(0..nzones as u64),
                         &root(),
@@ -178,11 +183,11 @@ impl LayoutReader for ZonedReader {
     }
 
     fn dtype(&self) -> &DType {
-        self.data_child.dtype()
+        self.layout.dtype()
     }
 
     fn row_count(&self) -> Precision<u64> {
-        self.data_child.row_count()
+        Precision::exact(self.layout.row_count())
     }
 
     fn register_splits(
@@ -191,7 +196,7 @@ impl LayoutReader for ZonedReader {
         row_offset: u64,
         splits: &mut BTreeSet<u64>,
     ) -> VortexResult<()> {
-        self.data_child
+        self.data_child()?
             .register_splits(field_mask, row_offset, splits)
     }
 
@@ -203,7 +208,7 @@ impl LayoutReader for ZonedReader {
     ) -> VortexResult<MaskFuture> {
         log::debug!("Stats pruning evaluation: {} - {}", &self.name, expr);
         let data_eval = self
-            .data_child
+            .data_child()?
             .pruning_evaluation(row_range, expr, mask.clone())?;
 
         let Some(pruning_mask_future) = self.pruning_mask_future(expr.clone()) else {
@@ -236,10 +241,10 @@ impl LayoutReader for ZonedReader {
         Ok(MaskFuture::new(mask.len(), async move {
             log::debug!("Invoking stats pruning evaluation {}: {}", name, expr);
 
-            let pruning_mask = pruning_mask_future.clone().await?.mask()?;
+            let pruning_mask = pruning_mask_future.await?.mask()?;
 
             let mut builder = BooleanBufferBuilder::new(mask.len());
-            for (zone_idx, &zone_length) in zone_range.clone().zip_eq(&zone_lengths) {
+            for (zone_idx, &zone_length) in zone_range.zip_eq(&zone_lengths) {
                 builder.append_n(zone_length, !pruning_mask.value(usize::try_from(zone_idx)?));
             }
 
@@ -273,7 +278,7 @@ impl LayoutReader for ZonedReader {
         expr: &ExprRef,
         mask: MaskFuture,
     ) -> VortexResult<MaskFuture> {
-        self.data_child.filter_evaluation(row_range, expr, mask)
+        self.data_child()?.filter_evaluation(row_range, expr, mask)
     }
 
     fn projection_evaluation(
@@ -284,7 +289,8 @@ impl LayoutReader for ZonedReader {
     ) -> VortexResult<BoxFuture<'static, VortexResult<ArrayRef>>> {
         // TODO(ngates): there are some projection expressions that we may also be able to
         //  short-circuit with statistics.
-        self.data_child.projection_evaluation(row_range, expr, mask)
+        self.data_child()?
+            .projection_evaluation(row_range, expr, mask)
     }
 }