Skip to content

Commit 6ee8276

Browse files
joseph-isaacsjoseph-isaacs
committed
fix
Signed-off-by: Joe Isaacs <joe.isaacs@live.co.uk>
1 parent 8557f46 commit 6ee8276

2 files changed

Lines changed: 103 additions & 30 deletions

File tree

‎encodings/sparse/public-api.lock‎

Lines changed: 71 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,14 +32,16 @@ pub fn vortex_sparse::Sparse::buffer_name(_array: vortex_array::array::view::Arr
3232

3333
pub fn vortex_sparse::Sparse::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::array::typed::ArrayParts<Self>>
3434

35-
pub fn vortex_sparse::Sparse::execute(array: vortex_array::array::typed::Array<Self>, ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<vortex_array::executor::ExecutionResult>
35+
pub fn vortex_sparse::Sparse::execute(array: vortex_array::array::typed::Array<Self>, _ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<vortex_array::executor::ExecutionResult>
3636

3737
pub fn vortex_sparse::Sparse::execute_parent(array: vortex_array::array::view::ArrayView<'_, Self>, parent: &vortex_array::array::erased::ArrayRef, child_idx: usize, ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<core::option::Option<vortex_array::array::erased::ArrayRef>>
3838

3939
pub fn vortex_sparse::Sparse::id(&self) -> vortex_array::array::ArrayId
4040

4141
pub fn vortex_sparse::Sparse::nbuffers(_array: vortex_array::array::view::ArrayView<'_, Self>) -> usize
4242

43+
pub fn vortex_sparse::Sparse::reduce(array: vortex_array::array::view::ArrayView<'_, Self>) -> vortex_error::VortexResult<core::option::Option<vortex_array::array::erased::ArrayRef>>
44+
4345
pub fn vortex_sparse::Sparse::reduce_parent(array: vortex_array::array::view::ArrayView<'_, Self>, parent: &vortex_array::array::erased::ArrayRef, child_idx: usize) -> vortex_error::VortexResult<core::option::Option<vortex_array::array::erased::ArrayRef>>
4446

4547
pub fn vortex_sparse::Sparse::serialize(array: vortex_array::array::view::ArrayView<'_, Self>, _session: &vortex_session::VortexSession) -> vortex_error::VortexResult<core::option::Option<alloc::vec::Vec<u8>>>
@@ -138,4 +140,72 @@ pub fn vortex_sparse::SparseMetadata::clear(&mut self)
138140

139141
pub fn vortex_sparse::SparseMetadata::encoded_len(&self) -> usize
140142

143+
pub struct vortex_sparse::SparseSlots
144+
145+
pub vortex_sparse::SparseSlots::patch_chunk_offsets: core::option::Option<vortex_array::array::erased::ArrayRef>
146+
147+
pub vortex_sparse::SparseSlots::patch_indices: vortex_array::array::erased::ArrayRef
148+
149+
pub vortex_sparse::SparseSlots::patch_values: vortex_array::array::erased::ArrayRef
150+
151+
impl vortex_sparse::SparseSlots
152+
153+
pub const vortex_sparse::SparseSlots::COUNT: usize
154+
155+
pub const vortex_sparse::SparseSlots::NAMES: [&'static str; 3]
156+
157+
pub const vortex_sparse::SparseSlots::PATCH_CHUNK_OFFSETS: usize
158+
159+
pub const vortex_sparse::SparseSlots::PATCH_INDICES: usize
160+
161+
pub const vortex_sparse::SparseSlots::PATCH_VALUES: usize
162+
163+
pub fn vortex_sparse::SparseSlots::from_slots(slots: alloc::vec::Vec<core::option::Option<vortex_array::array::erased::ArrayRef>>) -> Self
164+
165+
pub fn vortex_sparse::SparseSlots::into_slots(self) -> alloc::vec::Vec<core::option::Option<vortex_array::array::erased::ArrayRef>>
166+
167+
pub struct vortex_sparse::SparseSlotsView<'a>
168+
169+
pub vortex_sparse::SparseSlotsView::patch_chunk_offsets: core::option::Option<&'a vortex_array::array::erased::ArrayRef>
170+
171+
pub vortex_sparse::SparseSlotsView::patch_indices: &'a vortex_array::array::erased::ArrayRef
172+
173+
pub vortex_sparse::SparseSlotsView::patch_values: &'a vortex_array::array::erased::ArrayRef
174+
175+
impl<'a> vortex_sparse::SparseSlotsView<'a>
176+
177+
pub fn vortex_sparse::SparseSlotsView<'a>::from_slots(slots: &'a [core::option::Option<vortex_array::array::erased::ArrayRef>]) -> Self
178+
179+
pub fn vortex_sparse::SparseSlotsView<'a>::to_owned(&self) -> vortex_sparse::SparseSlots
180+
181+
impl<'a> core::clone::Clone for vortex_sparse::SparseSlotsView<'a>
182+
183+
pub fn vortex_sparse::SparseSlotsView<'a>::clone(&self) -> vortex_sparse::SparseSlotsView<'a>
184+
185+
impl<'a> core::fmt::Debug for vortex_sparse::SparseSlotsView<'a>
186+
187+
pub fn vortex_sparse::SparseSlotsView<'a>::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result
188+
189+
impl<'a> core::marker::Copy for vortex_sparse::SparseSlotsView<'a>
190+
191+
pub trait vortex_sparse::SparseArraySlotsExt: vortex_array::array::typed::TypedArrayRef<vortex_sparse::Sparse>
192+
193+
pub fn vortex_sparse::SparseArraySlotsExt::patch_chunk_offsets(&self) -> core::option::Option<&vortex_array::array::erased::ArrayRef>
194+
195+
pub fn vortex_sparse::SparseArraySlotsExt::patch_indices(&self) -> &vortex_array::array::erased::ArrayRef
196+
197+
pub fn vortex_sparse::SparseArraySlotsExt::patch_values(&self) -> &vortex_array::array::erased::ArrayRef
198+
199+
pub fn vortex_sparse::SparseArraySlotsExt::slots_view(&self) -> vortex_sparse::SparseSlotsView<'_>
200+
201+
impl<T: vortex_array::array::typed::TypedArrayRef<vortex_sparse::Sparse>> vortex_sparse::SparseArraySlotsExt for T
202+
203+
pub fn T::patch_chunk_offsets(&self) -> core::option::Option<&vortex_array::array::erased::ArrayRef>
204+
205+
pub fn T::patch_indices(&self) -> &vortex_array::array::erased::ArrayRef
206+
207+
pub fn T::patch_values(&self) -> &vortex_array::array::erased::ArrayRef
208+
209+
pub fn T::slots_view(&self) -> vortex_sparse::SparseSlotsView<'_>
210+
141211
pub type vortex_sparse::SparseArray = vortex_array::array::typed::Array<vortex_sparse::Sparse>

‎encodings/sparse/src/lib.rs‎

Lines changed: 32 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -9,22 +9,22 @@ use std::hash::Hasher;
99

1010
use kernel::PARENT_KERNELS;
1111
use prost::Message as _;
12+
use vortex_array::AnyCanonical;
1213
use vortex_array::Array;
1314
use vortex_array::ArrayEq;
1415
use vortex_array::ArrayHash;
1516
use vortex_array::ArrayId;
1617
use vortex_array::ArrayParts;
17-
use vortex_array::AnyCanonical;
1818
use vortex_array::ArrayRef;
1919
use vortex_array::ArrayView;
2020
use vortex_array::Canonical;
2121
use vortex_array::ExecutionCtx;
2222
use vortex_array::ExecutionResult;
2323
use vortex_array::IntoArray;
2424
use vortex_array::Precision;
25-
use vortex_array::arrays::Primitive;
2625
use vortex_array::arrays::BoolArray;
2726
use vortex_array::arrays::ConstantArray;
27+
use vortex_array::arrays::Primitive;
2828
use vortex_array::arrays::PrimitiveArray;
2929
use vortex_array::arrays::bool::BoolArrayExt;
3030
use vortex_array::buffer::BufferHandle;
@@ -33,6 +33,8 @@ use vortex_array::dtype::DType;
3333
use vortex_array::dtype::Nullability;
3434
use vortex_array::patches::Patches;
3535
use vortex_array::patches::PatchesMetadata;
36+
use vortex_array::require_child;
37+
use vortex_array::require_opt_child;
3638
use vortex_array::scalar::Scalar;
3739
use vortex_array::scalar::ScalarValue;
3840
use vortex_array::scalar_fn::fns::operators::Operator;
@@ -81,28 +83,6 @@ pub(crate) struct SparseParts {
8183
pub len: usize,
8284
}
8385

84-
impl SparseParts {
85-
/// Resolve patches by subtracting the offset from indices.
86-
pub fn resolve_patches(mut self) -> VortexResult<Self> {
87-
if self.patches.offset() != 0 {
88-
let offset_scalar =
89-
Scalar::from(self.patches.offset()).cast(self.patches.indices().dtype())?;
90-
let indices = self.patches.indices().binary(
91-
ConstantArray::new(offset_scalar, self.patches.indices().len()).into_array(),
92-
Operator::Sub,
93-
)?;
94-
self.patches = Patches::new(
95-
self.patches.array_len(),
96-
0,
97-
indices,
98-
self.patches.values().clone(),
99-
None,
100-
)?;
101-
}
102-
Ok(self)
103-
}
104-
}
105-
10686
pub(crate) trait SparseOwnedExt {
10787
fn into_parts(self) -> VortexResult<SparseParts>;
10888
}
@@ -272,26 +252,49 @@ impl VTable for Sparse {
272252
}
273253

274254
fn execute(array: Array<Self>, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
255+
// Resolve offset first: wrap indices in Binary(indices, offset, Sub) and rebuild
256+
// with offset=0. Uses slot children (not data) since the executor may have
257+
// updated slots via reduce_parent/execute_parent.
258+
let array = if array.patches().offset() != 0 {
259+
let offset = array.patches().offset();
260+
let indices = array.patch_indices();
261+
let offset_scalar = Scalar::from(offset).cast(indices.dtype())?;
262+
let resolved_indices = indices.binary(
263+
ConstantArray::new(offset_scalar, indices.len()).into_array(),
264+
Operator::Sub,
265+
)?;
266+
let patches = Patches::new(
267+
array.len(),
268+
0,
269+
resolved_indices,
270+
array.patch_values().clone(),
271+
None,
272+
)?;
273+
Sparse::try_new_from_patches(patches, array.fill_scalar().clone())?
274+
} else {
275+
array
276+
};
277+
275278
// Require children to be executed through the scheduler,
276279
// enabling cross-step optimization via reduce_parent rules.
277-
let array = vortex_array::require_child!(
280+
let array = require_child!(
278281
array, array.patch_indices(), SparseSlots::PATCH_INDICES => Primitive
279282
);
280-
let array = vortex_array::require_child!(
283+
let array = require_child!(
281284
array, array.patch_values(), SparseSlots::PATCH_VALUES => AnyCanonical
282285
);
283-
vortex_array::require_opt_child!(
286+
require_opt_child!(
284287
array,
285288
array.patch_chunk_offsets(),
286289
SparseSlots::PATCH_CHUNK_OFFSETS => Primitive
287290
);
288291

289-
let parts = array.into_parts()?.resolve_patches()?;
292+
let parts = array.into_parts()?;
293+
// TODO(joe): remove ctx from execute_sparse since all slots should be canonical.
290294
execute_sparse(parts, ctx).map(ExecutionResult::done)
291295
}
292296
}
293297

294-
295298
#[derive(Clone, Debug)]
296299
pub struct SparseData {
297300
patches: Patches,

0 commit comments

Comments
 (0)