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

vortex_array/array/
mod.rs
