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",
@@ -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 }

Cargo.toml

@@ -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,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

@@ -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;
@@ -44,9 +47,12 @@ 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;
+pub use plans::plan_row_idx_expression;
+pub use plans::row_idx_dtype;
 pub use typed::DynPlan;
 pub use typed::Plan;
 pub use typed::PlanParts;

vortex-layout/src/plan/mod.rs

@@ -19,12 +19,10 @@ use super::Concat;
 use super::Eval;
 use super::Pack;
 use super::PlanRef;
-use super::RowIdx;
 use super::Take;
 use super::plans::EvalIdentityRule;
 use super::plans::ExpressionConcatRule;
 use super::plans::ExpressionPackRule;
-use super::plans::ExpressionRowIdxRule;
 use super::plans::ExpressionTakeRule;
 
 static EVAL_IDENTITY_RULE: PlanReduceRuleAdapter<Eval, EvalIdentityRule> =
@@ -36,15 +34,12 @@ static EXPRESSION_CONCAT_RULE: PlanParentReduceRuleAdapter<Concat, ExpressionCon
     PlanParentReduceRuleAdapter::new(ExpressionConcatRule);
 static EXPRESSION_TAKE_RULE: PlanParentReduceRuleAdapter<Take, ExpressionTakeRule> =
     PlanParentReduceRuleAdapter::new(ExpressionTakeRule);
-static EXPRESSION_ROW_IDX_RULE: PlanParentReduceRuleAdapter<RowIdx, ExpressionRowIdxRule> =
-    PlanParentReduceRuleAdapter::new(ExpressionRowIdxRule);
 static EXPRESSION_PACK_RULE: PlanParentReduceRuleAdapter<Pack, ExpressionPackRule> =
     PlanParentReduceRuleAdapter::new(ExpressionPackRule);
 
 static PARENT_RULES: PlanParentRuleSet = PlanParentRuleSet::new(&[
     &EXPRESSION_CONCAT_RULE,
     &EXPRESSION_TAKE_RULE,
-    &EXPRESSION_ROW_IDX_RULE,
     &EXPRESSION_PACK_RULE,
 ]);
 

vortex-layout/src/plan/optimizer/mod.rs

@@ -2,21 +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_array::expr::ExactBoundExpr;
-use vortex_array::expr::label_bound_tree;
+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::layouts::row_idx::RowIdx as RowIdxFn;
 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;
@@ -142,6 +153,57 @@ 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)?;
+                chunk_futures.push(chunk.execute(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}]"))
     }
@@ -161,23 +223,6 @@ impl PlanParentReduceRule<Concat> for ExpressionConcatRule {
         _child_idx: usize,
     ) -> VortexResult<Option<PlanRef>> {
         let expression = parent.expression();
-        // Row-index expressions are relative to the whole row domain, so they cannot be evaluated
-        // chunk by chunk.
-        let references_row_idx = label_bound_tree(
-            expression,
-            |node| {
-                node.as_scalar()
-                    .is_some_and(|scalar_fn| scalar_fn.is::<RowIdxFn>())
-            },
-            |acc, &child| acc | child,
-        )
-        .get(&ExactBoundExpr(expression.clone()))
-        .copied()
-        .unwrap_or(false);
-        if references_row_idx {
-            return Ok(None);
-        }
-
         let chunks = child
             .children()
             .iter()

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")