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

vortex_array/array/
mod.rs
