Skip to content

Commit c4dd28a

Browse files
committed
wip: chunked iter exe
Signed-off-by: Joe Isaacs <joe.isaacs@live.co.uk>
1 parent 25a2f01 commit c4dd28a

8 files changed

Lines changed: 12144 additions & 11683 deletions

File tree

vortex-array/public-api.lock

Lines changed: 12086 additions & 11652 deletions
Large diffs are not rendered by default.

vortex-array/src/arrays/chunked/vtable/mod.rs

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -251,9 +251,7 @@ impl VTable for Chunked {
251251
}
252252

253253
// All chunks are now canonical — combine them.
254-
Ok(ExecutionResult::done(
255-
_canonicalize(array.as_view(), ctx)?,
256-
))
254+
Ok(ExecutionResult::done(_canonicalize(array.as_view(), ctx)?))
257255
}
258256

259257
fn reduce(array: ArrayView<'_, Self>) -> VortexResult<Option<ArrayRef>> {

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

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ use vortex_error::vortex_panic;
1111
use vortex_mask::Mask;
1212
use vortex_session::VortexSession;
1313

14+
use crate::AnyCanonical;
1415
use crate::ArrayEq;
1516
use crate::ArrayHash;
1617
use crate::ArrayRef;
@@ -34,6 +35,7 @@ use crate::buffer::BufferHandle;
3435
use crate::dtype::DType;
3536
use crate::executor::ExecutionCtx;
3637
use crate::executor::ExecutionResult;
38+
use crate::require_child;
3739
use crate::scalar::Scalar;
3840
use crate::serde::ArrayChildren;
3941
use crate::validity::Validity;
@@ -142,14 +144,18 @@ impl VTable for Filter {
142144
if let Some(canonical) = execute_filter_fast_paths(array.as_view(), ctx)? {
143145
return Ok(ExecutionResult::done(canonical));
144146
}
147+
148+
let array = require_child!(array, array.child(), CHILD_SLOT => AnyCanonical);
149+
145150
let Mask::Values(mask_values) = &array.mask else {
146151
unreachable!("`execute_filter_fast_paths` handles AllTrue and AllFalse")
147152
};
148153

149-
// We rely on the optimization pass that runs prior to this execution for filter pushdown,
150-
// so now we can just execute the filter without worrying.
154+
// Child is pre-canonicalized — apply the filter directly.
155+
debug_assert!(array.child().is_canonical());
156+
let child = array.child().to_canonical()?;
151157
Ok(ExecutionResult::done(
152-
execute_filter(array.child().clone().execute(ctx)?, mask_values).into_array(),
158+
execute_filter(child, mask_values).into_array(),
153159
))
154160
}
155161

vortex-array/src/arrays/list/vtable/mod.rs

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,9 +23,11 @@ use crate::array::Array;
2323
use crate::array::ArrayId;
2424
use crate::array::ArrayView;
2525
use crate::array::VTable;
26+
use crate::arrays::Primitive;
2627
use crate::arrays::list::ListArrayExt;
2728
use crate::arrays::list::ListData;
2829
use crate::arrays::list::array::NUM_SLOTS;
30+
use crate::arrays::list::array::OFFSETS_SLOT;
2931
use crate::arrays::list::array::SLOT_NAMES;
3032
use crate::arrays::list::compute::PARENT_KERNELS;
3133
use crate::arrays::list::compute::rules::PARENT_RULES;
@@ -34,6 +36,7 @@ use crate::buffer::BufferHandle;
3436
use crate::dtype::DType;
3537
use crate::dtype::Nullability;
3638
use crate::dtype::PType;
39+
use crate::require_child;
3740
use crate::serde::ArrayChildren;
3841
use crate::validity::Validity;
3942
use crate::vtable;
@@ -114,7 +117,7 @@ impl VTable for List {
114117
let elements = slots[crate::arrays::list::array::ELEMENTS_SLOT]
115118
.as_ref()
116119
.vortex_expect("ListArray elements slot");
117-
let offsets = slots[crate::arrays::list::array::OFFSETS_SLOT]
120+
let offsets = slots[OFFSETS_SLOT]
118121
.as_ref()
119122
.vortex_expect("ListArray offsets slot");
120123
vortex_ensure!(
@@ -180,6 +183,7 @@ impl VTable for List {
180183
}
181184

182185
fn execute(array: Array<Self>, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
186+
let array = require_child!(array, array.offsets(), OFFSETS_SLOT => Primitive);
183187
Ok(ExecutionResult::done(
184188
list_view_from_list(array, ctx)?.into_array(),
185189
))

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

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,10 +13,10 @@ use vortex_error::vortex_ensure;
1313
use vortex_error::vortex_panic;
1414
use vortex_session::VortexSession;
1515

16+
use crate::AnyCanonical;
1617
use crate::ArrayEq;
1718
use crate::ArrayHash;
1819
use crate::ArrayRef;
19-
use crate::Canonical;
2020
use crate::IntoArray;
2121
use crate::Precision;
2222
use crate::array::Array;
@@ -35,6 +35,7 @@ use crate::buffer::BufferHandle;
3535
use crate::dtype::DType;
3636
use crate::executor::ExecutionCtx;
3737
use crate::executor::ExecutionResult;
38+
use crate::require_child;
3839
use crate::scalar::Scalar;
3940
use crate::serde::ArrayChildren;
4041
use crate::validity::Validity;
@@ -166,7 +167,10 @@ impl VTable for Masked {
166167
// While we could manually convert the dtype, `mask_validity_canonical` is already O(1) for
167168
// `AllTrue` masks (no data copying), so there's no benefit.
168169

169-
let child = array.child().clone().execute::<Canonical>(ctx)?;
170+
let array = require_child!(array, array.child(), CHILD_SLOT => AnyCanonical);
171+
172+
debug_assert!(array.child().is_canonical());
173+
let child = array.child().to_canonical()?;
170174
Ok(ExecutionResult::done(
171175
mask_validity_canonical(child, &validity_mask, ctx)?.into_array(),
172176
))

vortex-array/src/arrays/patched/vtable/mod.rs

Lines changed: 24 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -31,10 +31,15 @@ use crate::array::ArrayView;
3131
use crate::array::VTable;
3232
use crate::array::ValidityChild;
3333
use crate::array::ValidityVTableFromChild;
34+
use crate::arrays::Primitive;
3435
use crate::arrays::PrimitiveArray;
3536
use crate::arrays::patched::PatchedArrayExt;
3637
use crate::arrays::patched::PatchedData;
38+
use crate::arrays::patched::array::INDICES_SLOT;
39+
use crate::arrays::patched::array::INNER_SLOT;
40+
use crate::arrays::patched::array::LANE_OFFSETS_SLOT;
3741
use crate::arrays::patched::array::SLOT_NAMES;
42+
use crate::arrays::patched::array::VALUES_SLOT;
3843
use crate::arrays::patched::compute::rules::PARENT_RULES;
3944
use crate::arrays::patched::vtable::kernels::PARENT_KERNELS;
4045
use crate::arrays::primitive::PrimitiveDataParts;
@@ -45,6 +50,7 @@ use crate::dtype::DType;
4550
use crate::dtype::NativePType;
4651
use crate::dtype::PType;
4752
use crate::match_each_native_ptype;
53+
use crate::require_child;
4854
use crate::serde::ArrayChildren;
4955
use crate::vtable;
5056

@@ -242,12 +248,18 @@ impl VTable for Patched {
242248
SLOT_NAMES[idx].to_string()
243249
}
244250

245-
fn execute(array: Array<Self>, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
251+
fn execute(array: Array<Self>, _ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
252+
let array = require_child!(array, array.base_array(), INNER_SLOT => Primitive);
253+
let array = require_child!(array, array.lane_offsets(), LANE_OFFSETS_SLOT => Primitive);
254+
let array = require_child!(array, array.patch_indices(), INDICES_SLOT => Primitive);
255+
let array = require_child!(array, array.patch_values(), VALUES_SLOT => Primitive);
256+
246257
let inner = array
247258
.base_array()
248259
.clone()
249-
.execute::<Canonical>(ctx)?
250-
.into_primitive();
260+
.try_downcast::<Primitive>()
261+
.ok()
262+
.vortex_expect("base_array pre-canonicalized to Primitive");
251263

252264
let PrimitiveDataParts {
253265
buffer,
@@ -258,17 +270,21 @@ impl VTable for Patched {
258270
let lane_offsets = array
259271
.lane_offsets()
260272
.clone()
261-
.execute::<PrimitiveArray>(ctx)?;
273+
.try_downcast::<Primitive>()
274+
.ok()
275+
.vortex_expect("lane_offsets pre-canonicalized to Primitive");
262276
let indices = array
263277
.patch_indices()
264278
.clone()
265-
.execute::<PrimitiveArray>(ctx)?;
266-
267-
// TODO(aduffy): add support for non-primitive PatchedArray patches application (?)
279+
.try_downcast::<Primitive>()
280+
.ok()
281+
.vortex_expect("patch_indices pre-canonicalized to Primitive");
268282
let values = array
269283
.patch_values()
270284
.clone()
271-
.execute::<PrimitiveArray>(ctx)?;
285+
.try_downcast::<Primitive>()
286+
.ok()
287+
.vortex_expect("patch_values pre-canonicalized to Primitive");
272288

273289
let patched_values = match_each_native_ptype!(values.ptype(), |V| {
274290
let offset = array.offset();

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

Lines changed: 8 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,6 @@ use crate::AnyCanonical;
1818
use crate::ArrayEq;
1919
use crate::ArrayHash;
2020
use crate::ArrayRef;
21-
use crate::Canonical;
2221
use crate::IntoArray;
2322
use crate::Precision;
2423
use crate::array::Array;
@@ -36,6 +35,7 @@ use crate::buffer::BufferHandle;
3635
use crate::dtype::DType;
3736
use crate::executor::ExecutionCtx;
3837
use crate::executor::ExecutionResult;
38+
use crate::require_child;
3939
use crate::scalar::Scalar;
4040
use crate::serde::ArrayChildren;
4141
use crate::validity::Validity;
@@ -141,20 +141,15 @@ impl VTable for Slice {
141141
vortex_bail!("Slice array is not serializable")
142142
}
143143

144-
fn execute(array: Array<Self>, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
145-
// Execute the child to get canonical form, then slice it
146-
let Some(canonical) = array.child().as_opt::<AnyCanonical>() else {
147-
// If the child is not canonical, recurse.
148-
return array
149-
.child()
150-
.clone()
151-
.execute::<ArrayRef>(ctx)?
152-
.slice(array.slice_range().clone())
153-
.map(ExecutionResult::done);
154-
};
144+
fn execute(array: Array<Self>, _ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
145+
let array = require_child!(array, array.child(), CHILD_SLOT => AnyCanonical);
155146

147+
// Child is now canonical — slice it.
156148
// TODO(ngates): we should inline canonical slice logic here.
157-
Canonical::from(canonical)
149+
debug_assert!(array.child().is_canonical());
150+
array
151+
.child()
152+
.to_canonical()?
158153
.into_array()
159154
.slice(array.range.clone())
160155
.map(ExecutionResult::done)

vortex-array/src/arrays/varbin/vtable/mod.rs

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,14 +18,17 @@ use crate::array::Array;
1818
use crate::array::ArrayId;
1919
use crate::array::ArrayView;
2020
use crate::array::VTable;
21+
use crate::arrays::Primitive;
2122
use crate::arrays::varbin::VarBinArrayExt;
2223
use crate::arrays::varbin::VarBinData;
2324
use crate::arrays::varbin::array::NUM_SLOTS;
25+
use crate::arrays::varbin::array::OFFSETS_SLOT;
2426
use crate::arrays::varbin::array::SLOT_NAMES;
2527
use crate::buffer::BufferHandle;
2628
use crate::dtype::DType;
2729
use crate::dtype::Nullability;
2830
use crate::dtype::PType;
31+
use crate::require_child;
2932
use crate::serde::ArrayChildren;
3033
use crate::validity::Validity;
3134
use crate::vtable;
@@ -89,7 +92,7 @@ impl VTable for VarBin {
8992
"VarBinArray expected {NUM_SLOTS} slots, found {}",
9093
slots.len()
9194
);
92-
let offsets = slots[crate::arrays::varbin::array::OFFSETS_SLOT]
95+
let offsets = slots[OFFSETS_SLOT]
9396
.as_ref()
9497
.vortex_expect("VarBinArray offsets slot");
9598
vortex_ensure!(
@@ -187,6 +190,7 @@ impl VTable for VarBin {
187190
}
188191

189192
fn execute(array: Array<Self>, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
193+
let array = require_child!(array, array.offsets(), OFFSETS_SLOT => Primitive);
190194
Ok(ExecutionResult::done(
191195
varbin_to_canonical(array.as_view(), ctx)?.into_array(),
192196
))

0 commit comments

Comments
 (0)