Skip to main content

vortex_array/arrays/interleave/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! The [`Interleave`] encoding: a lazy, random-access gather of `N` value arrays into one array,
5//! routed by a per-row `(array_index, row_index)` pair.
6//!
7//! # Specification
8//!
9//! An [`Interleave`] array has `N + 2` children: an `array_indices` selector and a `row_indices`
10//! selector followed by `N` *values*. The output has `array_indices.len()` rows, and output
11//! row `i` comes from `values[array_indices[i]][row_indices[i]]`.
12//!
13//! Unlike a `Merge`, which consumes each branch in order under a cursor, an [`Interleave`] is
14//! **random-access**: `row_indices` names an explicit position within the selected value array, so
15//! rows may be reordered, skipped, or repeated. A `Merge` is the special case where each value
16//! array is consumed front-to-back exactly once.
17//!
18//! Like a `Merge`, the value arrays are independent: each holds only its own rows, and the
19//! selectors stitch them back together. This distinguishes [`Interleave`] from an element-wise
20//! select such as `zip`, whose arguments are all full-length.
21//!
22//! ## Invariants
23//!
24//! - Both selectors are **non-nullable** and equal in length, which is the output length. They
25//!   record *where* each output row comes from, which is always a definite decision. Predicate
26//!   nullability must be resolved into definite indices by the caller *before* the interleave is
27//!   built.
28//! - `array_indices[i] < values.len()` and `row_indices[i] < values[array_indices[i]].len()` for
29//!   every `i`. These per-row bounds depend on the selector *values* and so are a runtime
30//!   precondition of the caller, checked in the execution kernels rather than at construction.
31//! - All values share a logical type up to nullability. The output type is that shared type with
32//!   the union of the values' nullabilities. This is orthogonal to the selectors: a row's *value*
33//!   may be null even though its `(array_index, row_index)` is definite.
34//! - The output length equals `array_indices.len()` (`== row_indices.len()`).
35//!
36//! ## Selector types
37//!
38//! `array_indices` encodes the value array per row as a non-nullable **unsigned integer**
39//! (`array_indices[i]` is the index into `values`). `row_indices` is likewise a non-nullable
40//! **unsigned integer** naming the position within the selected value array.
41
42mod 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
81/// An [`Interleave`]-encoded Vortex array. See the [module docs](self) for the specification.
82pub type InterleaveArray = Array<Interleave>;
83
84/// The [`Interleave`] encoding. See the [module docs](self).
85#[derive(Clone, Debug)]
86pub struct Interleave;
87
88/// Per-array metadata for an [`InterleaveArray`].
89///
90/// The selectors and values live in the array's slots; the selectors occupy `slots[0]` and
91/// `slots[1]`, and the `num_values` values follow at `slots[2..2 + num_values]`.
92#[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
115/// Accessors for the values and selectors of an [`InterleaveArray`].
116pub trait InterleaveArrayExt: TypedArrayRef<Interleave> {
117    /// The number of value arrays (two fewer than the number of children).
118    fn num_values(&self) -> usize {
119        self.num_values
120    }
121
122    /// The `idx`-th value array (holding the rows that `array_indices` routes to it).
123    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    /// The selector routing each output row to a value array.
130    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    /// The selector naming each output row's position within its value array.
137    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    /// The single source of truth for [`InterleaveArray`] invariants.
147    ///
148    /// Validates `values`, `array_indices`, and `row_indices` against the [specification](self) and
149    /// returns the output [`DType`] (the shared value type with the union of value nullabilities).
150    /// Both the public constructor and the [`VTable::validate`] hook funnel through here.
151    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        // Both selectors are non-nullable unsigned integers: `array_indices` indexes the values and
163        // `row_indices` names a position within the selected value.
164        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    /// Constructs a new [`InterleaveArray`] from `values` and the `array_indices` / `row_indices`
207    /// selectors.
208    ///
209    /// See the [module docs](self) for the full specification and invariants. The selectors must be
210    /// non-nullable: they record a definite `(array_index, row_index)` per row, so null-predicate
211    /// handling is the caller's responsibility, resolved before the interleave is constructed. The
212    /// per-row bounds on the selector values are a runtime precondition checked during execution.
213    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        // SAFETY: `check` just validated every invariant and computed the matching `dtype`.
220        Ok(unsafe { Self::new_unchecked(values, array_indices, row_indices, dtype) })
221    }
222
223    /// Constructs an [`InterleaveArray`] without re-validating the spec invariants.
224    ///
225    /// This is the assembly half of [`try_new`](Self::try_new): it lays the selectors into
226    /// `slots[0]`/`slots[1]` and the values into `slots[2..]` and stamps `dtype`, skipping the
227    /// `Interleave::check` pass. It exists for hot internal paths — notably building the
228    /// pushed-down validity interleave in [`ValidityVTable::validity`] — where the inputs are known
229    /// to satisfy the invariants by construction and re-checking them is pure overhead.
230    ///
231    /// # Safety
232    ///
233    /// The caller must uphold every [module invariant](self): at least two `values` sharing a dtype
234    /// up to nullability, both selectors non-nullable unsigned integers of equal length, and `dtype`
235    /// equal to the value returned by `Interleave::check` for these arguments (the shared value
236    /// type with the union of the values' nullabilities).
237    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        // SAFETY: the caller of `new_unchecked` upholds every invariant; here we only assemble the
249        // canonical slot layout (`array_indices`, `row_indices`, then values) that follows.
250        unsafe { Self::new_unchecked_slots(slots, dtype) }
251    }
252
253    /// Constructs an [`InterleaveArray`] from a pre-assembled `slots` buffer, skipping both the
254    /// spec re-check and the slot copy that [`new_unchecked`](Self::new_unchecked) performs.
255    ///
256    /// This is the lowest-level assembly path: it stamps `dtype` onto `slots` as-is. Callers that
257    /// already own a correctly laid-out [`ArraySlots`] — for example reusing a validated parent's
258    /// selectors while swapping its values — avoid materializing an intermediate `Vec<ArrayRef>`
259    /// and re-pushing it into a fresh `ArraySlots`.
260    ///
261    /// # Safety
262    ///
263    /// The caller must uphold every [module invariant](self) *and* the slot layout: `slots[0]` is
264    /// the `array_indices` selector, `slots[1]` is the `row_indices` selector, and `slots[2..]` are
265    /// the value arrays. Every slot must be `Some`, there must be at least two values
266    /// (`slots.len() >= 4`), and `dtype` must equal the value returned by `Interleave::check` for
267    /// these arguments.
268    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        // All semantic invariants live in `check`; here we only confirm the array's cached `dtype`
322        // and `len` agree with what the children imply.
323        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        // Random-access gather: read the routing pair for `index` directly, then pull that row from
398        // the selected value array. No cursor walk is required.
399        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        // The value may be non-nullable while the interleaved output is nullable; align the dtype.
414        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        // SAFETY: `value_validity_array` yields a non-nullable boolean array per value,
435        // the selectors already validated.
436        let interleaved = unsafe {
437            InterleaveArray::new_unchecked_slots(slots, DType::Bool(Nullability::NonNullable))
438        };
439        Ok(Validity::Array(interleaved.into_array()))
440    }
441}
442
443/// Materializes a value's validity as a non-nullable boolean array of the value's length, where
444/// `true` marks a valid (non-null) row.
445fn 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    use crate::dtype::PType;
467
468    /// Reference (oracle) implementation of the interleave spec, used only to validate the optimized
469    /// [execute](super::execute) path. It is intentionally simple and slow: it pulls each output
470    /// element one [`Scalar`] at a time via [`ArrayRef::execute_scalar`] and never touches raw bits.
471    ///
472    /// This is deliberately *not* wired into the array execution path — it exists purely as a
473    /// trustworthy comparison point in tests.
474    fn interleave_reference(
475        values: &[ArrayRef],
476        array_indices: &ArrayRef,
477        row_indices: &ArrayRef,
478        ctx: &mut ExecutionCtx,
479    ) -> VortexResult<ArrayRef> {
480        let len = array_indices.len();
481        let nullable = values.iter().any(|v| v.dtype().is_nullable());
482        let mut out: Vec<Option<bool>> = Vec::with_capacity(len);
483
484        for i in 0..len {
485            let j = array_indices
486                .execute_scalar(i, ctx)?
487                .as_primitive()
488                .as_::<usize>()
489                .vortex_expect("array_indices is non-nullable");
490            let row = row_indices
491                .execute_scalar(i, ctx)?
492                .as_primitive()
493                .as_::<usize>()
494                .vortex_expect("row_indices is non-nullable");
495            out.push(values[j].execute_scalar(row, ctx)?.as_bool().value());
496        }
497
498        Ok(if nullable {
499            BoolArray::from_iter(out).into_array()
500        } else {
501            BoolArray::from_iter(
502                out.into_iter()
503                    .map(|v| v.vortex_expect("non-nullable value produced a null")),
504            )
505            .into_array()
506        })
507    }
508
509    /// Builds the compact value arrays and the unsigned `(array_indices, row_indices)` selectors for
510    /// a gather described by per-output `(array_index, row_index)` pairs over `branches`.
511    fn build(
512        branches: &[&[Option<bool>]],
513        indices: &[(usize, usize)],
514    ) -> (Vec<ArrayRef>, ArrayRef, ArrayRef) {
515        let nullable = branches.iter().flat_map(|b| b.iter()).any(Option::is_none);
516        let to_value = |vals: &[Option<bool>]| -> ArrayRef {
517            if nullable {
518                BoolArray::from_iter(vals.iter().copied()).into_array()
519            } else {
520                BoolArray::from_iter(
521                    vals.iter()
522                        .map(|v| v.vortex_expect("non-nullable value produced a null")),
523                )
524                .into_array()
525            }
526        };
527
528        let values = branches.iter().map(|b| to_value(b)).collect();
529        let array_indices = PrimitiveArray::from_iter(
530            indices
531                .iter()
532                .map(|&(a, _)| u32::try_from(a).vortex_expect("array index fits in u32")),
533        )
534        .into_array();
535        let row_indices = PrimitiveArray::from_iter(
536            indices
537                .iter()
538                .map(|&(_, r)| u32::try_from(r).vortex_expect("row index fits in u32")),
539        )
540        .into_array();
541        (values, array_indices, row_indices)
542    }
543
544    /// Asserts that the optimized execute path and the reference implementation agree, exercising
545    /// `InterleaveArray` construction, `execute`, `scalar_at`, and `validity` (via
546    /// `assert_arrays_eq`).
547    fn check(branches: &[&[Option<bool>]], indices: &[(usize, usize)]) -> VortexResult<()> {
548        let (values, array_indices, row_indices) = build(branches, indices);
549
550        let interleaved =
551            InterleaveArray::try_new(values.clone(), array_indices.clone(), row_indices.clone())?
552                .into_array();
553
554        let mut ctx = array_session().create_execution_ctx();
555        let reference = interleave_reference(&values, &array_indices, &row_indices, &mut ctx)?;
556
557        assert_arrays_eq!(interleaved, reference, &mut ctx);
558        Ok(())
559    }
560
561    #[test]
562    fn interleave_reorders_and_repeats() -> VortexResult<()> {
563        // Random access: rows are pulled out of order and branch 0 row 0 is repeated.
564        check(
565            &[&[Some(true), Some(false)], &[Some(false), Some(true)]],
566            &[(0, 1), (1, 0), (0, 0), (1, 1), (0, 0)],
567        )
568    }
569
570    #[test]
571    fn interleave_skips_rows() -> VortexResult<()> {
572        // Branch 0 row 1 and branch 1 row 0 are never gathered.
573        check(
574            &[
575                &[Some(true), Some(false), Some(true)],
576                &[Some(false), Some(true)],
577            ],
578            &[(0, 0), (1, 1), (0, 2)],
579        )
580    }
581
582    #[test]
583    fn interleave_binary_spans_word_boundary() -> VortexResult<()> {
584        // Exercises a two-value gather across more than one 64-bit packing word, with out-of-order
585        // routing into both branches so neither buffer is consumed contiguously.
586        let branch0: Vec<Option<bool>> = (0..100).map(|i| Some(i % 3 == 0)).collect();
587        let branch1: Vec<Option<bool>> = (0..100).map(|i| Some(i % 5 == 0)).collect();
588        let indices: Vec<(usize, usize)> = (0..200).map(|i| (i % 2, (i * 7) % 100)).collect();
589        check(&[&branch0, &branch1], &indices)
590    }
591
592    #[test]
593    fn interleave_binary_nullable_spans_word_boundary() -> VortexResult<()> {
594        // Same two-value gather, now with nulls so the validity packing is exercised too.
595        let branch0: Vec<Option<bool>> = (0..70)
596            .map(|i| (i % 4 != 0).then_some(i % 2 == 0))
597            .collect();
598        let branch1: Vec<Option<bool>> = (0..70)
599            .map(|i| (i % 3 != 0).then_some(i % 2 == 1))
600            .collect();
601        let indices: Vec<(usize, usize)> = (0..150).map(|i| (i % 2, (i * 11) % 70)).collect();
602        check(&[&branch0, &branch1], &indices)
603    }
604
605    #[test]
606    fn interleave_three_values() -> VortexResult<()> {
607        // An unsigned `array_indices` routes among three values with full random access.
608        check(
609            &[
610                &[Some(true), Some(false)],
611                &[Some(false)],
612                &[Some(true), Some(true), Some(false)],
613            ],
614            &[(2, 1), (0, 0), (1, 0), (2, 2), (0, 1), (2, 0)],
615        )
616    }
617
618    #[test]
619    fn interleave_only_one_branch() -> VortexResult<()> {
620        check(
621            &[&[Some(true), Some(false), Some(true)], &[Some(false)]],
622            &[(0, 2), (0, 0), (0, 1)],
623        )
624    }
625
626    #[test]
627    fn interleave_nullable_with_nulls_in_values() -> VortexResult<()> {
628        check(
629            &[&[None, Some(true), None], &[Some(false), None]],
630            &[(1, 1), (0, 0), (1, 0), (0, 2), (0, 1)],
631        )
632    }
633
634    #[test]
635    fn interleave_empty() -> VortexResult<()> {
636        check(&[&[Some(true)], &[Some(false)]], &[])
637    }
638
639    #[test]
640    fn rejects_boolean_array_indices() {
641        let value = BoolArray::from_iter([true, false]).into_array();
642        let array_indices = BoolArray::from_iter([true, false]).into_array();
643        let row_indices = PrimitiveArray::from_iter([0u32, 1]).into_array();
644        let err = InterleaveArray::try_new(vec![value.clone(), value], array_indices, row_indices)
645            .err()
646            .vortex_expect("expected interleave to reject a boolean array_indices");
647        assert!(err.to_string().contains("unsigned integer"), "{err}");
648    }
649
650    #[test]
651    fn rejects_signed_integer_array_indices() {
652        let value = BoolArray::from_iter([true]).into_array();
653        let array_indices = PrimitiveArray::from_iter([0i32, 1]).into_array();
654        let row_indices = PrimitiveArray::from_iter([0u32, 0]).into_array();
655        let err = InterleaveArray::try_new(vec![value.clone(), value], array_indices, row_indices)
656            .err()
657            .vortex_expect("expected interleave to reject a signed integer array_indices");
658        assert!(err.to_string().contains("unsigned integer"), "{err}");
659    }
660
661    #[test]
662    fn rejects_nullable_row_indices() {
663        let value = BoolArray::from_iter([true, false]).into_array();
664        let array_indices = PrimitiveArray::from_iter([0u32, 1]).into_array();
665        let row_indices = PrimitiveArray::from_option_iter([Some(0u32), Some(1)]).into_array();
666        let err = InterleaveArray::try_new(vec![value.clone(), value], array_indices, row_indices)
667            .err()
668            .vortex_expect("expected interleave to reject nullable row_indices");
669        assert!(err.to_string().contains("non-nullable"), "{err}");
670    }
671
672    #[test]
673    fn rejects_mismatched_selector_lengths() {
674        let value = BoolArray::from_iter([true, false]).into_array();
675        let array_indices = PrimitiveArray::from_iter([0u32, 1]).into_array();
676        let row_indices = PrimitiveArray::from_iter([0u32]).into_array();
677        let err = InterleaveArray::try_new(vec![value.clone(), value], array_indices, row_indices)
678            .err()
679            .vortex_expect("expected interleave to reject mismatched selector lengths");
680        assert!(err.to_string().contains("equal length"), "{err}");
681    }
682
683    #[test]
684    fn execute_rejects_out_of_bounds_array_index() -> VortexResult<()> {
685        let value = BoolArray::from_iter([true]).into_array();
686        let array_indices = PrimitiveArray::from_iter([2u32]).into_array();
687        let row_indices = PrimitiveArray::from_iter([0u32]).into_array();
688        let interleaved =
689            InterleaveArray::try_new(vec![value.clone(), value], array_indices, row_indices)?
690                .into_array();
691
692        let mut ctx = array_session().create_execution_ctx();
693        let err = interleaved
694            .execute::<Canonical>(&mut ctx)
695            .err()
696            .vortex_expect("expected execution to reject out-of-bounds array index");
697        assert!(
698            err.to_string().contains("array index out of bounds"),
699            "{err}"
700        );
701        Ok(())
702    }
703
704    #[test]
705    fn execute_rejects_out_of_bounds_row_index() -> VortexResult<()> {
706        let value = BoolArray::from_iter([true]).into_array();
707        let array_indices = PrimitiveArray::from_iter([0u32]).into_array();
708        let row_indices = PrimitiveArray::from_iter([1u32]).into_array();
709        let interleaved =
710            InterleaveArray::try_new(vec![value.clone(), value], array_indices, row_indices)?
711                .into_array();
712
713        let mut ctx = array_session().create_execution_ctx();
714        let err = interleaved
715            .execute::<Canonical>(&mut ctx)
716            .err()
717            .vortex_expect("expected execution to reject out-of-bounds row index");
718        assert!(err.to_string().contains("row index out of bounds"), "{err}");
719        Ok(())
720    }
721
722    #[test]
723    fn executes_primitive_values() -> VortexResult<()> {
724        let v0 = PrimitiveArray::from_iter([1.0f64, 2.0]).into_array();
725        let v1 = PrimitiveArray::from_option_iter([Some(10.0f64), None]).into_array();
726        let array_indices = PrimitiveArray::from_iter([0u8, 1, 0, 1]).into_array();
727        let row_indices = PrimitiveArray::from_iter([0u32, 0, 1, 1]).into_array();
728        let interleaved =
729            InterleaveArray::try_new(vec![v0, v1], array_indices, row_indices)?.into_array();
730        let expected =
731            PrimitiveArray::from_option_iter([Some(1.0f64), Some(10.0), Some(2.0), None])
732                .into_array();
733        let mut ctx = array_session().create_execution_ctx();
734        assert_arrays_eq!(interleaved, expected, &mut ctx);
735        Ok(())
736    }
737
738    #[test]
739    fn executes_primitive_constant_values() -> VortexResult<()> {
740        let constant = ConstantArray::new(1.0f64, 2).into_array();
741        let column = PrimitiveArray::from_iter([10.0f64, 20.0]).into_array();
742        let array_indices = PrimitiveArray::from_iter([0u8, 1, 0, 1]).into_array();
743        let row_indices = PrimitiveArray::from_iter([0u32, 0, 1, 1]).into_array();
744        let interleaved =
745            InterleaveArray::try_new(vec![constant, column], array_indices, row_indices)?
746                .into_array();
747        let expected = PrimitiveArray::from_iter([1.0f64, 10.0, 1.0, 20.0]).into_array();
748        let mut ctx = array_session().create_execution_ctx();
749
750        assert_arrays_eq!(interleaved, expected, &mut ctx);
751        Ok(())
752    }
753
754    #[test]
755    fn executes_null_primitive_constant_values() -> VortexResult<()> {
756        let constant = ConstantArray::new(
757            Scalar::null(DType::Primitive(PType::F64, Nullability::Nullable)),
758            2,
759        )
760        .into_array();
761        let column = PrimitiveArray::from_iter([10.0f64, 20.0]).into_array();
762        let array_indices = PrimitiveArray::from_iter([0u8, 1, 0, 1]).into_array();
763        let row_indices = PrimitiveArray::from_iter([0u32, 0, 1, 1]).into_array();
764        let interleaved =
765            InterleaveArray::try_new(vec![constant, column], array_indices, row_indices)?
766                .into_array();
767        let expected =
768            PrimitiveArray::from_option_iter([None, Some(10.0f64), None, Some(20.0)]).into_array();
769        let mut ctx = array_session().create_execution_ctx();
770
771        assert_arrays_eq!(interleaved, expected, &mut ctx);
772        Ok(())
773    }
774}