Skip to content

Commit 5bccfa1

Browse files
joseph-isaacsjoseph-isaacs
committed
u
Signed-off-by: Joe Isaacs <joe.isaacs@live.co.uk>
1 parent 16bb338 commit 5bccfa1

5 files changed

Lines changed: 100 additions & 71 deletions

File tree

‎encodings/alp/src/alp/array.rs‎

Lines changed: 8 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ use vortex_array::dtype::PType;
2323
use vortex_array::patches::Patches;
2424
use vortex_array::patches::PatchesMetadata;
2525
use vortex_array::require_child;
26+
use vortex_array::require_patches;
2627
use vortex_array::serde::ArrayChildren;
2728
use vortex_array::stats::ArrayStats;
2829
use vortex_array::stats::StatsSetRef;
@@ -205,34 +206,13 @@ impl VTable for ALP {
205206

206207
fn execute(array: Arc<Array<Self>>, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
207208
let array = require_child!(array, array.encoded(), ENCODED_SLOT => Primitive);
208-
if array
209-
.patches()
210-
.is_some_and(|p| !p.indices().is::<Primitive>())
211-
{
212-
return Ok(ExecutionResult::execute_slot::<Primitive>(
213-
array,
214-
PATCH_INDICES_SLOT,
215-
));
216-
}
217-
if array
218-
.patches()
219-
.is_some_and(|p| !p.values().is::<Primitive>())
220-
{
221-
return Ok(ExecutionResult::execute_slot::<Primitive>(
222-
array,
223-
PATCH_VALUES_SLOT,
224-
));
225-
}
226-
if array.patches().is_some_and(|p| {
227-
p.chunk_offsets()
228-
.as_ref()
229-
.is_some_and(|co| !co.is::<Primitive>())
230-
}) {
231-
return Ok(ExecutionResult::execute_slot::<Primitive>(
232-
array,
233-
PATCH_CHUNK_OFFSETS_SLOT,
234-
));
235-
}
209+
require_patches!(
210+
array,
211+
array.patches(),
212+
PATCH_INDICES_SLOT,
213+
PATCH_VALUES_SLOT,
214+
PATCH_CHUNK_OFFSETS_SLOT
215+
);
236216

237217
Ok(ExecutionResult::done(
238218
execute_decompress(Arc::unwrap_or_clone(array).into_inner(), ctx)?.into_array(),

‎encodings/alp/src/alp_rd/array.rs‎

Lines changed: 8 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ use vortex_array::dtype::PType;
2626
use vortex_array::patches::Patches;
2727
use vortex_array::patches::PatchesMetadata;
2828
use vortex_array::require_child;
29+
use vortex_array::require_patches;
2930
use vortex_array::serde::ArrayChildren;
3031
use vortex_array::stats::ArrayStats;
3132
use vortex_array::stats::StatsSetRef;
@@ -271,24 +272,13 @@ impl VTable for ALPRD {
271272
fn execute(array: Arc<Array<Self>>, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
272273
let array = require_child!(array, array.left_parts(), 0 => Primitive);
273274
let array = require_child!(array, array.right_parts(), 1 => Primitive);
274-
if array
275-
.left_parts_patches()
276-
.is_some_and(|p| !p.indices().is::<Primitive>())
277-
{
278-
return Ok(ExecutionResult::execute_slot::<Primitive>(
279-
array,
280-
LP_PATCH_INDICES_SLOT,
281-
));
282-
}
283-
if array
284-
.left_parts_patches()
285-
.is_some_and(|p| !p.values().is::<Primitive>())
286-
{
287-
return Ok(ExecutionResult::execute_slot::<Primitive>(
288-
array,
289-
LP_PATCH_VALUES_SLOT,
290-
));
291-
}
275+
require_patches!(
276+
array,
277+
array.left_parts_patches(),
278+
LP_PATCH_INDICES_SLOT,
279+
LP_PATCH_VALUES_SLOT,
280+
LP_PATCH_CHUNK_OFFSETS_SLOT
281+
);
292282

293283
let right_bit_width = array.right_bit_width();
294284
let ALPRDArrayParts {

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

Lines changed: 10 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -15,14 +15,15 @@ use vortex_array::IntoArray;
1515
use vortex_array::Precision;
1616
use vortex_array::ProstMetadata;
1717
use vortex_array::SerializeMetadata;
18-
use vortex_array::arrays::Primitive;
1918
use vortex_array::buffer::BufferHandle;
2019
use vortex_array::builders::ArrayBuilder;
2120
use vortex_array::dtype::DType;
2221
use vortex_array::dtype::PType;
2322
use vortex_array::match_each_integer_ptype;
2423
use vortex_array::patches::Patches;
2524
use vortex_array::patches::PatchesMetadata;
25+
use vortex_array::require_patches;
26+
use vortex_array::require_validity;
2627
use vortex_array::serde::ArrayChildren;
2728
use vortex_array::stats::StatsSetRef;
2829
use vortex_array::validity::Validity;
@@ -312,30 +313,14 @@ impl VTable for BitPacked {
312313
}
313314

314315
fn execute(array: Arc<Array<Self>>, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
315-
if array
316-
.patches()
317-
.is_some_and(|p| !p.indices().is::<Primitive>())
318-
{
319-
return Ok(ExecutionResult::execute_slot::<Primitive>(
320-
array,
321-
PATCH_INDICES_SLOT,
322-
));
323-
}
324-
if array
325-
.patches()
326-
.is_some_and(|p| !p.values().is::<Primitive>())
327-
{
328-
return Ok(ExecutionResult::execute_slot::<Primitive>(
329-
array,
330-
PATCH_VALUES_SLOT,
331-
));
332-
}
333-
if matches!(&array.validity, Validity::Array(v) if !v.is::<AnyCanonical>()) {
334-
return Ok(ExecutionResult::execute_slot::<AnyCanonical>(
335-
array,
336-
VALIDITY_SLOT,
337-
));
338-
}
316+
require_patches!(
317+
array,
318+
array.patches(),
319+
PATCH_INDICES_SLOT,
320+
PATCH_VALUES_SLOT,
321+
PATCH_CHUNK_OFFSETS_SLOT
322+
);
323+
require_validity!(array, &array.validity, VALIDITY_SLOT => AnyCanonical);
339324

340325
Ok(ExecutionResult::done(
341326
unpack_array(&array, ctx)?.into_array(),

‎vortex-array/public-api.lock‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23886,6 +23886,12 @@ pub macro vortex_array::register_kernel!
2388623886

2388723887
pub macro vortex_array::require_child!
2388823888

23889+
pub macro vortex_array::require_opt_child!
23890+
23891+
pub macro vortex_array::require_patches!
23892+
23893+
pub macro vortex_array::require_validity!
23894+
2388923895
pub macro vortex_array::vtable!
2389023896

2389123897
pub enum vortex_array::Canonical

‎vortex-array/src/executor.rs‎

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -476,6 +476,74 @@ macro_rules! require_child {
476476
}};
477477
}
478478

479+
/// Like [`require_child!`], but for optional children. If the child is `None`, this is a no-op.
480+
/// If the child is `Some` but does not match `$M`, early-returns an [`ExecutionResult`] requesting
481+
/// execution of child `$idx`.
482+
///
483+
/// Unlike `require_child!`, this is a statement macro (no value produced) and does not clone
484+
/// `$parent` — it is moved into the early-return path.
485+
///
486+
/// ```ignore
487+
/// require_opt_child!(array, array.patches().map(|p| p.indices()), 1 => Primitive);
488+
/// ```
489+
#[macro_export]
490+
macro_rules! require_opt_child {
491+
($parent:expr, $child_opt:expr, $idx:expr => $M:ty) => {
492+
if $child_opt.is_some_and(|child| !child.is::<$M>()) {
493+
return Ok($crate::ExecutionResult::execute_slot::<$M>($parent, $idx));
494+
}
495+
};
496+
}
497+
498+
/// Require that all children of a [`Patches`](crate::patches::Patches) (indices, values, and
499+
/// optionally chunk_offsets) are `Primitive`. If no patches are present, this is a no-op.
500+
///
501+
/// Like [`require_opt_child!`], `$parent` is moved (not cloned) into the early-return path.
502+
///
503+
/// ```ignore
504+
/// require_patches!(array, array.patches(), PATCH_INDICES_SLOT, PATCH_VALUES_SLOT, PATCH_CHUNK_OFFSETS_SLOT);
505+
/// ```
506+
#[macro_export]
507+
macro_rules! require_patches {
508+
($parent:expr, $patches:expr, $indices_slot:expr, $values_slot:expr, $chunk_offsets_slot:expr) => {
509+
$crate::require_opt_child!(
510+
$parent,
511+
$patches.map(|p| p.indices()),
512+
$indices_slot => $crate::arrays::Primitive
513+
);
514+
$crate::require_opt_child!(
515+
$parent,
516+
$patches.map(|p| p.values()),
517+
$values_slot => $crate::arrays::Primitive
518+
);
519+
$crate::require_opt_child!(
520+
$parent,
521+
$patches.and_then(|p| p.chunk_offsets().as_ref()),
522+
$chunk_offsets_slot => $crate::arrays::Primitive
523+
);
524+
};
525+
}
526+
527+
/// Require that a [`Validity::Array`] child matches `$M`. If validity is not array-backed
528+
/// (e.g. `NonNullable` or `AllValid`), this is a no-op. If it is array-backed but does not
529+
/// match `$M`, early-returns an [`ExecutionResult`] requesting execution of the validity slot.
530+
///
531+
/// Like [`require_opt_child!`], `$parent` is moved (not cloned) into the early-return path.
532+
///
533+
/// ```ignore
534+
/// require_validity!(array, &array.validity, VALIDITY_SLOT => AnyCanonical);
535+
/// ```
536+
#[macro_export]
537+
macro_rules! require_validity {
538+
($parent:expr, $validity:expr, $idx:expr => $M:ty) => {
539+
if let $crate::validity::Validity::Array(v) = $validity {
540+
if !v.is::<$M>() {
541+
return Ok($crate::ExecutionResult::execute_slot::<$M>($parent, $idx));
542+
}
543+
}
544+
};
545+
}
546+
479547
/// Extension trait for creating an execution context from a session.
480548
pub trait VortexSessionExecute {
481549
/// Create a new execution context from this session.

0 commit comments

Comments
 (0)