vortex_array/array/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::any::Any;
5use std::fmt::Debug;
6use std::fmt::Formatter;
7use std::hash::Hasher;
8use std::sync::Arc;
9
10use vortex_buffer::ByteBuffer;
11use vortex_error::VortexExpect;
12use vortex_error::VortexResult;
13use vortex_error::vortex_ensure;
14use vortex_error::vortex_err;
15use vortex_error::vortex_panic;
16use vortex_session::registry::Id;
17
18use crate::ExecutionCtx;
19use crate::buffer::BufferHandle;
20use crate::builders::ArrayBuilder;
21use crate::dtype::DType;
22use crate::dtype::Nullability;
23use crate::executor::ExecutionResult;
24use crate::executor::ExecutionStep;
25use crate::scalar::Scalar;
26use crate::validity::Validity;
27
28mod erased;
29pub use erased::*;
30
31mod plugin;
32pub use plugin::*;
33mod probe;
34pub use probe::*;
35
36mod foreign;
37pub(crate) use foreign::*;
38
39mod typed;
40pub use typed::*;
41
42pub mod vtable;
43pub use vtable::*;
44
45mod view;
46use smallvec::SmallVec;
47pub use view::*;
48
49use crate::hash::ArrayEq;
50use crate::hash::ArrayHash;
51
52/// The slots of an array: a collection of optional child arrays.
53///
54/// Most encodings have 4 or fewer slots, so we use a `SmallVec` to avoid
55/// heap allocation in the common case.
56pub type ArraySlots = SmallVec<[Option<ArrayRef>; 4]>;
57
58/// A borrowed run of required slots, e.g. the variadic tail of a slot layout.
59///
60/// Wraps a `&[Option<ArrayRef>]` whose entries are guaranteed present by encoding
61/// validation, exposing them as `&ArrayRef` without per-call-site unwrapping.
62#[derive(Clone, Copy, Debug)]
63pub struct SlotSlice<'a> {
64    slots: &'a [Option<ArrayRef>],
65    expect: &'static str,
66}
67
68impl<'a> SlotSlice<'a> {
69    /// Wrap a slice of slots that validation guarantees are all present.
70    ///
71    /// `expect` names the slot run in the panic message if a slot is unexpectedly absent.
72    pub fn new(slots: &'a [Option<ArrayRef>], expect: &'static str) -> Self {
73        Self { slots, expect }
74    }
75
76    /// The number of slots in the run.
77    pub fn len(&self) -> usize {
78        self.slots.len()
79    }
80
81    /// Returns `true` if the run contains no slots.
82    pub fn is_empty(&self) -> bool {
83        self.slots.is_empty()
84    }
85
86    /// Returns the slot at `idx`, or `None` if out of bounds.
87    pub fn get(&self, idx: usize) -> Option<&'a ArrayRef> {
88        self.slots
89            .get(idx)
90            .map(|slot| slot.as_ref().vortex_expect(self.expect))
91    }
92
93    /// Iterate the slots in order.
94    pub fn iter(&self) -> impl ExactSizeIterator<Item = &'a ArrayRef> + use<'a> {
95        let expect = self.expect;
96        self.slots
97            .iter()
98            .map(move |slot| slot.as_ref().vortex_expect(expect))
99    }
100
101    /// Clone every slot into an owned `Vec`.
102    pub fn to_vec(&self) -> Vec<ArrayRef> {
103        self.iter().cloned().collect()
104    }
105}
106
107impl std::ops::Index<usize> for SlotSlice<'_> {
108    type Output = ArrayRef;
109
110    fn index(&self, idx: usize) -> &Self::Output {
111        self.slots[idx].as_ref().vortex_expect(self.expect)
112    }
113}
114
115/// The public API trait for all Vortex arrays.
116///
117/// This trait is sealed and cannot be implemented outside of `vortex-array`.
118/// Use [`ArrayRef`] as the primary handle for working with arrays.
119#[doc(hidden)]
120pub(crate) trait DynArrayData: 'static + private::Sealed + Send + Sync + Debug {
121    /// Returns the array as a reference to a generic [`Any`] trait object.
122    fn as_any(&self) -> &dyn Any;
123
124    /// Returns the array as a mutable reference to a generic [`Any`] trait object.
125    fn as_any_mut(&mut self) -> &mut dyn Any;
126
127    /// Returns the [`Validity`] of the array.
128    fn validity(&self, this: &ArrayRef) -> VortexResult<Validity>;
129
130    /// Writes the array into the canonical builder.
131    ///
132    /// The [`DType`] of the builder must match that of the array.
133    fn append_to_builder(
134        &self,
135        this: &ArrayRef,
136        builder: &mut dyn ArrayBuilder,
137        ctx: &mut ExecutionCtx,
138    ) -> VortexResult<()>;
139
140    // --- Visitor methods (formerly in ArrayVisitor) ---
141
142    /// Returns the buffers of the array.
143    fn buffers(&self, this: &ArrayRef) -> Vec<ByteBuffer>;
144
145    /// Returns the buffer handles of the array.
146    fn buffer_handles(&self, this: &ArrayRef) -> Vec<BufferHandle>;
147
148    /// Returns the names of the buffers of the array.
149    fn buffer_names(&self, this: &ArrayRef) -> Vec<String>;
150
151    /// Returns the array's buffers with their names.
152    fn named_buffers(&self, this: &ArrayRef) -> Vec<(String, BufferHandle)>;
153
154    /// Returns the number of buffers of the array.
155    fn nbuffers(&self, this: &ArrayRef) -> usize;
156
157    /// Returns the name of the slot at the given index.
158    fn slot_name(&self, this: &ArrayRef, idx: usize) -> String;
159
160    /// Formats a human-readable metadata description.
161    fn metadata_fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result;
162
163    /// Hashes the array contents including len, dtype, and encoding id.
164    fn dyn_array_hash(&self, state: &mut dyn Hasher, accuracy: crate::EqMode);
165
166    /// Compares two arrays of the same concrete type for equality.
167    fn dyn_array_eq(&self, other: &ArrayRef, accuracy: crate::EqMode) -> bool;
168
169    /// Returns a new array with the given slots.
170    fn with_slots(&self, this: &ArrayRef, slots: ArraySlots) -> VortexResult<ArrayRef>;
171
172    /// Returns a new array with the given buffers.
173    fn with_buffers(&self, this: &ArrayRef, buffers: Vec<BufferHandle>) -> VortexResult<ArrayRef>;
174
175    /// Returns a new array with the given slots, bypassing encoding-level validation.
176    ///
177    /// Used by the executor to temporarily carry an array that has had one of its child slots
178    /// taken out (leaving `None`) without panicking `V::validate`. The caller must ensure the
179    /// missing slot is filled back in (via `put_slot_unchecked`) or driven to completion by the
180    /// builder path before the array becomes externally observable.
181    ///
182    /// # Safety
183    ///
184    /// The array returned may have slots whose content does not match the encoding's normal
185    /// invariants. Callers must re-establish those invariants before handing the array to
186    /// anything outside the executor.
187    unsafe fn with_slots_unchecked(&self, this: &ArrayRef, slots: ArraySlots) -> ArrayRef;
188
189    /// Attempt to reduce the array to a simpler representation.
190    fn reduce(&self, this: &ArrayRef) -> VortexResult<Option<ArrayRef>>;
191
192    /// Attempt to reduce the parent of this array.
193    fn reduce_parent(
194        &self,
195        this: &ArrayRef,
196        parent: &ArrayRef,
197        child_idx: usize,
198    ) -> VortexResult<Option<ArrayRef>>;
199
200    /// Execute the array by taking a single encoding-specific execution step.
201    ///
202    /// This is the checked entry point. If the encoding reports
203    /// [`ExecutionStep::Done`](ExecutionStep::Done), implementations must validate that the
204    /// returned array preserves this array's logical `len` and `dtype`, and must transfer this
205    /// array's statistics to the returned array.
206    fn execute(&self, this: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult>;
207
208    /// Execute the array by taking a single encoding-specific execution step without applying
209    /// `Done`-result postconditions.
210    ///
211    /// This exists for the iterative executor, which may call into `execute` on suspended
212    /// executor-private arrays whose slots temporarily contain `None`. In that mode the executor
213    /// itself is responsible for deciding when a `Done` result represents a real logical array,
214    /// enforcing any `len`/`dtype` invariants, and transferring statistics.
215    ///
216    /// # Safety
217    /// The `array` returned should have it's `DType` and len checked
218    /// (optionally it should have its stats propagated from `this`).
219    unsafe fn execute_unchecked(
220        &self,
221        this: ArrayRef,
222        ctx: &mut ExecutionCtx,
223    ) -> VortexResult<ExecutionResult>;
224
225    /// Read the scalar at `index`, including its nullness, retaining nothing.
226    ///
227    /// Caller must guarantee `index < len`.
228    ///
229    /// Kept apart from [`Self::probe_scalar_retained`] so this entry is a bare trampoline into
230    /// the encoding: sharing one function made the one-off path pay the retained branch's
231    /// register saves before its tail call.
232    fn probe_scalar_once(
233        &self,
234        this: &ArrayRef,
235        index: usize,
236        ctx: &mut ExecutionCtx,
237    ) -> VortexResult<Scalar>;
238
239    /// Read the scalar at `index`, including its nullness, keeping preparation in `state`.
240    ///
241    /// Caller must guarantee `index < len`.
242    fn probe_scalar_retained(
243        &self,
244        this: &ArrayRef,
245        index: usize,
246        state: &mut Option<Box<dyn Any>>,
247        ctx: &mut ExecutionCtx,
248    ) -> VortexResult<Scalar>;
249
250    /// Whether the row at `index` is valid, using the validity kept in `state`.
251    ///
252    /// Caller must guarantee `index < len`.
253    fn probe_is_valid_retained(
254        &self,
255        this: &ArrayRef,
256        index: usize,
257        state: &mut Option<Box<dyn Any>>,
258        ctx: &mut ExecutionCtx,
259    ) -> VortexResult<bool>;
260}
261
262/// Trait for converting a type into a Vortex [`ArrayRef`].
263pub trait IntoArray {
264    /// Convert this value into the erased array handle used by generic APIs.
265    fn into_array(self) -> ArrayRef;
266}
267
268mod private {
269    use super::*;
270
271    pub trait Sealed {}
272
273    impl<V: VTable> Sealed for ArrayData<V> {}
274}
275
276// =============================================================================
277// New path: DynArrayData and supporting trait impls for ArrayData<V>
278// =============================================================================
279
280/// DynArrayData implementation for [`ArrayData<V>`].
281///
282/// This is self-contained: identity methods use `ArrayData<V>`'s own fields (dtype, len, stats),
283/// while data-access methods delegate to VTable methods on the inner `V::TypedArrayData`.
284impl<V: VTable> DynArrayData for ArrayData<V> {
285    fn as_any(&self) -> &dyn Any {
286        self
287    }
288
289    fn as_any_mut(&mut self) -> &mut dyn Any {
290        self
291    }
292
293    fn validity(&self, this: &ArrayRef) -> VortexResult<Validity> {
294        if this.dtype().is_nullable() {
295            let view = unsafe { ArrayView::new_unchecked(this, &self.data) };
296            let validity = <V::ValidityVTable as ValidityVTable<V>>::validity(view)?;
297            if let Validity::Array(array) = &validity {
298                vortex_ensure!(array.len() == this.len(), "Validity array length mismatch");
299                vortex_ensure!(
300                    matches!(array.dtype(), DType::Bool(Nullability::NonNullable)),
301                    "Validity array is not non-nullable boolean: {}",
302                    this.encoding_id(),
303                );
304            }
305            Ok(validity)
306        } else {
307            Ok(Validity::NonNullable)
308        }
309    }
310
311    fn append_to_builder(
312        &self,
313        this: &ArrayRef,
314        builder: &mut dyn ArrayBuilder,
315        ctx: &mut ExecutionCtx,
316    ) -> VortexResult<()> {
317        if builder.dtype() != this.dtype() {
318            vortex_panic!(
319                "Builder dtype mismatch: expected {}, got {}",
320                this.dtype(),
321                builder.dtype(),
322            );
323        }
324        let len = builder.len();
325
326        let view = unsafe { ArrayView::new_unchecked(this, &self.data) };
327        V::append_to_builder(view, builder, ctx)?;
328
329        assert_eq!(
330            len + this.len(),
331            builder.len(),
332            "Builder length mismatch after writing array for encoding {}",
333            this.encoding_id(),
334        );
335        Ok(())
336    }
337
338    fn buffers(&self, this: &ArrayRef) -> Vec<ByteBuffer> {
339        let view = unsafe { ArrayView::new_unchecked(this, &self.data) };
340        (0..V::nbuffers(view))
341            .map(|i| V::buffer(view, i).to_host_sync())
342            .collect()
343    }
344
345    fn buffer_handles(&self, this: &ArrayRef) -> Vec<BufferHandle> {
346        let view = unsafe { ArrayView::new_unchecked(this, &self.data) };
347        (0..V::nbuffers(view)).map(|i| V::buffer(view, i)).collect()
348    }
349
350    fn buffer_names(&self, this: &ArrayRef) -> Vec<String> {
351        let view = unsafe { ArrayView::new_unchecked(this, &self.data) };
352        (0..V::nbuffers(view))
353            .filter_map(|i| V::buffer_name(view, i))
354            .collect()
355    }
356
357    fn named_buffers(&self, this: &ArrayRef) -> Vec<(String, BufferHandle)> {
358        let view = unsafe { ArrayView::new_unchecked(this, &self.data) };
359        (0..V::nbuffers(view))
360            .filter_map(|i| V::buffer_name(view, i).map(|name| (name, V::buffer(view, i))))
361            .collect()
362    }
363
364    fn nbuffers(&self, this: &ArrayRef) -> usize {
365        let view = unsafe { ArrayView::new_unchecked(this, &self.data) };
366        V::nbuffers(view)
367    }
368
369    fn slot_name(&self, this: &ArrayRef, idx: usize) -> String {
370        let view = unsafe { ArrayView::new_unchecked(this, &self.data) };
371        V::slot_name(view, idx)
372    }
373
374    fn metadata_fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
375        std::fmt::Display::fmt(&self.data, f)
376    }
377
378    fn dyn_array_hash(&self, state: &mut dyn Hasher, accuracy: crate::EqMode) {
379        let mut wrapper = HasherWrapper(state);
380        // Note: metadata (len, dtype, encoding_id) and slots are hashed by ArrayRef.
381        self.data.array_hash(&mut wrapper, accuracy);
382    }
383
384    fn dyn_array_eq(&self, other: &ArrayRef, accuracy: crate::EqMode) -> bool {
385        // Note: metadata (len, dtype, encoding_id) and slots are compared by ArrayRef.
386        other
387            .dyn_array()
388            .as_any()
389            .downcast_ref::<Self>()
390            .is_some_and(|other_inner| self.data.array_eq(&other_inner.data, accuracy))
391    }
392
393    fn with_slots(&self, this: &ArrayRef, slots: ArraySlots) -> VortexResult<ArrayRef> {
394        let stats = this.statistics().to_owned();
395        Ok(Array::<V>::try_from_parts(
396            ArrayParts::new(
397                self.vtable.clone(),
398                this.dtype().clone(),
399                this.len(),
400                self.data.clone(),
401            )
402            .with_slots(slots),
403        )?
404        .with_stats_set(stats)
405        .into_array())
406    }
407
408    fn with_buffers(&self, this: &ArrayRef, buffers: Vec<BufferHandle>) -> VortexResult<ArrayRef> {
409        let view = unsafe { ArrayView::new_unchecked(this, &self.data) };
410        let stats = this.statistics().to_owned();
411        Ok(
412            Array::<V>::try_from_parts(V::with_buffers(&self.vtable, view, &buffers)?)?
413                .with_stats_set(stats)
414                .into_array(),
415        )
416    }
417
418    unsafe fn with_slots_unchecked(&self, this: &ArrayRef, slots: ArraySlots) -> ArrayRef {
419        // SAFETY: we intentionally skip `V::validate` here. Caller guarantees that the resulting
420        // array is either repaired or not externally observed.
421        let store = unsafe {
422            ArrayInner::<ArrayData<V>>::new_unchecked(
423                self.vtable.clone(),
424                this.len(),
425                this.dtype().clone(),
426                self.data.clone(),
427                slots,
428                this.statistics().to_array_stats(),
429            )
430        };
431        ArrayRef::from_inner(Arc::new(store))
432    }
433
434    fn reduce(&self, this: &ArrayRef) -> VortexResult<Option<ArrayRef>> {
435        let view = unsafe { ArrayView::new_unchecked(this, &self.data) };
436        let Some(reduced) = V::reduce(view)? else {
437            return Ok(None);
438        };
439        vortex_ensure!(
440            reduced.len() == this.len(),
441            "Reduced array length mismatch from {} to {}",
442            this.encoding_id(),
443            reduced.encoding_id()
444        );
445        // Rules may reduce to a non-nullable constant
446        vortex_ensure!(
447            reduced.dtype() == this.dtype()
448                || (reduced.dtype().eq_ignore_nullability(this.dtype())
449                    && !reduced.dtype().is_nullable()),
450            "Reduced array dtype mismatch from {} ({}) to {} ({})",
451            this.encoding_id(),
452            this.dtype(),
453            reduced.encoding_id(),
454            reduced.dtype()
455        );
456        Ok(Some(reduced))
457    }
458
459    fn reduce_parent(
460        &self,
461        this: &ArrayRef,
462        parent: &ArrayRef,
463        child_idx: usize,
464    ) -> VortexResult<Option<ArrayRef>> {
465        let view = unsafe { ArrayView::new_unchecked(this, &self.data) };
466        let Some(reduced) = V::reduce_parent(view, parent, child_idx)? else {
467            return Ok(None);
468        };
469
470        vortex_ensure!(
471            reduced.len() == parent.len(),
472            "Reduced array length mismatch from {} to {}",
473            parent.encoding_id(),
474            reduced.encoding_id()
475        );
476        vortex_ensure!(
477            reduced.dtype() == parent.dtype(),
478            "Reduced array dtype mismatch from {} to {}",
479            parent.encoding_id(),
480            reduced.encoding_id()
481        );
482
483        Ok(Some(reduced))
484    }
485
486    fn execute(&self, this: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
487        let len = this.len();
488        let dtype = this.dtype().clone();
489        let stats = this.statistics().to_array_stats();
490        let result = unsafe { self.execute_unchecked(this, ctx)? };
491
492        if matches!(result.step(), ExecutionStep::Done) {
493            if cfg!(debug_assertions) {
494                vortex_ensure!(
495                    result.array().len() == len,
496                    "Result length mismatch for {:?}",
497                    self.vtable
498                );
499                vortex_ensure!(
500                    result.array().dtype() == &dtype,
501                    "Executed canonical dtype mismatch for {:?}",
502                    self.vtable
503                );
504            }
505
506            result
507                .array()
508                .statistics()
509                .set_iter(crate::stats::StatsSet::from(stats).into_iter());
510        }
511
512        Ok(result)
513    }
514
515    unsafe fn execute_unchecked(
516        &self,
517        this: ArrayRef,
518        ctx: &mut ExecutionCtx,
519    ) -> VortexResult<ExecutionResult> {
520        let typed = Array::<V>::try_from_array_ref(this)
521            .map_err(|_| vortex_err!("Failed to downcast array for execute"))
522            .vortex_expect("Failed to downcast array for execute");
523        V::execute(typed, ctx)
524    }
525
526    fn probe_scalar_once(
527        &self,
528        this: &ArrayRef,
529        index: usize,
530        ctx: &mut ExecutionCtx,
531    ) -> VortexResult<Scalar> {
532        // SAFETY: this adapter belongs to the ArrayData<V> stored in `this`.
533        let view = unsafe { ArrayView::new_unchecked(this, &self.data) };
534        <V::OperationsVTable as OperationsVTable<V>>::probe_scalar(
535            &mut ProbeState::once(view),
536            index,
537            ctx,
538        )
539    }
540
541    fn probe_scalar_retained(
542        &self,
543        this: &ArrayRef,
544        index: usize,
545        state: &mut Option<Box<dyn Any>>,
546        ctx: &mut ExecutionCtx,
547    ) -> VortexResult<Scalar> {
548        // SAFETY: this adapter belongs to the ArrayData<V> stored in `this`.
549        let view = unsafe { ArrayView::new_unchecked(this, &self.data) };
550        let mut state = ProbeState::repeated(view, repeated_state(state)?);
551        <V::OperationsVTable as OperationsVTable<V>>::probe_scalar(&mut state, index, ctx)
552    }
553
554    fn probe_is_valid_retained(
555        &self,
556        this: &ArrayRef,
557        index: usize,
558        state: &mut Option<Box<dyn Any>>,
559        ctx: &mut ExecutionCtx,
560    ) -> VortexResult<bool> {
561        repeated_state::<EncodingProbeState<V>>(state)?.is_valid(this, index, ctx)
562    }
563}
564
565/// Wrapper around `&mut dyn Hasher` that implements `Hasher` (and is `Sized`).
566struct HasherWrapper<'a>(&'a mut dyn Hasher);
567
568impl Hasher for HasherWrapper<'_> {
569    fn finish(&self) -> u64 {
570        self.0.finish()
571    }
572
573    fn write(&mut self, bytes: &[u8]) {
574        self.0.write(bytes);
575    }
576}
577
578/// ArrayId is a globally unique name for the array's vtable.
579pub type ArrayId = Id;
vortex_array/array/mod.rs

vortex_array/array/
mod.rs
