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",
@@ -325,6 +326,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

@@ -18,7 +18,7 @@ on demand.
 | `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 |
+| `RowIdx` | generate row numbers for the current execution 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,
@@ -55,9 +55,14 @@ ID, dtype, row count, and lazy children. Only the unsized tail containing the vt
 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
 
-Plans currently stop at construction and optimization. Still to come: a plan registry and foreign
-operator placeholder so third-party operators survive a round trip, a serialization envelope, and
-an execution stage that walks an optimized plan, reads the referenced segments, and returns the
-query result.
+Still to come: a plan registry and foreign operator placeholder so third-party operators survive
+a round trip, and a serialization envelope.

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,67 @@
+// 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,
+    row_offset: u64,
+}
+
+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,
+            row_offset: 0,
+        }
+    }
+
+    /// Sets the global row index of the first row in the root plan's row domain.
+    pub fn with_row_offset(mut self, row_offset: u64) -> Self {
+        self.row_offset = row_offset;
+        self
+    }
+
+    /// Returns the global row index of the first row in the current plan's row domain.
+    pub fn row_offset(&self) -> u64 {
+        self.row_offset
+    }
+
+    /// Derives the execution context for a child whose row domain starts within this one.
+    pub(crate) fn child_row_domain(&self, relative_row_offset: u64) -> VortexResult<Self> {
+        let row_offset = self
+            .row_offset
+            .checked_add(relative_row_offset)
+            .ok_or_else(|| vortex_error::vortex_err!("Plan row-domain offset overflow"))?;
+        Ok(Self {
+            segment_source: Arc::clone(&self.segment_source),
+            session: self.session.clone(),
+            row_offset,
+        })
+    }
+
+    /// 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

@@ -126,6 +126,7 @@ fn lower_dict(layout: &DictLayout) -> VortexResult<TakePlan> {
         TakePlan::from_children_unchecked(
             layout.dtype().clone(),
             layout.row_count(),
+            layout.has_all_values_referenced(),
             lazy_children(layout.to_layout(), vec![1, 0]),
         )
     })

vortex-layout/src/plan/lower.rs

@@ -9,6 +9,7 @@
 
 mod children;
 mod display;
+mod execution;
 mod lower;
 mod optimize;
 pub mod optimizer;
@@ -22,6 +23,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 lower::lower;
 pub use optimize::optimize;
 pub use plans::Concat;
@@ -43,6 +46,7 @@ pub use plans::SegmentScan;
 pub use plans::SegmentScanData;
 pub use plans::SegmentScanPlan;
 pub use plans::Take;
+pub use plans::TakeData;
 pub use plans::TakePlan;
 pub use plans::Zoned;
 pub use plans::ZonedPlan;

vortex-layout/src/plan/mod.rs

@@ -2,18 +2,32 @@
 // 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::EmptyMetadata;
+use vortex_array::IntoArray;
+use vortex_array::MaskFuture;
+use vortex_array::arrays::ChunkedArray;
 use vortex_array::dtype::DType;
+use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
+use vortex_error::vortex_ensure;
+use vortex_error::vortex_err;
 use vortex_session::registry::CachedId;
 
 use crate::plan::Eval;
 use crate::plan::EvalPlan;
 use crate::plan::Plan;
+use crate::plan::PlanArrayFuture;
 use crate::plan::PlanChildren;
+use crate::plan::PlanExecutionContext;
 use crate::plan::PlanId;
 use crate::plan::PlanParts;
 use crate::plan::PlanRef;
@@ -139,6 +153,62 @@ impl PlanVTable for Concat {
         Ok(())
     }
 
+    fn execute(
+        plan: &Plan<Self>,
+        ctx: &PlanExecutionContext,
+        row_range: &Range<u64>,
+        mask: MaskFuture,
+    ) -> VortexResult<PlanArrayFuture> {
+        vortex_ensure!(
+            row_range.start <= row_range.end && row_range.end <= plan.row_count(),
+            "Concat row range {:?} is outside 0..{}",
+            row_range,
+            plan.row_count()
+        );
+        vortex_ensure!(
+            mask.len() == usize::try_from(row_range.end - row_range.start)?,
+            "Concat mask length mismatch"
+        );
+        if row_range.is_empty() {
+            let empty = Canonical::empty(plan.dtype()).into_array();
+            return Ok(future::ready(Ok(empty)).boxed());
+        }
+
+        let mut chunk_futures = Vec::new();
+        for (chunk, &chunk_offset) in plan.children().iter().zip(plan.row_offsets()) {
+            let chunk = chunk?;
+            let chunk_end = chunk_offset
+                .checked_add(chunk.row_count())
+                .ok_or_else(|| 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)?;
+                let child_ctx = ctx.child_row_domain(chunk_offset)?;
+                chunk_futures.push(chunk.execute(
+                    &child_ctx,
+                    &child_range,
+                    mask.slice(mask_range),
+                )?);
+            }
+        }
+
+        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 child_name(_plan: &Plan<Self>, index: usize) -> Cow<'_, str> {
         Cow::Owned(format!("chunks[{index}]"))
     }

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

@@ -3,15 +3,20 @@
 
 use std::borrow::Cow;
 use std::fmt;
+use std::ops::Range;
 
+use futures::FutureExt;
 use vortex_array::EmptyMetadata;
+use vortex_array::MaskFuture;
 use vortex_array::expr::BoundExpression;
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
 use vortex_session::registry::CachedId;
 
 use crate::plan::Plan;
+use crate::plan::PlanArrayFuture;
 use crate::plan::PlanChildren;
+use crate::plan::PlanExecutionContext;
 use crate::plan::PlanId;
 use crate::plan::PlanParts;
 use crate::plan::PlanRef;
@@ -106,6 +111,17 @@ impl PlanVTable for Eval {
         Ok(())
     }
 
+    fn execute(
+        plan: &Plan<Self>,
+        ctx: &PlanExecutionContext,
+        row_range: &Range<u64>,
+        mask: MaskFuture,
+    ) -> VortexResult<PlanArrayFuture> {
+        let child = plan.child_plan()?.execute(ctx, row_range, mask)?;
+        let expression = plan.expression().clone();
+        Ok(async move { child.await?.apply_bound(&expression) }.boxed())
+    }
+
     fn child_name(_plan: &Plan<Self>, index: usize) -> Cow<'_, str> {
         if index == 0 {
             Cow::Borrowed("child")