Defer chunk reader construction until after pruning

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

@@ -41,7 +41,7 @@ use crate::segments::SegmentSource;
 pub struct ChunkedReader {
     layout: ChunkedLayout,
     name: Arc<str>,
-    lazy_children: LazyReaderChildren,
+    lazy_children: Arc<LazyReaderChildren>,
     /// Lazily computed classification of which chunks register no interior splits, letting
     /// [`ChunkedReader::register_splits`] avoid materializing those chunks' readers.
     chunk_skips: OnceCell<ChunkSkips>,
@@ -98,7 +98,7 @@ impl ChunkedReader {
         Self {
             layout,
             name,
-            lazy_children,
+            lazy_children: Arc::new(lazy_children),
             chunk_skips: OnceCell::new(),
         }
     }
@@ -305,9 +305,16 @@ impl LayoutReader for ChunkedReader {
         let mut chunk_evals = vec![];
 
         for (chunk_idx, _, chunk_range, mask_range) in self.ranges(row_range) {
+            let chunk_mask = mask.slice(mask_range);
+            if chunk_mask.all_false() {
+                // Preserve already-excluded rows without materializing the child reader.
+                chunk_evals.push(MaskFuture::ready(chunk_mask));
+                continue;
+            }
+
             let chunk_reader = self.chunk_reader(chunk_idx)?;
             let chunk_eval = chunk_reader
-                .pruning_evaluation(&chunk_range, expr, mask.slice(mask_range))
+                .pruning_evaluation(&chunk_range, expr, chunk_mask)
                 .map_err(|err| {
                     err.with_context(format!(
                         "While evaluating pruning filter on chunk {chunk_idx}"
@@ -351,12 +358,29 @@ impl LayoutReader for ChunkedReader {
         let mut chunk_evals = vec![];
 
         for (chunk_idx, _, chunk_range, mask_range) in self.ranges(row_range) {
-            let chunk_reader = self.chunk_reader(chunk_idx)?;
-            let chunk_eval = chunk_reader
-                .filter_evaluation(&chunk_range, expr, mask.slice(mask_range))
-                .map_err(|err| {
-                    err.with_context(format!("While evaluating filter on chunk {chunk_idx}"))
-                })?;
+            let lazy_children = Arc::clone(&self.lazy_children);
+            let expr = expr.clone();
+            let chunk_mask = mask.slice(mask_range);
+            let chunk_len = chunk_mask.len();
+
+            let chunk_eval = MaskFuture::new(chunk_len, async move {
+                let chunk_mask = chunk_mask.await?;
+                if chunk_mask.all_false() {
+                    return Ok(chunk_mask);
+                }
+
+                let chunk_reader = Arc::clone(lazy_children.get(chunk_idx)?);
+                let chunk_eval = chunk_reader
+                    .filter_evaluation(
+                        &chunk_range,
+                        &expr,
+                        MaskFuture::ready(chunk_mask),
+                    )
+                    .map_err(|err| {
+                        err.with_context(format!("While evaluating filter on chunk {chunk_idx}"))
+                    })?;
+                chunk_eval.await
+            });
             chunk_evals.push(chunk_eval);
         }
 
@@ -442,6 +466,7 @@ mod test {
     use vortex_buffer::buffer;
     use vortex_io::runtime::single::block_on;
     use vortex_io::session::RuntimeSessionExt;
+    use vortex_mask::Mask;
     use vortex_session::registry::ReadContext;
 
     use crate::LayoutRef;
@@ -546,6 +571,35 @@ mod test {
         assert_eq!(splits, expected.into_iter().collect::<Vec<_>>());
     }
 
+    #[test]
+    fn test_filter_skips_fully_masked_chunks() {
+        let layout = nested_chunked_layout();
+        block_on(|_handle| async {
+            let reader = layout
+                .new_reader(
+                    "".into(),
+                    Arc::new(TestSegments::default()),
+                    &SESSION,
+                    &Default::default(),
+                )
+                .unwrap();
+            let expr = root().bind(reader.dtype()).unwrap();
+            let row_count = usize::try_from(layout.row_count()).unwrap();
+
+            let result = reader
+                .filter_evaluation(
+                    &(0..layout.row_count()),
+                    &expr,
+                    MaskFuture::ready(Mask::new_false(row_count)),
+                )
+                .unwrap()
+                .await
+                .unwrap();
+
+            assert!(result.all_false());
+        })
+    }
+
     #[rstest]
     fn test_chunked_evaluator(
         #[from(chunked_layout)] (segments, layout): (Arc<dyn SegmentSource>, LayoutRef),

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

@@ -140,13 +140,18 @@ impl LayoutReader for ZonedReader {
         mask: Mask,
     ) -> VortexResult<MaskFuture> {
         trace!("Stats pruning evaluation: {} - {}", &self.name, expr);
-        let data_eval = self
-            .data_child()?
-            .pruning_evaluation(row_range, expr, mask.clone())?;
 
         let Some(pruning_mask_future) = self.pruning.pruning_mask_future(expr.clone()) else {
             trace!("Stats pruning evaluation: not prune-able {expr}");
-            return Ok(data_eval);
+            let lazy_children = Arc::clone(&self.lazy_children);
+            let row_range = row_range.clone();
+            let expr = expr.clone();
+            return Ok(MaskFuture::new(mask.len(), async move {
+                let data_child = Arc::clone(lazy_children.get(0)?);
+                data_child
+                    .pruning_evaluation(&row_range, &expr, mask)?
+                    .await
+            }));
         };
 
         let row_count = row_range.end - row_range.start;
@@ -169,6 +174,8 @@ impl LayoutReader for ZonedReader {
             .try_collect()?;
 
         let name = Arc::clone(&self.name);
+        let lazy_children = Arc::clone(&self.lazy_children);
+        let row_range = row_range.clone();
         let expr = expr.clone();
 
         Ok(MaskFuture::new(mask.len(), async move {
@@ -188,8 +195,15 @@ impl LayoutReader for ZonedReader {
             let mask_density = mask.density();
             let mut stats_mask = mask.bitand(&stats_mask);
 
-            // Forward to data child for further pruning.
+            // Only materialize the data child after statistics have pruned the input, and pass
+            // that sparse mask down so chunked readers can skip fully excluded chunks.
             if !stats_mask.all_false() {
+                let data_child = Arc::clone(lazy_children.get(0)?);
+                let data_eval = data_child.pruning_evaluation(
+                    &row_range,
+                    &expr,
+                    stats_mask.clone(),
+                )?;
                 let data_mask = data_eval.await?;
                 stats_mask = stats_mask.bitand(&data_mask);
             }