Skip to main content

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