Skip to content

Commit 3436405

Browse files
committed
implement append_to_builder for Patched
We call append_to_builder on the inner first, then just do a single pass and overwrite what it just wrote. Signed-off-by: Andrew Duffy <andrew@a10y.dev>
1 parent d40a607 commit 3436405

File tree

2 files changed

+58
-0
lines changed

2 files changed

+58
-0
lines changed

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

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ use std::hash::Hash;
99
use std::hash::Hasher;
1010

1111
use vortex_buffer::Buffer;
12+
use vortex_error::VortexExpect;
1213
use vortex_error::VortexResult;
1314
use vortex_error::vortex_bail;
1415
use vortex_error::vortex_ensure;
@@ -33,6 +34,8 @@ use crate::arrays::patched::patch_lanes;
3334
use crate::arrays::patched::vtable::kernels::PARENT_KERNELS;
3435
use crate::arrays::primitive::PrimitiveArrayParts;
3536
use crate::buffer::BufferHandle;
37+
use crate::builders::ArrayBuilder;
38+
use crate::builders::PrimitiveBuilder;
3639
use crate::dtype::DType;
3740
use crate::dtype::NativePType;
3841
use crate::match_each_native_ptype;
@@ -164,6 +167,56 @@ impl VTable for Patched {
164167
Ok(ProstMetadata(inner))
165168
}
166169

170+
fn append_to_builder(
171+
array: &Self::Array,
172+
builder: &mut dyn ArrayBuilder,
173+
ctx: &mut ExecutionCtx,
174+
) -> VortexResult<()> {
175+
let dtype = array.dtype();
176+
177+
if !dtype.is_primitive() {
178+
// Default pathway: canonicalize and propagate.
179+
let canonical = array.to_array().execute::<Canonical>(ctx)?.into_array();
180+
builder.extend_from_array(&canonical);
181+
return Ok(());
182+
}
183+
184+
let ptype = dtype.as_ptype();
185+
186+
let len = array.len();
187+
array.inner.append_to_builder(builder, ctx)?;
188+
189+
let offset = array.offset;
190+
let lane_offsets: Buffer<u32> =
191+
Buffer::from_byte_buffer(array.lane_offsets.clone().unwrap_host());
192+
let indices: Buffer<u16> = Buffer::from_byte_buffer(array.indices.clone().unwrap_host());
193+
let values = array.values.clone().execute::<PrimitiveArray>(ctx)?;
194+
195+
match_each_native_ptype!(ptype, |V| {
196+
let typed_builder = builder
197+
.as_any_mut()
198+
.downcast_mut::<PrimitiveBuilder<V>>()
199+
.vortex_expect("correctly typed builder");
200+
201+
// Overwrite the last `len` elements of the builder. These would have been
202+
// populated by the inner.append_to_builder() call above.
203+
let mut output = typed_builder.values_mut();
204+
205+
apply_patches_primitive::<V>(
206+
&mut output[(values.len() - len)..],
207+
offset,
208+
len,
209+
array.n_chunks,
210+
array.n_lanes,
211+
&lane_offsets,
212+
&indices,
213+
values.as_slice::<V>(),
214+
);
215+
});
216+
217+
Ok(())
218+
}
219+
167220
fn build(
168221
dtype: &DType,
169222
len: usize,

vortex-array/src/builders/primitive.rs

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,11 @@ impl<T: NativePType> PrimitiveBuilder<T> {
6262
self.values.as_ref()
6363
}
6464

65+
/// Returns the raw primitive values in this builder as a mutable slice.
66+
pub fn values_mut(&mut self) -> &[T] {
67+
self.values.as_mut()
68+
}
69+
6570
/// Create a new handle to the next `len` uninitialized values in the builder.
6671
///
6772
/// All reads/writes through the handle to the values buffer or the validity buffer will operate

0 commit comments

Comments
 (0)