Skip to content

Commit d5fc0af

Browse files
joseph-isaacsclaude
joseph-isaacs
and
claude
committed
feat(layout): add Eval, RowIdx, and Take exec nodes
- `EvalNode` applies its bound expression to each array its child produces as it streams through. The child runs over the same selection, so the expression never sees an unselected row. - `RowIdxNode` emits the global row index of every selected row at start, offset by the graph's row offset. - `TakeNode` runs the codes over the selection and the values over their whole domain. It is `Ready::Closed(values)`: it does not run until the values have closed, joins them once, and then wraps each codes array as a dictionary over them as it arrives. Signed-off-by: Joe Isaacs <[email protected]> Co-Authored-By: Claude Fable 5.1 <[email protected]> Claude-Session: https://claude.ai/code/session_01NWYmfHTFtaEdhdyd3u5hv3 Signed-off-by: Joe Isaacs <[email protected]>
1 parent 2651b89 commit d5fc0af

7 files changed

Lines changed: 337 additions & 0 deletions

File tree

‎vortex-layout/src/plan/physical/operators/eval.rs‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
// SPDX-License-Identifier: Apache-2.0
2+
// SPDX-FileCopyrightText: Copyright the Vortex contributors
3+
4+
//! Applying an expression to each batch.
5+
6+
use vortex_array::expr::BoundExpression;
7+
use vortex_error::VortexResult;
8+
9+
use crate::plan::physical::operator::Cx;
10+
use crate::plan::physical::operator::Input;
11+
use crate::plan::physical::operator::Operator;
12+
use crate::plan::physical::operator::Step;
13+
14+
/// Applies an expression to each batch.
15+
pub(crate) struct EvalStage {
16+
expression: BoundExpression,
17+
}
18+
19+
impl EvalStage {
20+
pub(crate) fn new(expression: BoundExpression) -> Self {
21+
Self { expression }
22+
}
23+
}
24+
25+
impl Operator for EvalStage {
26+
fn compute(&mut self, input: Input, _cx: &mut Cx<'_>) -> VortexResult<Step> {
27+
let Input::Batch(batch) = input else {
28+
return Ok(Step::End);
29+
};
30+
Ok(Step::Final(batch.apply_bound(&self.expression)?))
31+
}
32+
}

‎vortex-layout/src/plan/physical/operators/mod.rs‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,25 +7,35 @@
77
//! batch twice, holds more than the batches it must join, or scans inlets it does not read.
88
99
mod concat;
10+
mod eval;
1011
mod filter;
1112
mod once;
1213
mod pack;
1314
mod scan;
15+
mod take;
1416

1517
pub(crate) use concat::*;
18+
pub(crate) use eval::*;
1619
pub(crate) use filter::*;
1720
pub(crate) use once::*;
1821
pub(crate) use pack::*;
1922
pub(crate) use scan::*;
2023
use smallvec::SmallVec;
24+
pub(crate) use take::*;
2125
use vortex_array::ArrayRef;
26+
use vortex_array::Canonical;
2227
use vortex_array::IntoArray;
2328
use vortex_array::arrays::ChunkedArray;
29+
use vortex_array::dtype::DType;
2430
use vortex_error::VortexResult;
2531
use vortex_error::vortex_err;
2632

33+
use crate::plan::physical::operator::Cx;
2734
use crate::plan::physical::pipeline::Inlet;
2835

36+
/// An inlet capacity that never blocks the sink, for a source that needs its inlet whole.
37+
pub(crate) const UNBOUNDED: usize = usize::MAX;
38+
2939
/// Takes the first `len` rows of the inlet, across as many batches as hold them, slicing the last
3040
/// and leaving its rest in place. Rows spanning batches come back chunked, not copied.
3141
fn take_rows(inlet: &mut Inlet<'_>, len: usize) -> VortexResult<ArrayRef> {
@@ -51,3 +61,26 @@ fn take_rows(inlet: &mut Inlet<'_>, len: usize) -> VortexResult<ArrayRef> {
5161
// SAFETY: the pieces come from one inlet, whose batches share its sink's dtype.
5262
Ok(unsafe { ChunkedArray::new_unchecked(pieces, dtype) }.into_array())
5363
}
64+
65+
/// Joins arrays covering consecutive rows into one.
66+
pub(crate) fn join(dtype: &DType, mut arrays: Vec<ArrayRef>) -> VortexResult<ArrayRef> {
67+
match arrays.len() {
68+
0 => Ok(Canonical::empty(dtype).into_array()),
69+
1 => Ok(arrays.remove(0)),
70+
_ => Ok(ChunkedArray::try_new(arrays, dtype.clone())?.into_array()),
71+
}
72+
}
73+
74+
/// Takes every batch of a closed inlet.
75+
fn drain(inlet: &mut Inlet<'_>) -> Vec<ArrayRef> {
76+
let mut batches = Vec::with_capacity(inlet.len());
77+
while let Some(batch) = inlet.take() {
78+
batches.push(batch);
79+
}
80+
batches
81+
}
82+
83+
/// Joins the whole of closed inlet `index` into one array of `dtype`.
84+
fn join_inlet(cx: &mut Cx<'_>, index: usize, dtype: &DType) -> VortexResult<ArrayRef> {
85+
join(dtype, drain(&mut cx.inlet(index)))
86+
}

‎vortex-layout/src/plan/physical/operators/take.rs‎

Lines changed: 89 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,89 @@
1+
// SPDX-License-Identifier: Apache-2.0
2+
// SPDX-FileCopyrightText: Copyright the Vortex contributors
3+
4+
//! Looking codes up in dictionary values.
5+
6+
use vortex_array::ArrayRef;
7+
use vortex_array::IntoArray;
8+
use vortex_array::arrays::DictArray;
9+
use vortex_array::arrays::Shared;
10+
use vortex_array::arrays::SharedArray;
11+
use vortex_error::VortexResult;
12+
13+
use super::UNBOUNDED;
14+
use super::join_inlet;
15+
use crate::plan::physical::TakePlan;
16+
use crate::plan::physical::operator::Cx;
17+
use crate::plan::physical::operator::Input;
18+
use crate::plan::physical::operator::Operator;
19+
use crate::plan::physical::operator::Source;
20+
use crate::plan::physical::operator::SourceBlocked;
21+
use crate::plan::physical::operator::SourceReady;
22+
use crate::plan::physical::operator::Step;
23+
use crate::plan::physical::pipeline::DEFAULT_CAPACITY;
24+
25+
/// Wraps each batch of codes as a dictionary over the values, once the values inlet has ended.
26+
pub(crate) struct TakeSource {
27+
plan: TakePlan,
28+
/// The values, once their inlet has ended.
29+
values: Option<ArrayRef>,
30+
}
31+
32+
const CODES: usize = 0;
33+
const VALUES: usize = 1;
34+
35+
impl TakeSource {
36+
/// A take over `values`, or over what its values inlet produces when `None`.
37+
pub(crate) fn new(plan: TakePlan, values: Option<ArrayRef>) -> Self {
38+
Self { plan, values }
39+
}
40+
41+
/// The batches inlet `inlet` holds: the values are taken whole, so they never block.
42+
pub(crate) fn capacity(inlet: usize) -> usize {
43+
if inlet == VALUES {
44+
UNBOUNDED
45+
} else {
46+
DEFAULT_CAPACITY
47+
}
48+
}
49+
}
50+
51+
impl Source for TakeSource {
52+
fn ready(&mut self, cx: &mut Cx<'_>) -> VortexResult<SourceReady> {
53+
if self.values.is_none() && !cx.inlet(VALUES).closed() {
54+
return Ok(SourceReady::Blocked(SourceBlocked::Inlet(VALUES)));
55+
}
56+
let codes = cx.inlet(CODES);
57+
Ok(if codes.is_empty() && !codes.closed() {
58+
SourceReady::Blocked(SourceBlocked::Inlet(CODES))
59+
} else {
60+
SourceReady::Ready
61+
})
62+
}
63+
}
64+
65+
impl Operator for TakeSource {
66+
fn compute(&mut self, _input: Input, cx: &mut Cx<'_>) -> VortexResult<Step> {
67+
let values = match &self.values {
68+
Some(values) => values.clone(),
69+
None => {
70+
let values = shared(join_inlet(cx, VALUES, self.plan.values()?.dtype())?);
71+
self.values = Some(values.clone());
72+
values
73+
}
74+
};
75+
Ok(match cx.inlet(CODES).take() {
76+
Some(codes) => Step::Final(DictArray::try_new(codes, values)?.into_array()),
77+
None => Step::End,
78+
})
79+
}
80+
}
81+
82+
/// `array` as a [`SharedArray`], so however many dictionaries use it, it is canonicalized once.
83+
fn shared(array: ArrayRef) -> ArrayRef {
84+
if array.is::<Shared>() {
85+
array
86+
} else {
87+
SharedArray::new(array).into_array()
88+
}
89+
}

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

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,10 @@ use vortex_array::arrays::VarBinViewArray;
2020
use vortex_array::assert_arrays_eq;
2121
use vortex_array::buffer::BufferHandle;
2222
use vortex_array::dtype::StructFields;
23+
use vortex_array::expr::get_item;
24+
use vortex_array::expr::gt;
25+
use vortex_array::expr::lit;
26+
use vortex_array::expr::root;
2327
use vortex_array::serde::SerializeOptions;
2428
use vortex_buffer::Alignment;
2529
use vortex_buffer::Buffer;
@@ -36,12 +40,15 @@ use super::*;
3640
use crate::LayoutRef;
3741
use crate::OwnedLayoutChildren;
3842
use crate::layouts::chunked::ChunkedLayout;
43+
use crate::layouts::dict::DictLayout;
3944
use crate::layouts::flat::FlatLayout;
4045
use crate::layouts::struct_::StructLayout;
46+
use crate::plan::physical::EvalPlan;
4147
use crate::plan::physical::FilterPlan;
4248
use crate::plan::physical::PackPlan;
4349
use crate::plan::physical::PlanRef;
4450
use crate::plan::physical::SegmentScan;
51+
use crate::plan::physical::Take;
4552
use crate::plan::physical::lower;
4653
use crate::plan::physical::operator::Cx;
4754
use crate::plan::physical::operator::InletSet;
@@ -52,6 +59,7 @@ use crate::plan::physical::operator::SourceBlocked;
5259
use crate::plan::physical::operator::SourceReady;
5360
use crate::plan::physical::operator::Step;
5461
use crate::plan::physical::operators::PackSource;
62+
use crate::plan::physical::optimize;
5563
use crate::segments::SegmentId;
5664
use crate::test::SESSION;
5765

@@ -581,6 +589,68 @@ fn scan_keeps_the_selection_with_or_without_a_filter(#[case] sel: Sel) -> Vortex
581589
Ok(())
582590
}
583591

592+
/// A take reads its values over their whole domain and its codes over the selection, including
593+
/// when a predicate has been pushed onto the values, and emits nothing before the values are
594+
/// whole. Later scans of the plan reuse the values and read only the codes.
595+
#[rstest]
596+
#[case::values(false)]
597+
#[case::predicate(true)]
598+
fn take_waits_for_whole_values_and_keeps_them(#[case] predicate: bool) -> VortexResult<()> {
599+
let mut store = Store::default();
600+
let values = VarBinViewArray::from_iter_str(["a", "b", "c"]).into_array();
601+
let codes = PrimitiveArray::from_iter((0..ROWS).map(|v| (v % 3) as u8)).into_array();
602+
let layout = DictLayout::new(store.flat(&values)?, store.flat(&codes)?).into_layout();
603+
let mut plan = lower(&layout)?;
604+
let mut expected = values.take(codes)?;
605+
if predicate {
606+
let expression = gt(root(), lit("a"))
607+
.bind(plan.dtype())?
608+
.optimize_recursive()?;
609+
expected = expected.apply_bound(&expression)?;
610+
plan = optimize(EvalPlan::try_new(expression, plan)?.into_plan())?;
611+
}
612+
assert!(plan.is::<Take>());
613+
614+
let mask = Sel::EveryOther.mask(10);
615+
// Codes (segment 1) land first; nothing comes out until the values (segment 0) do.
616+
let first = run(&store, &plan, 0..10, mask.clone(), scripted(&[1, 0]))?;
617+
assert_eq!(reads(&first.events), 2);
618+
let first_piece = first
619+
.events
620+
.iter()
621+
.position(|e| matches!(e, Event::Piece(_)))
622+
.vortex_expect("a piece");
623+
assert!(
624+
first.events[first_piece..]
625+
.iter()
626+
.all(|e| !matches!(e, Event::Delivered(_))),
627+
"the only array must follow both deliveries"
628+
);
629+
assert_view(&expected, &(0..10), &mask, first.arrays)?;
630+
631+
Ok(())
632+
}
633+
634+
#[test]
635+
fn eval_applies_expression_to_selected_rows() -> VortexResult<()> {
636+
let mut store = Store::default();
637+
let (plan, expected) = two_columns(&mut store)?;
638+
let expression = gt(get_item("a", root()), lit(4_i32)).bind(plan.dtype())?;
639+
let expected = expected.apply_bound(&expression)?;
640+
let plan = EvalPlan::try_new(expression, plan)?.into_plan();
641+
642+
let rows = 2..18;
643+
let mask = Sel::EveryOther.mask(16);
644+
let run = run(
645+
&store,
646+
&plan,
647+
rows.clone(),
648+
mask.clone(),
649+
delivery(Delivery::Lifo),
650+
)?;
651+
assert_view(&expected, &rows, &mask, run.arrays)
652+
}
653+
584654
/// A column read by two fields of one struct is read and decoded once, and both fields get it.
585655
#[test]
586656
fn a_segment_two_readers_need_is_read_once() -> VortexResult<()> {

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

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,11 +3,13 @@
33

44
use std::borrow::Cow;
55
use std::fmt;
6+
use std::ops::Range;
67

78
use vortex_array::EmptyMetadata;
89
use vortex_array::expr::BoundExpression;
910
use vortex_error::VortexResult;
1011
use vortex_error::vortex_bail;
12+
use vortex_mask::Mask;
1113
use vortex_session::registry::CachedId;
1214

1315
use crate::plan::physical::Plan;
@@ -17,7 +19,12 @@ use crate::plan::physical::PlanParts;
1719
use crate::plan::physical::PlanRef;
1820
use crate::plan::physical::PlanVTable;
1921
use crate::plan::physical::check_child_count;
22+
use crate::plan::physical::operators::EvalStage;
2023
use crate::plan::physical::optimizer::PlanReduceRule;
24+
use crate::plan::physical::pipeline::Chain;
25+
use crate::plan::physical::pipeline::Compiler;
26+
use crate::plan::physical::pipeline::Reach;
27+
use crate::segments::SegmentId;
2128

2229
/// Applies an expression to the output of its child.
2330
#[derive(Clone, Debug)]
@@ -113,6 +120,25 @@ impl PlanVTable for Eval {
113120
Cow::Owned(format!("child[{index}]"))
114121
}
115122
}
123+
124+
fn compile(
125+
plan: &Plan<Self>,
126+
rows: Range<u64>,
127+
mask: &Mask,
128+
compiler: &mut Compiler<'_>,
129+
) -> VortexResult<Option<Chain>> {
130+
let chain = compiler.compile(&plan.child_plan()?, rows, mask)?;
131+
Ok(chain.map(|chain| chain.with(EvalStage::new(plan.expression().clone()))))
132+
}
133+
134+
fn reach(
135+
plan: &Plan<Self>,
136+
rows: Range<u64>,
137+
at: &Reach,
138+
visit: &mut dyn FnMut(SegmentId, Range<u64>),
139+
) -> VortexResult<()> {
140+
plan.child_plan()?.reach(rows, at, visit)
141+
}
116142
}
117143

118144
fn validate_expression_child(expression: &BoundExpression, child: &PlanRef) -> VortexResult<()> {

0 commit comments

Comments
 (0)