diff --git a/vortex-layout/src/plan/pipeline/ops/eval.rs b/vortex-layout/src/plan/pipeline/ops/eval.rs new file mode 100644 index 00000000000..27f18ce2e5b --- /dev/null +++ b/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 { + match input { + Input::Chunk(batch) => Ok(Step::Last(batch.apply_bound(&self.expression)?)), + Input::End => Ok(Step::Finished), + Input::None => Ok(Step::Consumed), + } + } +} diff --git a/vortex-layout/src/plan/pipeline/ops/mod.rs b/vortex-layout/src/plan/pipeline/ops/mod.rs index f5506551a3d..9cd076fc5c9 100644 --- a/vortex-layout/src/plan/pipeline/ops/mod.rs +++ b/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 { @@ -51,3 +60,21 @@ pub(crate) fn take_rows(inlet: &mut Inlet<'_>, len: usize) -> VortexResult) -> VortexResult { + 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 { + let mut batches = Vec::with_capacity(inlet.len()); + while let Some(batch) = inlet.take() { + batches.push(batch); + } + batches +} diff --git a/vortex-layout/src/plan/pipeline/ops/take.rs b/vortex-layout/src/plan/pipeline/ops/take.rs new file mode 100644 index 00000000000..e7de9fca232 --- /dev/null +++ b/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, +} + +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) -> Self { + Self { plan, values } + } +} + +impl Operator for TakeSource { + fn compute(&mut self, _input: Input, cx: &mut Cx<'_>) -> VortexResult { + 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::() { + 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 + } + } +} diff --git a/vortex-layout/src/plan/pipeline/tests.rs b/vortex-layout/src/plan/pipeline/tests.rs index 3d8d4845bd1..61ded4bd2aa 100644 --- a/vortex-layout/src/plan/pipeline/tests.rs +++ b/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::()); + + 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<()> { diff --git a/vortex-layout/src/plan/plans/eval.rs b/vortex-layout/src/plan/plans/eval.rs index f1edb967e41..b0ae10850ee 100644 --- a/vortex-layout/src/plan/plans/eval.rs +++ b/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, + rows: Range, + mask: &Mask, + compiler: &mut Compiler<'_>, + ) -> VortexResult> { + let chain = compiler.compile(&plan.child_plan()?, rows, mask)?; + Ok(chain.map(|chain| chain.with(EvalStage::new(plan.expression().clone())))) + } + + fn reach( + plan: &Plan, + rows: Range, + at: &Reach, + visit: &mut dyn FnMut(SegmentId, Range), + ) -> VortexResult<()> { + plan.child_plan()?.reach(rows, at, visit) + } } fn validate_expression_child(expression: &BoundExpression, child: &PlanRef) -> VortexResult<()> { diff --git a/vortex-layout/src/plan/plans/row_idx.rs b/vortex-layout/src/plan/plans/row_idx.rs index 6b1edbbc433..db8a8911038 100644 --- a/vortex-layout/src/plan/plans/row_idx.rs +++ b/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, + rows: Range, + mask: &Mask, + compiler: &mut Compiler<'_>, + ) -> VortexResult> { + 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, + rows: Range, + at: &Reach, + visit: &mut dyn FnMut(SegmentId, Range), + ) -> VortexResult<()> { + let _ = (plan, rows, at, visit); + Ok(()) + } } /// Plans an expression over a data source and its global row-index domain. diff --git a/vortex-layout/src/plan/plans/take.rs b/vortex-layout/src/plan/plans/take.rs index 95919090f3c..7ab2569d88f 100644 --- a/vortex-layout/src/plan/plans/take.rs +++ b/vortex-layout/src/plan/plans/take.rs @@ -2,12 +2,16 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::borrow::Cow; +use std::ops::Range; +use vortex_array::Canonical; use vortex_array::EmptyMetadata; +use vortex_array::IntoArray; use vortex_array::dtype::DType; use vortex_array::expr::ExactBoundExpr; use vortex_array::expr::label_bound_tree; use vortex_error::VortexResult; +use vortex_mask::Mask; use vortex_session::registry::CachedId; use crate::plan::Eval; @@ -20,6 +24,12 @@ use crate::plan::PlanRef; use crate::plan::PlanVTable; use crate::plan::check_child_count; use crate::plan::optimizer::PlanParentReduceRule; +use crate::plan::pipeline::Chain; +use crate::plan::pipeline::Compiler; +use crate::plan::pipeline::Reach; +use crate::plan::pipeline::ops::OnceSource; +use crate::plan::pipeline::ops::TakeSource; +use crate::segments::SegmentId; const CODES: usize = 0; const VALUES: usize = 1; @@ -117,6 +127,40 @@ impl PlanVTable for Take { _ => Cow::Owned(format!("child[{index}]")), } } + + fn compile( + plan: &Plan, + rows: Range, + mask: &Mask, + compiler: &mut Compiler<'_>, + ) -> VortexResult> { + let values = plan.values()?; + let values_rows = 0..values.row_count(); + let Some(codes) = compiler.compile(&plan.codes()?, rows, mask)? else { + return Ok(None); + }; + let len = usize::try_from(values_rows.end)?; + let values = match compiler.compile(&values, values_rows, &Mask::new_true(len))? { + Some(values) => values, + None => Chain::new(OnceSource::new( + Canonical::empty(values.dtype()).into_array(), + )), + }; + let source = TakeSource::new(plan.clone(), None); + Ok(Some(compiler.join(vec![codes, values], source))) + } + + fn reach( + plan: &Plan, + rows: Range, + at: &Reach, + visit: &mut dyn FnMut(SegmentId, Range), + ) -> VortexResult<()> { + let values = plan.values()?; + let len = values.row_count(); + values.reach(0..len, &at.fixed(&rows), visit)?; + plan.codes()?.reach(rows, at, visit) + } } /// Pushes a strict, infallible boolean expression onto the dictionary values of a [`Take`].