diff --git a/vortex-layout/src/layouts/chunked/reader.rs b/vortex-layout/src/layouts/chunked/reader.rs index b5f4f5b438d..7eed0e1ca24 100644 --- a/vortex-layout/src/layouts/chunked/reader.rs +++ b/vortex-layout/src/layouts/chunked/reader.rs @@ -41,7 +41,7 @@ use crate::segments::SegmentSource; pub struct ChunkedReader { layout: ChunkedLayout, name: Arc, - lazy_children: LazyReaderChildren, + lazy_children: Arc, /// Lazily computed classification of which chunks register no interior splits, letting /// [`ChunkedReader::register_splits`] avoid materializing those chunks' readers. chunk_skips: OnceCell, @@ -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::>()); } + #[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, LayoutRef), diff --git a/vortex-layout/src/layouts/zoned/reader.rs b/vortex-layout/src/layouts/zoned/reader.rs index 52a57c1b92b..162898aa382 100644 --- a/vortex-layout/src/layouts/zoned/reader.rs +++ b/vortex-layout/src/layouts/zoned/reader.rs @@ -140,13 +140,18 @@ impl LayoutReader for ZonedReader { mask: Mask, ) -> VortexResult { 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); }