vortex_array/array/
erased.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::any::type_name;
5use std::fmt::Debug;
6use std::fmt::Formatter;
7use std::hash::Hash;
8use std::hash::Hasher;
9use std::ops::Range;
10use std::sync::Arc;
11
12use vortex_buffer::ByteBuffer;
13use vortex_error::VortexExpect;
14use vortex_error::VortexResult;
15use vortex_error::vortex_ensure;
16use vortex_error::vortex_err;
17use vortex_error::vortex_panic;
18use vortex_mask::Mask;
19
20use crate::AnyCanonical;
21use crate::Array;
22use crate::ArrayEq;
23use crate::ArrayHash;
24use crate::ArrayView;
25use crate::Canonical;
26use crate::ExecutionCtx;
27use crate::ExecutionResult;
28use crate::IntoArray;
29use crate::VTable;
30use crate::VortexSessionExecute;
31use crate::aggregate_fn::fns::sum::sum;
32use crate::array::ArrayData;
33use crate::array::ArrayId;
34use crate::array::ArrayInner;
35use crate::array::ArraySlots;
36use crate::array::DynArrayData;
37use crate::arrays::Constant;
38use crate::arrays::DictArray;
39use crate::arrays::FilterArray;
40use crate::arrays::SliceArray;
41use crate::buffer::BufferHandle;
42use crate::builders::ArrayBuilder;
43use crate::dtype::DType;
44use crate::expr::stats::Precision;
45use crate::expr::stats::Stat;
46use crate::expr::stats::StatsProviderExt;
47use crate::legacy_session;
48use crate::matcher::Matcher;
49use crate::optimizer::ArrayOptimizer;
50use crate::scalar::Scalar;
51use crate::scalar::ScalarValue;
52use crate::stats::StatsSetRef;
53use crate::validity::Validity;
54
55/// A depth-first pre-order iterator over an Array.
56pub struct DepthFirstArrayIterator {
57    stack: Vec<ArrayRef>,
58}
59
60impl Iterator for DepthFirstArrayIterator {
61    type Item = ArrayRef;
62
63    fn next(&mut self) -> Option<Self::Item> {
64        let next = self.stack.pop()?;
65        for child in next.children().into_iter().rev() {
66            self.stack.push(child);
67        }
68        Some(next)
69    }
70}
71
72/// A reference-counted pointer to a type-erased array.
73///
74/// Wraps `Arc<ArrayInner<dyn DynArrayData>>` — a single 16-byte fat pointer.
75/// Metadata (`len`, `dtype`, `encoding_id`) lives in `ArrayInner::meta` and is
76/// accessed as a normal struct field read — no vtable dispatch, no extra allocation.
77#[derive(Clone)]
78pub struct ArrayRef(Arc<ArrayInner<dyn DynArrayData>>);
79
80impl ArrayRef {
81    /// Create from an `Arc<ArrayInner<dyn DynArrayData>>`.
82    pub(crate) fn from_inner<D: DynArrayData>(inner: Arc<ArrayInner<D>>) -> Self {
83        Self(inner)
84    }
85
86    /// Returns a reference to the `dyn DynArrayData` inside the inner.
87    #[inline(always)]
88    pub(crate) fn dyn_array(&self) -> &dyn DynArrayData {
89        &self.0.data
90    }
91
92    /// Returns a mutable reference to the inner if this is the sole owner.
93    #[inline(always)]
94    pub(crate) fn inner_mut(&mut self) -> Option<&mut ArrayInner<dyn DynArrayData>> {
95        Arc::get_mut(&mut self.0)
96    }
97
98    /// Returns the Arc::as_ptr().addr() of the underlying array.
99    /// This function is used in a couple of places, and we should migrate them to using array_eq.
100    #[doc(hidden)]
101    pub fn addr(&self) -> usize {
102        Arc::as_ptr(&self.0).addr()
103    }
104
105    /// Downcast the inner to a concrete `ArrayInner<ArrayData<V>>`.
106    ///
107    /// Uses the same raw-pointer technique as `Arc::downcast`.
108    #[allow(dead_code)]
109    pub(crate) fn downcast_inner<V: VTable>(self) -> Result<Arc<ArrayInner<ArrayData<V>>>, Self> {
110        // TODO(joe): can we use encoding id here?
111        if self.0.data.as_any().is::<ArrayData<V>>() {
112            Ok(unsafe { self.downcast_inner_unchecked() })
113        } else {
114            Err(self)
115        }
116    }
117
118    /// Downcast without a runtime type check.
119    ///
120    /// # Safety
121    /// The caller must guarantee the concrete type behind `dyn DynArrayData` is `ArrayData<V>`.
122    #[inline(always)]
123    pub(crate) unsafe fn downcast_inner_unchecked<V: VTable>(
124        self,
125    ) -> Arc<ArrayInner<ArrayData<V>>> {
126        debug_assert!(self.0.data.as_any().is::<ArrayData<V>>());
127        // Recover the original concrete Arc. The fat pointer's data pointer is the
128        // same allocation that was originally `Arc<ArrayInner<ArrayData<V>>>` before
129        // unsized coercion to `Arc<ArrayInner<dyn DynArrayData>>`.
130        let raw = Arc::into_raw(self.0);
131        // # Safety all arrays are constructed in this way and type aliased.
132        unsafe { Arc::from_raw(raw.cast::<ArrayInner<ArrayData<V>>>()) }
133    }
134
135    /// Returns true if the two ArrayRefs point to the same allocation.
136    pub fn ptr_eq(this: &ArrayRef, other: &ArrayRef) -> bool {
137        Arc::ptr_eq(&this.0, &other.0)
138    }
139}
140
141impl Debug for ArrayRef {
142    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
143        f.debug_struct("Array")
144            .field("encoding", &self.0.encoding_id)
145            .field("dtype", &self.0.dtype)
146            .field("len", &self.0.len)
147            .field("data", &self.0.data)
148            .finish()
149    }
150}
151
152impl ArrayHash for ArrayRef {
153    fn array_hash<H: Hasher>(&self, state: &mut H, accuracy: crate::EqMode) {
154        self.0.len.hash(state);
155        self.0.dtype.hash(state);
156        self.0.encoding_id.hash(state);
157        self.0.slots.len().hash(state);
158        for slot in &self.0.slots {
159            slot.array_hash(state, accuracy);
160        }
161        self.0
162            .data
163            .dyn_array_hash(state as &mut dyn Hasher, accuracy);
164    }
165}
166
167impl ArrayEq for ArrayRef {
168    fn array_eq(&self, other: &Self, accuracy: crate::EqMode) -> bool {
169        self.0.len == other.0.len
170            && self.0.dtype == other.0.dtype
171            && self.0.encoding_id == other.0.encoding_id
172            && self.0.slots.len() == other.0.slots.len()
173            && self
174                .0
175                .slots
176                .iter()
177                .zip(other.0.slots.iter())
178                .all(|(slot, other_slot)| slot.array_eq(other_slot, accuracy))
179            && self.0.data.dyn_array_eq(other, accuracy)
180    }
181}
182impl ArrayRef {
183    /// Returns the length of the array.
184    #[inline]
185    pub fn len(&self) -> usize {
186        self.0.len
187    }
188
189    /// Returns whether the array is empty (has zero rows).
190    #[inline]
191    pub fn is_empty(&self) -> bool {
192        self.0.len == 0
193    }
194
195    /// Returns the logical Vortex [`DType`] of the array.
196    #[inline]
197    pub fn dtype(&self) -> &DType {
198        &self.0.dtype
199    }
200
201    /// Returns the encoding ID of the array.
202    #[inline]
203    pub fn encoding_id(&self) -> ArrayId {
204        self.0.encoding_id
205    }
206
207    /// Performs a constant-time slice of the array.
208    pub fn slice(&self, range: Range<usize>) -> VortexResult<ArrayRef> {
209        let len = self.len();
210        let start = range.start;
211        let stop = range.end;
212
213        if start == 0 && stop == len {
214            return Ok(self.clone());
215        }
216
217        vortex_ensure!(start <= len, "OutOfBounds: start {start} > length {}", len);
218        vortex_ensure!(stop <= len, "OutOfBounds: stop {stop} > length {}", len);
219
220        vortex_ensure!(start <= stop, "start ({start}) must be <= stop ({stop})");
221
222        if start == stop {
223            return Ok(Canonical::empty(self.dtype()).into_array());
224        }
225
226        let sliced = SliceArray::try_new(self.clone(), range)?
227            .into_array()
228            .optimize()?;
229
230        // Propagate some stats from the original array to the sliced array.
231        if !sliced.is::<Constant>() {
232            self.statistics().with_iter(|iter| {
233                sliced.statistics().inherit(iter.filter(|(stat, value)| {
234                    matches!(
235                        stat,
236                        Stat::IsConstant | Stat::IsSorted | Stat::IsStrictSorted
237                    ) && value
238                        .as_ref()
239                        .as_exact()
240                        .is_some_and(|v| matches!(v, ScalarValue::Bool(true)))
241                }));
242            });
243        }
244
245        Ok(sliced)
246    }
247
248    /// Wraps the array in a [`FilterArray`] such that it is logically filtered by the given mask.
249    pub fn filter(&self, mask: Mask) -> VortexResult<ArrayRef> {
250        FilterArray::try_new(self.clone(), mask)?
251            .into_array()
252            .optimize()
253    }
254
255    /// Wraps the array in a [`DictArray`] such that it is logically taken by the given indices.
256    pub fn take(&self, indices: ArrayRef) -> VortexResult<ArrayRef> {
257        DictArray::try_new(indices, self.clone())?
258            .into_array()
259            .optimize()
260    }
261
262    /// Fetch the scalar at the given index.
263    #[deprecated(
264        note = "Use `execute_scalar` instead, which allows passing an execution context for more \
265        efficient execution when fetching multiple scalars from the same array."
266    )]
267    #[allow(clippy::disallowed_methods)]
268    pub fn scalar_at(&self, index: usize) -> VortexResult<Scalar> {
269        self.execute_scalar(index, &mut legacy_session().create_execution_ctx())
270    }
271
272    /// Execute the array to extract a scalar at the given index.
273    pub fn execute_scalar(&self, index: usize, ctx: &mut ExecutionCtx) -> VortexResult<Scalar> {
274        vortex_ensure!(index < self.len(), OutOfBounds: index, 0, self.len());
275        if self.dtype().is_nullable() && self.is_invalid(index, ctx)? {
276            return Ok(Scalar::null(self.dtype().clone()));
277        }
278        let scalar = self.0.data.execute_scalar(self, index, ctx)?;
279        debug_assert_eq!(self.dtype(), scalar.dtype(), "Scalar dtype mismatch");
280        Ok(scalar)
281    }
282
283    /// Returns whether the item at `index` is valid.
284    pub fn is_valid(&self, index: usize, ctx: &mut ExecutionCtx) -> VortexResult<bool> {
285        vortex_ensure!(index < self.len(), OutOfBounds: index, 0, self.len());
286        match self.validity()? {
287            Validity::NonNullable | Validity::AllValid => Ok(true),
288            Validity::AllInvalid => Ok(false),
289            Validity::Array(a) => a
290                .execute_scalar(index, ctx)?
291                .as_bool()
292                .value()
293                .ok_or_else(|| vortex_err!("validity value at index {} is null", index)),
294        }
295    }
296
297    /// Returns whether the item at `index` is invalid.
298    pub fn is_invalid(&self, index: usize, ctx: &mut ExecutionCtx) -> VortexResult<bool> {
299        Ok(!self.is_valid(index, ctx)?)
300    }
301
302    /// Returns whether all items in the array are valid.
303    pub fn all_valid(&self, ctx: &mut ExecutionCtx) -> VortexResult<bool> {
304        if self.is_empty() {
305            return Ok(true);
306        }
307
308        match self.validity()? {
309            Validity::NonNullable | Validity::AllValid => Ok(true),
310            Validity::AllInvalid => Ok(false),
311            Validity::Array(a) => Ok(a.statistics().compute_min::<bool>(ctx).unwrap_or(false)),
312        }
313    }
314
315    /// Returns whether the array is all invalid.
316    pub fn all_invalid(&self, ctx: &mut ExecutionCtx) -> VortexResult<bool> {
317        if self.is_empty() {
318            return Ok(true);
319        }
320
321        match self.validity()? {
322            Validity::NonNullable | Validity::AllValid => Ok(false),
323            Validity::AllInvalid => Ok(true),
324            Validity::Array(a) => Ok(!a.statistics().compute_max::<bool>(ctx).unwrap_or(true)),
325        }
326    }
327
328    /// Returns the number of valid elements in the array.
329    pub fn valid_count(&self, ctx: &mut ExecutionCtx) -> VortexResult<usize> {
330        let len = self.len();
331        if let Precision::Exact(invalid_count) = self.statistics().get_as::<usize>(Stat::NullCount)
332        {
333            return Ok(len - invalid_count);
334        }
335
336        let count = match self.validity()? {
337            Validity::NonNullable | Validity::AllValid => len,
338            Validity::AllInvalid => 0,
339            Validity::Array(a) => {
340                let array_sum = sum(&a, ctx)?;
341                array_sum
342                    .as_primitive()
343                    .as_::<usize>()
344                    .ok_or_else(|| vortex_err!("sum of validity array is null"))?
345            }
346        };
347        vortex_ensure!(count <= len, "Valid count exceeds array length");
348
349        self.statistics()
350            .set(Stat::NullCount, Precision::exact(len - count));
351
352        Ok(count)
353    }
354
355    /// Returns the number of invalid elements in the array.
356    pub fn invalid_count(&self, ctx: &mut ExecutionCtx) -> VortexResult<usize> {
357        Ok(self.len() - self.valid_count(ctx)?)
358    }
359
360    /// Returns the [`Validity`] of the array.
361    pub fn validity(&self) -> VortexResult<Validity> {
362        self.0.data.validity(self)
363    }
364
365    /// Returns the canonical representation of the array.
366    #[deprecated(note = "use `array.execute::<Canonical>(ctx)` instead")]
367    #[allow(clippy::disallowed_methods)]
368    pub fn into_canonical(self) -> VortexResult<Canonical> {
369        self.execute(&mut legacy_session().create_execution_ctx())
370    }
371
372    /// Returns the canonical representation of the array.
373    #[deprecated(note = "use `array.execute::<Canonical>(ctx)` instead")]
374    pub fn to_canonical(&self) -> VortexResult<Canonical> {
375        #[expect(deprecated)]
376        let result = self.clone().into_canonical();
377        result
378    }
379
380    /// Writes the array into the canonical builder.
381    pub fn append_to_builder(
382        &self,
383        builder: &mut dyn ArrayBuilder,
384        ctx: &mut ExecutionCtx,
385    ) -> VortexResult<()> {
386        self.0.data.append_to_builder(self, builder, ctx)
387    }
388
389    /// Returns the statistics of the array.
390    pub fn statistics(&self) -> StatsSetRef<'_> {
391        self.0.stats.to_ref(self)
392    }
393
394    /// Does the array match the given matcher.
395    #[inline]
396    pub fn is<M: Matcher>(&self) -> bool {
397        M::matches(self)
398    }
399
400    /// Returns the array downcast by the given matcher.
401    #[inline]
402    pub fn as_<M: Matcher>(&self) -> M::Match<'_> {
403        self.as_opt::<M>().vortex_expect("Failed to downcast")
404    }
405
406    /// Returns the array downcast by the given matcher.
407    #[inline]
408    pub fn as_opt<M: Matcher>(&self) -> Option<M::Match<'_>> {
409        M::try_match(self)
410    }
411
412    /// Returns the array downcast to the given `Array<V>` as an owned typed handle.
413    pub fn try_downcast<V: VTable>(self) -> Result<Array<V>, ArrayRef> {
414        Array::<V>::try_from_array_ref(self)
415    }
416
417    /// Returns the array downcast to the given `Array<V>` as an owned typed handle.
418    ///
419    /// # Panics
420    ///
421    /// Panics if the array is not of the given type.
422    pub fn downcast<V: VTable>(self) -> Array<V> {
423        Self::try_downcast(self)
424            .unwrap_or_else(|_| vortex_panic!("Failed to downcast to {}", type_name::<V>()))
425    }
426
427    /// Returns a reference to the typed `ArrayData<V>` if this array matches the given vtable type.
428    pub fn as_typed<V: VTable>(&self) -> Option<ArrayView<'_, V>> {
429        let inner = self.0.data.as_any().downcast_ref::<ArrayData<V>>()?;
430        Some(unsafe { ArrayView::new_unchecked(self, &inner.data) })
431    }
432
433    /// Returns the constant scalar if this is a constant array.
434    pub fn as_constant(&self) -> Option<Scalar> {
435        self.as_opt::<Constant>().map(|a| a.scalar().clone())
436    }
437
438    /// Total size of the array in bytes, including all children and buffers.
439    pub fn nbytes(&self) -> u64 {
440        let mut nbytes = 0;
441        for array in self.depth_first_traversal() {
442            for buffer in array.buffers() {
443                nbytes += buffer.len() as u64;
444            }
445        }
446        nbytes
447    }
448
449    /// Whether the array is of a canonical encoding.
450    pub fn is_canonical(&self) -> bool {
451        self.is::<AnyCanonical>()
452    }
453
454    /// Returns a new array with the slot at `slot_idx` replaced by `replacement`.
455    ///
456    /// This is only valid for physical rewrites: the replacement must have the same logical
457    /// `DType` and `len` as the existing slot.
458    ///
459    /// # Safety
460    ///
461    /// If this returns `Ok`, the caller must guarantee that the replacement slot represents the
462    /// same logical values as the original slot. Only the physical representation may change.
463    /// Existing parent statistics are preserved and must remain valid.
464    ///
465    /// Takes ownership to allow in-place mutation when the refcount is 1.
466    pub unsafe fn with_slot(
467        self,
468        slot_idx: usize,
469        replacement: ArrayRef,
470    ) -> VortexResult<ArrayRef> {
471        let mut slots: ArraySlots = self.slots().iter().cloned().collect();
472        let nslots = slots.len();
473        vortex_ensure!(
474            slot_idx < nslots,
475            "slot index {} out of bounds for array with {} slots",
476            slot_idx,
477            nslots
478        );
479        let existing = slots[slot_idx]
480            .as_ref()
481            .vortex_expect("with_slot cannot replace an absent slot");
482        vortex_ensure!(
483            existing.dtype() == replacement.dtype(),
484            "slot {} dtype changed from {} to {} during physical rewrite",
485            slot_idx,
486            existing.dtype(),
487            replacement.dtype()
488        );
489        vortex_ensure!(
490            existing.len() == replacement.len(),
491            "slot {} len changed from {} to {} during physical rewrite",
492            slot_idx,
493            existing.len(),
494            replacement.len()
495        );
496        slots[slot_idx] = Some(replacement);
497        // SAFETY: upheld by the caller of this unsafe API.
498        unsafe { self.with_slots(slots) }
499    }
500
501    /// Take a slot for executor-owned physical rewrites.
502    ///
503    /// On return the produced parent has the taken slot set to `None`
504    /// callers must put the slot back (typically via [`Self::put_slot_unchecked`]) before the parent is
505    /// returned from the execution loop.
506    ///
507    /// When the `Arc` was shared this allocates a fresh parent.
508    ///
509    /// # Safety
510    /// The caller must put back a slot with the same logical dtype and length before exposing the
511    /// parent array, and must only use this for physical rewrites.
512    pub(crate) unsafe fn take_slot_unchecked(
513        mut self,
514        slot_idx: usize,
515    ) -> VortexResult<(ArrayRef, ArrayRef)> {
516        if let Some(inner) = Arc::get_mut(&mut self.0) {
517            let child = inner.slots[slot_idx]
518                .take()
519                .vortex_expect("take_slot_unchecked cannot take an absent slot");
520            return Ok((self, child));
521        }
522
523        // Arc is shared: clone the child out and build a fresh parent with slot_idx = None,
524        // bypassing encoding-level validation so the absent slot does not panic `V::validate`.
525        let child = self.slots()[slot_idx]
526            .as_ref()
527            .vortex_expect("take_slot_unchecked cannot take an absent slot")
528            .clone();
529
530        let mut new_slots: ArraySlots = self.slots().iter().cloned().collect();
531        new_slots[slot_idx] = None;
532
533        // SAFETY: ensured by the caller — the None slot is either put back or driven to completion
534        // via the builder path before the parent escapes the executor.
535        let new_parent = unsafe { self.0.data.with_slots_unchecked(&self, new_slots) };
536        Ok((new_parent, child))
537    }
538
539    /// Puts an array into `slot_idx` by either, cloning the inner array if the Arc is not exclusive
540    /// or replacing the slot in this `ArrayRef`.
541    /// This is the mirror of [`Self::take_slot_unchecked`].
542    ///
543    /// # Safety
544    /// The replacement must have the same logical dtype and length as the taken slot, and this
545    /// must only be used for physical rewrites.
546    pub(crate) unsafe fn put_slot_unchecked(
547        mut self,
548        slot_idx: usize,
549        replacement: ArrayRef,
550    ) -> VortexResult<ArrayRef> {
551        if let Some(inner) = Arc::get_mut(&mut self.0) {
552            inner.slots[slot_idx] = Some(replacement);
553            return Ok(self);
554        }
555
556        let mut slots: ArraySlots = self.slots().iter().cloned().collect();
557        slots[slot_idx] = Some(replacement);
558        self.0.data.with_slots(&self, slots)
559    }
560
561    /// Returns a new array with the provided slots.
562    ///
563    /// This is only valid for physical rewrites: slot count, presence, logical `DType`, and
564    /// logical `len` must remain unchanged.
565    ///
566    /// # Safety
567    ///
568    /// If this returns `Ok`, the caller must guarantee that each replacement slot represents the
569    /// same logical values as the original slot. Only physical representation may change. Existing
570    /// parent statistics are preserved and must remain valid.
571    pub unsafe fn with_slots(self, slots: ArraySlots) -> VortexResult<ArrayRef> {
572        let old_slots = self.slots();
573        vortex_ensure!(
574            old_slots.len() == slots.len(),
575            "slot count changed from {} to {} during physical rewrite",
576            old_slots.len(),
577            slots.len()
578        );
579        for (idx, (old_slot, new_slot)) in old_slots.iter().zip(slots.iter()).enumerate() {
580            vortex_ensure!(
581                old_slot.is_some() == new_slot.is_some(),
582                "slot {} presence changed during physical rewrite",
583                idx
584            );
585            if let (Some(old_slot), Some(new_slot)) = (old_slot.as_ref(), new_slot.as_ref()) {
586                vortex_ensure!(
587                    old_slot.dtype() == new_slot.dtype(),
588                    "slot {} dtype changed from {} to {} during physical rewrite",
589                    idx,
590                    old_slot.dtype(),
591                    new_slot.dtype()
592                );
593                vortex_ensure!(
594                    old_slot.len() == new_slot.len(),
595                    "slot {} len changed from {} to {} during physical rewrite",
596                    idx,
597                    old_slot.len(),
598                    new_slot.len()
599                );
600            }
601        }
602        self.0.data.with_slots(&self, slots)
603    }
604
605    /// Returns a new array with the provided top-level buffer handles.
606    ///
607    /// This is only valid for physical rewrites: buffer count, logical `DType`, logical `len`, and
608    /// child slots must remain unchanged. Encoding-specific validation checks buffer shape,
609    /// alignment, and metadata consistency.
610    ///
611    /// # Safety
612    ///
613    /// If this returns `Ok`, the caller must guarantee that the replacement buffers represent the
614    /// same logical values as the original buffers. Only the buffer handle implementation,
615    /// placement, or backing storage may change. Existing statistics are preserved and must remain
616    /// valid.
617    pub unsafe fn with_buffers(
618        self,
619        buffers: impl IntoIterator<Item = BufferHandle>,
620    ) -> VortexResult<ArrayRef> {
621        let buffers = buffers.into_iter().collect::<Vec<_>>();
622        let nbuffers = self.nbuffers();
623        vortex_ensure!(
624            nbuffers == buffers.len(),
625            "buffer count changed from {} to {} during physical rewrite",
626            nbuffers,
627            buffers.len()
628        );
629        for (idx, (old_buffer, new_buffer)) in self
630            .buffer_handles()
631            .into_iter()
632            .zip(buffers.iter())
633            .enumerate()
634        {
635            vortex_ensure!(
636                old_buffer.len() == new_buffer.len(),
637                "buffer {} length changed from {} to {} during physical rewrite",
638                idx,
639                old_buffer.len(),
640                new_buffer.len()
641            );
642        }
643        self.0.data.with_buffers(&self, buffers)
644    }
645
646    pub fn reduce(&self) -> VortexResult<Option<ArrayRef>> {
647        self.0.data.reduce(self)
648    }
649
650    pub fn reduce_parent(
651        &self,
652        parent: &ArrayRef,
653        child_idx: usize,
654    ) -> VortexResult<Option<ArrayRef>> {
655        self.0.data.reduce_parent(self, parent, child_idx)
656    }
657
658    pub(crate) fn execute_encoding(self, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
659        let inner = Arc::as_ptr(&self.0);
660        // SAFETY: the Arc outlives the DynArrayData function call
661        unsafe { (&*inner).data.execute(self, ctx) }
662    }
663
664    /// Execute a single encoding step without applying `Done`-result postconditions.
665    ///
666    /// This is for the iterative executor only. It may operate on suspended executor-private
667    /// arrays whose slots temporarily contain `None`, so the executor itself must interpret
668    /// `Done`, enforce any `len`/`dtype` invariants, and transfer statistics.
669    pub(crate) fn execute_encoding_unchecked(
670        self,
671        ctx: &mut ExecutionCtx,
672    ) -> VortexResult<ExecutionResult> {
673        let inner = Arc::as_ptr(&self.0);
674        // SAFETY: `inner` points at the allocation owned by `self.0`. `self` stays alive for the
675        // duration of the call, so the pointee remains valid. Avoiding an extra `Arc` clone here
676        // preserves uniqueness so execute-time metadata cursors can use `Arc::get_mut`.
677        unsafe { (&*inner).data.execute_unchecked(self, ctx) }
678    }
679
680    // ArrayVisitor delegation methods
681
682    /// Returns an iterator over the children of the array: its non-None slots in order.
683    pub fn children_iter(&self) -> impl Iterator<Item = &ArrayRef> {
684        self.0.slots.iter().filter_map(|s| s.as_ref())
685    }
686
687    /// Returns the children of the array.
688    pub fn children(&self) -> Vec<ArrayRef> {
689        self.children_iter().cloned().collect()
690    }
691
692    /// Returns the number of children of the array.
693    pub fn nchildren(&self) -> usize {
694        self.children_iter().count()
695    }
696
697    /// Returns the nth child of the array without allocating a Vec.
698    ///
699    /// Returns `None` if the index is out of bounds.
700    pub fn nth_child(&self, idx: usize) -> Option<ArrayRef> {
701        self.children_iter().nth(idx).cloned()
702    }
703
704    /// Returns the names of the children of the array: the slot names of the non-None slots
705    /// in order.
706    pub fn children_names(&self) -> Vec<String> {
707        self.0
708            .slots
709            .iter()
710            .enumerate()
711            .filter(|(_, s)| s.is_some())
712            .map(|(slot_idx, _)| self.slot_name(slot_idx))
713            .collect()
714    }
715
716    /// Returns the array's children with their names.
717    pub fn named_children(&self) -> Vec<(String, ArrayRef)> {
718        self.children_names()
719            .into_iter()
720            .zip(self.children_iter().cloned())
721            .collect()
722    }
723
724    /// Returns the data buffers of the array.
725    pub fn buffers(&self) -> Vec<ByteBuffer> {
726        self.0.data.buffers(self)
727    }
728
729    /// Returns the buffer handles of the array.
730    pub fn buffer_handles(&self) -> Vec<BufferHandle> {
731        self.0.data.buffer_handles(self)
732    }
733
734    /// Returns the names of the buffers of the array.
735    pub fn buffer_names(&self) -> Vec<String> {
736        self.0.data.buffer_names(self)
737    }
738
739    /// Returns the array's buffers with their names.
740    pub fn named_buffers(&self) -> Vec<(String, BufferHandle)> {
741        self.0.data.named_buffers(self)
742    }
743
744    /// Returns the number of data buffers of the array.
745    pub fn nbuffers(&self) -> usize {
746        self.0.data.nbuffers(self)
747    }
748
749    /// Returns the slots of the array.
750    pub fn slots(&self) -> &[Option<ArrayRef>] {
751        &self.0.slots
752    }
753
754    /// Returns the name of the slot at the given index.
755    pub fn slot_name(&self, idx: usize) -> String {
756        self.0.data.slot_name(self, idx)
757    }
758
759    /// Formats a human-readable metadata description.
760    pub fn metadata_fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
761        self.0.data.metadata_fmt(f)
762    }
763
764    /// Returns whether all buffers are host-resident.
765    pub fn is_host(&self) -> bool {
766        for array in self.depth_first_traversal() {
767            if !array.buffer_handles().iter().all(BufferHandle::is_on_host) {
768                return false;
769            }
770        }
771        true
772    }
773
774    // ArrayVisitorExt delegation methods
775
776    /// Count the number of buffers encoded by self and all child arrays.
777    pub fn nbuffers_recursive(&self) -> usize {
778        self.children()
779            .iter()
780            .map(|c| c.nbuffers_recursive())
781            .sum::<usize>()
782            + self.nbuffers()
783    }
784
785    /// Depth-first traversal of the array and its children.
786    pub fn depth_first_traversal(&self) -> DepthFirstArrayIterator {
787        DepthFirstArrayIterator {
788            stack: vec![self.clone()],
789        }
790    }
791}
792
793impl IntoArray for ArrayRef {
794    #[inline(always)]
795    fn into_array(self) -> ArrayRef {
796        self
797    }
798}
799
800impl<V: VTable> Matcher for V {
801    type Match<'a> = ArrayView<'a, V>;
802
803    #[inline]
804    fn matches(array: &ArrayRef) -> bool {
805        array.0.data.as_any().is::<ArrayData<V>>()
806    }
807
808    #[inline]
809    fn try_match(array: &'_ ArrayRef) -> Option<ArrayView<'_, V>> {
810        let inner = array.0.data.as_any().downcast_ref::<ArrayData<V>>()?;
811        // # Safety checked by `downcast_ref`.
812        Some(unsafe { ArrayView::new_unchecked(array, &inner.data) })
813    }
814}
vortex_array/array/erased.rs

vortex_array/array/
erased.rs
