Add plan-native scan execution

Signed-off-by: Joe Isaacs <joe.isaacs@live.co.uk>

Author: joseph-isaacs

Cargo.lock

@@ -23,6 +23,7 @@ members = [
     "vortex-btrblocks",
     "vortex-layout",
     "vortex-scan",
+    "vortex-scan-v2",
     "vortex-file",
     "vortex-ipc",
     "vortex",
@@ -322,6 +323,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 }

Cargo.toml

@@ -25,9 +25,10 @@ general logical query plan.
 Every rewrite must preserve the query result, including its dtype, row domain, row order, row
 identity, null behavior, and observable errors.
 
-## Future execution
+## Execution
 
-Plans currently stop at construction and optimization. A future PR will add a method for executing
-an optimized plan. That method will walk the physical plan, read the referenced layout data,
-evaluate its expressions, and return the result of the query. The execution API and return type
-will be defined as part of that integration rather than fixed by the planning IR today.
+Each plan node can execute a row range and selection mask. Leaf plans read their referenced
+segments, structural plans combine their children, and expression plans evaluate the remaining
+derived work. The separate `vortex-scan-v2` crate copies the existing scan orchestration around
+this API so the original `LayoutReader` scanner remains unchanged while the plan-native path is
+developed.

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

@@ -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/row_idx/mod.rs

@@ -0,0 +1,41 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use std::sync::Arc;
+
+use futures::future::BoxFuture;
+use vortex_array::ArrayRef;
+use vortex_error::VortexResult;
+use vortex_session::VortexSession;
+
+use crate::segments::SegmentSource;
+
+/// Future resolving to the array produced by a physical plan.
+pub type PlanArrayFuture = BoxFuture<'static, VortexResult<ArrayRef>>;
+
+/// Runtime dependencies shared by every node in a plan execution.
+#[derive(Clone)]
+pub struct PlanExecutionContext {
+    segment_source: Arc<dyn SegmentSource>,
+    session: VortexSession,
+}
+
+impl PlanExecutionContext {
+    /// Creates an execution context over a segment source and Vortex session.
+    pub fn new(segment_source: Arc<dyn SegmentSource>, session: VortexSession) -> Self {
+        Self {
+            segment_source,
+            session,
+        }
+    }
+
+    /// Returns the segment source used to satisfy leaf reads.
+    pub fn segment_source(&self) -> &Arc<dyn SegmentSource> {
+        &self.segment_source
+    }
+
+    /// Returns the Vortex session used for array decoding and expression execution.
+    pub fn session(&self) -> &VortexSession {
+        &self.session
+    }
+}

vortex-layout/src/plan/execution.rs

@@ -5,6 +5,7 @@
 
 mod children;
 mod display;
+mod execution;
 pub mod optimizer;
 mod plans;
 
@@ -19,6 +20,8 @@ pub use display::PlanSummaryExtractor;
 pub use display::PlanTreeContext;
 pub use display::PlanTreeDisplay;
 pub use display::PlanTreeExtractor;
+pub use execution::PlanArrayFuture;
+pub use execution::PlanExecutionContext;
 pub use plans::ChunkedPlan;
 pub use plans::DictPlan;
 pub use plans::ExpressionPlan;
@@ -64,6 +67,19 @@ pub trait Plan: Any + Send + Sync {
     /// domain.
     fn optimize(&self) -> VortexResult<PlanRef>;
 
+    /// Executes this plan for `row_range`, returning values selected by `mask`.
+    ///
+    /// The row range is expressed in this plan's row domain. The returned array has one row for
+    /// every true value in `mask`.
+    fn execute(
+        &self,
+        _ctx: &PlanExecutionContext,
+        _row_range: &std::ops::Range<u64>,
+        _mask: vortex_array::MaskFuture,
+    ) -> VortexResult<PlanArrayFuture> {
+        vortex_bail!("Plan execution is not implemented for '{}'", self.name())
+    }
+
     /// Returns the dtype produced by this plan.
     fn dtype(&self) -> &DType;
 
@@ -86,7 +102,7 @@ pub trait Plan: Any + Send + Sync {
     }
 }
 
-/// Constructs a physical plan without changing the layout or scan APIs.
+/// Constructs a physical plan for a stored layout tree.
 ///
 /// Known layouts are represented by optimizer-visible plan nodes, which may defer constructing
 /// their children. Unsupported layout kinds return an error when their plan is requested.

vortex-layout/src/plan/mod.rs

@@ -2,18 +2,31 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 use std::borrow::Cow;
+use std::future;
+use std::ops::Range;
 use std::sync::Arc;
 
+use futures::FutureExt;
+use futures::TryStreamExt;
+use futures::stream::FuturesOrdered;
+use vortex_array::Canonical;
+use vortex_array::IntoArray;
+use vortex_array::MaskFuture;
+use vortex_array::arrays::ChunkedArray;
 use vortex_array::dtype::DType;
 use vortex_array::expr::ExactBoundExpr;
 use vortex_array::expr::label_bound_tree;
+use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
+use vortex_error::vortex_ensure;
 
 use crate::layouts::chunked::ChunkedLayout;
 use crate::layouts::row_idx::RowIdx;
 use crate::plan::ExpressionPlan;
 use crate::plan::LazyPlanChildren;
 use crate::plan::Plan;
+use crate::plan::PlanArrayFuture;
+use crate::plan::PlanExecutionContext;
 use crate::plan::PlanRef;
 use crate::plan::new_plan;
 use crate::plan::optimizer::PlanParentReduceRule;
@@ -60,6 +73,62 @@ impl Plan for ChunkedPlan {
         Ok(Arc::new(self.with_chunks(self.dtype.clone(), chunks)))
     }
 
+    fn execute(
+        &self,
+        ctx: &PlanExecutionContext,
+        row_range: &Range<u64>,
+        mask: MaskFuture,
+    ) -> VortexResult<PlanArrayFuture> {
+        vortex_ensure!(
+            row_range.start <= row_range.end && row_range.end <= self.row_count(),
+            "Chunked plan row range {:?} is outside 0..{}",
+            row_range,
+            self.row_count()
+        );
+        vortex_ensure!(
+            mask.len() == usize::try_from(row_range.end - row_range.start)?,
+            "Chunked plan mask length mismatch"
+        );
+        if row_range.is_empty() {
+            let empty = Canonical::empty(&self.dtype).into_array();
+            return Ok(future::ready(Ok(empty)).boxed());
+        }
+
+        let mut chunk_futures = Vec::new();
+        let mut chunk_offset = 0_u64;
+        for chunk_index in 0..self.chunks.len() {
+            let chunk = self
+                .chunks
+                .get(chunk_index)?
+                .ok_or_else(|| vortex_error::vortex_err!("Chunk {chunk_index} has no plan"))?;
+            let chunk_end = chunk_offset
+                .checked_add(chunk.row_count())
+                .ok_or_else(|| vortex_error::vortex_err!("Chunk row offset overflow"))?;
+            let start = row_range.start.max(chunk_offset);
+            let end = row_range.end.min(chunk_end);
+            if start < end {
+                let child_range = start - chunk_offset..end - chunk_offset;
+                let mask_range = usize::try_from(start - row_range.start)?
+                    ..usize::try_from(end - row_range.start)?;
+                chunk_futures.push(chunk.execute(ctx, &child_range, mask.slice(mask_range))?);
+            }
+            chunk_offset = chunk_end;
+        }
+
+        Ok(async move {
+            let chunks: Vec<_> = FuturesOrdered::from_iter(chunk_futures)
+                .try_collect()
+                .await?;
+            vortex_ensure!(!chunks.is_empty(), "Non-empty row range selected no chunks");
+            if chunks.len() == 1 {
+                return Ok(chunks.into_iter().next().vortex_expect("one chunk"));
+            }
+            let dtype = chunks[0].dtype().clone();
+            Ok(ChunkedArray::try_new(chunks, dtype)?.into_array())
+        }
+        .boxed())
+    }
+
     fn dtype(&self) -> &DType {
         &self.dtype
     }

vortex-layout/src/plan/plans/chunked.rs

@@ -2,16 +2,25 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 use std::borrow::Cow;
+use std::ops::Range;
 use std::sync::Arc;
 
+use futures::FutureExt;
+use futures::try_join;
+use vortex_array::IntoArray;
+use vortex_array::MaskFuture;
+use vortex_array::arrays::DictArray;
 use vortex_array::expr::ExactBoundExpr;
 use vortex_array::expr::label_bound_tree;
+use vortex_array::optimizer::ArrayOptimizer;
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
 
 use crate::layouts::dict::DictLayout;
 use crate::plan::ExpressionPlan;
 use crate::plan::Plan;
+use crate::plan::PlanArrayFuture;
+use crate::plan::PlanExecutionContext;
 use crate::plan::PlanRef;
 use crate::plan::new_plan;
 use crate::plan::optimizer::PlanParentReduceRule;
@@ -67,6 +76,35 @@ impl Plan for DictPlan {
         Ok(Arc::new(self.with_children(codes, values)))
     }
 
+    fn execute(
+        &self,
+        ctx: &PlanExecutionContext,
+        row_range: &Range<u64>,
+        mask: MaskFuture,
+    ) -> VortexResult<PlanArrayFuture> {
+        let codes = self.codes.execute(ctx, row_range, mask)?;
+        let values_len = usize::try_from(self.values.row_count())?;
+        let values = self.values.execute(
+            ctx,
+            &(0..self.values.row_count()),
+            MaskFuture::new_true(values_len),
+        )?;
+        let all_values_referenced = self.layout.has_all_values_referenced();
+
+        Ok(async move {
+            let (codes, values) = try_join!(codes, values)?;
+            // SAFETY: DictLayout validation guarantees integer codes and matching child dtypes.
+            let dictionary = unsafe {
+                DictArray::new_unchecked(codes, values)
+                    .set_all_values_referenced(all_values_referenced)
+            }
+            .into_array()
+            .optimize()?;
+            Ok(dictionary)
+        }
+        .boxed())
+    }
+
     fn dtype(&self) -> &vortex_array::dtype::DType {
         &self.dtype
     }

vortex-layout/src/plan/plans/dict.rs

@@ -3,8 +3,11 @@
 
 use std::any::TypeId;
 use std::borrow::Cow;
+use std::ops::Range;
 use std::sync::Arc;
 
+use futures::FutureExt;
+use vortex_array::MaskFuture;
 use vortex_array::dtype::DType;
 use vortex_array::dtype::FieldName;
 use vortex_array::expr::BoundExpression;
@@ -17,6 +20,8 @@ use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
 
 use crate::plan::Plan;
+use crate::plan::PlanArrayFuture;
+use crate::plan::PlanExecutionContext;
 use crate::plan::PlanRef;
 use crate::plan::optimizer::reduce_parent;
 
@@ -100,6 +105,17 @@ impl Plan for ExpressionPlan {
         self.optimize_top_down(None)
     }
 
+    fn execute(
+        &self,
+        ctx: &PlanExecutionContext,
+        row_range: &Range<u64>,
+        mask: MaskFuture,
+    ) -> VortexResult<PlanArrayFuture> {
+        let child = self.child.execute(ctx, row_range, mask)?;
+        let expression = self.expression.clone();
+        Ok(async move { child.await?.apply_bound(&expression) }.boxed())
+    }
+
     fn dtype(&self) -> &DType {
         self.expression.dtype()
     }