vortex_array/arrays/interleave/
array.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::fmt::Display;
5use std::fmt::Formatter;
6use std::hash::Hasher;
7
8use vortex_error::VortexExpect;
9use vortex_error::VortexResult;
10
11use super::Interleave;
12use crate::ArrayEq;
13use crate::ArrayHash;
14use crate::ArrayRef;
15use crate::EqMode;
16use crate::array::Array;
17use crate::array::ArrayParts;
18use crate::array::ArraySlots;
19use crate::array::TypedArrayRef;
20use crate::array_slots;
21use crate::dtype::DType;
22
23/// Slot layout of an interleave: two selectors followed by its value arrays.
24#[array_slots(Interleave)]
25pub(super) struct InterleaveSlots {
26    /// The value-array index for each output row.
27    #[slot(0)]
28    pub array_indices: ArrayRef,
29    /// The row index within the selected value array for each output row.
30    #[slot(1)]
31    pub row_indices: ArrayRef,
32    /// The independent source arrays selected by the two selectors.
33    #[slot(2..)]
34    pub values: Vec<ArrayRef>,
35}
36
37/// An [`Interleave`]-encoded Vortex array. See the [module docs](super) for the specification.
38pub type InterleaveArray = Array<Interleave>;
39
40/// Per-array metadata for an [`InterleaveArray`].
41///
42/// The selectors occupy the leading slots, followed by the `num_values` value arrays.
43#[derive(Clone, Debug)]
44pub struct InterleaveData {
45    pub(crate) num_values: usize,
46}
47
48impl Display for InterleaveData {
49    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
50        write!(f, "num_values: {}", self.num_values)
51    }
52}
53
54impl ArrayHash for InterleaveData {
55    fn array_hash<H: Hasher>(&self, state: &mut H, _accuracy: EqMode) {
56        state.write_usize(self.num_values);
57    }
58}
59
60impl ArrayEq for InterleaveData {
61    fn array_eq(&self, other: &Self, _accuracy: EqMode) -> bool {
62        self.num_values == other.num_values
63    }
64}
65
66/// Accessors for the values and selectors of an [`InterleaveArray`].
67pub trait InterleaveArrayExt: TypedArrayRef<Interleave> {
68    /// The number of value arrays (two fewer than the number of children).
69    fn num_values(&self) -> usize {
70        self.num_values
71    }
72
73    /// The `idx`-th value array (holding the rows that `array_indices` routes to it).
74    fn value(&self, idx: usize) -> &ArrayRef {
75        self.slots()[InterleaveSlots::VALUES_OFFSET + idx]
76            .as_ref()
77            .vortex_expect("validated interleave value slot")
78    }
79
80    /// The selector routing each output row to a value array.
81    fn array_indices(&self) -> &ArrayRef {
82        self.slots()[InterleaveSlots::ARRAY_INDICES]
83            .as_ref()
84            .vortex_expect("validated interleave array_indices slot")
85    }
86
87    /// The selector naming each output row's position within its value array.
88    fn row_indices(&self) -> &ArrayRef {
89        self.slots()[InterleaveSlots::ROW_INDICES]
90            .as_ref()
91            .vortex_expect("validated interleave row_indices slot")
92    }
93}
94impl<T: TypedArrayRef<Interleave>> InterleaveArrayExt for T {}
95
96impl Array<Interleave> {
97    /// Constructs a new [`InterleaveArray`] from `values` and the `array_indices` / `row_indices`
98    /// selectors.
99    ///
100    /// See the [module docs](super) for the full specification and invariants. The selectors must be
101    /// non-nullable: they record a definite `(array_index, row_index)` per row, so null-predicate
102    /// handling is the caller's responsibility, resolved before the interleave is constructed. The
103    /// per-row bounds on the selector values are a runtime precondition checked during execution.
104    pub fn try_new(
105        values: Vec<ArrayRef>,
106        array_indices: ArrayRef,
107        row_indices: ArrayRef,
108    ) -> VortexResult<Self> {
109        let dtype = Interleave::check(&values, &array_indices, &row_indices)?;
110        // SAFETY: `check` just validated every invariant and computed the matching `dtype`.
111        Ok(unsafe { Self::new_unchecked(values, array_indices, row_indices, dtype) })
112    }
113
114    /// Constructs an [`InterleaveArray`] without re-validating the spec invariants.
115    ///
116    /// This is the assembly half of [`try_new`](Self::try_new): it lays out the selectors followed
117    /// by the value arrays and stamps `dtype`, skipping the `Interleave::check` pass for callers
118    /// that already uphold the invariants.
119    ///
120    /// # Safety
121    ///
122    /// The caller must uphold every [module invariant](super): at least two `values` sharing a dtype
123    /// up to nullability, both selectors non-nullable unsigned integers of equal length, and `dtype`
124    /// equal to the value returned by `Interleave::check` for these arguments (the shared value
125    /// type with the union of the values' nullabilities).
126    pub unsafe fn new_unchecked(
127        values: Vec<ArrayRef>,
128        array_indices: ArrayRef,
129        row_indices: ArrayRef,
130        dtype: DType,
131    ) -> Self {
132        let slots = InterleaveSlots {
133            array_indices,
134            row_indices,
135            values,
136        }
137        .into_slots();
138
139        // SAFETY: the caller of `new_unchecked` upholds every invariant; here we only assemble the
140        // canonical slot layout (`array_indices`, `row_indices`, then values) that follows.
141        unsafe { Self::new_unchecked_slots(slots, dtype) }
142    }
143
144    /// Constructs an [`InterleaveArray`] from a pre-assembled `slots` buffer, skipping both the
145    /// spec re-check and the slot copy that [`new_unchecked`](Self::new_unchecked) performs.
146    ///
147    /// This is the lowest-level assembly path: it stamps `dtype` onto `slots` as-is. Callers that
148    /// already own a correctly laid-out [`ArraySlots`] — for example reusing a validated parent's
149    /// selectors while swapping its values — avoid materializing an intermediate `Vec<ArrayRef>`
150    /// and re-pushing it into a fresh `ArraySlots`.
151    ///
152    /// # Safety
153    ///
154    /// The caller must uphold every [module invariant](super) *and* the slot layout: the
155    /// `array_indices` selector comes first, followed by the `row_indices` selector and then the
156    /// value arrays. Every slot must be `Some`, there must be at least two values, and `dtype`
157    /// must equal the value returned by `Interleave::check` for these arguments.
158    pub unsafe fn new_unchecked_slots(slots: ArraySlots, dtype: DType) -> Self {
159        let num_values = slots.len() - InterleaveSlots::VALUES_OFFSET;
160        let len = slots[InterleaveSlots::ARRAY_INDICES]
161            .as_ref()
162            .vortex_expect("interleave array_indices slot present")
163            .len();
164
165        unsafe {
166            Array::from_parts_unchecked(ArrayParts::new(
167                Interleave,
168                dtype,
169                len,
170                InterleaveData { num_values },
171                slots,
172            ))
173        }
174    }
175}
vortex_array/arrays/interleave/array.rs

vortex_array/arrays/interleave/
array.rs
