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::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
55pub 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#[derive(Clone)]
78pub struct ArrayRef(Arc<ArrayInner<dyn DynArrayData>>);
79
80impl ArrayRef {
81 pub(crate) fn from_inner<D: DynArrayData>(inner: Arc<ArrayInner<D>>) -> Self {
83 Self(inner)
84 }
85
86 #[allow(clippy::inline_always)]
88 #[inline(always)]
89 pub(crate) fn dyn_array(&self) -> &dyn DynArrayData {
90 &self.0.data
91 }
92
93 #[allow(clippy::inline_always)]
95 #[inline(always)]
96 pub(crate) fn inner_mut(&mut self) -> Option<&mut ArrayInner<dyn DynArrayData>> {
97 Arc::get_mut(&mut self.0)
98 }
99
100 #[doc(hidden)]
103 pub fn addr(&self) -> usize {
104 Arc::as_ptr(&self.0).addr()
105 }
106
107 #[allow(dead_code)]
111 pub(crate) fn downcast_inner<V: VTable>(self) -> Result<Arc<ArrayInner<ArrayData<V>>>, Self> {
112 if self.0.data.as_any().is::<ArrayData<V>>() {
114 Ok(unsafe { self.downcast_inner_unchecked() })
115 } else {
116 Err(self)
117 }
118 }
119
120 #[allow(clippy::inline_always)]
125 #[inline(always)]
126 pub(crate) unsafe fn downcast_inner_unchecked<V: VTable>(
127 self,
128 ) -> Arc<ArrayInner<ArrayData<V>>> {
129 debug_assert!(self.0.data.as_any().is::<ArrayData<V>>());
130 let raw = Arc::into_raw(self.0);
134 unsafe { Arc::from_raw(raw.cast::<ArrayInner<ArrayData<V>>>()) }
136 }
137
138 pub fn ptr_eq(this: &ArrayRef, other: &ArrayRef) -> bool {
140 Arc::ptr_eq(&this.0, &other.0)
141 }
142}
143
144impl Debug for ArrayRef {
145 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
146 f.debug_struct("Array")
147 .field("encoding", &self.0.encoding_id)
148 .field("dtype", &self.0.dtype)
149 .field("len", &self.0.len)
150 .field("data", &self.0.data)
151 .finish()
152 }
153}
154
155impl ArrayHash for ArrayRef {
156 fn array_hash<H: Hasher>(&self, state: &mut H, accuracy: crate::EqMode) {
157 self.0.len.hash(state);
158 self.0.dtype.hash(state);
159 self.0.encoding_id.hash(state);
160 self.0.slots.len().hash(state);
161 for slot in &self.0.slots {
162 slot.array_hash(state, accuracy);
163 }
164 self.0
165 .data
166 .dyn_array_hash(state as &mut dyn Hasher, accuracy);
167 }
168}
169
170impl ArrayEq for ArrayRef {
171 fn array_eq(&self, other: &Self, accuracy: crate::EqMode) -> bool {
172 self.0.len == other.0.len
173 && self.0.dtype == other.0.dtype
174 && self.0.encoding_id == other.0.encoding_id
175 && self.0.slots.len() == other.0.slots.len()
176 && self
177 .0
178 .slots
179 .iter()
180 .zip(other.0.slots.iter())
181 .all(|(slot, other_slot)| slot.array_eq(other_slot, accuracy))
182 && self.0.data.dyn_array_eq(other, accuracy)
183 }
184}
185impl ArrayRef {
186 #[inline]
188 pub fn len(&self) -> usize {
189 self.0.len
190 }
191
192 #[inline]
194 pub fn is_empty(&self) -> bool {
195 self.0.len == 0
196 }
197
198 #[inline]
200 pub fn dtype(&self) -> &DType {
201 &self.0.dtype
202 }
203
204 #[inline]
206 pub fn encoding_id(&self) -> ArrayId {
207 self.0.encoding_id
208 }
209
210 pub fn slice(&self, range: Range<usize>) -> VortexResult<ArrayRef> {
212 let len = self.len();
213 let start = range.start;
214 let stop = range.end;
215
216 if start == 0 && stop == len {
217 return Ok(self.clone());
218 }
219
220 vortex_ensure!(start <= len, "OutOfBounds: start {start} > length {}", len);
221 vortex_ensure!(stop <= len, "OutOfBounds: stop {stop} > length {}", len);
222
223 vortex_ensure!(start <= stop, "start ({start}) must be <= stop ({stop})");
224
225 if start == stop {
226 return Ok(Canonical::empty(self.dtype()).into_array());
227 }
228
229 let sliced = SliceArray::try_new(self.clone(), range)?
230 .into_array()
231 .optimize()?;
232
233 if !sliced.is::<Constant>() {
235 self.statistics().with_iter(|iter| {
236 sliced.statistics().inherit(iter.filter(|(stat, value)| {
237 matches!(
238 stat,
239 Stat::IsConstant | Stat::IsSorted | Stat::IsStrictSorted
240 ) && value
241 .as_ref()
242 .as_exact()
243 .is_some_and(|v| matches!(v, ScalarValue::Bool(true)))
244 }));
245 });
246 }
247
248 Ok(sliced)
249 }
250
251 pub fn filter(&self, mask: Mask) -> VortexResult<ArrayRef> {
253 FilterArray::try_new(self.clone(), mask)?
254 .into_array()
255 .optimize()
256 }
257
258 pub fn take(&self, indices: ArrayRef) -> VortexResult<ArrayRef> {
260 DictArray::try_new(indices, self.clone())?
261 .into_array()
262 .optimize()
263 }
264
265 #[deprecated(
267 note = "Use `execute_scalar` instead, which allows passing an execution context for more \
268 efficient execution when fetching multiple scalars from the same array."
269 )]
270 #[allow(clippy::disallowed_methods)]
271 pub fn scalar_at(&self, index: usize) -> VortexResult<Scalar> {
272 self.execute_scalar(index, &mut legacy_session().create_execution_ctx())
273 }
274
275 pub fn execute_scalar(&self, index: usize, ctx: &mut ExecutionCtx) -> VortexResult<Scalar> {
277 vortex_ensure!(index < self.len(), OutOfBounds: index, 0, self.len());
278 if self.dtype().is_nullable() && self.is_invalid(index, ctx)? {
279 return Ok(Scalar::null(self.dtype().clone()));
280 }
281 let scalar = self.0.data.execute_scalar(self, index, ctx)?;
282 debug_assert_eq!(self.dtype(), scalar.dtype(), "Scalar dtype mismatch");
283 Ok(scalar)
284 }
285
286 pub fn is_valid(&self, index: usize, ctx: &mut ExecutionCtx) -> VortexResult<bool> {
288 vortex_ensure!(index < self.len(), OutOfBounds: index, 0, self.len());
289 match self.validity()? {
290 Validity::NonNullable | Validity::AllValid => Ok(true),
291 Validity::AllInvalid => Ok(false),
292 Validity::Array(a) => a
293 .execute_scalar(index, ctx)?
294 .as_bool()
295 .value()
296 .ok_or_else(|| vortex_err!("validity value at index {} is null", index)),
297 }
298 }
299
300 pub fn is_invalid(&self, index: usize, ctx: &mut ExecutionCtx) -> VortexResult<bool> {
302 Ok(!self.is_valid(index, ctx)?)
303 }
304
305 pub fn all_valid(&self, ctx: &mut ExecutionCtx) -> VortexResult<bool> {
307 if self.is_empty() {
308 return Ok(true);
309 }
310
311 match self.validity()? {
312 Validity::NonNullable | Validity::AllValid => Ok(true),
313 Validity::AllInvalid => Ok(false),
314 Validity::Array(a) => Ok(a.statistics().compute_min::<bool>(ctx).unwrap_or(false)),
315 }
316 }
317
318 pub fn all_invalid(&self, ctx: &mut ExecutionCtx) -> VortexResult<bool> {
320 if self.is_empty() {
321 return Ok(true);
322 }
323
324 match self.validity()? {
325 Validity::NonNullable | Validity::AllValid => Ok(false),
326 Validity::AllInvalid => Ok(true),
327 Validity::Array(a) => Ok(!a.statistics().compute_max::<bool>(ctx).unwrap_or(true)),
328 }
329 }
330
331 pub fn valid_count(&self, ctx: &mut ExecutionCtx) -> VortexResult<usize> {
333 let len = self.len();
334 if let Precision::Exact(invalid_count) = self.statistics().get_as::<usize>(Stat::NullCount)
335 {
336 return Ok(len - invalid_count);
337 }
338
339 let count = match self.validity()? {
340 Validity::NonNullable | Validity::AllValid => len,
341 Validity::AllInvalid => 0,
342 Validity::Array(a) => {
343 let array_sum = sum(&a, ctx)?;
344 array_sum
345 .as_primitive()
346 .as_::<usize>()
347 .ok_or_else(|| vortex_err!("sum of validity array is null"))?
348 }
349 };
350 vortex_ensure!(count <= len, "Valid count exceeds array length");
351
352 self.statistics()
353 .set(Stat::NullCount, Precision::exact(len - count));
354
355 Ok(count)
356 }
357
358 pub fn invalid_count(&self, ctx: &mut ExecutionCtx) -> VortexResult<usize> {
360 Ok(self.len() - self.valid_count(ctx)?)
361 }
362
363 pub fn validity(&self) -> VortexResult<Validity> {
365 self.0.data.validity(self)
366 }
367
368 #[deprecated(note = "use `array.execute::<Canonical>(ctx)` instead")]
370 #[allow(clippy::disallowed_methods)]
371 pub fn into_canonical(self) -> VortexResult<Canonical> {
372 self.execute(&mut legacy_session().create_execution_ctx())
373 }
374
375 #[deprecated(note = "use `array.execute::<Canonical>(ctx)` instead")]
377 pub fn to_canonical(&self) -> VortexResult<Canonical> {
378 #[expect(deprecated)]
379 let result = self.clone().into_canonical();
380 result
381 }
382
383 pub fn append_to_builder(
385 &self,
386 builder: &mut dyn ArrayBuilder,
387 ctx: &mut ExecutionCtx,
388 ) -> VortexResult<()> {
389 self.0.data.append_to_builder(self, builder, ctx)
390 }
391
392 pub fn statistics(&self) -> StatsSetRef<'_> {
394 self.0.stats.to_ref(self)
395 }
396
397 #[inline]
399 pub fn is<M: Matcher>(&self) -> bool {
400 M::matches(self)
401 }
402
403 #[inline]
405 pub fn as_<M: Matcher>(&self) -> M::Match<'_> {
406 self.as_opt::<M>().vortex_expect("Failed to downcast")
407 }
408
409 #[inline]
411 pub fn as_opt<M: Matcher>(&self) -> Option<M::Match<'_>> {
412 M::try_match(self)
413 }
414
415 pub fn try_downcast<V: VTable>(self) -> Result<Array<V>, ArrayRef> {
417 Array::<V>::try_from_array_ref(self)
418 }
419
420 pub fn downcast<V: VTable>(self) -> Array<V> {
426 Self::try_downcast(self)
427 .unwrap_or_else(|_| vortex_panic!("Failed to downcast to {}", type_name::<V>()))
428 }
429
430 pub fn as_typed<V: VTable>(&self) -> Option<ArrayView<'_, V>> {
432 let inner = self.0.data.as_any().downcast_ref::<ArrayData<V>>()?;
433 Some(unsafe { ArrayView::new_unchecked(self, &inner.data) })
434 }
435
436 pub fn as_constant(&self) -> Option<Scalar> {
438 self.as_opt::<Constant>().map(|a| a.scalar().clone())
439 }
440
441 pub fn nbytes(&self) -> u64 {
443 let mut nbytes = 0;
444 for array in self.depth_first_traversal() {
445 for buffer in array.buffers() {
446 nbytes += buffer.len() as u64;
447 }
448 }
449 nbytes
450 }
451
452 pub fn is_canonical(&self) -> bool {
454 self.is::<AnyCanonical>()
455 }
456
457 pub unsafe fn with_slot(
470 self,
471 slot_idx: usize,
472 replacement: ArrayRef,
473 ) -> VortexResult<ArrayRef> {
474 let mut slots: ArraySlots = self.slots().iter().cloned().collect();
475 let nslots = slots.len();
476 vortex_ensure!(
477 slot_idx < nslots,
478 "slot index {} out of bounds for array with {} slots",
479 slot_idx,
480 nslots
481 );
482 let existing = slots[slot_idx]
483 .as_ref()
484 .vortex_expect("with_slot cannot replace an absent slot");
485 vortex_ensure!(
486 existing.dtype() == replacement.dtype(),
487 "slot {} dtype changed from {} to {} during physical rewrite",
488 slot_idx,
489 existing.dtype(),
490 replacement.dtype()
491 );
492 vortex_ensure!(
493 existing.len() == replacement.len(),
494 "slot {} len changed from {} to {} during physical rewrite",
495 slot_idx,
496 existing.len(),
497 replacement.len()
498 );
499 slots[slot_idx] = Some(replacement);
500 unsafe { self.with_slots(slots) }
502 }
503
504 pub(crate) unsafe fn take_slot_unchecked(
516 mut self,
517 slot_idx: usize,
518 ) -> VortexResult<(ArrayRef, ArrayRef)> {
519 if let Some(inner) = Arc::get_mut(&mut self.0) {
520 let child = inner.slots[slot_idx]
521 .take()
522 .vortex_expect("take_slot_unchecked cannot take an absent slot");
523 return Ok((self, child));
524 }
525
526 let child = self.slots()[slot_idx]
529 .as_ref()
530 .vortex_expect("take_slot_unchecked cannot take an absent slot")
531 .clone();
532
533 let mut new_slots: ArraySlots = self.slots().iter().cloned().collect();
534 new_slots[slot_idx] = None;
535
536 let new_parent = unsafe { self.0.data.with_slots_unchecked(&self, new_slots) };
539 Ok((new_parent, child))
540 }
541
542 pub(crate) unsafe fn put_slot_unchecked(
550 mut self,
551 slot_idx: usize,
552 replacement: ArrayRef,
553 ) -> VortexResult<ArrayRef> {
554 if let Some(inner) = Arc::get_mut(&mut self.0) {
555 inner.slots[slot_idx] = Some(replacement);
556 return Ok(self);
557 }
558
559 let mut slots: ArraySlots = self.slots().iter().cloned().collect();
560 slots[slot_idx] = Some(replacement);
561 self.0.data.with_slots(&self, slots)
562 }
563
564 pub unsafe fn with_slots(self, slots: ArraySlots) -> VortexResult<ArrayRef> {
575 let old_slots = self.slots();
576 vortex_ensure!(
577 old_slots.len() == slots.len(),
578 "slot count changed from {} to {} during physical rewrite",
579 old_slots.len(),
580 slots.len()
581 );
582 for (idx, (old_slot, new_slot)) in old_slots.iter().zip(slots.iter()).enumerate() {
583 vortex_ensure!(
584 old_slot.is_some() == new_slot.is_some(),
585 "slot {} presence changed during physical rewrite",
586 idx
587 );
588 if let (Some(old_slot), Some(new_slot)) = (old_slot.as_ref(), new_slot.as_ref()) {
589 vortex_ensure!(
590 old_slot.dtype() == new_slot.dtype(),
591 "slot {} dtype changed from {} to {} during physical rewrite",
592 idx,
593 old_slot.dtype(),
594 new_slot.dtype()
595 );
596 vortex_ensure!(
597 old_slot.len() == new_slot.len(),
598 "slot {} len changed from {} to {} during physical rewrite",
599 idx,
600 old_slot.len(),
601 new_slot.len()
602 );
603 }
604 }
605 self.0.data.with_slots(&self, slots)
606 }
607
608 pub unsafe fn with_buffers(
621 self,
622 buffers: impl IntoIterator<Item = BufferHandle>,
623 ) -> VortexResult<ArrayRef> {
624 let buffers = buffers.into_iter().collect::<Vec<_>>();
625 let nbuffers = self.nbuffers();
626 vortex_ensure!(
627 nbuffers == buffers.len(),
628 "buffer count changed from {} to {} during physical rewrite",
629 nbuffers,
630 buffers.len()
631 );
632 for (idx, (old_buffer, new_buffer)) in self
633 .buffer_handles()
634 .into_iter()
635 .zip(buffers.iter())
636 .enumerate()
637 {
638 vortex_ensure!(
639 old_buffer.len() == new_buffer.len(),
640 "buffer {} length changed from {} to {} during physical rewrite",
641 idx,
642 old_buffer.len(),
643 new_buffer.len()
644 );
645 }
646 self.0.data.with_buffers(&self, buffers)
647 }
648
649 pub fn reduce(&self) -> VortexResult<Option<ArrayRef>> {
650 self.0.data.reduce(self)
651 }
652
653 pub fn reduce_parent(
654 &self,
655 parent: &ArrayRef,
656 child_idx: usize,
657 ) -> VortexResult<Option<ArrayRef>> {
658 self.0.data.reduce_parent(self, parent, child_idx)
659 }
660
661 pub(crate) fn execute_encoding(self, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
662 let inner = Arc::as_ptr(&self.0);
663 unsafe { (&*inner).data.execute(self, ctx) }
665 }
666
667 pub(crate) fn execute_encoding_unchecked(
673 self,
674 ctx: &mut ExecutionCtx,
675 ) -> VortexResult<ExecutionResult> {
676 let inner = Arc::as_ptr(&self.0);
677 unsafe { (&*inner).data.execute_unchecked(self, ctx) }
681 }
682
683 pub fn children_iter(&self) -> impl Iterator<Item = &ArrayRef> {
687 self.0.slots.iter().filter_map(|s| s.as_ref())
688 }
689
690 pub fn children(&self) -> Vec<ArrayRef> {
692 self.children_iter().cloned().collect()
693 }
694
695 pub fn nchildren(&self) -> usize {
697 self.children_iter().count()
698 }
699
700 pub fn nth_child(&self, idx: usize) -> Option<ArrayRef> {
704 self.children_iter().nth(idx).cloned()
705 }
706
707 pub fn children_names(&self) -> Vec<String> {
710 self.0
711 .slots
712 .iter()
713 .enumerate()
714 .filter(|(_, s)| s.is_some())
715 .map(|(slot_idx, _)| self.slot_name(slot_idx))
716 .collect()
717 }
718
719 pub fn named_children(&self) -> Vec<(String, ArrayRef)> {
721 self.children_names()
722 .into_iter()
723 .zip(self.children_iter().cloned())
724 .collect()
725 }
726
727 pub fn buffers(&self) -> Vec<ByteBuffer> {
729 self.0.data.buffers(self)
730 }
731
732 pub fn buffer_handles(&self) -> Vec<BufferHandle> {
734 self.0.data.buffer_handles(self)
735 }
736
737 pub fn buffer_names(&self) -> Vec<String> {
739 self.0.data.buffer_names(self)
740 }
741
742 pub fn named_buffers(&self) -> Vec<(String, BufferHandle)> {
744 self.0.data.named_buffers(self)
745 }
746
747 pub fn nbuffers(&self) -> usize {
749 self.0.data.nbuffers(self)
750 }
751
752 pub fn slots(&self) -> &[Option<ArrayRef>] {
754 &self.0.slots
755 }
756
757 pub fn slot_name(&self, idx: usize) -> String {
759 self.0.data.slot_name(self, idx)
760 }
761
762 pub fn metadata_fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
764 self.0.data.metadata_fmt(f)
765 }
766
767 pub fn is_host(&self) -> bool {
769 for array in self.depth_first_traversal() {
770 if !array.buffer_handles().iter().all(BufferHandle::is_on_host) {
771 return false;
772 }
773 }
774 true
775 }
776
777 pub fn nbuffers_recursive(&self) -> usize {
781 self.children()
782 .iter()
783 .map(|c| c.nbuffers_recursive())
784 .sum::<usize>()
785 + self.nbuffers()
786 }
787
788 pub fn depth_first_traversal(&self) -> DepthFirstArrayIterator {
790 DepthFirstArrayIterator {
791 stack: vec![self.clone()],
792 }
793 }
794}
795
796impl IntoArray for ArrayRef {
797 #[allow(clippy::inline_always)]
798 #[inline(always)]
799 fn into_array(self) -> ArrayRef {
800 self
801 }
802}
803
804impl<V: VTable> Matcher for V {
805 type Match<'a> = ArrayView<'a, V>;
806
807 #[inline]
808 fn matches(array: &ArrayRef) -> bool {
809 array.0.data.as_any().is::<ArrayData<V>>()
810 }
811
812 #[inline]
813 fn try_match(array: &'_ ArrayRef) -> Option<ArrayView<'_, V>> {
814 let inner = array.0.data.as_any().downcast_ref::<ArrayData<V>>()?;
815 Some(unsafe { ArrayView::new_unchecked(array, &inner.data) })
817 }
818}