Skip to content

Commit 77c8e95

Browse files
joseph-isaacsjoseph-isaacs
authored
feat: iterative execution patched array (#7345)
Use iterative execute require_child to execute its children. We should look into a better API to access slots. I am working on this --------- Signed-off-by: Joe Isaacs <joe.isaacs@live.co.uk>
1 parent 72d39fb commit 77c8e95

2 files changed

Lines changed: 49 additions & 31 deletions

File tree

‎vortex-array/public-api.lock‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -3350,7 +3350,7 @@ pub fn vortex_array::arrays::patched::Patched::child_name(array: vortex_array::A
33503350

33513351
pub fn vortex_array::arrays::patched::Patched::deserialize(&self, dtype: &vortex_array::dtype::DType, len: usize, metadata: &[u8], _buffers: &[vortex_array::buffer::BufferHandle], children: &dyn vortex_array::serde::ArrayChildren, _session: &vortex_session::VortexSession) -> vortex_error::VortexResult<vortex_array::ArrayParts<Self>>
33523352

3353-
pub fn vortex_array::arrays::patched::Patched::execute(array: vortex_array::Array<Self>, ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult<vortex_array::ExecutionResult>
3353+
pub fn vortex_array::arrays::patched::Patched::execute(array: vortex_array::Array<Self>, _ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult<vortex_array::ExecutionResult>
33543354

33553355
pub fn vortex_array::arrays::patched::Patched::execute_parent(array: vortex_array::ArrayView<'_, Self>, parent: &vortex_array::ArrayRef, child_idx: usize, ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult<core::option::Option<vortex_array::ArrayRef>>
33563356

@@ -6136,7 +6136,7 @@ pub fn vortex_array::arrays::patched::Patched::child_name(array: vortex_array::A
61366136

61376137
pub fn vortex_array::arrays::patched::Patched::deserialize(&self, dtype: &vortex_array::dtype::DType, len: usize, metadata: &[u8], _buffers: &[vortex_array::buffer::BufferHandle], children: &dyn vortex_array::serde::ArrayChildren, _session: &vortex_session::VortexSession) -> vortex_error::VortexResult<vortex_array::ArrayParts<Self>>
61386138

6139-
pub fn vortex_array::arrays::patched::Patched::execute(array: vortex_array::Array<Self>, ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult<vortex_array::ExecutionResult>
6139+
pub fn vortex_array::arrays::patched::Patched::execute(array: vortex_array::Array<Self>, _ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult<vortex_array::ExecutionResult>
61406140

61416141
pub fn vortex_array::arrays::patched::Patched::execute_parent(array: vortex_array::ArrayView<'_, Self>, parent: &vortex_array::ArrayRef, child_idx: usize, ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult<core::option::Option<vortex_array::ArrayRef>>
61426142

@@ -20000,7 +20000,7 @@ pub fn vortex_array::arrays::patched::Patched::child_name(array: vortex_array::A
2000020000

2000120001
pub fn vortex_array::arrays::patched::Patched::deserialize(&self, dtype: &vortex_array::dtype::DType, len: usize, metadata: &[u8], _buffers: &[vortex_array::buffer::BufferHandle], children: &dyn vortex_array::serde::ArrayChildren, _session: &vortex_session::VortexSession) -> vortex_error::VortexResult<vortex_array::ArrayParts<Self>>
2000220002

20003-
pub fn vortex_array::arrays::patched::Patched::execute(array: vortex_array::Array<Self>, ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult<vortex_array::ExecutionResult>
20003+
pub fn vortex_array::arrays::patched::Patched::execute(array: vortex_array::Array<Self>, _ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult<vortex_array::ExecutionResult>
2000420004

2000520005
pub fn vortex_array::arrays::patched::Patched::execute_parent(array: vortex_array::ArrayView<'_, Self>, parent: &vortex_array::ArrayRef, child_idx: usize, ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult<core::option::Option<vortex_array::ArrayRef>>
2000620006

@@ -20972,7 +20972,7 @@ pub fn vortex_array::arrays::patched::Patched::child_name(array: vortex_array::A
2097220972

2097320973
pub fn vortex_array::arrays::patched::Patched::deserialize(&self, dtype: &vortex_array::dtype::DType, len: usize, metadata: &[u8], _buffers: &[vortex_array::buffer::BufferHandle], children: &dyn vortex_array::serde::ArrayChildren, _session: &vortex_session::VortexSession) -> vortex_error::VortexResult<vortex_array::ArrayParts<Self>>
2097420974

20975-
pub fn vortex_array::arrays::patched::Patched::execute(array: vortex_array::Array<Self>, ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult<vortex_array::ExecutionResult>
20975+
pub fn vortex_array::arrays::patched::Patched::execute(array: vortex_array::Array<Self>, _ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult<vortex_array::ExecutionResult>
2097620976

2097720977
pub fn vortex_array::arrays::patched::Patched::execute_parent(array: vortex_array::ArrayView<'_, Self>, parent: &vortex_array::ArrayRef, child_idx: usize, ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult<core::option::Option<vortex_array::ArrayRef>>
2097820978

@@ -23652,7 +23652,7 @@ pub fn vortex_array::arrays::patched::Patched::child_name(array: vortex_array::A
2365223652

2365323653
pub fn vortex_array::arrays::patched::Patched::deserialize(&self, dtype: &vortex_array::dtype::DType, len: usize, metadata: &[u8], _buffers: &[vortex_array::buffer::BufferHandle], children: &dyn vortex_array::serde::ArrayChildren, _session: &vortex_session::VortexSession) -> vortex_error::VortexResult<vortex_array::ArrayParts<Self>>
2365423654

23655-
pub fn vortex_array::arrays::patched::Patched::execute(array: vortex_array::Array<Self>, ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult<vortex_array::ExecutionResult>
23655+
pub fn vortex_array::arrays::patched::Patched::execute(array: vortex_array::Array<Self>, _ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult<vortex_array::ExecutionResult>
2365623656

2365723657
pub fn vortex_array::arrays::patched::Patched::execute_parent(array: vortex_array::ArrayView<'_, Self>, parent: &vortex_array::ArrayRef, child_idx: usize, ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult<core::option::Option<vortex_array::ArrayRef>>
2365823658

@@ -24872,7 +24872,7 @@ pub fn vortex_array::arrays::patched::Patched::child_name(array: vortex_array::A
2487224872

2487324873
pub fn vortex_array::arrays::patched::Patched::deserialize(&self, dtype: &vortex_array::dtype::DType, len: usize, metadata: &[u8], _buffers: &[vortex_array::buffer::BufferHandle], children: &dyn vortex_array::serde::ArrayChildren, _session: &vortex_session::VortexSession) -> vortex_error::VortexResult<vortex_array::ArrayParts<Self>>
2487424874

24875-
pub fn vortex_array::arrays::patched::Patched::execute(array: vortex_array::Array<Self>, ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult<vortex_array::ExecutionResult>
24875+
pub fn vortex_array::arrays::patched::Patched::execute(array: vortex_array::Array<Self>, _ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult<vortex_array::ExecutionResult>
2487624876

2487724877
pub fn vortex_array::arrays::patched::Patched::execute_parent(array: vortex_array::ArrayView<'_, Self>, parent: &vortex_array::ArrayRef, child_idx: usize, ctx: &mut vortex_array::ExecutionCtx) -> vortex_error::VortexResult<core::option::Option<vortex_array::ArrayRef>>
2487824878

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

Lines changed: 43 additions & 25 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

5056
/// A [`Patched`]-encoded Vortex array.
@@ -245,12 +251,42 @@ impl VTable for Patched {
245251
SLOT_NAMES[idx].to_string()
246252
}
247253

248-
fn execute(array: Array<Self>, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
249-
let inner = array
250-
.base_array()
251-
.clone()
252-
.execute::<Canonical>(ctx)?
253-
.into_primitive();
254+
fn execute(array: Array<Self>, _ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
255+
let array = require_child!(array, array.base_array(), INNER_SLOT => Primitive);
256+
let array = require_child!(array, array.lane_offsets(), LANE_OFFSETS_SLOT => Primitive);
257+
let array = require_child!(array, array.patch_indices(), INDICES_SLOT => Primitive);
258+
let array = require_child!(array, array.patch_values(), VALUES_SLOT => Primitive);
259+
260+
let len = array.len();
261+
262+
fn take_slot(slots: &mut [Option<ArrayRef>], idx: usize) -> ArrayRef {
263+
slots[idx].take().vortex_expect("slot must be present")
264+
}
265+
266+
fn downcast_slot(slot: ArrayRef) -> PrimitiveArray {
267+
slot.try_downcast::<Primitive>()
268+
.ok()
269+
.vortex_expect("slot must be primitive")
270+
}
271+
272+
let (n_lanes, offset, inner, lane_offsets, indices, values) = match array.try_into_parts() {
273+
Ok(mut parts) => {
274+
let PatchedData { n_lanes, offset } = parts.data;
275+
let inner = downcast_slot(take_slot(&mut parts.slots, INNER_SLOT));
276+
let lane_offsets = downcast_slot(take_slot(&mut parts.slots, LANE_OFFSETS_SLOT));
277+
let indices = downcast_slot(take_slot(&mut parts.slots, INDICES_SLOT));
278+
let values = downcast_slot(take_slot(&mut parts.slots, VALUES_SLOT));
279+
(n_lanes, offset, inner, lane_offsets, indices, values)
280+
}
281+
Err(array) => {
282+
let PatchedData { n_lanes, offset } = array.data().clone();
283+
let inner = downcast_slot(array.base_array().clone());
284+
let lane_offsets = downcast_slot(array.lane_offsets().clone());
285+
let indices = downcast_slot(array.patch_indices().clone());
286+
let values = downcast_slot(array.patch_values().clone());
287+
(n_lanes, offset, inner, lane_offsets, indices, values)
288+
}
289+
};
254290

255291
// TODO(joe): use iterative execution
256292
let PrimitiveDataParts {
@@ -259,32 +295,14 @@ impl VTable for Patched {
259295
validity,
260296
} = inner.into_data_parts();
261297

262-
let lane_offsets = array
263-
.lane_offsets()
264-
.clone()
265-
.execute::<PrimitiveArray>(ctx)?;
266-
let indices = array
267-
.patch_indices()
268-
.clone()
269-
.execute::<PrimitiveArray>(ctx)?;
270-
271-
// TODO(aduffy): add support for non-primitive PatchedArray patches application (?)
272-
let values = array
273-
.patch_values()
274-
.clone()
275-
.execute::<PrimitiveArray>(ctx)?;
276-
277298
let patched_values = match_each_native_ptype!(values.ptype(), |V| {
278-
let offset = array.offset();
279-
let len = array.len();
280-
281299
let mut output = Buffer::<V>::from_byte_buffer(buffer.unwrap_host()).into_mut();
282300

283301
apply_patches_primitive::<V>(
284302
&mut output,
285303
offset,
286304
len,
287-
array.n_lanes(),
305+
n_lanes,
288306
lane_offsets.as_slice::<u32>(),
289307
indices.as_slice::<u16>(),
290308
values.as_slice::<V>(),

0 commit comments

Comments
 (0)