Prune plan scans with zoned statistics

Cargo.toml

@@ -23,6 +23,7 @@ members = [
     "vortex-btrblocks",
     "vortex-layout",
     "vortex-scan",
+    "vortex-scan-v2",
     "vortex-file",
     "vortex-ipc",
     "vortex",
@@ -323,6 +324,7 @@ vortex-proto = { version = "0.1.0", path = "./vortex-proto", default-features =
 vortex-row = { version = "0.1.0", path = "./vortex-row", default-features = false }
 vortex-runend = { version = "0.1.0", path = "./encodings/runend", default-features = false }
 vortex-scan = { version = "0.1.0", path = "./vortex-scan", default-features = false }
+vortex-scan-v2 = { version = "0.1.0", path = "./vortex-scan-v2", default-features = false }
 vortex-sequence = { version = "0.1.0", path = "encodings/sequence", default-features = false }
 vortex-session = { version = "0.1.0", path = "./vortex-session", default-features = false }
 vortex-sparse = { version = "0.1.0", path = "./encodings/sparse", default-features = false }

docs/developer-guide/index.md

@@ -23,6 +23,7 @@ internals/session
 internals/async-runtime
 internals/vtables
 internals/execution
+internals/scan-planning
 internals/stats-pruning
 internals/io
 internals/serialization

docs/developer-guide/internals/scan-planning.md

@@ -0,0 +1,68 @@
+# Scan Plans
+
+A scan plan is the physical plan for satisfying one scan query. It is a tree of physical operators
+over a row domain, describing the reads and derived work needed to produce that query's result.
+
+## Operators, not layout mirrors
+
+Plan operators describe *what work happens*, not *which layout produced it*. Their identity and
+operator-specific state are independent of the source layout kind. The complete plan node is not:
+its common lazy-child container can own hidden source state used to materialize individual children
+on demand.
+
+| Operator | Work |
+| --- | --- |
+| `SegmentScan` | read one segment and decode it to an array |
+| `Concat` | concatenate its children row-wise |
+| `Pack` | assemble a struct from one child per field, plus optional validity |
+| `Take` | index `values` by `codes` |
+| `ListPack` | assemble a list from elements and offsets, plus optional validity |
+| `Eval` | apply an expression to its child |
+| `RowIdx` | offset row numbers into the file's row domain |
+
+Naming operators for what they compute is what lets one rule cover every case. `Concat` of
+`Concat` flattens on shape alone, and `Take` over `SegmentScan` is the dictionary pushdown,
+regardless of the source layout.
+
+The stored layout tree describes all physical data in a file. A plan is query-specific: it is built
+from that tree for one projection, filter, and row domain. Different queries over the same file can
+therefore produce different plans.
+
+## Optimization
+
+Child replacement is implemented by the common plan container rather than by every operator. It
+replaces the external child container, clones `PlanData`, then invokes the operator's
+`PlanVTable::with_children` callback to validate the new children and refresh derived caches such
+as `Concat` row offsets. Rules therefore rewrite the generic tree without reconstructing common
+plan fields inside each operator.
+
+Optimization rewrites the initial tree so that each expression is evaluated as close as possible to
+the physical data that can satisfy it. Every rewrite must preserve the query result, including its
+dtype, row domain, row order, row identity, null behavior, and observable errors.
+
+Planning does not read segment data. It constructs and optimizes a description of the work that a
+later execution stage will perform.
+
+## Vtables
+
+Each operator is a small vtable type implementing `PlanVTable`, paired with a `Plan<V>` container
+over a shared `PlanRef`. `PlanRef` points to one allocation whose ordinary fields hold the operator
+ID, dtype, row count, and lazy children. Only the unsized tail containing the vtable and
+`V::PlanData` is erased behind `dyn DynPlan`, so common-field reads do not use dynamic dispatch.
+`Plan<V>` provides typed access to that operator data through `Deref`.
+
+`PlanVTable` also carries `id` and a `Metadata` codec. Operators with no unrecoverable state
+already serialize their metadata; the ones holding a read context or a bound expression return
+`None` until those codecs exist.
+
+## Execution
+
+Each operator executes over a row range and selection mask. `SegmentScan` reads its segment,
+structural operators combine their children, and `Eval` applies the remaining derived work.
+`vortex-scan-v2` copies the existing scan orchestration around this API, so the original
+`LayoutReader` scanner is untouched while the plan-native path is developed.
+
+## Future work
+
+Still to come: a plan registry and foreign operator placeholder so third-party operators survive
+a round trip, and a serialization envelope.

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

@@ -297,7 +297,7 @@ fn row_idx_dtype() -> DType {
 }
 
 // Returns a SequenceArray representing the row indices for the given row range,
-fn idx_array(row_offset: u64, row_range: &Range<u64>) -> SequenceArray {
+pub(crate) fn idx_array(row_offset: u64, row_range: &Range<u64>) -> SequenceArray {
     Sequence::try_new(
         PValue::U64(row_offset + row_range.start),
         PValue::U64(1),

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

@@ -358,7 +358,11 @@ impl ZonedLayout {
 }
 
 impl ZonedData {
-    fn aggregate_fns(&self) -> Arc<[AggregateFnRef]> {
+    pub(crate) fn zone_len(&self) -> usize {
+        self.zone_len
+    }
+
+    pub(crate) fn aggregate_fns(&self) -> Arc<[AggregateFnRef]> {
         match &self.zone_map_schema {
             ZoneMapSchema::LegacyStats(stats) => stats
                 .iter()

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

@@ -18,6 +18,7 @@ use vortex_array::aggregate_fn::fns::bounded_max::BOUNDED_MAX_BOUND;
 use vortex_array::aggregate_fn::fns::bounded_max::BoundedMax;
 use vortex_array::aggregate_fn::fns::sum::Sum;
 use vortex_array::aggregate_fn::fns::sum::normalize_legacy_partial_array;
+use vortex_array::arrays::BoolArray;
 use vortex_array::arrays::ConstantArray;
 use vortex_array::arrays::PrimitiveArray;
 use vortex_array::arrays::StructArray;
@@ -92,7 +93,7 @@ impl ZoneMap {
         Ok(unsafe { Self::new_unchecked(column_dtype, array, aggregate_fns, zone_len, row_count) })
     }
 
-    pub(super) unsafe fn new_unchecked(
+    pub(crate) unsafe fn new_unchecked(
         column_dtype: DType,
         array: StructArray,
         aggregate_fns: Arc<[AggregateFnRef]>,
@@ -152,19 +153,32 @@ impl ZoneMap {
         session: &VortexSession,
     ) -> VortexResult<Mask> {
         let mut ctx = session.create_execution_ctx();
-        let num_zones = self.array.len();
-        let predicate = self.lower_stats(predicate.clone())?;
+        self.applied_predicate(predicate)?
+            .null_as_false()
+            .execute(&mut ctx)
+    }
 
-        let array = self.array.clone().into_array();
-        let applied = array.apply_bound(&predicate)?;
+    /// Evaluates a pruning predicate while preserving unknown (null) proof values.
+    pub(crate) fn evaluate(
+        &self,
+        predicate: &BoundExpression,
+        session: &VortexSession,
+    ) -> VortexResult<BoolArray> {
+        let mut ctx = session.create_execution_ctx();
+        self.applied_predicate(predicate)?
+            .execute::<BoolArray>(&mut ctx)
+    }
 
+    fn applied_predicate(&self, predicate: &BoundExpression) -> VortexResult<ArrayRef> {
+        let num_zones = self.array.len();
+        let predicate = self.lower_stats(predicate.clone())?;
+        let applied = self.array.clone().into_array().apply_bound(&predicate)?;
         if !contains_row_count(&applied) {
-            return applied.null_as_false().execute(&mut ctx);
+            return Ok(applied);
         }
 
         let row_count_array = row_count_array(self.zone_len, self.row_count, num_zones)?;
-        let substituted = substitute_row_count(applied, &row_count_array)?;
-        substituted.null_as_false().execute(&mut ctx)
+        substitute_row_count(applied, &row_count_array)
     }
 
     fn lower_stats(&self, predicate: BoundExpression) -> VortexResult<BoundExpression> {

vortex-layout/src/lib.rs

@@ -15,6 +15,7 @@
 //! optional bound filter, optional row range, [`Selection`](vortex_scan::selection::Selection),
 //! split strategy, and task concurrency settings, then produces array streams or iterators.
 pub mod layouts;
+pub mod plan;
 
 pub use children::*;
 pub use encoding::*;

vortex-layout/src/plan/children.rs

@@ -0,0 +1,136 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use std::fmt;
+use std::sync::Arc;
+
+use once_cell::sync::OnceCell;
+use vortex_error::VortexResult;
+use vortex_error::vortex_bail;
+use vortex_error::vortex_err;
+
+use crate::plan::PlanRef;
+
+type ChildInitializer = dyn Fn(usize) -> VortexResult<PlanRef> + 'static + Send + Sync;
+
+/// Ordered plan children that may be initialized one slot at a time.
+///
+/// Eagerly constructed operators store already-filled slots. Layout lowering instead installs an
+/// initializer that owns the source layout and lowers each child on first access.
+#[derive(Clone)]
+pub struct PlanChildren {
+    initializer: Option<Arc<ChildInitializer>>,
+    cache: Arc<[OnceCell<PlanRef>]>,
+}
+
+impl PlanChildren {
+    /// Creates lazy child slots backed by `initializer`.
+    pub(crate) fn lazy(
+        len: usize,
+        initializer: impl Fn(usize) -> VortexResult<PlanRef> + 'static + Send + Sync,
+    ) -> Self {
+        Self {
+            initializer: Some(Arc::new(initializer)),
+            cache: (0..len).map(|_| OnceCell::new()).collect::<Vec<_>>().into(),
+        }
+    }
+
+    /// Returns the number of children without initializing any slot.
+    pub fn len(&self) -> usize {
+        self.cache.len()
+    }
+
+    /// Returns whether there are no children.
+    pub fn is_empty(&self) -> bool {
+        self.cache.is_empty()
+    }
+
+    /// Returns a child, initializing and caching its slot on first access.
+    pub fn get(&self, index: usize) -> VortexResult<Option<PlanRef>> {
+        let Some(cell) = self.cache.get(index) else {
+            return Ok(None);
+        };
+        if let Some(child) = cell.get() {
+            return Ok(Some(child.clone()));
+        }
+
+        let initializer = self
+            .initializer
+            .as_ref()
+            .ok_or_else(|| vortex_err!("Plan child {index} was not initialized"))?;
+        Ok(Some(cell.get_or_try_init(|| initializer(index))?.clone()))
+    }
+
+    /// Iterates over the children in logical order, initializing slots as they are visited.
+    pub fn iter(&self) -> impl ExactSizeIterator<Item = VortexResult<PlanRef>> + '_ {
+        (0..self.len()).map(|index| {
+            self.get(index)?
+                .ok_or_else(|| vortex_err!("Plan child {index} is absent"))
+        })
+    }
+
+    /// Materializes all children into an eager vector.
+    pub fn to_vec(&self) -> VortexResult<Vec<PlanRef>> {
+        self.iter().collect()
+    }
+
+    /// Returns a child collection with one slot replaced.
+    pub fn with_child(&self, index: usize, child: PlanRef) -> VortexResult<Self> {
+        if index >= self.len() {
+            vortex_bail!("Plan child index out of bounds: {index} of {}", self.len());
+        }
+
+        let source = self.clone();
+        Ok(Self::lazy(source.len(), move |child_index| {
+            if child_index == index {
+                return Ok(child.clone());
+            }
+            source
+                .get(child_index)?
+                .ok_or_else(|| vortex_err!("Plan child {child_index} is absent"))
+        }))
+    }
+}
+
+impl From<Vec<PlanRef>> for PlanChildren {
+    fn from(children: Vec<PlanRef>) -> Self {
+        let cache = children
+            .into_iter()
+            .map(OnceCell::with_value)
+            .collect::<Vec<_>>()
+            .into();
+        Self {
+            initializer: None,
+            cache,
+        }
+    }
+}
+
+impl<const N: usize> From<[PlanRef; N]> for PlanChildren {
+    fn from(children: [PlanRef; N]) -> Self {
+        Vec::from(children).into()
+    }
+}
+
+impl Default for PlanChildren {
+    fn default() -> Self {
+        Vec::new().into()
+    }
+}
+
+impl fmt::Debug for PlanChildren {
+    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
+        formatter
+            .debug_struct("PlanChildren")
+            .field("len", &self.len())
+            .field(
+                "initialized",
+                &self
+                    .cache
+                    .iter()
+                    .filter(|slot| slot.get().is_some())
+                    .count(),
+            )
+            .finish()
+    }
+}