Vortex Layouts - chunk pruning (#1824)

gatesn · web-flow · commit 98ac114a7702 · 2025-01-06T14:30:07.000Z
diff --git a/vortex-expr/src/pruning.rs b/vortex-expr/src/pruning.rs
@@ -127,7 +127,7 @@ impl PruningPredicate {
     ///
     /// Returns Ok(None) if any of the required statistics are not present in metadata.
     /// If it returns Ok(Some(array)), the array is a boolean array with the same length as the
-    /// metadata, and the values indicate whether the corresponding chunk can be pruned.
+    /// metadata, and a true value means the chunk _can_ be pruned.
     pub fn evaluate(&self, metadata: &ArrayData) -> VortexResult<Option<ArrayData>> {
         let known_stats = HashSet::from_iter(
             metadata
diff --git a/vortex-layout/src/layouts/chunked/scan.rs b/vortex-layout/src/layouts/chunked/scan.rs
@@ -1,11 +1,13 @@
 use std::sync::{Arc, OnceLock, RwLock};
 
+use arrow_buffer::BooleanBuffer;
 use vortex_array::array::{ChunkedArray, StructArray};
 use vortex_array::stats::{stats_from_bitset_bytes, Stat};
 use vortex_array::validity::Validity;
-use vortex_array::{ArrayData, ContextRef, IntoArrayData};
+use vortex_array::{ArrayData, ContextRef, IntoArrayData, IntoArrayVariant};
 use vortex_dtype::DType;
-use vortex_error::{vortex_err, vortex_panic, VortexExpect, VortexResult};
+use vortex_error::{vortex_err, vortex_panic, VortexError, VortexExpect, VortexResult};
+use vortex_expr::pruning::PruningPredicate;
 
 use crate::layouts::chunked::stats::StatsTable;
 use crate::layouts::chunked::ChunkedLayout;
@@ -29,13 +31,19 @@ use crate::{LayoutData, LayoutEncoding, RowMask};
 /// boundary we will read the chunk twice). However, if we have overlapping row ranges, as can
 /// happen if the parent is performing multiple scans (filter + projection), then we may read the
 /// same chunk many times.
+#[derive(Clone, Debug)]
 pub struct ChunkedScan {
     layout: LayoutData,
     scan: Scan,
     dtype: DType,
     ctx: ContextRef,
+    // Shared stats table scanner
     stats_scanner: Arc<RwLock<Box<dyn Scanner>>>,
+    // Cached pruning mask for the scan
+    pruning_mask: Arc<OnceLock<Option<BooleanBuffer>>>,
+    // Shared lazy chunk scanners
     chunk_scans: Arc<[OnceLock<Box<dyn LayoutScan>>]>,
+    // The stats that are present in the layout
     present_stats: Arc<[Stat]>,
 }
 
@@ -66,7 +74,6 @@ impl ChunkedScan {
             .unwrap_or_default()
             .into();
 
-        // Construct a scanner for the stats array
         let stats_scanner = layout
             .metadata()
             .is_some()
@@ -100,6 +107,7 @@ impl ChunkedScan {
             dtype,
             ctx,
             stats_scanner: Arc::new(RwLock::new(stats_scanner)),
+            pruning_mask: Arc::new(OnceLock::new()),
             chunk_scans,
             present_stats,
         })
@@ -117,34 +125,23 @@ impl LayoutScan for ChunkedScan {
 
     fn create_scanner(&self, mask: RowMask) -> VortexResult<Box<dyn Scanner>> {
         Ok(Box::new(ChunkedScanner {
-            scan: self.scan.clone(),
-            layout: self.layout.clone(),
-            ctx: self.ctx.clone(),
-            scan_dtype: self.dtype.clone(),
-            present_stats: self.present_stats.clone(),
-            stats_scanner: self.stats_scanner.clone(),
-            chunk_scans: self.chunk_scans.clone(),
+            chunked_scan: self.clone(),
             mask,
             chunk_states: None,
         }) as _)
     }
 }
 
 /// A scanner for a chunked layout.
+#[derive(Debug)]
 struct ChunkedScanner {
-    scan: Scan,
-    layout: LayoutData,
-    ctx: ContextRef,
-    scan_dtype: DType,
-    present_stats: Arc<[Stat]>,
-    stats_scanner: Arc<RwLock<Box<dyn Scanner>>>,
-    chunk_scans: Arc<[OnceLock<Box<dyn LayoutScan>>]>,
+    chunked_scan: ChunkedScan,
     mask: RowMask,
-
     // State for each chunk in the layout
     chunk_states: Option<Vec<ChunkState>>,
 }
 
+#[derive(Debug)]
 enum ChunkState {
     Pending(Box<dyn Scanner>),
     Resolved(Option<ArrayData>),
@@ -155,33 +152,55 @@ impl Scanner for ChunkedScanner {
         // If we haven't set up our chunk state yet, then fetch the stats table and do so.
         if self.chunk_states.is_none() {
             // First, we grab the stats table
-            let _stats_table = match self
+            let stats_table = match self
+                .chunked_scan
                 .stats_scanner
                 .write()
                 .map_err(|_| vortex_err!("poisoned"))?
                 .poll(segments)?
             {
                 Poll::Some(stats_array) => StatsTable::try_new(
-                    self.layout.dtype().clone(),
+                    self.chunked_scan.layout.dtype().clone(),
                     stats_array,
-                    self.present_stats.clone(),
+                    self.chunked_scan.present_stats.clone(),
                 )?,
                 Poll::NeedMore(segments) => {
                     // Otherwise, we need more segments to read the stats table.
                     return Ok(Poll::NeedMore(segments));
                 }
             };
 
+            // And compute the pruning predicate
+            let pruning_mask = self.chunked_scan.pruning_mask.get_or_try_init(|| {
+                Ok::<_, VortexError>(
+                    self.chunked_scan
+                        .scan
+                        .filter
+                        .as_ref()
+                        .and_then(PruningPredicate::try_new)
+                        .and_then(|predicate| predicate.evaluate(stats_table.array()).transpose())
+                        .transpose()?
+                        .map(|mask| mask.into_bool())
+                        .transpose()?
+                        .map(|mask| mask.boolean_buffer()),
+                )
+            })?;
+
             // Now we can set up the chunk state.
-            let mut chunks = Vec::with_capacity(self.chunk_scans.len());
+            let mut chunks = Vec::with_capacity(self.chunked_scan.chunk_scans.len());
             let mut row_offset = 0;
-            for chunk_idx in 0..self.chunk_scans.len() {
-                let chunk_scan = self.chunk_scans[chunk_idx].get_or_try_init(|| {
-                    self.layout
-                        .child(chunk_idx, self.layout.dtype().clone())
-                        .vortex_expect("Child index out of bounds")
-                        .new_scan(self.scan.clone(), self.ctx.clone())
-                })?;
+            for chunk_idx in 0..self.chunked_scan.chunk_scans.len() {
+                let chunk_scan =
+                    self.chunked_scan.chunk_scans[chunk_idx].get_or_try_init(|| {
+                        self.chunked_scan
+                            .layout
+                            .child(chunk_idx, self.chunked_scan.layout.dtype().clone())
+                            .vortex_expect("Child index out of bounds")
+                            .new_scan(
+                                self.chunked_scan.scan.clone(),
+                                self.chunked_scan.ctx.clone(),
+                            )
+                    })?;
 
                 // Figure out the row range of the chunk
                 let chunk_len = chunk_scan.layout().row_count();
@@ -194,8 +213,13 @@ impl Scanner for ChunkedScanner {
                     continue;
                 }
 
-                // Try to skip the chunk based on the stats table
-                // TODO
+                // Try to skip the chunk based on the pruning predicate
+                if let Some(pruning_mask) = pruning_mask {
+                    if pruning_mask.value(chunk_idx) {
+                        chunks.push(ChunkState::Resolved(None));
+                        continue;
+                    }
+                }
 
                 // Otherwise, we need to read it. So we set up a mask for the chunk range.
                 let chunk_mask = self
@@ -250,7 +274,7 @@ impl Scanner for ChunkedScanner {
             .collect::<Vec<_>>();
 
         Ok(Poll::Some(
-            ChunkedArray::try_new(chunks, self.scan_dtype.clone())?.into_array(),
+            ChunkedArray::try_new(chunks, self.chunked_scan.dtype.clone())?.into_array(),
         ))
     }
 }
@@ -263,25 +287,37 @@ enum PollStats {
 
 #[cfg(test)]
 mod test {
-    use vortex_array::{ArrayDType, ArrayLen, IntoArrayData, IntoArrayVariant};
+    use std::assert_matches::assert_matches;
+
+    use vortex_array::{ArrayLen, IntoArrayData, IntoArrayVariant};
     use vortex_buffer::buffer;
+    use vortex_dtype::Nullability::NonNullable;
+    use vortex_dtype::{DType, PType};
+    use vortex_expr::{BinaryExpr, Identity, Literal, Operator};
 
+    use crate::layouts::chunked::scan::{ChunkState, ChunkedScan, ChunkedScanner};
     use crate::layouts::chunked::writer::ChunkedLayoutWriter;
-    use crate::scanner::Scan;
+    use crate::scanner::{Poll, Scan, Scanner};
     use crate::segments::test::TestSegments;
     use crate::strategies::LayoutWriterExt;
-    use crate::LayoutData;
+    use crate::{LayoutData, RowMask};
 
-    /// Create a chunked layout with three chunks of `1..=3` primitive arrays.
+    /// Create a chunked layout with three chunks of primitive arrays.
     fn chunked_layout() -> (TestSegments, LayoutData) {
-        let arr = buffer![1, 2, 3].into_array();
         let mut segments = TestSegments::default();
-        let layout = ChunkedLayoutWriter::new(arr.dtype(), Default::default())
-            .push_all(
-                &mut segments,
-                [Ok(arr.clone()), Ok(arr.clone()), Ok(arr.clone())],
-            )
-            .unwrap();
+        let layout = ChunkedLayoutWriter::new(
+            &DType::Primitive(PType::I32, NonNullable),
+            Default::default(),
+        )
+        .push_all(
+            &mut segments,
+            [
+                Ok(buffer![1, 2, 3].into_array()),
+                Ok(buffer![4, 5, 6].into_array()),
+                Ok(buffer![7, 8, 9].into_array()),
+            ],
+        )
+        .unwrap();
         (segments, layout)
     }
 
@@ -293,6 +329,46 @@ mod test {
         let result = segments.do_scan(scan.as_ref()).into_primitive().unwrap();
 
         assert_eq!(result.len(), 9);
-        assert_eq!(result.as_slice::<i32>(), &[1, 2, 3, 1, 2, 3, 1, 2, 3]);
+        assert_eq!(result.as_slice::<i32>(), &[1, 2, 3, 4, 5, 6, 7, 8, 9]);
+    }
+
+    #[test]
+    fn test_chunked_pruning_mask() {
+        let (segments, layout) = chunked_layout();
+
+        let scan = ChunkedScan::try_new(
+            layout,
+            Scan {
+                projection: Identity::new_expr(),
+                filter: Some(BinaryExpr::new_expr(
+                    Identity::new_expr(),
+                    Operator::Gt,
+                    Literal::new_expr(6.into()),
+                )),
+            },
+            Default::default(),
+        )
+        .unwrap();
+
+        // Populate the stats scanner so that we can compute the pruning mask
+        _ = scan.stats_scanner.write().unwrap().poll(&segments).unwrap();
+
+        let mut scanner = ChunkedScanner {
+            chunked_scan: scan,
+            mask: RowMask::new_valid_between(0, 9),
+            chunk_states: None,
+        };
+
+        // Then we poll the chunked scanner without any segments so _only_ the stats were
+        // available.
+        let Poll::NeedMore(_segments) = scanner.poll(&TestSegments::default()).unwrap() else {
+            unreachable!()
+        };
+
+        // Now we validate that based on the pruning mask, we have excluded the first two chunks
+        let chunk_states = scanner.chunk_states.as_ref().unwrap().as_slice();
+        assert_matches!(chunk_states[0], ChunkState::Resolved(None));
+        assert_matches!(chunk_states[1], ChunkState::Resolved(None));
+        assert_matches!(chunk_states[2], ChunkState::Pending(_));
     }
 }
diff --git a/vortex-layout/src/layouts/flat/scan.rs b/vortex-layout/src/layouts/flat/scan.rs
@@ -9,17 +9,12 @@ use crate::scanner::{LayoutScan, Poll, Scan, Scanner};
 use crate::segments::{SegmentId, SegmentReader};
 use crate::{LayoutData, LayoutEncoding, RowMask};
 
-#[derive(Clone, Eq, PartialEq)]
-enum State {
-    Initial,
-}
-
+#[derive(Debug)]
 pub struct FlatScan {
     layout: LayoutData,
     scan: Scan,
     dtype: DType,
     ctx: ContextRef,
-    state: State,
 }
 
 impl FlatScan {
@@ -33,7 +28,6 @@ impl FlatScan {
             scan,
             dtype,
             ctx,
-            state: State::Initial,
         })
     }
 }
@@ -67,6 +61,9 @@ impl LayoutScan for FlatScan {
     }
 }
 
+// TODO(ngates): this needs to move into a shared Scanner inside the Scan. Then each scanner can
+//  share work.
+#[derive(Debug)]
 struct FlatScanner {
     segment_id: SegmentId,
     dtype: DType,
@@ -175,7 +172,7 @@ mod test {
             filter: Some(BinaryExpr::new_expr(
                 Arc::new(Identity),
                 Operator::Gt,
-                Literal::new_expr(3.into()),
+                Literal::new_expr(3i32.into()),
             )),
         };
 
diff --git a/vortex-layout/src/layouts/struct_/scan.rs b/vortex-layout/src/layouts/struct_/scan.rs
@@ -7,6 +7,7 @@ use crate::scanner::{LayoutScan, Poll, Scan, Scanner};
 use crate::segments::SegmentReader;
 use crate::{LayoutData, LayoutEncoding, RowMask};
 
+#[derive(Debug)]
 pub struct StructScan {
     layout: LayoutData,
     scan: Scan,
@@ -49,11 +50,12 @@ impl LayoutScan for StructScan {
     }
 }
 
-#[derive(Clone)]
+#[derive(Clone, Debug)]
 enum State {
     Initial,
 }
 
+#[derive(Debug)]
 struct StructScanner {
     layout: LayoutData,
     scan: Scan,
diff --git a/vortex-layout/src/lib.rs b/vortex-layout/src/lib.rs
@@ -1,4 +1,5 @@
 #![feature(once_cell_try)]
+#![feature(assert_matches)]
 #![allow(dead_code)]
 mod data;
 pub mod scanner;
diff --git a/vortex-layout/src/scanner/mod.rs b/vortex-layout/src/scanner/mod.rs
@@ -1,5 +1,7 @@
 mod scan;
 
+use std::fmt::Debug;
+
 pub use scan::*;
 use vortex_array::ArrayData;
 use vortex_dtype::DType;
@@ -9,7 +11,7 @@ use crate::segments::{SegmentId, SegmentReader};
 use crate::{LayoutData, RowMask};
 
 /// A [`LayoutScan`] provides an encapsulation of an invocation of a scan operation.
-pub trait LayoutScan: 'static + Send + Sync {
+pub trait LayoutScan: 'static + Send + Sync + Debug {
     /// Returns the [`LayoutData`] that this scan is operating on.
     fn layout(&self) -> &LayoutData;
 
@@ -45,7 +47,7 @@ pub enum Poll {
 }
 
 /// A trait for scanning a single row range of a layout.
-pub trait Scanner: 'static + Send + Sync {
+pub trait Scanner: 'static + Send + Sync + Debug {
     /// Attempts to return the [`ArrayData`] result of this ranged scan. If the scanner cannot
     /// make progress, it can return a vec of additional data segments using [`Poll::NeedMore`].
     ///
@@ -67,6 +69,7 @@ pub trait ScannerExt: Scanner {
 impl<S: Scanner> ScannerExt for S {}
 
 /// A scanner with an [`ArrayData`] that is always returned.
+#[derive(Debug)]
 pub struct ResolvedScanner(pub ArrayData);
 
 impl Scanner for ResolvedScanner {
diff --git a/vortex-layout/src/scanner/scan.rs b/vortex-layout/src/scanner/scan.rs

Original file line number	Diff line number	Diff line change
`@@ -7,6 +7,7 @@ use crate::scanner::{LayoutScan, Poll, Scan, Scanner};`
`7`	`7`	`use crate::segments::SegmentReader;`
`8`	`8`	`use crate::{LayoutData, LayoutEncoding, RowMask};`
`9`	`9`
	`10`	`+#[derive(Debug)]`
`10`	`11`	`pub struct StructScan {`
`11`	`12`	`layout: LayoutData,`
`12`	`13`	`scan: Scan,`
`@@ -49,11 +50,12 @@ impl LayoutScan for StructScan {`
`49`	`50`	`}`
`50`	`51`	`}`
`51`	`52`
`52`		`-#[derive(Clone)]`
	`53`	`+#[derive(Clone, Debug)]`
`53`	`54`	`enum State {`
`54`	`55`	`Initial,`
`55`	`56`	`}`
`56`	`57`
	`58`	`+#[derive(Debug)]`
`57`	`59`	`struct StructScanner {`
`58`	`60`	`layout: LayoutData,`
`59`	`61`	`scan: Scan,`
Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,5 @@`
`1`	`1`	`#![feature(once_cell_try)]`
	`2`	`+#![feature(assert_matches)]`
`2`	`3`	`#![allow(dead_code)]`
`3`	`4`	`mod data;`
`4`	`5`	`pub mod scanner;`