1mod execute;
43
44use std::fmt::Display;
45use std::fmt::Formatter;
46use std::hash::Hasher;
47
48use vortex_error::VortexExpect;
49use vortex_error::VortexResult;
50use vortex_error::vortex_bail;
51use vortex_error::vortex_ensure;
52use vortex_error::vortex_panic;
53use vortex_session::VortexSession;
54use vortex_session::registry::CachedId;
55
56use crate::ArrayEq;
57use crate::ArrayHash;
58use crate::ArrayRef;
59use crate::EqMode;
60use crate::ExecutionCtx;
61use crate::IntoArray;
62use crate::array::Array;
63use crate::array::ArrayId;
64use crate::array::ArrayParts;
65use crate::array::ArraySlots;
66use crate::array::ArrayView;
67use crate::array::OperationsVTable;
68use crate::array::TypedArrayRef;
69use crate::array::VTable;
70use crate::array::ValidityVTable;
71use crate::array::with_empty_buffers;
72use crate::arrays::ConstantArray;
73use crate::buffer::BufferHandle;
74use crate::dtype::DType;
75use crate::dtype::Nullability;
76use crate::executor::ExecutionResult;
77use crate::scalar::Scalar;
78use crate::serde::ArrayChildren;
79use crate::validity::Validity;
80
81pub type InterleaveArray = Array<Interleave>;
83
84#[derive(Clone, Debug)]
86pub struct Interleave;
87
88#[derive(Clone, Debug)]
93pub struct InterleaveData {
94 pub(crate) num_values: usize,
95}
96
97impl Display for InterleaveData {
98 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
99 write!(f, "num_values: {}", self.num_values)
100 }
101}
102
103impl ArrayHash for InterleaveData {
104 fn array_hash<H: Hasher>(&self, state: &mut H, _accuracy: EqMode) {
105 state.write_usize(self.num_values);
106 }
107}
108
109impl ArrayEq for InterleaveData {
110 fn array_eq(&self, other: &Self, _accuracy: EqMode) -> bool {
111 self.num_values == other.num_values
112 }
113}
114
115pub trait InterleaveArrayExt: TypedArrayRef<Interleave> {
117 fn num_values(&self) -> usize {
119 self.num_values
120 }
121
122 fn value(&self, idx: usize) -> &ArrayRef {
124 self.as_ref().slots()[idx + 2]
125 .as_ref()
126 .vortex_expect("validated interleave value slot")
127 }
128
129 fn array_indices(&self) -> &ArrayRef {
131 self.as_ref().slots()[0]
132 .as_ref()
133 .vortex_expect("validated interleave array_indices slot")
134 }
135
136 fn row_indices(&self) -> &ArrayRef {
138 self.as_ref().slots()[1]
139 .as_ref()
140 .vortex_expect("validated interleave row_indices slot")
141 }
142}
143impl<T: TypedArrayRef<Interleave>> InterleaveArrayExt for T {}
144
145impl Interleave {
146 fn check(
152 values: &[ArrayRef],
153 array_indices: &ArrayRef,
154 row_indices: &ArrayRef,
155 ) -> VortexResult<DType> {
156 vortex_ensure!(
157 values.len() >= 2,
158 "interleave requires at least 2 values, got {}",
159 values.len()
160 );
161
162 for (name, selector) in [
165 ("array_indices", array_indices),
166 ("row_indices", row_indices),
167 ] {
168 match selector.dtype() {
169 DType::Primitive(ptype, nullability) if ptype.is_unsigned_int() => {
170 vortex_ensure!(
171 !nullability.is_nullable(),
172 "interleave {name} must be non-nullable, got {}",
173 selector.dtype()
174 );
175 }
176 other => vortex_bail!(
177 "interleave {name} must be a non-nullable unsigned integer, got {other}"
178 ),
179 }
180 }
181
182 vortex_ensure!(
183 array_indices.len() == row_indices.len(),
184 "interleave selectors must have equal length, got array_indices {} and row_indices {}",
185 array_indices.len(),
186 row_indices.len()
187 );
188
189 let base_dtype = values[0].dtype();
190 let mut nullability = Nullability::NonNullable;
191 for value in values {
192 vortex_ensure!(
193 value.dtype().eq_ignore_nullability(base_dtype),
194 "interleave values must share a dtype up to nullability: {} vs {}",
195 base_dtype,
196 value.dtype()
197 );
198 nullability |= value.dtype().nullability();
199 }
200
201 Ok(base_dtype.with_nullability(nullability))
202 }
203}
204
205impl Array<Interleave> {
206 pub fn try_new(
214 values: Vec<ArrayRef>,
215 array_indices: ArrayRef,
216 row_indices: ArrayRef,
217 ) -> VortexResult<Self> {
218 let dtype = Interleave::check(&values, &array_indices, &row_indices)?;
219 Ok(unsafe { Self::new_unchecked(values, array_indices, row_indices, dtype) })
221 }
222
223 pub unsafe fn new_unchecked(
238 values: Vec<ArrayRef>,
239 array_indices: ArrayRef,
240 row_indices: ArrayRef,
241 dtype: DType,
242 ) -> Self {
243 let mut slots: ArraySlots = ArraySlots::with_capacity(values.len() + 2);
244 slots.push(Some(array_indices));
245 slots.push(Some(row_indices));
246 slots.extend(values.into_iter().map(Some));
247
248 unsafe { Self::new_unchecked_slots(slots, dtype) }
251 }
252
253 pub unsafe fn new_unchecked_slots(slots: ArraySlots, dtype: DType) -> Self {
269 let num_values = slots.len() - 2;
270 let len = slots[0]
271 .as_ref()
272 .vortex_expect("interleave array_indices slot present")
273 .len();
274
275 unsafe {
276 Array::from_parts_unchecked(
277 ArrayParts::new(Interleave, dtype, len, InterleaveData { num_values })
278 .with_slots(slots),
279 )
280 }
281 }
282}
283
284impl VTable for Interleave {
285 type TypedArrayData = InterleaveData;
286 type OperationsVTable = Self;
287 type ValidityVTable = Self;
288
289 fn id(&self) -> ArrayId {
290 static ID: CachedId = CachedId::new("vortex.interleave");
291 *ID
292 }
293
294 fn validate(
295 &self,
296 data: &Self::TypedArrayData,
297 dtype: &DType,
298 len: usize,
299 slots: &[Option<ArrayRef>],
300 ) -> VortexResult<()> {
301 vortex_ensure!(
302 slots.len() == data.num_values + 2,
303 "InterleaveArray expected {} slots (values + array_indices + row_indices), got {}",
304 data.num_values + 2,
305 slots.len()
306 );
307 vortex_ensure!(
308 slots.iter().all(|s| s.is_some()),
309 "InterleaveArray slots must all be present"
310 );
311
312 let array_indices = slots[0]
313 .clone()
314 .vortex_expect("validated array_indices slot");
315 let row_indices = slots[1].clone().vortex_expect("validated row_indices slot");
316 let values: Vec<ArrayRef> = slots[2..]
317 .iter()
318 .map(|s| s.clone().vortex_expect("validated value slot"))
319 .collect();
320
321 let expected_dtype = Interleave::check(&values, &array_indices, &row_indices)?;
324 vortex_ensure!(
325 dtype == &expected_dtype,
326 "InterleaveArray dtype {} does not match the dtype implied by its children {}",
327 dtype,
328 expected_dtype
329 );
330 vortex_ensure!(
331 len == array_indices.len(),
332 "InterleaveArray length {} does not match array_indices length {}",
333 len,
334 array_indices.len()
335 );
336 Ok(())
337 }
338
339 fn nbuffers(_array: ArrayView<'_, Self>) -> usize {
340 0
341 }
342
343 fn buffer(_array: ArrayView<'_, Self>, _idx: usize) -> BufferHandle {
344 vortex_panic!("InterleaveArray has no buffers")
345 }
346
347 fn buffer_name(_array: ArrayView<'_, Self>, _idx: usize) -> Option<String> {
348 None
349 }
350
351 fn with_buffers(
352 &self,
353 array: ArrayView<'_, Self>,
354 buffers: &[BufferHandle],
355 ) -> VortexResult<ArrayParts<Self>> {
356 with_empty_buffers(self, array, buffers)
357 }
358
359 fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String {
360 match idx {
361 0 => "array_indices".to_string(),
362 1 => "row_indices".to_string(),
363 _ => format!("value_{}", idx - 2),
364 }
365 }
366
367 fn serialize(
368 _array: ArrayView<'_, Self>,
369 _session: &VortexSession,
370 ) -> VortexResult<Option<Vec<u8>>> {
371 vortex_bail!("Interleave array is not serializable")
372 }
373
374 fn deserialize(
375 &self,
376 _dtype: &DType,
377 _len: usize,
378 _metadata: &[u8],
379 _buffers: &[BufferHandle],
380 _children: &dyn ArrayChildren,
381 _session: &VortexSession,
382 ) -> VortexResult<ArrayParts<Self>> {
383 vortex_bail!("Interleave array is not serializable")
384 }
385
386 fn execute(array: Array<Self>, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
387 execute::execute(array, ctx)
388 }
389}
390
391impl OperationsVTable<Interleave> for Interleave {
392 fn scalar_at(
393 array: ArrayView<'_, Interleave>,
394 index: usize,
395 ctx: &mut ExecutionCtx,
396 ) -> VortexResult<Scalar> {
397 let branch_idx = array
400 .array_indices()
401 .execute_scalar(index, ctx)?
402 .as_primitive()
403 .as_::<usize>()
404 .vortex_expect("interleave array_indices is non-nullable");
405 let row = array
406 .row_indices()
407 .execute_scalar(index, ctx)?
408 .as_primitive()
409 .as_::<usize>()
410 .vortex_expect("interleave row_indices is non-nullable");
411
412 let scalar = array.value(branch_idx).execute_scalar(row, ctx)?;
413 Ok(if array.as_ref().dtype().is_nullable() {
415 scalar.into_nullable()
416 } else {
417 scalar
418 })
419 }
420}
421
422impl ValidityVTable<Interleave> for Interleave {
423 fn validity(array: ArrayView<'_, Interleave>) -> VortexResult<Validity> {
424 if !array.as_ref().dtype().is_nullable() {
425 return Ok(Validity::NonNullable);
426 }
427 let num_values = array.num_values();
428 let mut slots: ArraySlots = ArraySlots::with_capacity(num_values + 2);
429 slots.push(Some(array.array_indices().clone()));
430 slots.push(Some(array.row_indices().clone()));
431 for i in 0..num_values {
432 slots.push(Some(value_validity_array(array.value(i))?));
433 }
434 let interleaved = unsafe {
437 InterleaveArray::new_unchecked_slots(slots, DType::Bool(Nullability::NonNullable))
438 };
439 Ok(Validity::Array(interleaved.into_array()))
440 }
441}
442
443fn value_validity_array(value: &ArrayRef) -> VortexResult<ArrayRef> {
446 Ok(match value.validity()? {
447 Validity::NonNullable | Validity::AllValid => {
448 ConstantArray::new(true, value.len()).into_array()
449 }
450 Validity::AllInvalid => ConstantArray::new(false, value.len()).into_array(),
451 Validity::Array(array) => array,
452 })
453}
454
455#[cfg(test)]
456mod tests {
457 use vortex_error::VortexResult;
458
459 use super::*;
460 use crate::Canonical;
461 use crate::VortexSessionExecute;
462 use crate::array_session;
463 use crate::arrays::BoolArray;
464 use crate::arrays::PrimitiveArray;
465 use crate::assert_arrays_eq;
466
467 fn interleave_reference(
474 values: &[ArrayRef],
475 array_indices: &ArrayRef,
476 row_indices: &ArrayRef,
477 ctx: &mut ExecutionCtx,
478 ) -> VortexResult<ArrayRef> {
479 let len = array_indices.len();
480 let nullable = values.iter().any(|v| v.dtype().is_nullable());
481 let mut out: Vec<Option<bool>> = Vec::with_capacity(len);
482
483 for i in 0..len {
484 let j = array_indices
485 .execute_scalar(i, ctx)?
486 .as_primitive()
487 .as_::<usize>()
488 .vortex_expect("array_indices is non-nullable");
489 let row = row_indices
490 .execute_scalar(i, ctx)?
491 .as_primitive()
492 .as_::<usize>()
493 .vortex_expect("row_indices is non-nullable");
494 out.push(values[j].execute_scalar(row, ctx)?.as_bool().value());
495 }
496
497 Ok(if nullable {
498 BoolArray::from_iter(out).into_array()
499 } else {
500 BoolArray::from_iter(
501 out.into_iter()
502 .map(|v| v.vortex_expect("non-nullable value produced a null")),
503 )
504 .into_array()
505 })
506 }
507
508 fn build(
511 branches: &[&[Option<bool>]],
512 indices: &[(usize, usize)],
513 ) -> (Vec<ArrayRef>, ArrayRef, ArrayRef) {
514 let nullable = branches.iter().flat_map(|b| b.iter()).any(Option::is_none);
515 let to_value = |vals: &[Option<bool>]| -> ArrayRef {
516 if nullable {
517 BoolArray::from_iter(vals.iter().copied()).into_array()
518 } else {
519 BoolArray::from_iter(
520 vals.iter()
521 .map(|v| v.vortex_expect("non-nullable value produced a null")),
522 )
523 .into_array()
524 }
525 };
526
527 let values = branches.iter().map(|b| to_value(b)).collect();
528 let array_indices = PrimitiveArray::from_iter(
529 indices
530 .iter()
531 .map(|&(a, _)| u32::try_from(a).vortex_expect("array index fits in u32")),
532 )
533 .into_array();
534 let row_indices = PrimitiveArray::from_iter(
535 indices
536 .iter()
537 .map(|&(_, r)| u32::try_from(r).vortex_expect("row index fits in u32")),
538 )
539 .into_array();
540 (values, array_indices, row_indices)
541 }
542
543 fn check(branches: &[&[Option<bool>]], indices: &[(usize, usize)]) -> VortexResult<()> {
547 let (values, array_indices, row_indices) = build(branches, indices);
548
549 let interleaved =
550 InterleaveArray::try_new(values.clone(), array_indices.clone(), row_indices.clone())?
551 .into_array();
552
553 let mut ctx = array_session().create_execution_ctx();
554 let reference = interleave_reference(&values, &array_indices, &row_indices, &mut ctx)?;
555
556 assert_arrays_eq!(interleaved, reference, &mut ctx);
557 Ok(())
558 }
559
560 #[test]
561 fn interleave_reorders_and_repeats() -> VortexResult<()> {
562 check(
564 &[&[Some(true), Some(false)], &[Some(false), Some(true)]],
565 &[(0, 1), (1, 0), (0, 0), (1, 1), (0, 0)],
566 )
567 }
568
569 #[test]
570 fn interleave_skips_rows() -> VortexResult<()> {
571 check(
573 &[
574 &[Some(true), Some(false), Some(true)],
575 &[Some(false), Some(true)],
576 ],
577 &[(0, 0), (1, 1), (0, 2)],
578 )
579 }
580
581 #[test]
582 fn interleave_binary_spans_word_boundary() -> VortexResult<()> {
583 let branch0: Vec<Option<bool>> = (0..100).map(|i| Some(i % 3 == 0)).collect();
586 let branch1: Vec<Option<bool>> = (0..100).map(|i| Some(i % 5 == 0)).collect();
587 let indices: Vec<(usize, usize)> = (0..200).map(|i| (i % 2, (i * 7) % 100)).collect();
588 check(&[&branch0, &branch1], &indices)
589 }
590
591 #[test]
592 fn interleave_binary_nullable_spans_word_boundary() -> VortexResult<()> {
593 let branch0: Vec<Option<bool>> = (0..70)
595 .map(|i| (i % 4 != 0).then_some(i % 2 == 0))
596 .collect();
597 let branch1: Vec<Option<bool>> = (0..70)
598 .map(|i| (i % 3 != 0).then_some(i % 2 == 1))
599 .collect();
600 let indices: Vec<(usize, usize)> = (0..150).map(|i| (i % 2, (i * 11) % 70)).collect();
601 check(&[&branch0, &branch1], &indices)
602 }
603
604 #[test]
605 fn interleave_three_values() -> VortexResult<()> {
606 check(
608 &[
609 &[Some(true), Some(false)],
610 &[Some(false)],
611 &[Some(true), Some(true), Some(false)],
612 ],
613 &[(2, 1), (0, 0), (1, 0), (2, 2), (0, 1), (2, 0)],
614 )
615 }
616
617 #[test]
618 fn interleave_only_one_branch() -> VortexResult<()> {
619 check(
620 &[&[Some(true), Some(false), Some(true)], &[Some(false)]],
621 &[(0, 2), (0, 0), (0, 1)],
622 )
623 }
624
625 #[test]
626 fn interleave_nullable_with_nulls_in_values() -> VortexResult<()> {
627 check(
628 &[&[None, Some(true), None], &[Some(false), None]],
629 &[(1, 1), (0, 0), (1, 0), (0, 2), (0, 1)],
630 )
631 }
632
633 #[test]
634 fn interleave_empty() -> VortexResult<()> {
635 check(&[&[Some(true)], &[Some(false)]], &[])
636 }
637
638 #[test]
639 fn rejects_boolean_array_indices() {
640 let value = BoolArray::from_iter([true, false]).into_array();
641 let array_indices = BoolArray::from_iter([true, false]).into_array();
642 let row_indices = PrimitiveArray::from_iter([0u32, 1]).into_array();
643 let err = InterleaveArray::try_new(vec![value.clone(), value], array_indices, row_indices)
644 .err()
645 .vortex_expect("expected interleave to reject a boolean array_indices");
646 assert!(err.to_string().contains("unsigned integer"), "{err}");
647 }
648
649 #[test]
650 fn rejects_signed_integer_array_indices() {
651 let value = BoolArray::from_iter([true]).into_array();
652 let array_indices = PrimitiveArray::from_iter([0i32, 1]).into_array();
653 let row_indices = PrimitiveArray::from_iter([0u32, 0]).into_array();
654 let err = InterleaveArray::try_new(vec![value.clone(), value], array_indices, row_indices)
655 .err()
656 .vortex_expect("expected interleave to reject a signed integer array_indices");
657 assert!(err.to_string().contains("unsigned integer"), "{err}");
658 }
659
660 #[test]
661 fn rejects_nullable_row_indices() {
662 let value = BoolArray::from_iter([true, false]).into_array();
663 let array_indices = PrimitiveArray::from_iter([0u32, 1]).into_array();
664 let row_indices = PrimitiveArray::from_option_iter([Some(0u32), Some(1)]).into_array();
665 let err = InterleaveArray::try_new(vec![value.clone(), value], array_indices, row_indices)
666 .err()
667 .vortex_expect("expected interleave to reject nullable row_indices");
668 assert!(err.to_string().contains("non-nullable"), "{err}");
669 }
670
671 #[test]
672 fn rejects_mismatched_selector_lengths() {
673 let value = BoolArray::from_iter([true, false]).into_array();
674 let array_indices = PrimitiveArray::from_iter([0u32, 1]).into_array();
675 let row_indices = PrimitiveArray::from_iter([0u32]).into_array();
676 let err = InterleaveArray::try_new(vec![value.clone(), value], array_indices, row_indices)
677 .err()
678 .vortex_expect("expected interleave to reject mismatched selector lengths");
679 assert!(err.to_string().contains("equal length"), "{err}");
680 }
681
682 #[test]
683 fn execute_rejects_out_of_bounds_array_index() -> VortexResult<()> {
684 let value = BoolArray::from_iter([true]).into_array();
685 let array_indices = PrimitiveArray::from_iter([2u32]).into_array();
686 let row_indices = PrimitiveArray::from_iter([0u32]).into_array();
687 let interleaved =
688 InterleaveArray::try_new(vec![value.clone(), value], array_indices, row_indices)?
689 .into_array();
690
691 let mut ctx = array_session().create_execution_ctx();
692 let err = interleaved
693 .execute::<Canonical>(&mut ctx)
694 .err()
695 .vortex_expect("expected execution to reject out-of-bounds array index");
696 assert!(
697 err.to_string().contains("array index out of bounds"),
698 "{err}"
699 );
700 Ok(())
701 }
702
703 #[test]
704 fn execute_rejects_out_of_bounds_row_index() -> VortexResult<()> {
705 let value = BoolArray::from_iter([true]).into_array();
706 let array_indices = PrimitiveArray::from_iter([0u32]).into_array();
707 let row_indices = PrimitiveArray::from_iter([1u32]).into_array();
708 let interleaved =
709 InterleaveArray::try_new(vec![value.clone(), value], array_indices, row_indices)?
710 .into_array();
711
712 let mut ctx = array_session().create_execution_ctx();
713 let err = interleaved
714 .execute::<Canonical>(&mut ctx)
715 .err()
716 .vortex_expect("expected execution to reject out-of-bounds row index");
717 assert!(err.to_string().contains("row index out of bounds"), "{err}");
718 Ok(())
719 }
720
721 #[test]
722 #[should_panic(expected = "only implemented for boolean values")]
723 fn non_boolean_value_execution_panics() {
724 let v0 = PrimitiveArray::from_iter([1u32]).into_array();
726 let v1 = PrimitiveArray::from_iter([2u32]).into_array();
727 let array_indices = PrimitiveArray::from_iter([0u32, 1]).into_array();
728 let row_indices = PrimitiveArray::from_iter([0u32, 0]).into_array();
729 let interleaved = InterleaveArray::try_new(vec![v0, v1], array_indices, row_indices)
730 .vortex_expect("primitive values should construct")
731 .into_array();
732 let mut ctx = array_session().create_execution_ctx();
733 interleaved.execute::<Canonical>(&mut ctx).ok();
734 }
735}