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 #[inline(always)]
88 pub(crate) fn dyn_array(&self) -> &dyn DynArrayData {
89 &self.0.data
90 }
91
92 #[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 #[doc(hidden)]
101 pub fn addr(&self) -> usize {
102 Arc::as_ptr(&self.0).addr()
103 }
104
105 #[allow(dead_code)]
109 pub(crate) fn downcast_inner<V: VTable>(self) -> Result<Arc<ArrayInner<ArrayData<V>>>, Self> {
110 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 #[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 let raw = Arc::into_raw(self.0);
131 unsafe { Arc::from_raw(raw.cast::<ArrayInner<ArrayData<V>>>()) }
133 }
134
135 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 #[inline]
185 pub fn len(&self) -> usize {
186 self.0.len
187 }
188
189 #[inline]
191 pub fn is_empty(&self) -> bool {
192 self.0.len == 0
193 }
194
195 #[inline]
197 pub fn dtype(&self) -> &DType {
198 &self.0.dtype
199 }
200
201 #[inline]
203 pub fn encoding_id(&self) -> ArrayId {
204 self.0.encoding_id
205 }
206
207 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 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 pub fn filter(&self, mask: Mask) -> VortexResult<ArrayRef> {
250 FilterArray::try_new(self.clone(), mask)?
251 .into_array()
252 .optimize()
253 }
254
255 pub fn take(&self, indices: ArrayRef) -> VortexResult<ArrayRef> {
257 DictArray::try_new(indices, self.clone())?
258 .into_array()
259 .optimize()
260 }
261
262 #[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 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 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 pub fn is_invalid(&self, index: usize, ctx: &mut ExecutionCtx) -> VortexResult<bool> {
299 Ok(!self.is_valid(index, ctx)?)
300 }
301
302 pub fn all_valid(&self, ctx: &mut ExecutionCtx) -> VortexResult<bool> {
304 match self.validity()? {
305 Validity::NonNullable | Validity::AllValid => Ok(true),
306 Validity::AllInvalid => Ok(false),
307 Validity::Array(a) => Ok(a.statistics().compute_min::<bool>(ctx).unwrap_or(false)),
308 }
309 }
310
311 pub fn all_invalid(&self, ctx: &mut ExecutionCtx) -> VortexResult<bool> {
313 match self.validity()? {
314 Validity::NonNullable | Validity::AllValid => Ok(false),
315 Validity::AllInvalid => Ok(true),
316 Validity::Array(a) => Ok(!a.statistics().compute_max::<bool>(ctx).unwrap_or(true)),
317 }
318 }
319
320 pub fn valid_count(&self, ctx: &mut ExecutionCtx) -> VortexResult<usize> {
322 let len = self.len();
323 if let Precision::Exact(invalid_count) = self.statistics().get_as::<usize>(Stat::NullCount)
324 {
325 return Ok(len - invalid_count);
326 }
327
328 let count = match self.validity()? {
329 Validity::NonNullable | Validity::AllValid => len,
330 Validity::AllInvalid => 0,
331 Validity::Array(a) => {
332 let array_sum = sum(&a, ctx)?;
333 array_sum
334 .as_primitive()
335 .as_::<usize>()
336 .ok_or_else(|| vortex_err!("sum of validity array is null"))?
337 }
338 };
339 vortex_ensure!(count <= len, "Valid count exceeds array length");
340
341 self.statistics()
342 .set(Stat::NullCount, Precision::exact(len - count));
343
344 Ok(count)
345 }
346
347 pub fn invalid_count(&self, ctx: &mut ExecutionCtx) -> VortexResult<usize> {
349 Ok(self.len() - self.valid_count(ctx)?)
350 }
351
352 pub fn validity(&self) -> VortexResult<Validity> {
354 self.0.data.validity(self)
355 }
356
357 #[deprecated(note = "use `array.execute::<Canonical>(ctx)` instead")]
359 #[allow(clippy::disallowed_methods)]
360 pub fn into_canonical(self) -> VortexResult<Canonical> {
361 self.execute(&mut legacy_session().create_execution_ctx())
362 }
363
364 #[deprecated(note = "use `array.execute::<Canonical>(ctx)` instead")]
366 pub fn to_canonical(&self) -> VortexResult<Canonical> {
367 #[expect(deprecated)]
368 let result = self.clone().into_canonical();
369 result
370 }
371
372 pub fn append_to_builder(
374 &self,
375 builder: &mut dyn ArrayBuilder,
376 ctx: &mut ExecutionCtx,
377 ) -> VortexResult<()> {
378 self.0.data.append_to_builder(self, builder, ctx)
379 }
380
381 pub fn statistics(&self) -> StatsSetRef<'_> {
383 self.0.stats.to_ref(self)
384 }
385
386 #[inline]
388 pub fn is<M: Matcher>(&self) -> bool {
389 M::matches(self)
390 }
391
392 #[inline]
394 pub fn as_<M: Matcher>(&self) -> M::Match<'_> {
395 self.as_opt::<M>().vortex_expect("Failed to downcast")
396 }
397
398 #[inline]
400 pub fn as_opt<M: Matcher>(&self) -> Option<M::Match<'_>> {
401 M::try_match(self)
402 }
403
404 pub fn try_downcast<V: VTable>(self) -> Result<Array<V>, ArrayRef> {
406 Array::<V>::try_from_array_ref(self)
407 }
408
409 pub fn downcast<V: VTable>(self) -> Array<V> {
415 Self::try_downcast(self)
416 .unwrap_or_else(|_| vortex_panic!("Failed to downcast to {}", type_name::<V>()))
417 }
418
419 pub fn as_typed<V: VTable>(&self) -> Option<ArrayView<'_, V>> {
421 let inner = self.0.data.as_any().downcast_ref::<ArrayData<V>>()?;
422 Some(unsafe { ArrayView::new_unchecked(self, &inner.data) })
423 }
424
425 pub fn as_constant(&self) -> Option<Scalar> {
427 self.as_opt::<Constant>().map(|a| a.scalar().clone())
428 }
429
430 pub fn nbytes(&self) -> u64 {
432 let mut nbytes = 0;
433 for array in self.depth_first_traversal() {
434 for buffer in array.buffers() {
435 nbytes += buffer.len() as u64;
436 }
437 }
438 nbytes
439 }
440
441 pub fn is_canonical(&self) -> bool {
443 self.is::<AnyCanonical>()
444 }
445
446 pub unsafe fn with_slot(
459 self,
460 slot_idx: usize,
461 replacement: ArrayRef,
462 ) -> VortexResult<ArrayRef> {
463 let mut slots: ArraySlots = self.slots().iter().cloned().collect();
464 let nslots = slots.len();
465 vortex_ensure!(
466 slot_idx < nslots,
467 "slot index {} out of bounds for array with {} slots",
468 slot_idx,
469 nslots
470 );
471 let existing = slots[slot_idx]
472 .as_ref()
473 .vortex_expect("with_slot cannot replace an absent slot");
474 vortex_ensure!(
475 existing.dtype() == replacement.dtype(),
476 "slot {} dtype changed from {} to {} during physical rewrite",
477 slot_idx,
478 existing.dtype(),
479 replacement.dtype()
480 );
481 vortex_ensure!(
482 existing.len() == replacement.len(),
483 "slot {} len changed from {} to {} during physical rewrite",
484 slot_idx,
485 existing.len(),
486 replacement.len()
487 );
488 slots[slot_idx] = Some(replacement);
489 unsafe { self.with_slots(slots) }
491 }
492
493 pub(crate) unsafe fn take_slot_unchecked(
505 mut self,
506 slot_idx: usize,
507 ) -> VortexResult<(ArrayRef, ArrayRef)> {
508 if let Some(inner) = Arc::get_mut(&mut self.0) {
509 let child = inner.slots[slot_idx]
510 .take()
511 .vortex_expect("take_slot_unchecked cannot take an absent slot");
512 return Ok((self, child));
513 }
514
515 let child = self.slots()[slot_idx]
518 .as_ref()
519 .vortex_expect("take_slot_unchecked cannot take an absent slot")
520 .clone();
521
522 let mut new_slots: ArraySlots = self.slots().iter().cloned().collect();
523 new_slots[slot_idx] = None;
524
525 let new_parent = unsafe { self.0.data.with_slots_unchecked(&self, new_slots) };
528 Ok((new_parent, child))
529 }
530
531 pub(crate) unsafe fn put_slot_unchecked(
539 mut self,
540 slot_idx: usize,
541 replacement: ArrayRef,
542 ) -> VortexResult<ArrayRef> {
543 if let Some(inner) = Arc::get_mut(&mut self.0) {
544 inner.slots[slot_idx] = Some(replacement);
545 return Ok(self);
546 }
547
548 let mut slots: ArraySlots = self.slots().iter().cloned().collect();
549 slots[slot_idx] = Some(replacement);
550 self.0.data.with_slots(&self, slots)
551 }
552
553 pub unsafe fn with_slots(self, slots: ArraySlots) -> VortexResult<ArrayRef> {
564 let old_slots = self.slots();
565 vortex_ensure!(
566 old_slots.len() == slots.len(),
567 "slot count changed from {} to {} during physical rewrite",
568 old_slots.len(),
569 slots.len()
570 );
571 for (idx, (old_slot, new_slot)) in old_slots.iter().zip(slots.iter()).enumerate() {
572 vortex_ensure!(
573 old_slot.is_some() == new_slot.is_some(),
574 "slot {} presence changed during physical rewrite",
575 idx
576 );
577 if let (Some(old_slot), Some(new_slot)) = (old_slot.as_ref(), new_slot.as_ref()) {
578 vortex_ensure!(
579 old_slot.dtype() == new_slot.dtype(),
580 "slot {} dtype changed from {} to {} during physical rewrite",
581 idx,
582 old_slot.dtype(),
583 new_slot.dtype()
584 );
585 vortex_ensure!(
586 old_slot.len() == new_slot.len(),
587 "slot {} len changed from {} to {} during physical rewrite",
588 idx,
589 old_slot.len(),
590 new_slot.len()
591 );
592 }
593 }
594 self.0.data.with_slots(&self, slots)
595 }
596
597 pub unsafe fn with_buffers(
610 self,
611 buffers: impl IntoIterator<Item = BufferHandle>,
612 ) -> VortexResult<ArrayRef> {
613 let buffers = buffers.into_iter().collect::<Vec<_>>();
614 let nbuffers = self.nbuffers();
615 vortex_ensure!(
616 nbuffers == buffers.len(),
617 "buffer count changed from {} to {} during physical rewrite",
618 nbuffers,
619 buffers.len()
620 );
621 for (idx, (old_buffer, new_buffer)) in self
622 .buffer_handles()
623 .into_iter()
624 .zip(buffers.iter())
625 .enumerate()
626 {
627 vortex_ensure!(
628 old_buffer.len() == new_buffer.len(),
629 "buffer {} length changed from {} to {} during physical rewrite",
630 idx,
631 old_buffer.len(),
632 new_buffer.len()
633 );
634 }
635 self.0.data.with_buffers(&self, buffers)
636 }
637
638 pub fn reduce(&self) -> VortexResult<Option<ArrayRef>> {
639 self.0.data.reduce(self)
640 }
641
642 pub fn reduce_parent(
643 &self,
644 parent: &ArrayRef,
645 child_idx: usize,
646 ) -> VortexResult<Option<ArrayRef>> {
647 self.0.data.reduce_parent(self, parent, child_idx)
648 }
649
650 pub(crate) fn execute_encoding(self, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
651 let inner = Arc::as_ptr(&self.0);
652 unsafe { (&*inner).data.execute(self, ctx) }
654 }
655
656 pub(crate) fn execute_encoding_unchecked(
662 self,
663 ctx: &mut ExecutionCtx,
664 ) -> VortexResult<ExecutionResult> {
665 let inner = Arc::as_ptr(&self.0);
666 unsafe { (&*inner).data.execute_unchecked(self, ctx) }
670 }
671
672 pub fn children_iter(&self) -> impl Iterator<Item = &ArrayRef> {
676 self.0.slots.iter().filter_map(|s| s.as_ref())
677 }
678
679 pub fn children(&self) -> Vec<ArrayRef> {
681 self.children_iter().cloned().collect()
682 }
683
684 pub fn nchildren(&self) -> usize {
686 self.children_iter().count()
687 }
688
689 pub fn nth_child(&self, idx: usize) -> Option<ArrayRef> {
693 self.children_iter().nth(idx).cloned()
694 }
695
696 pub fn children_names(&self) -> Vec<String> {
699 self.0
700 .slots
701 .iter()
702 .enumerate()
703 .filter(|(_, s)| s.is_some())
704 .map(|(slot_idx, _)| self.slot_name(slot_idx))
705 .collect()
706 }
707
708 pub fn named_children(&self) -> Vec<(String, ArrayRef)> {
710 self.children_names()
711 .into_iter()
712 .zip(self.children_iter().cloned())
713 .collect()
714 }
715
716 pub fn buffers(&self) -> Vec<ByteBuffer> {
718 self.0.data.buffers(self)
719 }
720
721 pub fn buffer_handles(&self) -> Vec<BufferHandle> {
723 self.0.data.buffer_handles(self)
724 }
725
726 pub fn buffer_names(&self) -> Vec<String> {
728 self.0.data.buffer_names(self)
729 }
730
731 pub fn named_buffers(&self) -> Vec<(String, BufferHandle)> {
733 self.0.data.named_buffers(self)
734 }
735
736 pub fn nbuffers(&self) -> usize {
738 self.0.data.nbuffers(self)
739 }
740
741 pub fn slots(&self) -> &[Option<ArrayRef>] {
743 &self.0.slots
744 }
745
746 pub fn slot_name(&self, idx: usize) -> String {
748 self.0.data.slot_name(self, idx)
749 }
750
751 pub fn metadata_fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
753 self.0.data.metadata_fmt(f)
754 }
755
756 pub fn is_host(&self) -> bool {
758 for array in self.depth_first_traversal() {
759 if !array.buffer_handles().iter().all(BufferHandle::is_on_host) {
760 return false;
761 }
762 }
763 true
764 }
765
766 pub fn nbuffers_recursive(&self) -> usize {
770 self.children()
771 .iter()
772 .map(|c| c.nbuffers_recursive())
773 .sum::<usize>()
774 + self.nbuffers()
775 }
776
777 pub fn depth_first_traversal(&self) -> DepthFirstArrayIterator {
779 DepthFirstArrayIterator {
780 stack: vec![self.clone()],
781 }
782 }
783}
784
785impl IntoArray for ArrayRef {
786 #[inline(always)]
787 fn into_array(self) -> ArrayRef {
788 self
789 }
790}
791
792impl<V: VTable> Matcher for V {
793 type Match<'a> = ArrayView<'a, V>;
794
795 #[inline]
796 fn matches(array: &ArrayRef) -> bool {
797 array.0.data.as_any().is::<ArrayData<V>>()
798 }
799
800 #[inline]
801 fn try_match(array: &'_ ArrayRef) -> Option<ArrayView<'_, V>> {
802 let inner = array.0.data.as_any().downcast_ref::<ArrayData<V>>()?;
803 Some(unsafe { ArrayView::new_unchecked(array, &inner.data) })
805 }
806}