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_ensure_eq;
53use vortex_error::vortex_panic;
54use vortex_session::VortexSession;
55use vortex_session::registry::CachedId;
56
57use crate::ArrayEq;
58use crate::ArrayHash;
59use crate::ArrayRef;
60use crate::EqMode;
61use crate::ExecutionCtx;
62use crate::IntoArray;
63use crate::array::Array;
64use crate::array::ArrayId;
65use crate::array::ArrayParts;
66use crate::array::ArraySlots;
67use crate::array::ArrayView;
68use crate::array::OperationsVTable;
69use crate::array::TypedArrayRef;
70use crate::array::VTable;
71use crate::array::ValidityVTable;
72use crate::array::with_empty_buffers;
73use crate::arrays::ConstantArray;
74use crate::buffer::BufferHandle;
75use crate::dtype::DType;
76use crate::dtype::Nullability;
77use crate::executor::ExecutionResult;
78use crate::scalar::Scalar;
79use crate::serde::ArrayChildren;
80use crate::validity::Validity;
81
82pub type InterleaveArray = Array<Interleave>;
84
85#[derive(Clone, Debug)]
87pub struct Interleave;
88
89#[derive(Clone, Debug)]
94pub struct InterleaveData {
95 pub(crate) num_values: usize,
96}
97
98impl Display for InterleaveData {
99 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
100 write!(f, "num_values: {}", self.num_values)
101 }
102}
103
104impl ArrayHash for InterleaveData {
105 fn array_hash<H: Hasher>(&self, state: &mut H, _accuracy: EqMode) {
106 state.write_usize(self.num_values);
107 }
108}
109
110impl ArrayEq for InterleaveData {
111 fn array_eq(&self, other: &Self, _accuracy: EqMode) -> bool {
112 self.num_values == other.num_values
113 }
114}
115
116pub trait InterleaveArrayExt: TypedArrayRef<Interleave> {
118 fn num_values(&self) -> usize {
120 self.num_values
121 }
122
123 fn value(&self, idx: usize) -> &ArrayRef {
125 self.as_ref().slots()[idx + 2]
126 .as_ref()
127 .vortex_expect("validated interleave value slot")
128 }
129
130 fn array_indices(&self) -> &ArrayRef {
132 self.as_ref().slots()[0]
133 .as_ref()
134 .vortex_expect("validated interleave array_indices slot")
135 }
136
137 fn row_indices(&self) -> &ArrayRef {
139 self.as_ref().slots()[1]
140 .as_ref()
141 .vortex_expect("validated interleave row_indices slot")
142 }
143}
144impl<T: TypedArrayRef<Interleave>> InterleaveArrayExt for T {}
145
146impl Interleave {
147 fn check(
153 values: &[ArrayRef],
154 array_indices: &ArrayRef,
155 row_indices: &ArrayRef,
156 ) -> VortexResult<DType> {
157 vortex_ensure!(
158 values.len() >= 2,
159 "interleave requires at least 2 values, got {}",
160 values.len()
161 );
162
163 for (name, selector) in [
166 ("array_indices", array_indices),
167 ("row_indices", row_indices),
168 ] {
169 match selector.dtype() {
170 DType::Primitive(ptype, nullability) if ptype.is_unsigned_int() => {
171 vortex_ensure!(
172 !nullability.is_nullable(),
173 "interleave {name} must be non-nullable, got {}",
174 selector.dtype()
175 );
176 }
177 other => vortex_bail!(
178 "interleave {name} must be a non-nullable unsigned integer, got {other}"
179 ),
180 }
181 }
182
183 vortex_ensure_eq!(
184 array_indices.len(),
185 row_indices.len(),
186 "interleave selectors must have equal length",
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(ArrayParts::new(
277 Interleave,
278 dtype,
279 len,
280 InterleaveData { num_values },
281 slots,
282 ))
283 }
284 }
285}
286
287impl VTable for Interleave {
288 type TypedArrayData = InterleaveData;
289 type OperationsVTable = Self;
290 type ValidityVTable = Self;
291
292 fn id(&self) -> ArrayId {
293 static ID: CachedId = CachedId::new("vortex.interleave");
294 *ID
295 }
296
297 fn validate(
298 &self,
299 data: &Self::TypedArrayData,
300 dtype: &DType,
301 len: usize,
302 slots: &[Option<ArrayRef>],
303 ) -> VortexResult<()> {
304 vortex_ensure_eq!(
305 slots.len(),
306 data.num_values + 2,
307 "InterleaveArray expects values + array_indices + row_indices slots",
308 );
309 vortex_ensure!(
310 slots.iter().all(|s| s.is_some()),
311 "InterleaveArray slots must all be present"
312 );
313
314 let array_indices = slots[0]
315 .clone()
316 .vortex_expect("validated array_indices slot");
317 let row_indices = slots[1].clone().vortex_expect("validated row_indices slot");
318 let values: Vec<ArrayRef> = slots[2..]
319 .iter()
320 .map(|s| s.clone().vortex_expect("validated value slot"))
321 .collect();
322
323 let expected_dtype = Interleave::check(&values, &array_indices, &row_indices)?;
326 vortex_ensure_eq!(
327 dtype,
328 &expected_dtype,
329 "InterleaveArray dtype does not match the dtype implied by its children",
330 );
331 vortex_ensure_eq!(
332 len,
333 array_indices.len(),
334 "InterleaveArray length does not match array_indices length",
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 type ProbeState = ();
393
394 fn scalar_at(
395 array: ArrayView<'_, Interleave>,
396 index: usize,
397 ctx: &mut ExecutionCtx,
398 ) -> VortexResult<Scalar> {
399 let branch_idx = array
402 .array_indices()
403 .execute_scalar(index, ctx)?
404 .as_primitive()
405 .as_::<usize>()
406 .vortex_expect("interleave array_indices is non-nullable");
407 let row = array
408 .row_indices()
409 .execute_scalar(index, ctx)?
410 .as_primitive()
411 .as_::<usize>()
412 .vortex_expect("interleave row_indices is non-nullable");
413
414 let scalar = array.value(branch_idx).execute_scalar(row, ctx)?;
415 Ok(if array.as_ref().dtype().is_nullable() {
417 scalar.into_nullable()
418 } else {
419 scalar
420 })
421 }
422}
423
424impl ValidityVTable<Interleave> for Interleave {
425 fn validity(array: ArrayView<'_, Interleave>) -> VortexResult<Validity> {
426 if !array.as_ref().dtype().is_nullable() {
427 return Ok(Validity::NonNullable);
428 }
429 let num_values = array.num_values();
430 let mut slots: ArraySlots = ArraySlots::with_capacity(num_values + 2);
431 slots.push(Some(array.array_indices().clone()));
432 slots.push(Some(array.row_indices().clone()));
433 for i in 0..num_values {
434 slots.push(Some(value_validity_array(array.value(i))?));
435 }
436 let interleaved = unsafe {
439 InterleaveArray::new_unchecked_slots(slots, DType::Bool(Nullability::NonNullable))
440 };
441 Ok(Validity::Array(interleaved.into_array()))
442 }
443}
444
445fn value_validity_array(value: &ArrayRef) -> VortexResult<ArrayRef> {
448 Ok(match value.validity()? {
449 Validity::NonNullable | Validity::AllValid => {
450 ConstantArray::new(true, value.len()).into_array()
451 }
452 Validity::AllInvalid => ConstantArray::new(false, value.len()).into_array(),
453 Validity::Array(array) => array,
454 })
455}
456
457#[cfg(test)]
458mod tests {
459 use vortex_error::VortexResult;
460
461 use super::*;
462 use crate::Canonical;
463 use crate::VortexSessionExecute;
464 use crate::array_session;
465 use crate::arrays::BoolArray;
466 use crate::arrays::PrimitiveArray;
467 use crate::assert_arrays_eq;
468 use crate::dtype::PType;
469
470 fn interleave_reference(
477 values: &[ArrayRef],
478 array_indices: &ArrayRef,
479 row_indices: &ArrayRef,
480 ctx: &mut ExecutionCtx,
481 ) -> VortexResult<ArrayRef> {
482 let len = array_indices.len();
483 let nullable = values.iter().any(|v| v.dtype().is_nullable());
484 let mut out: Vec<Option<bool>> = Vec::with_capacity(len);
485
486 for i in 0..len {
487 let j = array_indices
488 .execute_scalar(i, ctx)?
489 .as_primitive()
490 .as_::<usize>()
491 .vortex_expect("array_indices is non-nullable");
492 let row = row_indices
493 .execute_scalar(i, ctx)?
494 .as_primitive()
495 .as_::<usize>()
496 .vortex_expect("row_indices is non-nullable");
497 out.push(values[j].execute_scalar(row, ctx)?.as_bool().value());
498 }
499
500 Ok(if nullable {
501 BoolArray::from_iter(out).into_array()
502 } else {
503 BoolArray::from_iter(
504 out.into_iter()
505 .map(|v| v.vortex_expect("non-nullable value produced a null")),
506 )
507 .into_array()
508 })
509 }
510
511 fn build(
514 branches: &[&[Option<bool>]],
515 indices: &[(usize, usize)],
516 ) -> (Vec<ArrayRef>, ArrayRef, ArrayRef) {
517 let nullable = branches.iter().flat_map(|b| b.iter()).any(Option::is_none);
518 let to_value = |vals: &[Option<bool>]| -> ArrayRef {
519 if nullable {
520 BoolArray::from_iter(vals.iter().copied()).into_array()
521 } else {
522 BoolArray::from_iter(
523 vals.iter()
524 .map(|v| v.vortex_expect("non-nullable value produced a null")),
525 )
526 .into_array()
527 }
528 };
529
530 let values = branches.iter().map(|b| to_value(b)).collect();
531 let array_indices = PrimitiveArray::from_iter(
532 indices
533 .iter()
534 .map(|&(a, _)| u32::try_from(a).vortex_expect("array index fits in u32")),
535 )
536 .into_array();
537 let row_indices = PrimitiveArray::from_iter(
538 indices
539 .iter()
540 .map(|&(_, r)| u32::try_from(r).vortex_expect("row index fits in u32")),
541 )
542 .into_array();
543 (values, array_indices, row_indices)
544 }
545
546 fn check(branches: &[&[Option<bool>]], indices: &[(usize, usize)]) -> VortexResult<()> {
550 let (values, array_indices, row_indices) = build(branches, indices);
551
552 let interleaved =
553 InterleaveArray::try_new(values.clone(), array_indices.clone(), row_indices.clone())?
554 .into_array();
555
556 let mut ctx = array_session().create_execution_ctx();
557 let reference = interleave_reference(&values, &array_indices, &row_indices, &mut ctx)?;
558
559 assert_arrays_eq!(interleaved, reference, &mut ctx);
560 Ok(())
561 }
562
563 #[test]
564 fn interleave_reorders_and_repeats() -> VortexResult<()> {
565 check(
567 &[&[Some(true), Some(false)], &[Some(false), Some(true)]],
568 &[(0, 1), (1, 0), (0, 0), (1, 1), (0, 0)],
569 )
570 }
571
572 #[test]
573 fn interleave_skips_rows() -> VortexResult<()> {
574 check(
576 &[
577 &[Some(true), Some(false), Some(true)],
578 &[Some(false), Some(true)],
579 ],
580 &[(0, 0), (1, 1), (0, 2)],
581 )
582 }
583
584 #[test]
585 fn interleave_binary_spans_word_boundary() -> VortexResult<()> {
586 let branch0: Vec<Option<bool>> = (0..100).map(|i| Some(i % 3 == 0)).collect();
589 let branch1: Vec<Option<bool>> = (0..100).map(|i| Some(i % 5 == 0)).collect();
590 let indices: Vec<(usize, usize)> = (0..200).map(|i| (i % 2, (i * 7) % 100)).collect();
591 check(&[&branch0, &branch1], &indices)
592 }
593
594 #[test]
595 fn interleave_binary_nullable_spans_word_boundary() -> VortexResult<()> {
596 let branch0: Vec<Option<bool>> = (0..70)
598 .map(|i| (i % 4 != 0).then_some(i % 2 == 0))
599 .collect();
600 let branch1: Vec<Option<bool>> = (0..70)
601 .map(|i| (i % 3 != 0).then_some(i % 2 == 1))
602 .collect();
603 let indices: Vec<(usize, usize)> = (0..150).map(|i| (i % 2, (i * 11) % 70)).collect();
604 check(&[&branch0, &branch1], &indices)
605 }
606
607 #[test]
608 fn interleave_three_values() -> VortexResult<()> {
609 check(
611 &[
612 &[Some(true), Some(false)],
613 &[Some(false)],
614 &[Some(true), Some(true), Some(false)],
615 ],
616 &[(2, 1), (0, 0), (1, 0), (2, 2), (0, 1), (2, 0)],
617 )
618 }
619
620 #[test]
621 fn interleave_only_one_branch() -> VortexResult<()> {
622 check(
623 &[&[Some(true), Some(false), Some(true)], &[Some(false)]],
624 &[(0, 2), (0, 0), (0, 1)],
625 )
626 }
627
628 #[test]
629 fn interleave_nullable_with_nulls_in_values() -> VortexResult<()> {
630 check(
631 &[&[None, Some(true), None], &[Some(false), None]],
632 &[(1, 1), (0, 0), (1, 0), (0, 2), (0, 1)],
633 )
634 }
635
636 #[test]
637 fn interleave_empty() -> VortexResult<()> {
638 check(&[&[Some(true)], &[Some(false)]], &[])
639 }
640
641 #[test]
642 fn rejects_boolean_array_indices() {
643 let value = BoolArray::from_iter([true, false]).into_array();
644 let array_indices = BoolArray::from_iter([true, false]).into_array();
645 let row_indices = PrimitiveArray::from_iter([0u32, 1]).into_array();
646 let err = InterleaveArray::try_new(vec![value.clone(), value], array_indices, row_indices)
647 .err()
648 .vortex_expect("expected interleave to reject a boolean array_indices");
649 assert!(err.to_string().contains("unsigned integer"), "{err}");
650 }
651
652 #[test]
653 fn rejects_signed_integer_array_indices() {
654 let value = BoolArray::from_iter([true]).into_array();
655 let array_indices = PrimitiveArray::from_iter([0i32, 1]).into_array();
656 let row_indices = PrimitiveArray::from_iter([0u32, 0]).into_array();
657 let err = InterleaveArray::try_new(vec![value.clone(), value], array_indices, row_indices)
658 .err()
659 .vortex_expect("expected interleave to reject a signed integer array_indices");
660 assert!(err.to_string().contains("unsigned integer"), "{err}");
661 }
662
663 #[test]
664 fn rejects_nullable_row_indices() {
665 let value = BoolArray::from_iter([true, false]).into_array();
666 let array_indices = PrimitiveArray::from_iter([0u32, 1]).into_array();
667 let row_indices = PrimitiveArray::from_option_iter([Some(0u32), Some(1)]).into_array();
668 let err = InterleaveArray::try_new(vec![value.clone(), value], array_indices, row_indices)
669 .err()
670 .vortex_expect("expected interleave to reject nullable row_indices");
671 assert!(err.to_string().contains("non-nullable"), "{err}");
672 }
673
674 #[test]
675 fn rejects_mismatched_selector_lengths() {
676 let value = BoolArray::from_iter([true, false]).into_array();
677 let array_indices = PrimitiveArray::from_iter([0u32, 1]).into_array();
678 let row_indices = PrimitiveArray::from_iter([0u32]).into_array();
679 let err = InterleaveArray::try_new(vec![value.clone(), value], array_indices, row_indices)
680 .err()
681 .vortex_expect("expected interleave to reject mismatched selector lengths");
682 assert!(err.to_string().contains("equal length"), "{err}");
683 }
684
685 #[test]
686 fn execute_rejects_out_of_bounds_array_index() -> VortexResult<()> {
687 let value = BoolArray::from_iter([true]).into_array();
688 let array_indices = PrimitiveArray::from_iter([2u32]).into_array();
689 let row_indices = PrimitiveArray::from_iter([0u32]).into_array();
690 let interleaved =
691 InterleaveArray::try_new(vec![value.clone(), value], array_indices, row_indices)?
692 .into_array();
693
694 let mut ctx = array_session().create_execution_ctx();
695 let err = interleaved
696 .execute::<Canonical>(&mut ctx)
697 .err()
698 .vortex_expect("expected execution to reject out-of-bounds array index");
699 assert!(
700 err.to_string().contains("array index out of bounds"),
701 "{err}"
702 );
703 Ok(())
704 }
705
706 #[test]
707 fn execute_rejects_out_of_bounds_row_index() -> VortexResult<()> {
708 let value = BoolArray::from_iter([true]).into_array();
709 let array_indices = PrimitiveArray::from_iter([0u32]).into_array();
710 let row_indices = PrimitiveArray::from_iter([1u32]).into_array();
711 let interleaved =
712 InterleaveArray::try_new(vec![value.clone(), value], array_indices, row_indices)?
713 .into_array();
714
715 let mut ctx = array_session().create_execution_ctx();
716 let err = interleaved
717 .execute::<Canonical>(&mut ctx)
718 .err()
719 .vortex_expect("expected execution to reject out-of-bounds row index");
720 assert!(err.to_string().contains("row index out of bounds"), "{err}");
721 Ok(())
722 }
723
724 #[test]
725 fn executes_primitive_values() -> VortexResult<()> {
726 let v0 = PrimitiveArray::from_iter([1.0f64, 2.0]).into_array();
727 let v1 = PrimitiveArray::from_option_iter([Some(10.0f64), None]).into_array();
728 let array_indices = PrimitiveArray::from_iter([0u8, 1, 0, 1]).into_array();
729 let row_indices = PrimitiveArray::from_iter([0u32, 0, 1, 1]).into_array();
730 let interleaved =
731 InterleaveArray::try_new(vec![v0, v1], array_indices, row_indices)?.into_array();
732 let expected =
733 PrimitiveArray::from_option_iter([Some(1.0f64), Some(10.0), Some(2.0), None])
734 .into_array();
735 let mut ctx = array_session().create_execution_ctx();
736 assert_arrays_eq!(interleaved, expected, &mut ctx);
737 Ok(())
738 }
739
740 #[test]
741 fn executes_primitive_constant_values() -> VortexResult<()> {
742 let constant = ConstantArray::new(1.0f64, 2).into_array();
743 let column = PrimitiveArray::from_iter([10.0f64, 20.0]).into_array();
744 let array_indices = PrimitiveArray::from_iter([0u8, 1, 0, 1]).into_array();
745 let row_indices = PrimitiveArray::from_iter([0u32, 0, 1, 1]).into_array();
746 let interleaved =
747 InterleaveArray::try_new(vec![constant, column], array_indices, row_indices)?
748 .into_array();
749 let expected = PrimitiveArray::from_iter([1.0f64, 10.0, 1.0, 20.0]).into_array();
750 let mut ctx = array_session().create_execution_ctx();
751
752 assert_arrays_eq!(interleaved, expected, &mut ctx);
753 Ok(())
754 }
755
756 #[test]
757 fn executes_null_primitive_constant_values() -> VortexResult<()> {
758 let constant = ConstantArray::new(
759 Scalar::null(DType::Primitive(PType::F64, Nullability::Nullable)),
760 2,
761 )
762 .into_array();
763 let column = PrimitiveArray::from_iter([10.0f64, 20.0]).into_array();
764 let array_indices = PrimitiveArray::from_iter([0u8, 1, 0, 1]).into_array();
765 let row_indices = PrimitiveArray::from_iter([0u32, 0, 1, 1]).into_array();
766 let interleaved =
767 InterleaveArray::try_new(vec![constant, column], array_indices, row_indices)?
768 .into_array();
769 let expected =
770 PrimitiveArray::from_option_iter([None, Some(10.0f64), None, Some(20.0)]).into_array();
771 let mut ctx = array_session().create_execution_ctx();
772
773 assert_arrays_eq!(interleaved, expected, &mut ctx);
774 Ok(())
775 }
776}