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

vortex_array/array/
mod.rs
