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::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
57pub 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#[derive(Clone)]
80pub struct ArrayRef(Arc<ArrayInner<dyn DynArrayData>>);
81
82impl ArrayRef {
83 pub(crate) fn from_inner<D: DynArrayData>(inner: Arc<ArrayInner<D>>) -> Self {
85 Self(inner)
86 }
87
88 #[allow(clippy::inline_always)]
90 #[inline(always)]
91 pub(crate) fn dyn_array(&self) -> &dyn DynArrayData {
92 &self.0.data
93 }
94
95 #[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 #[doc(hidden)]
105 pub fn addr(&self) -> usize {
106 Arc::as_ptr(&self.0).addr()
107 }
108
109 #[allow(dead_code)]
113 pub(crate) fn downcast_inner<V: VTable>(self) -> Result<Arc<ArrayInner<ArrayData<V>>>, Self> {
114 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 #[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 let raw = Arc::into_raw(self.0);
136 unsafe { Arc::from_raw(raw.cast::<ArrayInner<ArrayData<V>>>()) }
138 }
139
140 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 #[inline]
190 pub fn len(&self) -> usize {
191 self.0.len
192 }
193
194 #[inline]
196 pub fn is_empty(&self) -> bool {
197 self.0.len == 0
198 }
199
200 #[inline]
202 pub fn dtype(&self) -> &DType {
203 &self.0.dtype
204 }
205
206 #[inline]
208 pub fn encoding_id(&self) -> ArrayId {
209 self.0.encoding_id
210 }
211
212 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 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 pub fn filter(&self, mask: Mask) -> VortexResult<ArrayRef> {
255 FilterArray::try_new(self.clone(), mask)?
256 .into_array()
257 .optimize()
258 }
259
260 pub fn take(&self, indices: ArrayRef) -> VortexResult<ArrayRef> {
262 DictArray::try_new(indices, self.clone())?
263 .into_array()
264 .optimize()
265 }
266
267 #[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 #[inline]
282 pub fn execute_scalar(&self, index: usize, ctx: &mut ExecutionCtx) -> VortexResult<Scalar> {
283 self.probe().execute_scalar(index, ctx)
284 }
285
286 #[inline]
299 pub fn probe(&self) -> ArrayProbe<'_> {
300 ArrayProbe::Once(self)
301 }
302
303 pub fn repeated_probe(&self) -> RepeatedArrayProbe {
318 RepeatedArrayProbe::new(self.clone())
319 }
320
321 #[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 #[inline]
333 pub fn is_invalid(&self, index: usize, ctx: &mut ExecutionCtx) -> VortexResult<bool> {
334 Ok(!self.is_valid(index, ctx)?)
335 }
336
337 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 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 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 pub fn invalid_count(&self, ctx: &mut ExecutionCtx) -> VortexResult<usize> {
392 Ok(self.len() - self.valid_count(ctx)?)
393 }
394
395 pub fn validity(&self) -> VortexResult<Validity> {
397 self.0.data.validity(self)
398 }
399
400 #[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 #[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 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 pub fn statistics(&self) -> StatsSetRef<'_> {
426 self.0.stats.to_ref(self)
427 }
428
429 #[inline]
431 pub fn is<M: Matcher>(&self) -> bool {
432 M::matches(self)
433 }
434
435 #[inline]
437 pub fn as_<M: Matcher>(&self) -> M::Match<'_> {
438 self.as_opt::<M>().vortex_expect("Failed to downcast")
439 }
440
441 #[inline]
443 pub fn as_opt<M: Matcher>(&self) -> Option<M::Match<'_>> {
444 M::try_match(self)
445 }
446
447 pub fn try_downcast<V: VTable>(self) -> Result<Array<V>, ArrayRef> {
449 Array::<V>::try_from_array_ref(self)
450 }
451
452 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 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 pub fn as_constant(&self) -> Option<Scalar> {
470 self.as_opt::<Constant>().map(|a| a.scalar().clone())
471 }
472
473 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 pub fn is_canonical(&self) -> bool {
486 self.is::<AnyCanonical>()
487 }
488
489 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 unsafe { self.with_slots(slots) }
534 }
535
536 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 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 let new_parent = unsafe { self.0.data.with_slots_unchecked(&self, new_slots) };
571 Ok((new_parent, child))
572 }
573
574 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 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 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 unsafe { (&*inner).data.execute(self, ctx) }
697 }
698
699 pub(crate) fn execute_encoding_unchecked(
705 self,
706 ctx: &mut ExecutionCtx,
707 ) -> VortexResult<ExecutionResult> {
708 let inner = Arc::as_ptr(&self.0);
709 unsafe { (&*inner).data.execute_unchecked(self, ctx) }
713 }
714
715 pub fn children_iter(&self) -> impl Iterator<Item = &ArrayRef> {
719 self.0.slots.iter().filter_map(|s| s.as_ref())
720 }
721
722 pub fn children(&self) -> Vec<ArrayRef> {
724 self.children_iter().cloned().collect()
725 }
726
727 pub fn nchildren(&self) -> usize {
729 self.children_iter().count()
730 }
731
732 pub fn nth_child(&self, idx: usize) -> Option<&ArrayRef> {
736 self.children_iter().nth(idx)
737 }
738
739 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 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 pub fn buffers(&self) -> Vec<ByteBuffer> {
761 self.0.data.buffers(self)
762 }
763
764 pub fn buffer_handles(&self) -> Vec<BufferHandle> {
766 self.0.data.buffer_handles(self)
767 }
768
769 pub fn buffer_names(&self) -> Vec<String> {
771 self.0.data.buffer_names(self)
772 }
773
774 pub fn named_buffers(&self) -> Vec<(String, BufferHandle)> {
776 self.0.data.named_buffers(self)
777 }
778
779 pub fn nbuffers(&self) -> usize {
781 self.0.data.nbuffers(self)
782 }
783
784 pub fn slots(&self) -> &[Option<ArrayRef>] {
786 &self.0.slots
787 }
788
789 pub fn slot_name(&self, idx: usize) -> String {
791 self.0.data.slot_name(self, idx)
792 }
793
794 pub fn metadata_fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
796 self.0.data.metadata_fmt(f)
797 }
798
799 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 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 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 Some(unsafe { ArrayView::new_unchecked(array, &inner.data) })
849 }
850}