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_err;
15use vortex_mask::Mask;
16
17use crate::ArrayRef;
18use crate::ArraySlots;
19use crate::ExecutionCtx;
20use crate::VortexSessionExecute;
21use crate::aggregate_fn::NumericalAggregateOpts;
22use crate::aggregate_fn::fns::min_max::min_max;
23use crate::array::Array;
24use crate::array::ArrayParts;
25use crate::array::TypedArrayRef;
26use crate::array::child_to_validity;
27use crate::array::validity_to_child;
28use crate::array_slots;
29use crate::arrays::ListView;
30use crate::arrays::Primitive;
31use crate::arrays::PrimitiveArray;
32use crate::arrays::bool;
33use crate::arrays::primitive::PrimitiveArrayExt;
34use crate::builtins::ArrayBuiltins;
35use crate::dtype::DType;
36use crate::dtype::IntegerPType;
37use crate::dtype::PType;
38use crate::expr::stats::Stat;
39use crate::legacy_session;
40use crate::match_each_integer_ptype;
41use crate::match_each_unsigned_integer_ptype;
42use crate::scalar_fn::fns::operators::Operator;
43use crate::validity::Validity;
44
45#[array_slots(ListView)]
46pub struct ListViewSlots {
47 #[slot(0)]
50 pub elements: ArrayRef,
51 #[slot(1)]
56 pub offsets: ArrayRef,
57 #[slot(2)]
62 pub sizes: ArrayRef,
63 #[slot(3)]
65 pub validity: Option<ArrayRef>,
66}
67
68#[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 ) -> VortexResult<()> {
235 vortex_ensure!(
237 offsets.dtype().is_int() && !offsets.dtype().is_nullable(),
238 "offsets must be non-nullable integer array, got {}",
239 offsets.dtype()
240 );
241 vortex_ensure!(
242 sizes.dtype().is_int() && !sizes.dtype().is_nullable(),
243 "sizes must be non-nullable integer array, got {}",
244 sizes.dtype()
245 );
246
247 vortex_ensure!(
249 offsets.len() == sizes.len(),
250 "offsets and sizes must have the same length, got {} and {}",
251 offsets.len(),
252 sizes.len()
253 );
254
255 if let Some(validity_len) = validity.maybe_len() {
257 vortex_ensure!(
258 validity_len == offsets.len(),
259 "validity with size {validity_len} does not match array size {}",
260 offsets.len()
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 {
577 Array::from_parts_unchecked(
578 ArrayParts::new(ListView, dtype, len, data).with_slots(slots),
579 )
580 }
581 }
582
583 pub fn try_new(
585 elements: ArrayRef,
586 offsets: ArrayRef,
587 sizes: ArrayRef,
588 validity: Validity,
589 ) -> VortexResult<Self> {
590 let dtype = DType::List(Arc::new(elements.dtype().clone()), validity.nullability());
591 let len = offsets.len();
592 let slots = ListViewData::make_slots(&elements, &offsets, &sizes, &validity, len);
593 ListViewData::validate(&elements, &offsets, &sizes, &validity)?;
594 let data = ListViewData::try_new()?;
595 Ok(unsafe {
596 Array::from_parts_unchecked(
597 ArrayParts::new(ListView, dtype, len, data).with_slots(slots),
598 )
599 })
600 }
601
602 pub unsafe fn new_unchecked(
608 elements: ArrayRef,
609 offsets: ArrayRef,
610 sizes: ArrayRef,
611 validity: Validity,
612 ) -> Self {
613 let dtype = DType::List(Arc::new(elements.dtype().clone()), validity.nullability());
614 let len = offsets.len();
615 let slots = ListViewData::make_slots(&elements, &offsets, &sizes, &validity, len);
616 let data = unsafe { ListViewData::new_unchecked() };
617 unsafe {
618 Array::from_parts_unchecked(
619 ArrayParts::new(ListView, dtype, len, data).with_slots(slots),
620 )
621 }
622 }
623
624 pub unsafe fn with_zero_copy_to_list(self, is_zctl: bool) -> Self {
630 if cfg!(debug_assertions) && is_zctl {
631 #[allow(clippy::disallowed_methods)]
632 let mut ctx = legacy_session().create_execution_ctx();
633 let offsets_primitive = self
634 .offsets()
635 .clone()
636 .execute::<PrimitiveArray>(&mut ctx)
637 .vortex_expect("offsets must canonicalize to primitive");
638 let sizes_primitive = self
639 .sizes()
640 .clone()
641 .execute::<PrimitiveArray>(&mut ctx)
642 .vortex_expect("sizes must canonicalize to primitive");
643 validate_zctl(self.elements(), offsets_primitive, sizes_primitive)
644 .vortex_expect("Failed to validate zero-copy to list flag");
645 }
646 let dtype = self.dtype().clone();
647 let len = self.len();
648 let slots: ArraySlots = self.slots().iter().cloned().collect();
649 let data = unsafe { self.into_data().with_zero_copy_to_list(is_zctl) };
650 unsafe {
651 Array::from_parts_unchecked(
652 ArrayParts::new(ListView, dtype, len, data).with_slots(slots),
653 )
654 }
655 }
656
657 pub fn into_data_parts(self) -> ListViewDataParts {
658 let elements = self.slots()[ListViewSlots::ELEMENTS]
659 .clone()
660 .vortex_expect("ListViewArray elements slot");
661 let offsets = self.slots()[ListViewSlots::OFFSETS]
662 .clone()
663 .vortex_expect("ListViewArray offsets slot");
664 let sizes = self.slots()[ListViewSlots::SIZES]
665 .clone()
666 .vortex_expect("ListViewArray sizes slot");
667 let validity = self.listview_validity();
668 ListViewDataParts {
669 elements_dtype: Arc::new(elements.dtype().clone()),
670 elements,
671 offsets,
672 sizes,
673 validity,
674 }
675 }
676}
677
678fn validate_offsets_and_sizes<O, S>(
680 offsets_slice: &[O],
681 sizes_slice: &[S],
682 elements_len: u64,
683) -> VortexResult<()>
684where
685 O: IntegerPType,
686 S: IntegerPType,
687{
688 debug_assert_eq!(offsets_slice.len(), sizes_slice.len());
689
690 #[allow(clippy::absurd_extreme_comparisons, unused_comparisons)]
691 for i in 0..offsets_slice.len() {
692 let offset = offsets_slice[i];
693 let size = sizes_slice[i];
694
695 vortex_ensure!(offset >= O::zero(), "cannot have negative offsets");
696 vortex_ensure!(size >= S::zero(), "cannot have negative size");
697
698 let offset_u64 = offset
699 .to_u64()
700 .ok_or_else(|| vortex_err!("offset[{i}] = {offset:?} cannot be converted to u64"))?;
701
702 let size_u64 = size
703 .to_u64()
704 .ok_or_else(|| vortex_err!("size[{i}] = {size:?} cannot be converted to u64"))?;
705
706 let end = offset_u64.checked_add(size_u64).ok_or_else(|| {
708 vortex_err!("offset[{i}] ({offset_u64}) + size[{i}] ({size_u64}) would overflow u64")
709 })?;
710
711 if offset_u64 == elements_len {
712 vortex_ensure!(
713 size_u64 == 0,
714 "views to the end of the elements array (length {elements_len}) must have size 0 \
715 (had size {size_u64})"
716 );
717 }
718
719 vortex_ensure!(
720 end <= elements_len,
721 "offset[{i}] + size[{i}] = {offset_u64} + {size_u64} = {end} \
722 exceeds elements length {elements_len}",
723 );
724 }
725
726 Ok(())
727}
728
729#[allow(clippy::disallowed_methods)]
732fn validate_zctl(
733 elements: &ArrayRef,
734 offsets_primitive: PrimitiveArray,
735 sizes_primitive: PrimitiveArray,
736) -> VortexResult<()> {
737 let mut ctx = legacy_session().create_execution_ctx();
740 if let Some(is_sorted) = offsets_primitive.statistics().compute_is_sorted(&mut ctx) {
741 vortex_ensure!(is_sorted, "offsets must be sorted");
742 } else {
743 vortex_bail!("offsets must report is_sorted statistic");
744 }
745
746 fn validate_monotonic_ends<O: IntegerPType, S: IntegerPType>(
749 offsets_slice: &[O],
750 sizes_slice: &[S],
751 len: usize,
752 ) -> VortexResult<()> {
753 let mut max_end = 0usize;
754
755 for i in 0..len {
756 let offset = offsets_slice[i].to_usize().unwrap_or(usize::MAX);
757 let size = sizes_slice[i].to_usize().unwrap_or(usize::MAX);
758
759 vortex_ensure!(
761 offset >= max_end,
762 "Zero-copy-to-list requires views to be non-overlapping and ordered: \
763 view[{}] starts at {} but previous views extend to {}",
764 i,
765 offset,
766 max_end
767 );
768
769 let end = offset.saturating_add(size);
771 max_end = max_end.max(end);
772 }
773
774 Ok(())
775 }
776
777 let offsets_dtype = offsets_primitive.dtype();
778 let sizes_dtype = sizes_primitive.dtype();
779 let len = offsets_primitive.len();
780
781 let offsets_unsigned =
783 offsets_primitive.reinterpret_cast(offsets_dtype.as_ptype().to_unsigned());
784 let sizes_unsigned = sizes_primitive.reinterpret_cast(sizes_dtype.as_ptype().to_unsigned());
785
786 match_each_unsigned_integer_ptype!(offsets_unsigned.ptype(), |O| {
788 match_each_unsigned_integer_ptype!(sizes_unsigned.ptype(), |S| {
789 let offsets_slice = offsets_unsigned.as_slice::<O>();
790 let sizes_slice = sizes_unsigned.as_slice::<S>();
791
792 validate_monotonic_ends(offsets_slice, sizes_slice, len)?;
793 })
794 });
795
796 let mut element_references = vec![0u8; elements.len()];
801
802 fn count_references<O: IntegerPType, S: IntegerPType>(
803 element_references: &mut [u8],
804 offsets_primitive: PrimitiveArray,
805 sizes_primitive: PrimitiveArray,
806 ) {
807 let offsets_slice = offsets_primitive.as_slice::<O>();
808 let sizes_slice = sizes_primitive.as_slice::<S>();
809
810 for i in 0..offsets_slice.len() {
813 let offset: usize = offsets_slice[i].as_();
814 let size: usize = sizes_slice[i].as_();
815 for j in offset..offset + size {
816 element_references[j] = element_references[j].saturating_add(1);
817 }
818 }
819 }
820
821 match_each_unsigned_integer_ptype!(offsets_unsigned.ptype(), |O| {
822 match_each_unsigned_integer_ptype!(sizes_unsigned.ptype(), |S| {
823 count_references::<O, S>(&mut element_references, offsets_unsigned, sizes_unsigned);
824 })
825 });
826
827 let leftmost_used = element_references
829 .iter()
830 .position(|&references| references != 0);
831 let rightmost_used = element_references
832 .iter()
833 .rposition(|&references| references != 0);
834
835 if let (Some(first_ref), Some(last_ref)) = (leftmost_used, rightmost_used) {
836 vortex_ensure!(
837 element_references[first_ref..=last_ref]
838 .iter()
839 .all(|&references| references != 0),
840 "found gap in elements array between first and last referenced elements"
841 );
842 }
843
844 vortex_ensure!(element_references.iter().all(|&references| references <= 1));
845
846 Ok(())
847}