feat(layout): compile evals, row indices, and takes to pipelines

vortex-layout/src/plan/pipeline/ops/eval.rs

@@ -0,0 +1,33 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Applying an expression to each batch.
+
+use vortex_array::expr::BoundExpression;
+use vortex_error::VortexResult;
+
+use crate::plan::pipeline::Cx;
+use crate::plan::pipeline::Input;
+use crate::plan::pipeline::Operator;
+use crate::plan::pipeline::Step;
+
+/// Applies an expression to each batch.
+pub(crate) struct EvalStage {
+    expression: BoundExpression,
+}
+
+impl EvalStage {
+    pub(crate) fn new(expression: BoundExpression) -> Self {
+        Self { expression }
+    }
+}
+
+impl Operator for EvalStage {
+    fn compute(&mut self, input: Input, _cx: &mut Cx<'_>) -> VortexResult<Step> {
+        match input {
+            Input::Chunk(batch) => Ok(Step::Last(batch.apply_bound(&self.expression)?)),
+            Input::End => Ok(Step::Finished),
+            Input::None => Ok(Step::Consumed),
+        }
+    }
+}

vortex-layout/src/plan/pipeline/ops/mod.rs

@@ -7,25 +7,34 @@
 //! twice, holds more than the batches it must join, or scans inlets it does not read.
 
 mod concat;
+mod eval;
 mod filter;
 mod pack;
 mod port;
 mod scan;
+mod take;
 
 pub(crate) use concat::*;
+pub(crate) use eval::*;
 pub(crate) use filter::*;
 pub(crate) use pack::*;
 pub(crate) use port::*;
 pub(crate) use scan::*;
 use smallvec::SmallVec;
+pub(crate) use take::*;
 use vortex_array::ArrayRef;
+use vortex_array::Canonical;
 use vortex_array::IntoArray;
 use vortex_array::arrays::ChunkedArray;
+use vortex_array::dtype::DType;
 use vortex_error::VortexResult;
 use vortex_error::vortex_err;
 
 use crate::plan::pipeline::Inlet;
 
+/// An inlet capacity that never blocks the writer, for a reader that needs its inlet whole.
+const UNBOUNDED: usize = usize::MAX;
+
 /// Takes the first `len` rows of the inlet, across as many batches as hold them, slicing the last
 /// and leaving its rest in place. Rows spanning batches come back chunked, not copied.
 pub(crate) fn take_rows(inlet: &mut Inlet<'_>, len: usize) -> VortexResult<ArrayRef> {
@@ -51,3 +60,21 @@ pub(crate) fn take_rows(inlet: &mut Inlet<'_>, len: usize) -> VortexResult<Array
     // SAFETY: the pieces come from one inlet, whose batches share its writer's dtype.
     Ok(unsafe { ChunkedArray::new_unchecked(pieces, dtype) }.into_array())
 }
+
+/// Joins arrays covering consecutive rows into one.
+pub(crate) fn join(dtype: &DType, mut arrays: Vec<ArrayRef>) -> VortexResult<ArrayRef> {
+    match arrays.len() {
+        0 => Ok(Canonical::empty(dtype).into_array()),
+        1 => Ok(arrays.remove(0)),
+        _ => Ok(ChunkedArray::try_new(arrays, dtype.clone())?.into_array()),
+    }
+}
+
+/// Takes every batch of a closed inlet.
+pub(crate) fn drain(inlet: &mut Inlet<'_>) -> Vec<ArrayRef> {
+    let mut batches = Vec::with_capacity(inlet.len());
+    while let Some(batch) = inlet.take() {
+        batches.push(batch);
+    }
+    batches
+}

vortex-layout/src/plan/pipeline/ops/take.rs

@@ -0,0 +1,87 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Looking codes up in dictionary values.
+
+use vortex_array::ArrayRef;
+use vortex_array::IntoArray;
+use vortex_array::arrays::DictArray;
+use vortex_array::arrays::Shared;
+use vortex_array::arrays::SharedArray;
+use vortex_error::VortexResult;
+
+use super::*;
+use crate::plan::TakePlan;
+use crate::plan::pipeline::Blocked;
+use crate::plan::pipeline::Cx;
+use crate::plan::pipeline::DEFAULT_CAPACITY;
+use crate::plan::pipeline::Input;
+use crate::plan::pipeline::Operator;
+use crate::plan::pipeline::Source;
+use crate::plan::pipeline::Step;
+
+/// Wraps each batch of codes as a dictionary over the values, once the values inlet has ended.
+pub(crate) struct TakeSource {
+    plan: TakePlan,
+    values: Option<ArrayRef>,
+}
+
+const CODES: usize = 0;
+const VALUES: usize = 1;
+
+impl TakeSource {
+    /// A take over `values`, or over what its values inlet produces when `None`.
+    pub(crate) fn new(plan: TakePlan, values: Option<ArrayRef>) -> Self {
+        Self { plan, values }
+    }
+}
+
+impl Operator for TakeSource {
+    fn compute(&mut self, _input: Input, cx: &mut Cx<'_>) -> VortexResult<Step> {
+        let values = match &self.values {
+            Some(values) => values.clone(),
+            None => {
+                let mut inlet = cx.inlet(VALUES);
+                if !inlet.closed() {
+                    return Ok(Step::Blocked(Blocked::Inlet(VALUES)));
+                }
+                let values = join(self.plan.values()?.dtype(), drain(&mut inlet))?;
+                let values = if values.is::<Shared>() {
+                    values
+                } else {
+                    SharedArray::new(values).into_array()
+                };
+                self.values = Some(values.clone());
+                values
+            }
+        };
+        let mut codes = cx.inlet(CODES);
+        match codes.take() {
+            Some(batch) => {
+                let more = !codes.is_empty();
+                let array = DictArray::try_new(batch, values)?.into_array();
+                Ok(if more {
+                    Step::More(array)
+                } else {
+                    Step::Last(array)
+                })
+            }
+            None if codes.closed() => Ok(Step::Finished),
+            None => Ok(Step::Blocked(Blocked::Inlet(CODES))),
+        }
+    }
+}
+
+impl Source for TakeSource {
+    fn inlet_count(&self) -> usize {
+        if self.values.is_some() { 1 } else { 2 }
+    }
+
+    fn capacity(&self, inlet: usize) -> usize {
+        if inlet == VALUES {
+            UNBOUNDED
+        } else {
+            DEFAULT_CAPACITY
+        }
+    }
+}

vortex-layout/src/plan/pipeline/tests.rs

@@ -18,6 +18,10 @@ use vortex_array::arrays::StructArray;
 use vortex_array::arrays::VarBinViewArray;
 use vortex_array::assert_arrays_eq;
 use vortex_array::buffer::BufferHandle;
+use vortex_array::expr::get_item;
+use vortex_array::expr::gt;
+use vortex_array::expr::lit;
+use vortex_array::expr::root;
 use vortex_array::serde::SerializeOptions;
 use vortex_buffer::Alignment;
 use vortex_buffer::ByteBufferMut;
@@ -31,11 +35,15 @@ use super::*;
 use crate::LayoutRef;
 use crate::OwnedLayoutChildren;
 use crate::layouts::chunked::ChunkedLayout;
+use crate::layouts::dict::DictLayout;
 use crate::layouts::flat::FlatLayout;
 use crate::layouts::struct_::StructLayout;
+use crate::plan::EvalPlan;
 use crate::plan::PlanRef;
 use crate::plan::SegmentScan;
+use crate::plan::Take;
 use crate::plan::lower;
+use crate::plan::optimize;
 use crate::segments::SegmentId;
 use crate::test::SESSION;
 
@@ -569,6 +577,68 @@ fn scan_keeps_the_selection_with_or_without_a_filter(#[case] sel: Sel) -> Vortex
     Ok(())
 }
 
+/// A take reads its values over their whole domain and its codes over the selection, including
+/// when a predicate has been pushed onto the values, and emits nothing before the values are
+/// whole. Later scans of the plan reuse the values and read only the codes.
+#[rstest]
+#[case::values(false)]
+#[case::predicate(true)]
+fn take_waits_for_whole_values_and_keeps_them(#[case] predicate: bool) -> VortexResult<()> {
+    let mut store = Store::default();
+    let values = VarBinViewArray::from_iter_str(["a", "b", "c"]).into_array();
+    let codes = PrimitiveArray::from_iter((0..ROWS).map(|v| (v % 3) as u8)).into_array();
+    let layout = DictLayout::new(store.flat(&values)?, store.flat(&codes)?).into_layout();
+    let mut plan = lower(&layout)?;
+    let mut expected = values.take(codes)?;
+    if predicate {
+        let expression = gt(root(), lit("a"))
+            .bind(plan.dtype())?
+            .optimize_recursive()?;
+        expected = expected.apply_bound(&expression)?;
+        plan = optimize(EvalPlan::try_new(expression, plan)?.into_plan())?;
+    }
+    assert!(plan.is::<Take>());
+
+    let mask = Sel::EveryOther.mask(10);
+    // Codes (segment 1) land first; nothing comes out until the values (segment 0) do.
+    let first = run(&store, &plan, 0..10, mask.clone(), scripted(&[1, 0]))?;
+    assert_eq!(reads(&first.events), 2);
+    let first_piece = first
+        .events
+        .iter()
+        .position(|e| matches!(e, Event::Piece(_)))
+        .vortex_expect("a piece");
+    assert!(
+        first.events[first_piece..]
+            .iter()
+            .all(|e| !matches!(e, Event::Delivered(_))),
+        "the only array must follow both deliveries"
+    );
+    assert_view(&expected, &(0..10), &mask, first.arrays)?;
+
+    Ok(())
+}
+
+#[test]
+fn eval_applies_expression_to_selected_rows() -> VortexResult<()> {
+    let mut store = Store::default();
+    let (plan, expected) = two_columns(&mut store)?;
+    let expression = gt(get_item("a", root()), lit(4_i32)).bind(plan.dtype())?;
+    let expected = expected.apply_bound(&expression)?;
+    let plan = EvalPlan::try_new(expression, plan)?.into_plan();
+
+    let rows = 2..18;
+    let mask = Sel::EveryOther.mask(16);
+    let run = run(
+        &store,
+        &plan,
+        rows.clone(),
+        mask.clone(),
+        delivery(Delivery::Lifo),
+    )?;
+    assert_view(&expected, &rows, &mask, run.arrays)
+}
+
 /// A column read by two fields of one struct is read and decoded once, and both fields get it.
 #[test]
 fn a_segment_two_readers_need_is_read_once() -> VortexResult<()> {

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

@@ -3,11 +3,13 @@
 
 use std::borrow::Cow;
 use std::fmt;
+use std::ops::Range;
 
 use vortex_array::EmptyMetadata;
 use vortex_array::expr::BoundExpression;
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
+use vortex_mask::Mask;
 use vortex_session::registry::CachedId;
 
 use crate::plan::Plan;
@@ -18,6 +20,11 @@ use crate::plan::PlanRef;
 use crate::plan::PlanVTable;
 use crate::plan::check_child_count;
 use crate::plan::optimizer::PlanReduceRule;
+use crate::plan::pipeline::Chain;
+use crate::plan::pipeline::Compiler;
+use crate::plan::pipeline::Reach;
+use crate::plan::pipeline::ops::EvalStage;
+use crate::segments::SegmentId;
 
 /// Applies an expression to the output of its child.
 #[derive(Clone, Debug)]
@@ -113,6 +120,25 @@ impl PlanVTable for Eval {
             Cow::Owned(format!("child[{index}]"))
         }
     }
+
+    fn compile(
+        plan: &Plan<Self>,
+        rows: Range<u64>,
+        mask: &Mask,
+        compiler: &mut Compiler<'_>,
+    ) -> VortexResult<Option<Chain>> {
+        let chain = compiler.compile(&plan.child_plan()?, rows, mask)?;
+        Ok(chain.map(|chain| chain.with(EvalStage::new(plan.expression().clone()))))
+    }
+
+    fn reach(
+        plan: &Plan<Self>,
+        rows: Range<u64>,
+        at: &Reach,
+        visit: &mut dyn FnMut(SegmentId, Range<u64>),
+    ) -> VortexResult<()> {
+        plan.child_plan()?.reach(rows, at, visit)
+    }
 }
 
 fn validate_expression_child(expression: &BoundExpression, child: &PlanRef) -> VortexResult<()> {

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

@@ -3,8 +3,10 @@
 
 use std::fmt::Display;
 use std::fmt::Formatter;
+use std::ops::Range;
 
 use vortex_array::EmptyMetadata;
+use vortex_array::IntoArray;
 use vortex_array::dtype::DType;
 use vortex_array::dtype::FieldName;
 use vortex_array::dtype::Nullability;
@@ -16,9 +18,11 @@ use vortex_array::expr::traversal::NodeExt;
 use vortex_array::expr::traversal::Transformed;
 use vortex_array::expr::traversal::TraversalOrder;
 use vortex_array::scalar_fn::fns::pack::Pack as PackFn;
+use vortex_buffer::Buffer;
 use vortex_error::VortexResult;
 use vortex_error::vortex_ensure_eq;
 use vortex_error::vortex_err;
+use vortex_mask::Mask;
 use vortex_session::registry::CachedId;
 
 use crate::layouts::row_idx::RowIdx as RowIdxFn;
@@ -31,7 +35,12 @@ use crate::plan::PlanParts;
 use crate::plan::PlanRef;
 use crate::plan::PlanVTable;
 use crate::plan::check_child_count;
+use crate::plan::pipeline::Chain;
+use crate::plan::pipeline::Compiler;
+use crate::plan::pipeline::Reach;
+use crate::plan::pipeline::ops::OnceSource;
 use crate::plan::plans::pack::rewrite_partition_root;
+use crate::segments::SegmentId;
 
 const ROW_IDX_PARTITION_NAME: &str = "row_idx";
 const CHILD_PARTITION_NAME: &str = "child";
@@ -86,6 +95,36 @@ impl PlanVTable for RowIdx {
     ) -> VortexResult<()> {
         check_child_count("RowIdx", children, 0)
     }
+
+    fn compile(
+        plan: &Plan<Self>,
+        rows: Range<u64>,
+        mask: &Mask,
+        compiler: &mut Compiler<'_>,
+    ) -> VortexResult<Option<Chain>> {
+        let _ = plan;
+        if mask.all_false() {
+            return Ok(None);
+        }
+        let offset = compiler.row_offset();
+        let indices = Buffer::from_iter(rows.start + offset..rows.end + offset).into_array();
+        let indices = if mask.all_true() {
+            indices
+        } else {
+            indices.filter(mask.clone())?
+        };
+        Ok(Some(Chain::new(OnceSource::new(indices))))
+    }
+
+    fn reach(
+        plan: &Plan<Self>,
+        rows: Range<u64>,
+        at: &Reach,
+        visit: &mut dyn FnMut(SegmentId, Range<u64>),
+    ) -> VortexResult<()> {
+        let _ = (plan, rows, at, visit);
+        Ok(())
+    }
 }
 
 /// Plans an expression over a data source and its global row-index domain.