1use 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::ptr;
11use std::sync::Arc;
12
13use vortex_buffer::ByteBuffer;
14use vortex_error::VortexExpect;
15use vortex_error::VortexResult;
16use vortex_error::vortex_ensure;
17use vortex_error::vortex_ensure_eq;
18use vortex_error::vortex_err;
19use vortex_error::vortex_panic;
20use vortex_mask::Mask;
21
22use crate::AnyCanonical;
23use crate::Array;
24use crate::ArrayEq;
25use crate::ArrayHash;
26use crate::ArrayView;
27use crate::Canonical;
28use crate::ExecutionCtx;
29use crate::ExecutionResult;
30use crate::IntoArray;
31use crate::VTable;
32use crate::VortexSessionExecute;
33use crate::aggregate_fn::fns::sum::sum;
34use crate::array::ArrayData;
35use crate::array::ArrayId;
36use crate::array::ArrayInner;
37use crate::array::ArraySlots;
38use crate::array::DynArrayData;
39use crate::array::probe::ArrayProbe;
40use crate::array::probe::RepeatedArrayProbe;
41use crate::arrays::Constant;
42use crate::arrays::DictArray;
43use crate::arrays::FilterArray;
44use crate::arrays::SliceArray;
45use crate::buffer::BufferHandle;
46use crate::builders::ArrayBuilder;
47use crate::dtype::DType;
48use crate::expr::stats::Precision;
49use crate::expr::stats::Stat;
50use crate::expr::stats::StatsProviderExt;
51use crate::legacy_session;
52use crate::matcher::Matcher;
53use crate::optimizer::ArrayOptimizer;
54use crate::scalar::Scalar;
55use crate::scalar::ScalarValue;
56use crate::stats::StatsSetRef;
57use crate::validity::Validity;
58
59pub struct DepthFirstArrayIterator {
61 stack: Vec<ArrayRef>,
62}
63
64impl Iterator for DepthFirstArrayIterator {
65 type Item = ArrayRef;
66
67 fn next(&mut self) -> Option<Self::Item> {
68 let next = self.stack.pop()?;
69 for child in next.children().into_iter().rev() {
70 self.stack.push(child);
71 }
72 Some(next)
73 }
74}
75
76#[derive(Clone)]
82pub struct ArrayRef(Arc<ArrayInner<dyn DynArrayData>>);
83
84impl ArrayRef {
85 pub(crate) fn from_inner<D: DynArrayData>(inner: Arc<ArrayInner<D>>) -> Self {
87 Self(inner)
88 }
89
90 #[allow(clippy::inline_always)]
92 #[inline(always)]
93 pub(crate) fn dyn_array(&self) -> &dyn DynArrayData {
94 &self.0.data
95 }
96
97 #[allow(clippy::inline_always)]
99 #[inline(always)]
100 pub(crate) fn inner_mut(&mut self) -> Option<&mut ArrayInner<dyn DynArrayData>> {
101 Arc::get_mut(&mut self.0)
102 }
103
104 #[doc(hidden)]
107 pub fn addr(&self) -> usize {
108 Arc::as_ptr(&self.0).addr()
109 }
110
111 #[allow(dead_code)]
115 pub(crate) fn downcast_inner<V: VTable>(self) -> Result<Arc<ArrayInner<ArrayData<V>>>, Self> {
116 if self.0.data.as_any().is::<ArrayData<V>>() {
118 Ok(unsafe { self.downcast_inner_unchecked() })
119 } else {
120 Err(self)
121 }
122 }
123
124 #[allow(clippy::inline_always)]
129 #[inline(always)]
130 pub(crate) unsafe fn downcast_inner_unchecked<V: VTable>(
131 self,
132 ) -> Arc<ArrayInner<ArrayData<V>>> {
133 debug_assert!(self.0.data.as_any().is::<ArrayData<V>>());
134 let raw = Arc::into_raw(self.0);
138 unsafe { Arc::from_raw(raw.cast::<ArrayInner<ArrayData<V>>>()) }
140 }
141
142 pub fn ptr_eq(this: &ArrayRef, other: &ArrayRef) -> bool {
144 Arc::ptr_eq(&this.0, &other.0)
145 }
146}
147
148impl Debug for ArrayRef {
149 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
150 f.debug_struct("Array")
151 .field("encoding", &self.0.encoding_id)
152 .field("dtype", &self.0.dtype)
153 .field("len", &self.0.len)
154 .field("data", &self.0.data)
155 .finish()
156 }
157}
158
159impl ArrayHash for ArrayRef {
160 fn array_hash<H: Hasher>(&self, state: &mut H, accuracy: crate::EqMode) {
161 self.0.len.hash(state);
162 self.0.dtype.hash(state);
163 self.0.encoding_id.hash(state);
164 self.0.slots.len().hash(state);
165 for slot in &self.0.slots {
166 slot.array_hash(state, accuracy);
167 }
168 self.0
169 .data
170 .dyn_array_hash(state as &mut dyn Hasher, accuracy);
171 }
172}
173
174impl ArrayEq for ArrayRef {
175 fn array_eq(&self, other: &Self, accuracy: crate::EqMode) -> bool {
176 self.0.len == other.0.len
177 && self.0.dtype == other.0.dtype
178 && self.0.encoding_id == other.0.encoding_id
179 && self.0.slots.len() == other.0.slots.len()
180 && self
181 .0
182 .slots
183 .iter()
184 .zip(other.0.slots.iter())
185 .all(|(slot, other_slot)| slot.array_eq(other_slot, accuracy))
186 && self.0.data.dyn_array_eq(other, accuracy)
187 }
188}
189impl ArrayRef {
190 #[inline]
192 pub fn len(&self) -> usize {
193 self.0.len
194 }
195
196 #[inline]
198 pub fn is_empty(&self) -> bool {
199 self.0.len == 0
200 }
201
202 #[inline]
204 pub fn dtype(&self) -> &DType {
205 &self.0.dtype
206 }
207
208 #[inline]
210 pub fn encoding_id(&self) -> ArrayId {
211 self.0.encoding_id
212 }
213
214 pub fn slice(&self, range: Range<usize>) -> VortexResult<ArrayRef> {
216 let len = self.len();
217 let start = range.start;
218 let stop = range.end;
219
220 if start == 0 && stop == len {
221 return Ok(self.clone());
222 }
223
224 vortex_ensure!(start <= len, "OutOfBounds: start {start} > length {}", len);
225 vortex_ensure!(stop <= len, "OutOfBounds: stop {stop} > length {}", len);
226
227 vortex_ensure!(start <= stop, "start ({start}) must be <= stop ({stop})");
228
229 if start == stop {
230 return Ok(Canonical::empty(self.dtype()).into_array());
231 }
232
233 let sliced = SliceArray::try_new(self.clone(), range)?
234 .into_array()
235 .optimize()?;
236
237 if !sliced.is::<Constant>() {
239 self.statistics().with_iter(|iter| {
240 sliced.statistics().inherit(iter.filter(|(stat, value)| {
241 matches!(
242 stat,
243 Stat::IsConstant | Stat::IsSorted | Stat::IsStrictSorted
244 ) && value
245 .as_ref()
246 .as_exact()
247 .is_some_and(|v| matches!(v, ScalarValue::Bool(true)))
248 }));
249 });
250 }
251
252 Ok(sliced)
253 }
254
255 pub fn filter(&self, mask: Mask) -> VortexResult<ArrayRef> {
257 FilterArray::try_new(self.clone(), mask)?
258 .into_array()
259 .optimize()
260 }
261
262 pub fn take(&self, indices: ArrayRef) -> VortexResult<ArrayRef> {
264 DictArray::try_new(indices, self.clone())?
265 .into_array()
266 .optimize()
267 }
268
269 #[deprecated(
271 note = "Use `execute_scalar` instead, which allows passing an execution context for more \
272 efficient execution when fetching multiple scalars from the same array."
273 )]
274 #[allow(clippy::disallowed_methods)]
275 pub fn scalar_at(&self, index: usize) -> VortexResult<Scalar> {
276 self.execute_scalar(index, &mut legacy_session().create_execution_ctx())
277 }
278
279 #[inline]
284 pub fn execute_scalar(&self, index: usize, ctx: &mut ExecutionCtx) -> VortexResult<Scalar> {
285 self.probe().execute_scalar(index, ctx)
286 }
287
288 #[inline]
301 pub fn probe(&self) -> ArrayProbe<'_> {
302 ArrayProbe::Once(self)
303 }
304
305 pub fn repeated_probe(&self) -> RepeatedArrayProbe {
320 RepeatedArrayProbe::new(self.clone())
321 }
322
323 #[inline]
328 pub fn is_valid(&self, index: usize, ctx: &mut ExecutionCtx) -> VortexResult<bool> {
329 self.probe().execute_is_valid(index, ctx)
330 }
331
332 #[inline]
335 pub fn is_invalid(&self, index: usize, ctx: &mut ExecutionCtx) -> VortexResult<bool> {
336 Ok(!self.is_valid(index, ctx)?)
337 }
338
339 pub fn all_valid(&self, ctx: &mut ExecutionCtx) -> VortexResult<bool> {
341 if self.is_empty() {
342 return Ok(true);
343 }
344
345 match self.validity()? {
346 Validity::NonNullable | Validity::AllValid => Ok(true),
347 Validity::AllInvalid => Ok(false),
348 Validity::Array(a) => Ok(a.statistics().compute_min::<bool>(ctx).unwrap_or(false)),
349 }
350 }
351
352 pub fn all_invalid(&self, ctx: &mut ExecutionCtx) -> VortexResult<bool> {
354 if self.is_empty() {
355 return Ok(true);
356 }
357
358 match self.validity()? {
359 Validity::NonNullable | Validity::AllValid => Ok(false),
360 Validity::AllInvalid => Ok(true),
361 Validity::Array(a) => Ok(!a.statistics().compute_max::<bool>(ctx).unwrap_or(true)),
362 }
363 }
364
365 pub fn valid_count(&self, ctx: &mut ExecutionCtx) -> VortexResult<usize> {
367 let len = self.len();
368 if let Precision::Exact(invalid_count) = self.statistics().get_as::<usize>(Stat::NullCount)
369 {
370 return Ok(len - invalid_count);
371 }
372
373 let count = match self.validity()? {
374 Validity::NonNullable | Validity::AllValid => len,
375 Validity::AllInvalid => 0,
376 Validity::Array(a) => {
377 let array_sum = sum(&a, ctx)?;
378 array_sum
379 .as_primitive()
380 .as_::<usize>()
381 .ok_or_else(|| vortex_err!("sum of validity array is null"))?
382 }
383 };
384 vortex_ensure!(count <= len, "Valid count exceeds array length");
385
386 self.statistics()
387 .set(Stat::NullCount, Precision::exact(len - count));
388
389 Ok(count)
390 }
391
392 pub fn invalid_count(&self, ctx: &mut ExecutionCtx) -> VortexResult<usize> {
394 Ok(self.len() - self.valid_count(ctx)?)
395 }
396
397 pub fn validity(&self) -> VortexResult<Validity> {
399 self.0.data.validity(self)
400 }
401
402 #[deprecated(note = "use `array.execute::<Canonical>(ctx)` instead")]
404 #[allow(clippy::disallowed_methods)]
405 pub fn into_canonical(self) -> VortexResult<Canonical> {
406 self.execute(&mut legacy_session().create_execution_ctx())
407 }
408
409 #[deprecated(note = "use `array.execute::<Canonical>(ctx)` instead")]
411 pub fn to_canonical(&self) -> VortexResult<Canonical> {
412 #[expect(deprecated)]
413 let result = self.clone().into_canonical();
414 result
415 }
416
417 pub fn append_to_builder(
419 &self,
420 builder: &mut dyn ArrayBuilder,
421 ctx: &mut ExecutionCtx,
422 ) -> VortexResult<()> {
423 self.0.data.append_to_builder(self, builder, ctx)
424 }
425
426 pub fn statistics(&self) -> StatsSetRef<'_> {
428 self.0.stats.to_ref(self)
429 }
430
431 #[inline]
433 pub fn is<M: Matcher>(&self) -> bool {
434 M::matches(self)
435 }
436
437 #[inline]
439 pub fn as_<M: Matcher>(&self) -> M::Match<'_> {
440 self.as_opt::<M>().vortex_expect("Failed to downcast")
441 }
442
443 #[inline]
445 pub fn as_opt<M: Matcher>(&self) -> Option<M::Match<'_>> {
446 M::try_match(self)
447 }
448
449 pub fn try_downcast<V: VTable>(self) -> Result<Array<V>, ArrayRef> {
451 Array::<V>::try_from_array_ref(self)
452 }
453
454 pub fn downcast<V: VTable>(self) -> Array<V> {
460 Self::try_downcast(self)
461 .unwrap_or_else(|_| vortex_panic!("Failed to downcast to {}", type_name::<V>()))
462 }
463
464 pub fn as_typed<V: VTable>(&self) -> Option<ArrayView<'_, V>> {
466 let inner = self.0.data.as_any().downcast_ref::<ArrayData<V>>()?;
467 Some(unsafe { ArrayView::new_unchecked(self, &inner.data) })
468 }
469
470 #[inline]
475 pub(crate) unsafe fn as_typed_unchecked<V: VTable>(&self) -> ArrayView<'_, V> {
476 debug_assert!(self.is::<V>());
477 let inner = unsafe { &*ptr::from_ref(self.dyn_array()).cast::<ArrayData<V>>() };
480 unsafe { ArrayView::new_unchecked(self, &inner.data) }
482 }
483
484 pub fn as_constant(&self) -> Option<Scalar> {
486 self.as_opt::<Constant>().map(|a| a.scalar().clone())
487 }
488
489 pub fn nbytes(&self) -> u64 {
491 let mut nbytes = 0;
492 for array in self.depth_first_traversal() {
493 for buffer in array.buffers() {
494 nbytes += buffer.len() as u64;
495 }
496 }
497 nbytes
498 }
499
500 pub fn is_canonical(&self) -> bool {
502 self.is::<AnyCanonical>()
503 }
504
505 pub unsafe fn with_slot(
518 self,
519 slot_idx: usize,
520 replacement: ArrayRef,
521 ) -> VortexResult<ArrayRef> {
522 let mut slots: ArraySlots = self.slots().iter().cloned().collect();
523 let nslots = slots.len();
524 vortex_ensure!(
525 slot_idx < nslots,
526 "slot index {} out of bounds for array with {} slots",
527 slot_idx,
528 nslots
529 );
530 let existing = slots[slot_idx]
531 .as_ref()
532 .vortex_expect("with_slot cannot replace an absent slot");
533 vortex_ensure_eq!(
534 existing.dtype(),
535 replacement.dtype(),
536 "slot {} dtype changed during physical rewrite",
537 slot_idx
538 );
539 vortex_ensure_eq!(
540 existing.len(),
541 replacement.len(),
542 "slot {} len changed during physical rewrite",
543 slot_idx
544 );
545 slots[slot_idx] = Some(replacement);
546 unsafe { self.with_slots(slots) }
548 }
549
550 pub(crate) unsafe fn take_slot_unchecked(
562 mut self,
563 slot_idx: usize,
564 ) -> VortexResult<(ArrayRef, ArrayRef)> {
565 if let Some(inner) = Arc::get_mut(&mut self.0) {
566 let child = inner.slots[slot_idx]
567 .take()
568 .vortex_expect("take_slot_unchecked cannot take an absent slot");
569 return Ok((self, child));
570 }
571
572 let child = self.slots()[slot_idx]
575 .as_ref()
576 .vortex_expect("take_slot_unchecked cannot take an absent slot")
577 .clone();
578
579 let mut new_slots: ArraySlots = self.slots().iter().cloned().collect();
580 new_slots[slot_idx] = None;
581
582 let new_parent = unsafe { self.0.data.with_slots_unchecked(&self, new_slots) };
585 Ok((new_parent, child))
586 }
587
588 pub(crate) unsafe fn put_slot_unchecked(
596 mut self,
597 slot_idx: usize,
598 replacement: ArrayRef,
599 ) -> VortexResult<ArrayRef> {
600 if let Some(inner) = Arc::get_mut(&mut self.0) {
601 inner.slots[slot_idx] = Some(replacement);
602 return Ok(self);
603 }
604
605 let mut slots: ArraySlots = self.slots().iter().cloned().collect();
606 slots[slot_idx] = Some(replacement);
607 self.0.data.with_slots(&self, slots)
608 }
609
610 pub unsafe fn with_slots(self, slots: ArraySlots) -> VortexResult<ArrayRef> {
621 let old_slots = self.slots();
622 vortex_ensure_eq!(
623 old_slots.len(),
624 slots.len(),
625 "slot count changed during physical rewrite"
626 );
627 for (idx, (old_slot, new_slot)) in old_slots.iter().zip(slots.iter()).enumerate() {
628 vortex_ensure_eq!(
629 old_slot.is_some(),
630 new_slot.is_some(),
631 "slot {} presence changed during physical rewrite",
632 idx
633 );
634 if let (Some(old_slot), Some(new_slot)) = (old_slot.as_ref(), new_slot.as_ref()) {
635 vortex_ensure_eq!(
636 old_slot.dtype(),
637 new_slot.dtype(),
638 "slot {} dtype changed during physical rewrite",
639 idx
640 );
641 vortex_ensure_eq!(
642 old_slot.len(),
643 new_slot.len(),
644 "slot {} len changed during physical rewrite",
645 idx
646 );
647 }
648 }
649 self.0.data.with_slots(&self, slots)
650 }
651
652 pub unsafe fn with_buffers(
665 self,
666 buffers: impl IntoIterator<Item = BufferHandle>,
667 ) -> VortexResult<ArrayRef> {
668 let buffers = buffers.into_iter().collect::<Vec<_>>();
669 let nbuffers = self.nbuffers();
670 vortex_ensure_eq!(
671 nbuffers,
672 buffers.len(),
673 "buffer count changed during physical rewrite"
674 );
675 for (idx, (old_buffer, new_buffer)) in self
676 .buffer_handles()
677 .into_iter()
678 .zip(buffers.iter())
679 .enumerate()
680 {
681 vortex_ensure_eq!(
682 old_buffer.len(),
683 new_buffer.len(),
684 "buffer {} length changed during physical rewrite",
685 idx
686 );
687 }
688 self.0.data.with_buffers(&self, buffers)
689 }
690
691 pub fn reduce(&self) -> VortexResult<Option<ArrayRef>> {
692 self.0.data.reduce(self)
693 }
694
695 pub fn reduce_parent(
696 &self,
697 parent: &ArrayRef,
698 child_idx: usize,
699 ) -> VortexResult<Option<ArrayRef>> {
700 self.0.data.reduce_parent(self, parent, child_idx)
701 }
702
703 pub(crate) fn execute_encoding(self, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
704 let inner = Arc::as_ptr(&self.0);
705 unsafe { (&*inner).data.execute(self, ctx) }
707 }
708
709 pub(crate) fn execute_encoding_unchecked(
715 self,
716 ctx: &mut ExecutionCtx,
717 ) -> VortexResult<ExecutionResult> {
718 let inner = Arc::as_ptr(&self.0);
719 unsafe { (&*inner).data.execute_unchecked(self, ctx) }
723 }
724
725 pub fn children_iter(&self) -> impl Iterator<Item = &ArrayRef> {
729 self.0.slots.iter().filter_map(|s| s.as_ref())
730 }
731
732 pub fn children(&self) -> Vec<ArrayRef> {
734 self.children_iter().cloned().collect()
735 }
736
737 pub fn nchildren(&self) -> usize {
739 self.children_iter().count()
740 }
741
742 pub fn nth_child(&self, idx: usize) -> Option<&ArrayRef> {
746 self.children_iter().nth(idx)
747 }
748
749 pub fn children_names(&self) -> Vec<String> {
752 self.0
753 .slots
754 .iter()
755 .enumerate()
756 .filter(|(_, s)| s.is_some())
757 .map(|(slot_idx, _)| self.slot_name(slot_idx))
758 .collect()
759 }
760
761 pub fn named_children(&self) -> Vec<(String, ArrayRef)> {
763 self.children_names()
764 .into_iter()
765 .zip(self.children_iter().cloned())
766 .collect()
767 }
768
769 pub fn buffers(&self) -> Vec<ByteBuffer> {
771 self.0.data.buffers(self)
772 }
773
774 pub fn buffer_handles(&self) -> Vec<BufferHandle> {
776 self.0.data.buffer_handles(self)
777 }
778
779 pub fn buffer_names(&self) -> Vec<String> {
781 self.0.data.buffer_names(self)
782 }
783
784 pub fn named_buffers(&self) -> Vec<(String, BufferHandle)> {
786 self.0.data.named_buffers(self)
787 }
788
789 pub fn nbuffers(&self) -> usize {
791 self.0.data.nbuffers(self)
792 }
793
794 pub fn slots(&self) -> &[Option<ArrayRef>] {
796 &self.0.slots
797 }
798
799 pub fn slot_name(&self, idx: usize) -> String {
801 self.0.data.slot_name(self, idx)
802 }
803
804 pub fn metadata_fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
806 self.0.data.metadata_fmt(f)
807 }
808
809 pub fn is_host(&self) -> bool {
811 for array in self.depth_first_traversal() {
812 if !array.buffer_handles().iter().all(BufferHandle::is_on_host) {
813 return false;
814 }
815 }
816 true
817 }
818
819 pub fn nbuffers_recursive(&self) -> usize {
823 self.children()
824 .iter()
825 .map(|c| c.nbuffers_recursive())
826 .sum::<usize>()
827 + self.nbuffers()
828 }
829
830 pub fn depth_first_traversal(&self) -> DepthFirstArrayIterator {
832 DepthFirstArrayIterator {
833 stack: vec![self.clone()],
834 }
835 }
836}
837
838impl IntoArray for ArrayRef {
839 #[allow(clippy::inline_always)]
840 #[inline(always)]
841 fn into_array(self) -> ArrayRef {
842 self
843 }
844}
845
846impl<V: VTable> Matcher for V {
847 type Match<'a> = ArrayView<'a, V>;
848
849 #[inline]
850 fn matches(array: &ArrayRef) -> bool {
851 array.0.data.as_any().is::<ArrayData<V>>()
852 }
853
854 #[inline]
855 fn try_match(array: &'_ ArrayRef) -> Option<ArrayView<'_, V>> {
856 let inner = array.0.data.as_any().downcast_ref::<ArrayData<V>>()?;
857 Some(unsafe { ArrayView::new_unchecked(array, &inner.data) })
859 }
860}