1use std::fmt::Display;
5use std::fmt::Formatter;
6use std::sync::Arc;
7
8use num_traits::AsPrimitive;
9use vortex_buffer::BitBufferMut;
10use vortex_error::VortexExpect;
11use vortex_error::VortexResult;
12use vortex_error::vortex_bail;
13use vortex_error::vortex_ensure;
14use vortex_error::vortex_ensure_eq;
15use vortex_error::vortex_err;
16use vortex_mask::Mask;
17
18use crate::ArrayRef;
19use crate::ArraySlots;
20use crate::ExecutionCtx;
21use crate::VortexSessionExecute;
22use crate::aggregate_fn::NumericalAggregateOpts;
23use crate::aggregate_fn::fns::min_max::min_max;
24use crate::array::Array;
25use crate::array::ArrayParts;
26use crate::array::TypedArrayRef;
27use crate::array::child_to_validity;
28use crate::array::validity_to_child;
29use crate::array_slots;
30use crate::arrays::ListView;
31use crate::arrays::Primitive;
32use crate::arrays::PrimitiveArray;
33use crate::arrays::bool;
34use crate::arrays::primitive::PrimitiveArrayExt;
35use crate::builtins::ArrayBuiltins;
36use crate::dtype::DType;
37use crate::dtype::IntegerPType;
38use crate::dtype::PType;
39use crate::expr::stats::Stat;
40use crate::legacy_session;
41use crate::match_each_integer_ptype;
42use crate::match_each_unsigned_integer_ptype;
43use crate::scalar_fn::fns::operators::Operator;
44use crate::validity::Validity;
45
46#[array_slots(ListView)]
47pub struct ListViewSlots {
48 #[slot(0)]
51 pub elements: ArrayRef,
52 #[slot(1)]
57 pub offsets: ArrayRef,
58 #[slot(2)]
63 pub sizes: ArrayRef,
64 #[slot(3)]
66 pub validity: Option<ArrayRef>,
67}
68
69#[derive(Clone, Debug)]
133pub struct ListViewData {
134 is_zero_copy_to_list: bool,
143}
144
145impl Display for ListViewData {
146 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
147 write!(f, "is_zero_copy_to_list: {}", self.is_zero_copy_to_list)
148 }
149}
150
151pub struct ListViewDataParts {
152 pub elements_dtype: Arc<DType>,
153
154 pub elements: ArrayRef,
156
157 pub offsets: ArrayRef,
159
160 pub sizes: ArrayRef,
162
163 pub validity: Validity,
165}
166
167impl ListViewData {
168 pub(crate) fn make_slots(
169 elements: &ArrayRef,
170 offsets: &ArrayRef,
171 sizes: &ArrayRef,
172 validity: &Validity,
173 len: usize,
174 ) -> ArraySlots {
175 ListViewSlots {
176 elements: elements.clone(),
177 offsets: offsets.clone(),
178 sizes: sizes.clone(),
179 validity: validity_to_child(validity, len),
180 }
181 .into_slots()
182 }
183
184 pub fn new() -> Self {
191 Self {
192 is_zero_copy_to_list: false,
193 }
194 }
195
196 pub fn try_new() -> VortexResult<Self> {
203 Ok(Self::new())
204 }
205
206 pub unsafe fn new_unchecked() -> Self {
226 Self::new()
227 }
228
229 pub fn validate(
231 elements: &ArrayRef,
232 offsets: &ArrayRef,
233 sizes: &ArrayRef,
234 validity: &Validity,
235 ) -> VortexResult<()> {
236 vortex_ensure!(
238 offsets.dtype().is_int() && !offsets.dtype().is_nullable(),
239 "offsets must be non-nullable integer array, got {}",
240 offsets.dtype()
241 );
242 vortex_ensure!(
243 sizes.dtype().is_int() && !sizes.dtype().is_nullable(),
244 "sizes must be non-nullable integer array, got {}",
245 sizes.dtype()
246 );
247
248 vortex_ensure_eq!(
250 offsets.len(),
251 sizes.len(),
252 "offsets and sizes must have the same length",
253 );
254
255 if let Some(validity_len) = validity.maybe_len() {
257 vortex_ensure_eq!(
258 validity_len,
259 offsets.len(),
260 "validity length does not match array size",
261 );
262 }
263
264 if offsets.is_host() && sizes.is_host() {
266 #[allow(clippy::disallowed_methods)]
267 let mut ctx = legacy_session().create_execution_ctx();
268 let offsets_primitive = offsets.clone().execute::<PrimitiveArray>(&mut ctx)?;
269 let sizes_primitive = sizes.clone().execute::<PrimitiveArray>(&mut ctx)?;
270 let offsets_primitive =
273 offsets_primitive.reinterpret_cast(offsets_primitive.ptype().to_unsigned());
274 let sizes_primitive =
275 sizes_primitive.reinterpret_cast(sizes_primitive.ptype().to_unsigned());
276
277 match_each_unsigned_integer_ptype!(offsets_primitive.ptype(), |O| {
279 match_each_unsigned_integer_ptype!(sizes_primitive.ptype(), |S| {
280 let offsets_slice = offsets_primitive.as_slice::<O>();
281 let sizes_slice = sizes_primitive.as_slice::<S>();
282
283 validate_offsets_and_sizes::<O, S>(
284 offsets_slice,
285 sizes_slice,
286 elements.len() as u64,
287 )?;
288 })
289 });
290 }
291
292 Ok(())
293 }
294
295 pub unsafe fn with_zero_copy_to_list(mut self, is_zctl: bool) -> Self {
314 self.is_zero_copy_to_list = is_zctl;
315 self
316 }
317
318 pub fn is_zero_copy_to_list(&self) -> bool {
321 self.is_zero_copy_to_list
322 }
323}
324
325impl Default for ListViewData {
326 fn default() -> Self {
327 Self::new()
328 }
329}
330
331fn fill_referenced_mask<O: IntegerPType, S: IntegerPType>(
337 buf: &mut BitBufferMut,
338 offsets: &[O],
339 sizes: &[S],
340) {
341 let len = offsets.len();
342
343 assert_eq!(
344 len,
345 sizes.len(),
346 "offsets and sizes must be the same length"
347 );
348
349 for i in 0..len {
350 let start: usize = offsets[i].as_();
351 let size: usize = sizes[i].as_();
352 buf.fill_range(start, start + size, true);
353 }
354}
355
356pub trait ListViewArrayExt: ListViewArraySlotsExt {
357 fn nullability(&self) -> crate::dtype::Nullability {
358 match self.as_ref().dtype() {
359 DType::List(_, nullability) => *nullability,
360 _ => unreachable!("ListViewArrayExt requires a list dtype"),
361 }
362 }
363
364 fn listview_validity(&self) -> Validity {
365 child_to_validity(
366 self.as_ref().slots()[ListViewSlots::VALIDITY].as_ref(),
367 self.nullability(),
368 )
369 }
370
371 #[allow(clippy::disallowed_methods)]
372 fn offset_at(&self, index: usize) -> usize {
373 assert!(
374 index < self.as_ref().len(),
375 "Index {index} out of bounds 0..{}",
376 self.as_ref().len()
377 );
378 self.offsets()
379 .as_opt::<Primitive>()
380 .map(|p| match_each_integer_ptype!(p.ptype(), |P| { p.as_slice::<P>()[index].as_() }))
381 .unwrap_or_else(|| {
382 self.offsets()
383 .execute_scalar(index, &mut legacy_session().create_execution_ctx())
384 .vortex_expect("offsets must support execute_scalar")
385 .as_primitive()
386 .as_::<usize>()
387 .vortex_expect("offset must fit in usize")
388 })
389 }
390
391 #[allow(clippy::disallowed_methods)]
392 fn size_at(&self, index: usize) -> usize {
393 assert!(
394 index < self.as_ref().len(),
395 "Index {} out of bounds 0..{}",
396 index,
397 self.as_ref().len()
398 );
399 self.sizes()
400 .as_opt::<Primitive>()
401 .map(|p| match_each_integer_ptype!(p.ptype(), |P| { p.as_slice::<P>()[index].as_() }))
402 .unwrap_or_else(|| {
403 self.sizes()
404 .execute_scalar(index, &mut legacy_session().create_execution_ctx())
405 .vortex_expect("sizes must support execute_scalar")
406 .as_primitive()
407 .as_::<usize>()
408 .vortex_expect("size must fit in usize")
409 })
410 }
411
412 fn list_elements_at(&self, index: usize) -> VortexResult<ArrayRef> {
413 let offset = self.offset_at(index);
414 let size = self.size_at(index);
415 self.elements().slice(offset..offset + size)
416 }
417
418 fn compute_referenced_elements_mask(&self, ctx: &mut ExecutionCtx) -> VortexResult<Mask> {
429 assert!(!self.elements().is_empty());
430 let len = self.elements().len();
431
432 let offsets_primitive = self.offsets().clone().execute::<PrimitiveArray>(ctx)?;
433 let sizes_primitive = self.sizes().clone().execute::<PrimitiveArray>(ctx)?;
434
435 let mut buf = BitBufferMut::new_unset(len);
436
437 let offsets_primitive =
439 offsets_primitive.reinterpret_cast(offsets_primitive.ptype().to_unsigned());
440 let sizes_primitive =
441 sizes_primitive.reinterpret_cast(sizes_primitive.ptype().to_unsigned());
442 match_each_unsigned_integer_ptype!(offsets_primitive.ptype(), |O| {
443 match_each_unsigned_integer_ptype!(sizes_primitive.ptype(), |S| {
444 fill_referenced_mask::<O, S>(
445 &mut buf,
446 offsets_primitive.as_slice::<O>(),
447 sizes_primitive.as_slice::<S>(),
448 );
449 })
450 });
451
452 Ok(Mask::from_buffer(buf.freeze()))
453 }
454
455 fn compute_density(&self, ctx: &mut ExecutionCtx) -> VortexResult<f32> {
459 if self.elements().is_empty() {
460 return Ok(1.0);
461 }
462
463 if self.sizes().is_empty() {
464 return Ok(0.0);
465 }
466
467 let density = match self.compute_referenced_elements_mask(ctx)? {
468 Mask::AllTrue(_) => 1.0,
469 Mask::AllFalse(_) => 0.0,
470 Mask::Values(values) => values.true_count() as f32 / self.elements().len() as f32,
471 };
472
473 Ok(density)
474 }
475
476 fn upper_bound_density(&self, ctx: &mut ExecutionCtx) -> VortexResult<f32> {
483 let n_elts = self.elements().len();
484 if n_elts == 0 {
485 return Ok(1.0);
486 }
487
488 let sizes = self.sizes();
489 if sizes.is_empty() {
490 return Ok(0.0);
491 }
492
493 let sizes_sum = sizes
495 .statistics()
496 .compute_stat(Stat::Sum, ctx)?
497 .vortex_expect("sizes array has integer ptype elements")
498 .as_primitive()
499 .as_::<u64>()
500 .vortex_expect("integer ptypes can be upcast to u64");
501
502 let estimate = (sizes_sum as f32 / n_elts as f32).min(1.0);
505
506 debug_assert!(estimate >= 0.0);
507
508 Ok(estimate)
509 }
510
511 fn referenced_element_bounds(&self, ctx: &mut ExecutionCtx) -> VortexResult<(usize, usize)> {
524 let n_lists = self.as_ref().len();
525 vortex_ensure!(
526 n_lists > 0,
527 "referenced_element_bounds requires a non-empty array"
528 );
529
530 if self.is_zero_copy_to_list() {
531 let start = self.offset_at(0);
532 let end = self.offset_at(n_lists - 1) + self.size_at(n_lists - 1);
533 return Ok((start, end));
534 }
535
536 let start = self
537 .offsets()
538 .statistics()
539 .compute_min::<usize>(ctx)
540 .vortex_expect("offsets must report a usize min statistic");
541
542 let wide_dtype = DType::from(if self.offsets().dtype().as_ptype().is_unsigned_int() {
545 PType::U64
546 } else {
547 PType::I64
548 });
549 let offsets = self.offsets().cast(wide_dtype.clone())?;
550 let sizes = self.sizes().cast(wide_dtype)?;
551 let end = min_max(
552 &offsets.binary(sizes, Operator::Add)?,
553 ctx,
554 NumericalAggregateOpts::default(),
555 )?
556 .vortex_expect("non-empty array must report a min/max")
557 .max
558 .as_primitive()
559 .as_::<usize>()
560 .vortex_expect("max `offset + size` must fit in a usize");
561
562 Ok((start, end))
563 }
564}
565impl<T: TypedArrayRef<ListView>> ListViewArrayExt for T {}
566
567impl Array<ListView> {
568 pub fn new(elements: ArrayRef, offsets: ArrayRef, sizes: ArrayRef, validity: Validity) -> Self {
570 let dtype = DType::List(Arc::new(elements.dtype().clone()), validity.nullability());
571 let len = offsets.len();
572 let slots = ListViewData::make_slots(&elements, &offsets, &sizes, &validity, len);
573 ListViewData::validate(&elements, &offsets, &sizes, &validity)
574 .vortex_expect("`ListViewArray` construction failed");
575 let data = ListViewData::new();
576 unsafe { Array::from_parts_unchecked(ArrayParts::new(ListView, dtype, len, data, slots)) }
577 }
578
579 pub fn try_new(
581 elements: ArrayRef,
582 offsets: ArrayRef,
583 sizes: ArrayRef,
584 validity: Validity,
585 ) -> VortexResult<Self> {
586 let dtype = DType::List(Arc::new(elements.dtype().clone()), validity.nullability());
587 let len = offsets.len();
588 let slots = ListViewData::make_slots(&elements, &offsets, &sizes, &validity, len);
589 ListViewData::validate(&elements, &offsets, &sizes, &validity)?;
590 let data = ListViewData::try_new()?;
591 Ok(unsafe {
592 Array::from_parts_unchecked(ArrayParts::new(ListView, dtype, len, data, slots))
593 })
594 }
595
596 pub unsafe fn new_unchecked(
602 elements: ArrayRef,
603 offsets: ArrayRef,
604 sizes: ArrayRef,
605 validity: Validity,
606 ) -> Self {
607 let dtype = DType::List(Arc::new(elements.dtype().clone()), validity.nullability());
608 let len = offsets.len();
609 let slots = ListViewData::make_slots(&elements, &offsets, &sizes, &validity, len);
610 let data = unsafe { ListViewData::new_unchecked() };
611 unsafe { Array::from_parts_unchecked(ArrayParts::new(ListView, dtype, len, data, slots)) }
612 }
613
614 pub unsafe fn with_zero_copy_to_list(self, is_zctl: bool) -> Self {
620 if cfg!(debug_assertions) && is_zctl {
621 #[allow(clippy::disallowed_methods)]
622 let mut ctx = legacy_session().create_execution_ctx();
623 let offsets_primitive = self
624 .offsets()
625 .clone()
626 .execute::<PrimitiveArray>(&mut ctx)
627 .vortex_expect("offsets must canonicalize to primitive");
628 let sizes_primitive = self
629 .sizes()
630 .clone()
631 .execute::<PrimitiveArray>(&mut ctx)
632 .vortex_expect("sizes must canonicalize to primitive");
633 validate_zctl(self.elements(), offsets_primitive, sizes_primitive)
634 .vortex_expect("Failed to validate zero-copy to list flag");
635 }
636 let dtype = self.dtype().clone();
637 let len = self.len();
638 let slots: ArraySlots = self.slots().iter().cloned().collect();
639 let data = unsafe { self.into_data().with_zero_copy_to_list(is_zctl) };
640 unsafe { Array::from_parts_unchecked(ArrayParts::new(ListView, dtype, len, data, slots)) }
641 }
642
643 pub fn into_data_parts(self) -> ListViewDataParts {
644 let elements = self.slots()[ListViewSlots::ELEMENTS]
645 .clone()
646 .vortex_expect("ListViewArray elements slot");
647 let offsets = self.slots()[ListViewSlots::OFFSETS]
648 .clone()
649 .vortex_expect("ListViewArray offsets slot");
650 let sizes = self.slots()[ListViewSlots::SIZES]
651 .clone()
652 .vortex_expect("ListViewArray sizes slot");
653 let validity = self.listview_validity();
654 ListViewDataParts {
655 elements_dtype: Arc::new(elements.dtype().clone()),
656 elements,
657 offsets,
658 sizes,
659 validity,
660 }
661 }
662}
663
664fn validate_offsets_and_sizes<O, S>(
666 offsets_slice: &[O],
667 sizes_slice: &[S],
668 elements_len: u64,
669) -> VortexResult<()>
670where
671 O: IntegerPType,
672 S: IntegerPType,
673{
674 debug_assert_eq!(offsets_slice.len(), sizes_slice.len());
675
676 #[allow(clippy::absurd_extreme_comparisons, unused_comparisons)]
677 for i in 0..offsets_slice.len() {
678 let offset = offsets_slice[i];
679 let size = sizes_slice[i];
680
681 vortex_ensure!(offset >= O::zero(), "cannot have negative offsets");
682 vortex_ensure!(size >= S::zero(), "cannot have negative size");
683
684 let offset_u64 = offset
685 .to_u64()
686 .ok_or_else(|| vortex_err!("offset[{i}] = {offset:?} cannot be converted to u64"))?;
687
688 let size_u64 = size
689 .to_u64()
690 .ok_or_else(|| vortex_err!("size[{i}] = {size:?} cannot be converted to u64"))?;
691
692 let end = offset_u64.checked_add(size_u64).ok_or_else(|| {
694 vortex_err!("offset[{i}] ({offset_u64}) + size[{i}] ({size_u64}) would overflow u64")
695 })?;
696
697 if offset_u64 == elements_len {
698 vortex_ensure_eq!(
699 size_u64,
700 0,
701 "views to the end of the elements array (length {elements_len}) must have size 0",
702 );
703 }
704
705 vortex_ensure!(
706 end <= elements_len,
707 "offset[{i}] + size[{i}] = {offset_u64} + {size_u64} = {end} \
708 exceeds elements length {elements_len}",
709 );
710 }
711
712 Ok(())
713}
714
715#[allow(clippy::disallowed_methods)]
718fn validate_zctl(
719 elements: &ArrayRef,
720 offsets_primitive: PrimitiveArray,
721 sizes_primitive: PrimitiveArray,
722) -> VortexResult<()> {
723 let mut ctx = legacy_session().create_execution_ctx();
726 if let Some(is_sorted) = offsets_primitive.statistics().compute_is_sorted(&mut ctx) {
727 vortex_ensure!(is_sorted, "offsets must be sorted");
728 } else {
729 vortex_bail!("offsets must report is_sorted statistic");
730 }
731
732 fn validate_monotonic_ends<O: IntegerPType, S: IntegerPType>(
735 offsets_slice: &[O],
736 sizes_slice: &[S],
737 len: usize,
738 ) -> VortexResult<()> {
739 let mut max_end = 0usize;
740
741 for i in 0..len {
742 let offset = offsets_slice[i].to_usize().unwrap_or(usize::MAX);
743 let size = sizes_slice[i].to_usize().unwrap_or(usize::MAX);
744
745 vortex_ensure!(
747 offset >= max_end,
748 "Zero-copy-to-list requires views to be non-overlapping and ordered: \
749 view[{}] starts at {} but previous views extend to {}",
750 i,
751 offset,
752 max_end
753 );
754
755 let end = offset.saturating_add(size);
757 max_end = max_end.max(end);
758 }
759
760 Ok(())
761 }
762
763 let offsets_dtype = offsets_primitive.dtype();
764 let sizes_dtype = sizes_primitive.dtype();
765 let len = offsets_primitive.len();
766
767 let offsets_unsigned =
769 offsets_primitive.reinterpret_cast(offsets_dtype.as_ptype().to_unsigned());
770 let sizes_unsigned = sizes_primitive.reinterpret_cast(sizes_dtype.as_ptype().to_unsigned());
771
772 match_each_unsigned_integer_ptype!(offsets_unsigned.ptype(), |O| {
774 match_each_unsigned_integer_ptype!(sizes_unsigned.ptype(), |S| {
775 let offsets_slice = offsets_unsigned.as_slice::<O>();
776 let sizes_slice = sizes_unsigned.as_slice::<S>();
777
778 validate_monotonic_ends(offsets_slice, sizes_slice, len)?;
779 })
780 });
781
782 let mut element_references = vec![0u8; elements.len()];
787
788 fn count_references<O: IntegerPType, S: IntegerPType>(
789 element_references: &mut [u8],
790 offsets_primitive: PrimitiveArray,
791 sizes_primitive: PrimitiveArray,
792 ) {
793 let offsets_slice = offsets_primitive.as_slice::<O>();
794 let sizes_slice = sizes_primitive.as_slice::<S>();
795
796 for i in 0..offsets_slice.len() {
799 let offset: usize = offsets_slice[i].as_();
800 let size: usize = sizes_slice[i].as_();
801 for j in offset..offset + size {
802 element_references[j] = element_references[j].saturating_add(1);
803 }
804 }
805 }
806
807 match_each_unsigned_integer_ptype!(offsets_unsigned.ptype(), |O| {
808 match_each_unsigned_integer_ptype!(sizes_unsigned.ptype(), |S| {
809 count_references::<O, S>(&mut element_references, offsets_unsigned, sizes_unsigned);
810 })
811 });
812
813 let leftmost_used = element_references
815 .iter()
816 .position(|&references| references != 0);
817 let rightmost_used = element_references
818 .iter()
819 .rposition(|&references| references != 0);
820
821 if let (Some(first_ref), Some(last_ref)) = (leftmost_used, rightmost_used) {
822 vortex_ensure!(
823 element_references[first_ref..=last_ref]
824 .iter()
825 .all(|&references| references != 0),
826 "found gap in elements array between first and last referenced elements"
827 );
828 }
829
830 vortex_ensure!(element_references.iter().all(|&references| references <= 1));
831
832 Ok(())
833}