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_ensure;
13use vortex_error::vortex_ensure_eq;
14use vortex_error::vortex_err;
15use vortex_mask::Mask;
16
17use crate::ArrayRef;
18use crate::ArraySlots;
19use crate::ExecutionCtx;
20use crate::aggregate_fn::NumericalAggregateOpts;
21use crate::aggregate_fn::fns::min_max::min_max;
22use crate::array::Array;
23use crate::array::ArrayParts;
24use crate::array::TypedArrayRef;
25use crate::array::child_to_validity;
26use crate::array::validity_to_child;
27use crate::array_slots;
28use crate::arrays::ListView;
29use crate::arrays::Primitive;
30use crate::arrays::PrimitiveArray;
31use crate::arrays::bool;
32use crate::arrays::primitive::PrimitiveArrayExt;
33use crate::builtins::ArrayBuiltins;
34use crate::dtype::DType;
35use crate::dtype::IntegerPType;
36use crate::dtype::PType;
37use crate::expr::stats::Stat;
38use crate::match_each_integer_ptype;
39use crate::match_each_unsigned_integer_ptype;
40use crate::scalar_fn::fns::operators::Operator;
41use crate::validity::Validity;
42
43#[array_slots(ListView)]
44pub struct ListViewSlots {
45 #[slot(0)]
48 pub elements: ArrayRef,
49 #[slot(1)]
54 pub offsets: ArrayRef,
55 #[slot(2)]
60 pub sizes: ArrayRef,
61 #[slot(3)]
63 pub validity: Option<ArrayRef>,
64}
65
66#[derive(Clone, Debug)]
132pub struct ListViewData {
133 is_zero_copy_to_list: bool,
142}
143
144impl Display for ListViewData {
145 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
146 write!(f, "is_zero_copy_to_list: {}", self.is_zero_copy_to_list)
147 }
148}
149
150pub struct ListViewDataParts {
151 pub elements_dtype: Arc<DType>,
152
153 pub elements: ArrayRef,
155
156 pub offsets: ArrayRef,
158
159 pub sizes: ArrayRef,
161
162 pub validity: Validity,
164}
165
166impl ListViewData {
167 pub(crate) fn make_slots(
168 elements: &ArrayRef,
169 offsets: &ArrayRef,
170 sizes: &ArrayRef,
171 validity: &Validity,
172 len: usize,
173 ) -> ArraySlots {
174 ListViewSlots {
175 elements: elements.clone(),
176 offsets: offsets.clone(),
177 sizes: sizes.clone(),
178 validity: validity_to_child(validity, len),
179 }
180 .into_slots()
181 }
182
183 pub fn new() -> Self {
190 Self {
191 is_zero_copy_to_list: false,
192 }
193 }
194
195 pub fn try_new() -> VortexResult<Self> {
202 Ok(Self::new())
203 }
204
205 pub unsafe fn new_unchecked() -> Self {
225 Self::new()
226 }
227
228 pub fn validate(
230 elements: &ArrayRef,
231 offsets: &ArrayRef,
232 sizes: &ArrayRef,
233 validity: &Validity,
234 ctx: &mut ExecutionCtx,
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 let offsets_primitive = offsets.clone().execute::<PrimitiveArray>(ctx)?;
267 let sizes_primitive = sizes.clone().execute::<PrimitiveArray>(ctx)?;
268 let offsets_primitive =
271 offsets_primitive.reinterpret_cast(offsets_primitive.ptype().to_unsigned());
272 let sizes_primitive =
273 sizes_primitive.reinterpret_cast(sizes_primitive.ptype().to_unsigned());
274
275 match_each_unsigned_integer_ptype!(offsets_primitive.ptype(), |O| {
277 match_each_unsigned_integer_ptype!(sizes_primitive.ptype(), |S| {
278 let offsets_slice = offsets_primitive.as_slice::<O>();
279 let sizes_slice = sizes_primitive.as_slice::<S>();
280
281 validate_offsets_and_sizes::<O, S>(
282 offsets_slice,
283 sizes_slice,
284 elements.len() as u64,
285 )?;
286 })
287 });
288 }
289
290 Ok(())
291 }
292
293 pub unsafe fn with_zero_copy_to_list(mut self, is_zctl: bool) -> Self {
312 self.is_zero_copy_to_list = is_zctl;
313 self
314 }
315
316 pub fn is_zero_copy_to_list(&self) -> bool {
319 self.is_zero_copy_to_list
320 }
321}
322
323impl Default for ListViewData {
324 fn default() -> Self {
325 Self::new()
326 }
327}
328
329fn fill_referenced_mask<O: IntegerPType, S: IntegerPType>(
335 buf: &mut BitBufferMut,
336 offsets: &[O],
337 sizes: &[S],
338) {
339 let len = offsets.len();
340
341 assert_eq!(
342 len,
343 sizes.len(),
344 "offsets and sizes must be the same length"
345 );
346
347 for i in 0..len {
348 let start: usize = offsets[i].as_();
349 let size: usize = sizes[i].as_();
350 buf.fill_range(start, start + size, true);
351 }
352}
353
354pub trait ListViewArrayExt: ListViewArraySlotsExt {
355 fn nullability(&self) -> crate::dtype::Nullability {
356 match self.dtype() {
357 DType::List(_, nullability) => *nullability,
358 _ => unreachable!("ListViewArrayExt requires a list dtype"),
359 }
360 }
361
362 fn listview_validity(&self) -> Validity {
363 child_to_validity(
364 self.slots()[ListViewSlots::VALIDITY].as_ref(),
365 self.nullability(),
366 )
367 }
368
369 fn offset_at(&self, index: usize, ctx: &mut ExecutionCtx) -> usize {
370 assert!(
371 index < self.len(),
372 "Index {index} out of bounds 0..{}",
373 self.len()
374 );
375 self.offsets()
376 .as_opt::<Primitive>()
377 .map(|p| match_each_integer_ptype!(p.ptype(), |P| { p.as_slice::<P>()[index].as_() }))
378 .unwrap_or_else(|| {
379 self.offsets()
380 .execute_scalar(index, ctx)
381 .vortex_expect("offsets must support execute_scalar")
382 .as_primitive()
383 .as_::<usize>()
384 .vortex_expect("offset must fit in usize")
385 })
386 }
387
388 fn size_at(&self, index: usize, ctx: &mut ExecutionCtx) -> usize {
389 assert!(
390 index < self.len(),
391 "Index {} out of bounds 0..{}",
392 index,
393 self.len()
394 );
395 self.sizes()
396 .as_opt::<Primitive>()
397 .map(|p| match_each_integer_ptype!(p.ptype(), |P| { p.as_slice::<P>()[index].as_() }))
398 .unwrap_or_else(|| {
399 self.sizes()
400 .execute_scalar(index, ctx)
401 .vortex_expect("sizes must support execute_scalar")
402 .as_primitive()
403 .as_::<usize>()
404 .vortex_expect("size must fit in usize")
405 })
406 }
407
408 fn list_elements_at(&self, index: usize, ctx: &mut ExecutionCtx) -> VortexResult<ArrayRef> {
409 let offset = self.offset_at(index, ctx);
410 let size = self.size_at(index, ctx);
411 self.elements().slice(offset..offset + size)
412 }
413
414 fn compute_referenced_elements_mask(&self, ctx: &mut ExecutionCtx) -> VortexResult<Mask> {
425 assert!(!self.elements().is_empty());
426 let len = self.elements().len();
427
428 let offsets_primitive = self.offsets().clone().execute::<PrimitiveArray>(ctx)?;
429 let sizes_primitive = self.sizes().clone().execute::<PrimitiveArray>(ctx)?;
430
431 let mut buf = BitBufferMut::new_unset(len);
432
433 let offsets_primitive =
435 offsets_primitive.reinterpret_cast(offsets_primitive.ptype().to_unsigned());
436 let sizes_primitive =
437 sizes_primitive.reinterpret_cast(sizes_primitive.ptype().to_unsigned());
438 match_each_unsigned_integer_ptype!(offsets_primitive.ptype(), |O| {
439 match_each_unsigned_integer_ptype!(sizes_primitive.ptype(), |S| {
440 fill_referenced_mask::<O, S>(
441 &mut buf,
442 offsets_primitive.as_slice::<O>(),
443 sizes_primitive.as_slice::<S>(),
444 );
445 })
446 });
447
448 Ok(Mask::from_buffer(buf.freeze()))
449 }
450
451 fn compute_density(&self, ctx: &mut ExecutionCtx) -> VortexResult<f32> {
455 if self.elements().is_empty() {
456 return Ok(1.0);
457 }
458
459 if self.sizes().is_empty() {
460 return Ok(0.0);
461 }
462
463 let density = match self.compute_referenced_elements_mask(ctx)? {
464 Mask::AllTrue(_) => 1.0,
465 Mask::AllFalse(_) => 0.0,
466 Mask::Values(values) => values.true_count() as f32 / self.elements().len() as f32,
467 };
468
469 Ok(density)
470 }
471
472 fn upper_bound_density(&self, ctx: &mut ExecutionCtx) -> VortexResult<f32> {
479 let n_elts = self.elements().len();
480 if n_elts == 0 {
481 return Ok(1.0);
482 }
483
484 let sizes = self.sizes();
485 if sizes.is_empty() {
486 return Ok(0.0);
487 }
488
489 let sizes_sum = sizes
491 .statistics()
492 .compute_stat(Stat::Sum, ctx)?
493 .vortex_expect("sizes array has integer ptype elements")
494 .as_primitive()
495 .as_::<u64>()
496 .vortex_expect("integer ptypes can be upcast to u64");
497
498 let estimate = (sizes_sum as f32 / n_elts as f32).min(1.0);
501
502 debug_assert!(estimate >= 0.0);
503
504 Ok(estimate)
505 }
506
507 fn referenced_element_bounds(&self, ctx: &mut ExecutionCtx) -> VortexResult<(usize, usize)> {
520 let n_lists = self.len();
521 vortex_ensure!(
522 n_lists > 0,
523 "referenced_element_bounds requires a non-empty array"
524 );
525
526 if self.is_zero_copy_to_list() {
527 let start = self.offset_at(0, ctx);
528 let end = self.offset_at(n_lists - 1, ctx) + self.size_at(n_lists - 1, ctx);
529 return Ok((start, end));
530 }
531
532 let start = self
533 .offsets()
534 .statistics()
535 .compute_min::<usize>(ctx)
536 .vortex_expect("offsets must report a usize min statistic");
537
538 let wide_dtype = DType::from(if self.offsets().dtype().as_ptype().is_unsigned_int() {
541 PType::U64
542 } else {
543 PType::I64
544 });
545 let offsets = self.offsets().cast(wide_dtype.clone())?;
546 let sizes = self.sizes().cast(wide_dtype)?;
547 let end = min_max(
548 &offsets.binary(sizes, Operator::Add)?,
549 ctx,
550 NumericalAggregateOpts::default(),
551 )?
552 .vortex_expect("non-empty array must report a min/max")
553 .max
554 .as_primitive()
555 .as_::<usize>()
556 .vortex_expect("max `offset + size` must fit in a usize");
557
558 Ok((start, end))
559 }
560}
561impl<T: TypedArrayRef<ListView>> ListViewArrayExt for T {}
562
563impl Array<ListView> {
564 pub fn new(
566 elements: ArrayRef,
567 offsets: ArrayRef,
568 sizes: ArrayRef,
569 validity: Validity,
570 ctx: &mut ExecutionCtx,
571 ) -> Self {
572 let dtype = DType::List(Arc::new(elements.dtype().clone()), validity.nullability());
573 let len = offsets.len();
574 let slots = ListViewData::make_slots(&elements, &offsets, &sizes, &validity, len);
575 ListViewData::validate(&elements, &offsets, &sizes, &validity, ctx)
576 .vortex_expect("`ListViewArray` construction failed");
577 let data = ListViewData::new();
578 unsafe { Array::from_parts_unchecked(ArrayParts::new(ListView, dtype, len, data, slots)) }
579 }
580
581 pub fn try_new(
583 elements: ArrayRef,
584 offsets: ArrayRef,
585 sizes: ArrayRef,
586 validity: Validity,
587 ctx: &mut ExecutionCtx,
588 ) -> VortexResult<Self> {
589 let dtype = DType::List(Arc::new(elements.dtype().clone()), validity.nullability());
590 let len = offsets.len();
591 let slots = ListViewData::make_slots(&elements, &offsets, &sizes, &validity, len);
592 ListViewData::validate(&elements, &offsets, &sizes, &validity, ctx)?;
593 let data = ListViewData::try_new()?;
594 Ok(unsafe {
595 Array::from_parts_unchecked(ArrayParts::new(ListView, dtype, len, data, slots))
596 })
597 }
598
599 pub unsafe fn new_unchecked(
605 elements: ArrayRef,
606 offsets: ArrayRef,
607 sizes: ArrayRef,
608 validity: Validity,
609 ) -> Self {
610 let dtype = DType::List(Arc::new(elements.dtype().clone()), validity.nullability());
611 let len = offsets.len();
612 let slots = ListViewData::make_slots(&elements, &offsets, &sizes, &validity, len);
613 let data = unsafe { ListViewData::new_unchecked() };
614 unsafe { Array::from_parts_unchecked(ArrayParts::new(ListView, dtype, len, data, slots)) }
615 }
616
617 pub unsafe fn with_zero_copy_to_list(self, is_zctl: bool) -> Self {
623 if cfg!(debug_assertions)
626 && is_zctl
627 && self.offsets().is_host()
628 && self.sizes().is_host()
629 && let Some(offsets_primitive) = self.offsets().as_opt::<Primitive>()
630 && let Some(sizes_primitive) = self.sizes().as_opt::<Primitive>()
631 {
632 validate_zctl(
633 self.elements(),
634 offsets_primitive.into_owned(),
635 sizes_primitive.into_owned(),
636 )
637 .vortex_expect("Failed to validate zero-copy to list flag");
638 }
639 let dtype = self.dtype().clone();
640 let len = self.len();
641 let slots: ArraySlots = self.slots().iter().cloned().collect();
642 let data = unsafe { self.into_data().with_zero_copy_to_list(is_zctl) };
643 unsafe { Array::from_parts_unchecked(ArrayParts::new(ListView, dtype, len, data, slots)) }
644 }
645
646 pub fn into_data_parts(self) -> ListViewDataParts {
647 let elements = self.slots()[ListViewSlots::ELEMENTS]
648 .clone()
649 .vortex_expect("ListViewArray elements slot");
650 let offsets = self.slots()[ListViewSlots::OFFSETS]
651 .clone()
652 .vortex_expect("ListViewArray offsets slot");
653 let sizes = self.slots()[ListViewSlots::SIZES]
654 .clone()
655 .vortex_expect("ListViewArray sizes slot");
656 let validity = self.listview_validity();
657 ListViewDataParts {
658 elements_dtype: Arc::new(elements.dtype().clone()),
659 elements,
660 offsets,
661 sizes,
662 validity,
663 }
664 }
665}
666
667fn validate_offsets_and_sizes<O, S>(
669 offsets_slice: &[O],
670 sizes_slice: &[S],
671 elements_len: u64,
672) -> VortexResult<()>
673where
674 O: IntegerPType,
675 S: IntegerPType,
676{
677 debug_assert_eq!(offsets_slice.len(), sizes_slice.len());
678
679 #[allow(clippy::absurd_extreme_comparisons, unused_comparisons)]
680 for i in 0..offsets_slice.len() {
681 let offset = offsets_slice[i];
682 let size = sizes_slice[i];
683
684 vortex_ensure!(offset >= O::zero(), "cannot have negative offsets");
685 vortex_ensure!(size >= S::zero(), "cannot have negative size");
686
687 let offset_u64 = offset
688 .to_u64()
689 .ok_or_else(|| vortex_err!("offset[{i}] = {offset:?} cannot be converted to u64"))?;
690
691 let size_u64 = size
692 .to_u64()
693 .ok_or_else(|| vortex_err!("size[{i}] = {size:?} cannot be converted to u64"))?;
694
695 let end = offset_u64.checked_add(size_u64).ok_or_else(|| {
697 vortex_err!("offset[{i}] ({offset_u64}) + size[{i}] ({size_u64}) would overflow u64")
698 })?;
699
700 if offset_u64 == elements_len {
701 vortex_ensure_eq!(
702 size_u64,
703 0,
704 "views to the end of the elements array (length {elements_len}) must have size 0",
705 );
706 }
707
708 vortex_ensure!(
709 end <= elements_len,
710 "offset[{i}] + size[{i}] = {offset_u64} + {size_u64} = {end} \
711 exceeds elements length {elements_len}",
712 );
713 }
714
715 Ok(())
716}
717
718fn validate_zctl(
721 elements: &ArrayRef,
722 offsets_primitive: PrimitiveArray,
723 sizes_primitive: PrimitiveArray,
724) -> VortexResult<()> {
725 let offsets_sorted = match_each_integer_ptype!(offsets_primitive.ptype(), |P| {
728 offsets_primitive.as_slice::<P>().is_sorted()
729 });
730 vortex_ensure!(offsets_sorted, "offsets must be sorted");
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}