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

vortex_array/array/
mod.rs
