Skip to content

Commit 3a151c6

Browse files
gatesnclaude
gatesn
and
claude
committed
feat: migrate single-child wrappers to iterative execute (phase 4b)
Replace recursive child.execute() calls with ExecutionStep returns in Slice, Filter, Masked, FoR, and ZigZag vtables. Each now checks if its child is already in the needed form (canonical/primitive/constant) and returns Done directly, or returns ExecuteChild(0)/ColumnarizeChild(0) to let the scheduler handle child execution iteratively. Also handles ConstantArray children explicitly to prevent infinite loops in the scheduler (since constants are already columnar and won't be re-executed). FoR decompress is split into try_fused_decompress and apply_reference for reuse without recursive execution. Signed-off-by: Nicholas Gates <nick@nickgates.com> Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
1 parent a675c04 commit 3a151c6

7 files changed

Lines changed: 178 additions & 43 deletions

File tree

‎encodings/fastlanes/src/for/array/for_compress.rs‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -68,8 +68,9 @@ mod test {
6868

6969
use super::*;
7070
use crate::BitPackedArray;
71-
use crate::r#for::array::for_decompress::decompress;
71+
use crate::r#for::array::for_decompress::apply_reference;
7272
use crate::r#for::array::for_decompress::fused_decompress;
73+
use crate::r#for::array::for_decompress::try_fused_decompress;
7374

7475
static SESSION: LazyLock<VortexSession> =
7576
LazyLock::new(|| VortexSession::empty().with::<ArraySession>());
@@ -169,7 +170,13 @@ mod test {
169170
let expected_unsigned = PrimitiveArray::from_iter(unsigned);
170171
assert_arrays_eq!(encoded, expected_unsigned);
171172

172-
let decompressed = decompress(&compressed, &mut SESSION.create_execution_ctx())?;
173+
let mut ctx = SESSION.create_execution_ctx();
174+
let decompressed = if let Some(result) = try_fused_decompress(&compressed, &mut ctx)? {
175+
result
176+
} else {
177+
let encoded = compressed.encoded().to_primitive();
178+
apply_reference(&compressed, encoded)
179+
};
173180
array
174181
.as_slice::<i8>()
175182
.iter()

‎encodings/fastlanes/src/for/array/for_decompress.rs‎

Lines changed: 14 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -45,23 +45,28 @@ impl<T: PhysicalPType<Physical = T> + FoR> UnpackStrategy<T> for FoRStrategy<T>
4545
}
4646
}
4747

48-
pub fn decompress(array: &FoRArray, ctx: &mut ExecutionCtx) -> VortexResult<PrimitiveArray> {
49-
let ptype = array.ptype();
50-
51-
// Try to do fused unpack.
48+
/// Try the fused BitPacked decompression path. Returns `None` if the child is not BitPacked
49+
/// or the reference type is not unsigned.
50+
pub fn try_fused_decompress(
51+
array: &FoRArray,
52+
ctx: &mut ExecutionCtx,
53+
) -> VortexResult<Option<PrimitiveArray>> {
5254
if array.reference_scalar().dtype().is_unsigned_int()
5355
&& let Some(bp) = array.encoded().as_opt::<BitPackedVTable>()
5456
{
5557
return match_each_unsigned_integer_ptype!(array.ptype(), |T| {
56-
fused_decompress::<T>(array, bp, ctx)
58+
fused_decompress::<T>(array, bp, ctx).map(Some)
5759
});
5860
}
61+
Ok(None)
62+
}
5963

60-
// TODO(ngates): Do we need this to be into_encoded() somehow?
61-
let encoded = array.encoded().clone().execute::<PrimitiveArray>(ctx)?;
64+
/// Apply the FoR reference value to an already-decoded PrimitiveArray.
65+
pub fn apply_reference(array: &FoRArray, encoded: PrimitiveArray) -> PrimitiveArray {
66+
let ptype = array.ptype();
6267
let validity = encoded.validity().clone();
6368

64-
Ok(match_each_integer_ptype!(ptype, |T| {
69+
match_each_integer_ptype!(ptype, |T| {
6570
let min = array
6671
.reference_scalar()
6772
.as_primitive()
@@ -75,7 +80,7 @@ pub fn decompress(array: &FoRArray, ctx: &mut ExecutionCtx) -> VortexResult<Prim
7580
validity,
7681
)
7782
}
78-
}))
83+
})
7984
}
8085

8186
pub(crate) fn fused_decompress<

‎encodings/fastlanes/src/for/vtable/mod.rs‎

Lines changed: 55 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,8 +11,12 @@ use vortex_array::ExecutionCtx;
1111
use vortex_array::ExecutionStep;
1212
use vortex_array::IntoArray;
1313
use vortex_array::Precision;
14+
use vortex_array::arrays::ConstantArray;
15+
use vortex_array::arrays::ConstantVTable;
16+
use vortex_array::arrays::PrimitiveVTable;
1417
use vortex_array::buffer::BufferHandle;
1518
use vortex_array::dtype::DType;
19+
use vortex_array::match_each_integer_ptype;
1620
use vortex_array::scalar::Scalar;
1721
use vortex_array::scalar::ScalarValue;
1822
use vortex_array::serde::ArrayChildren;
@@ -21,14 +25,16 @@ use vortex_array::vtable;
2125
use vortex_array::vtable::ArrayId;
2226
use vortex_array::vtable::VTable;
2327
use vortex_array::vtable::ValidityVTableFromChild;
28+
use vortex_error::VortexExpect;
2429
use vortex_error::VortexResult;
2530
use vortex_error::vortex_bail;
2631
use vortex_error::vortex_ensure;
2732
use vortex_error::vortex_panic;
2833
use vortex_session::VortexSession;
2934

3035
use crate::FoRArray;
31-
use crate::r#for::array::for_decompress::decompress;
36+
use crate::r#for::array::for_decompress::apply_reference;
37+
use crate::r#for::array::for_decompress::try_fused_decompress;
3238
use crate::r#for::vtable::kernels::PARENT_KERNELS;
3339
use crate::r#for::vtable::rules::PARENT_RULES;
3440

@@ -167,7 +173,54 @@ impl VTable for FoRVTable {
167173
}
168174

169175
fn execute(array: &Self::Array, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionStep> {
170-
Ok(ExecutionStep::Done(decompress(array, ctx)?.into_array()))
176+
// Try fused decompress with BitPacked child (no child execution needed).
177+
if let Some(result) = try_fused_decompress(array, ctx)? {
178+
return Ok(ExecutionStep::Done(result.into_array()));
179+
}
180+
181+
// If child is already a PrimitiveArray, add the reference value.
182+
if array.encoded().is::<PrimitiveVTable>() {
183+
let encoded = array.encoded().as_::<PrimitiveVTable>().clone();
184+
return Ok(ExecutionStep::Done(
185+
apply_reference(array, encoded).into_array(),
186+
));
187+
}
188+
189+
// If child is a constant, compute the result as a constant.
190+
if let Some(constant) = array.encoded().as_opt::<ConstantVTable>() {
191+
let scalar = constant.scalar();
192+
if scalar.is_null() {
193+
return Ok(ExecutionStep::Done(
194+
ConstantArray::new(Scalar::null(array.dtype().clone()), array.len())
195+
.into_array(),
196+
));
197+
}
198+
return Ok(ExecutionStep::Done(match_each_integer_ptype!(
199+
array.ptype(),
200+
|T| {
201+
let enc_val = scalar
202+
.as_primitive()
203+
.typed_value::<T>()
204+
.vortex_expect("constant must be non-null after check");
205+
let ref_val = array
206+
.reference_scalar()
207+
.as_primitive()
208+
.typed_value::<T>()
209+
.vortex_expect("reference must be non-null");
210+
ConstantArray::new(
211+
Scalar::primitive(
212+
enc_val.wrapping_add(ref_val),
213+
scalar.dtype().nullability(),
214+
),
215+
array.len(),
216+
)
217+
.into_array()
218+
}
219+
)));
220+
}
221+
222+
// Otherwise, ask the scheduler to execute the child first.
223+
Ok(ExecutionStep::ExecuteChild(0))
171224
}
172225

173226
fn execute_parent(

‎encodings/zigzag/src/array.rs‎

Lines changed: 36 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,9 @@ use vortex_array::ExecutionCtx;
1212
use vortex_array::ExecutionStep;
1313
use vortex_array::IntoArray;
1414
use vortex_array::Precision;
15+
use vortex_array::arrays::ConstantArray;
16+
use vortex_array::arrays::ConstantVTable;
17+
use vortex_array::arrays::PrimitiveVTable;
1518
use vortex_array::buffer::BufferHandle;
1619
use vortex_array::dtype::DType;
1720
use vortex_array::dtype::PType;
@@ -149,10 +152,39 @@ impl VTable for ZigZagVTable {
149152
Ok(())
150153
}
151154

152-
fn execute(array: &Self::Array, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionStep> {
153-
Ok(ExecutionStep::Done(
154-
zigzag_decode(array.encoded().clone().execute(ctx)?).into_array(),
155-
))
155+
fn execute(array: &Self::Array, _ctx: &mut ExecutionCtx) -> VortexResult<ExecutionStep> {
156+
// If child is already a PrimitiveArray, decode it directly.
157+
if array.encoded().is::<PrimitiveVTable>() {
158+
let encoded = array.encoded().as_::<PrimitiveVTable>().clone();
159+
return Ok(ExecutionStep::Done(zigzag_decode(encoded).into_array()));
160+
}
161+
162+
// If child is a constant, decode the scalar value.
163+
if let Some(constant) = array.encoded().as_opt::<ConstantVTable>() {
164+
let scalar = constant.scalar();
165+
if scalar.is_null() {
166+
return Ok(ExecutionStep::Done(
167+
ConstantArray::new(Scalar::null(array.dtype().clone()), array.len())
168+
.into_array(),
169+
));
170+
}
171+
let result = match_each_unsigned_integer_ptype!(scalar.as_primitive().ptype(), |P| {
172+
let val = scalar
173+
.as_primitive()
174+
.typed_value::<P>()
175+
.vortex_expect("constant must be non-null after check");
176+
Scalar::primitive(
177+
<<P as ZigZagEncoded>::Int>::decode(val),
178+
array.dtype().nullability(),
179+
)
180+
});
181+
return Ok(ExecutionStep::Done(
182+
ConstantArray::new(result, array.len()).into_array(),
183+
));
184+
}
185+
186+
// Otherwise, ask the scheduler to execute the child first.
187+
Ok(ExecutionStep::ExecuteChild(0))
156188
}
157189

158190
fn reduce_parent(

‎vortex-array/src/arrays/filter/vtable.rs‎

Lines changed: 23 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -13,12 +13,16 @@ use vortex_error::vortex_panic;
1313
use vortex_mask::Mask;
1414
use vortex_session::VortexSession;
1515

16+
use crate::AnyCanonical;
1617
use crate::Array;
1718
use crate::ArrayEq;
1819
use crate::ArrayHash;
1920
use crate::ArrayRef;
21+
use crate::Canonical;
2022
use crate::IntoArray;
2123
use crate::Precision;
24+
use crate::arrays::ConstantArray;
25+
use crate::arrays::ConstantVTable;
2226
use crate::arrays::filter::array::FilterArray;
2327
use crate::arrays::filter::execute::execute_filter;
2428
use crate::arrays::filter::execute::execute_filter_fast_paths;
@@ -156,18 +160,30 @@ impl VTable for FilterVTable {
156160
}
157161

158162
fn execute(array: &Self::Array, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionStep> {
159-
if let Some(canonical) = execute_filter_fast_paths(array, ctx)? {
160-
return Ok(ExecutionStep::Done(canonical));
163+
if let Some(result) = execute_filter_fast_paths(array, ctx)? {
164+
return Ok(ExecutionStep::Done(result));
161165
}
162166
let Mask::Values(mask_values) = &array.mask else {
163167
unreachable!("`execute_filter_fast_paths` handles AllTrue and AllFalse")
164168
};
165169

166-
// We rely on the optimization pass that runs prior to this execution for filter pushdown,
167-
// so now we can just execute the filter without worrying.
168-
Ok(ExecutionStep::Done(
169-
execute_filter(array.child.clone().execute(ctx)?, mask_values).into_array(),
170-
))
170+
// If child is already canonical, filter it directly.
171+
if let Some(canonical) = array.child.as_opt::<AnyCanonical>() {
172+
return Ok(ExecutionStep::Done(
173+
execute_filter(Canonical::from(canonical), mask_values).into_array(),
174+
));
175+
}
176+
177+
// If child is a constant, filtering just changes the length.
178+
if let Some(constant) = array.child.as_opt::<ConstantVTable>() {
179+
return Ok(ExecutionStep::Done(
180+
ConstantArray::new(constant.scalar().clone(), mask_values.true_count())
181+
.into_array(),
182+
));
183+
}
184+
185+
// Otherwise, ask the scheduler to execute the child first.
186+
Ok(ExecutionStep::ExecuteChild(0))
171187
}
172188

173189
fn reduce_parent(

‎vortex-array/src/arrays/masked/vtable/mod.rs‎

Lines changed: 21 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -13,12 +13,15 @@ use vortex_error::vortex_ensure;
1313
use vortex_error::vortex_panic;
1414
use vortex_session::VortexSession;
1515

16+
use crate::AnyCanonical;
1617
use crate::ArrayRef;
1718
use crate::Canonical;
1819
use crate::EmptyMetadata;
1920
use crate::IntoArray;
2021
use crate::Precision;
2122
use crate::arrays::ConstantArray;
23+
use crate::arrays::ConstantVTable;
24+
use crate::arrays::constant::constant_canonicalize;
2225
use crate::arrays::masked::MaskedArray;
2326
use crate::arrays::masked::compute::rules::PARENT_RULES;
2427
use crate::arrays::masked::mask_validity_canonical;
@@ -178,10 +181,24 @@ impl VTable for MaskedVTable {
178181
// While we could manually convert the dtype, `mask_validity_canonical` is already O(1) for
179182
// `AllTrue` masks (no data copying), so there's no benefit.
180183

181-
let child = array.child().clone().execute::<Canonical>(ctx)?;
182-
Ok(ExecutionStep::Done(
183-
mask_validity_canonical(child, &validity_mask, ctx)?.into_array(),
184-
))
184+
// If child is already canonical, apply the validity mask directly.
185+
if let Some(canonical) = array.child().as_opt::<AnyCanonical>() {
186+
return Ok(ExecutionStep::Done(
187+
mask_validity_canonical(Canonical::from(canonical), &validity_mask, ctx)?
188+
.into_array(),
189+
));
190+
}
191+
192+
// If child is a constant, expand to canonical then apply the validity mask.
193+
if let Some(constant) = array.child().as_opt::<ConstantVTable>() {
194+
let canonical = constant_canonicalize(constant)?;
195+
return Ok(ExecutionStep::Done(
196+
mask_validity_canonical(canonical, &validity_mask, ctx)?.into_array(),
197+
));
198+
}
199+
200+
// Otherwise, ask the scheduler to execute the child first.
201+
Ok(ExecutionStep::ColumnarizeChild(0))
185202
}
186203

187204
fn reduce_parent(

‎vortex-array/src/arrays/slice/vtable.rs‎

Lines changed: 20 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,10 @@ use crate::ArrayEq;
2020
use crate::ArrayHash;
2121
use crate::ArrayRef;
2222
use crate::Canonical;
23+
use crate::IntoArray;
2324
use crate::Precision;
25+
use crate::arrays::ConstantArray;
26+
use crate::arrays::ConstantVTable;
2427
use crate::arrays::slice::array::SliceArray;
2528
use crate::arrays::slice::rules::PARENT_RULES;
2629
use crate::buffer::BufferHandle;
@@ -155,23 +158,25 @@ impl VTable for SliceVTable {
155158
Ok(())
156159
}
157160

158-
fn execute(array: &Self::Array, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionStep> {
159-
// Execute the child to get canonical form, then slice it
160-
let Some(canonical) = array.child.as_opt::<AnyCanonical>() else {
161-
// If the child is not canonical, recurse.
162-
return array
163-
.child
164-
.clone()
165-
.execute::<ArrayRef>(ctx)?
166-
.slice(array.slice_range().clone())
161+
fn execute(array: &Self::Array, _ctx: &mut ExecutionCtx) -> VortexResult<ExecutionStep> {
162+
// If child is already canonical, slice it directly.
163+
if let Some(canonical) = array.child.as_opt::<AnyCanonical>() {
164+
// TODO(ngates): we should inline canonical slice logic here.
165+
return Canonical::from(canonical)
166+
.as_ref()
167+
.slice(array.range.clone())
167168
.map(ExecutionStep::Done);
168-
};
169+
}
170+
171+
// If child is a constant, slicing just changes the length.
172+
if let Some(constant) = array.child.as_opt::<ConstantVTable>() {
173+
return Ok(ExecutionStep::Done(
174+
ConstantArray::new(constant.scalar().clone(), array.range.len()).into_array(),
175+
));
176+
}
169177

170-
// TODO(ngates): we should inline canonical slice logic here.
171-
Canonical::from(canonical)
172-
.as_ref()
173-
.slice(array.range.clone())
174-
.map(ExecutionStep::Done)
178+
// Otherwise, ask the scheduler to execute the child first.
179+
Ok(ExecutionStep::ExecuteChild(0))
175180
}
176181

177182
fn reduce_parent(

0 commit comments

Comments
 (0)