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 pub elements: ArrayRef,
50 pub offsets: ArrayRef,
55 pub sizes: ArrayRef,
60 pub validity: Option<ArrayRef>,
62}
63
64#[derive(Clone, Debug)]
128pub struct ListViewData {
129 is_zero_copy_to_list: bool,
138}
139
140impl Display for ListViewData {
141 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
142 write!(f, "is_zero_copy_to_list: {}", self.is_zero_copy_to_list)
143 }
144}
145
146pub struct ListViewDataParts {
147 pub elements_dtype: Arc<DType>,
148
149 pub elements: ArrayRef,
151
152 pub offsets: ArrayRef,
154
155 pub sizes: ArrayRef,
157
158 pub validity: Validity,
160}
161
162impl ListViewData {
163 pub(crate) fn make_slots(
164 elements: &ArrayRef,
165 offsets: &ArrayRef,
166 sizes: &ArrayRef,
167 validity: &Validity,
168 len: usize,
169 ) -> ArraySlots {
170 ListViewSlots {
171 elements: elements.clone(),
172 offsets: offsets.clone(),
173 sizes: sizes.clone(),
174 validity: validity_to_child(validity, len),
175 }
176 .into_slots()
177 }
178
179 pub fn new() -> Self {
186 Self {
187 is_zero_copy_to_list: false,
188 }
189 }
190
191 pub fn try_new() -> VortexResult<Self> {
198 Ok(Self::new())
199 }
200
201 pub unsafe fn new_unchecked() -> Self {
221 Self::new()
222 }
223
224 pub fn validate(
226 elements: &ArrayRef,
227 offsets: &ArrayRef,
228 sizes: &ArrayRef,
229 validity: &Validity,
230 ) -> VortexResult<()> {
231 vortex_ensure!(
233 offsets.dtype().is_int() && !offsets.dtype().is_nullable(),
234 "offsets must be non-nullable integer array, got {}",
235 offsets.dtype()
236 );
237 vortex_ensure!(
238 sizes.dtype().is_int() && !sizes.dtype().is_nullable(),
239 "sizes must be non-nullable integer array, got {}",
240 sizes.dtype()
241 );
242
243 vortex_ensure!(
245 offsets.len() == sizes.len(),
246 "offsets and sizes must have the same length, got {} and {}",
247 offsets.len(),
248 sizes.len()
249 );
250
251 if let Some(validity_len) = validity.maybe_len() {
253 vortex_ensure!(
254 validity_len == offsets.len(),
255 "validity with size {validity_len} does not match array size {}",
256 offsets.len()
257 );
258 }
259
260 if offsets.is_host() && sizes.is_host() {
262 #[allow(clippy::disallowed_methods)]
263 let mut ctx = legacy_session().create_execution_ctx();
264 let offsets_primitive = offsets.clone().execute::<PrimitiveArray>(&mut ctx)?;
265 let sizes_primitive = sizes.clone().execute::<PrimitiveArray>(&mut ctx)?;
266 let offsets_primitive =
269 offsets_primitive.reinterpret_cast(offsets_primitive.ptype().to_unsigned());
270 let sizes_primitive =
271 sizes_primitive.reinterpret_cast(sizes_primitive.ptype().to_unsigned());
272
273 match_each_unsigned_integer_ptype!(offsets_primitive.ptype(), |O| {
275 match_each_unsigned_integer_ptype!(sizes_primitive.ptype(), |S| {
276 let offsets_slice = offsets_primitive.as_slice::<O>();
277 let sizes_slice = sizes_primitive.as_slice::<S>();
278
279 validate_offsets_and_sizes::<O, S>(
280 offsets_slice,
281 sizes_slice,
282 elements.len() as u64,
283 )?;
284 })
285 });
286 }
287
288 Ok(())
289 }
290
291 pub unsafe fn with_zero_copy_to_list(mut self, is_zctl: bool) -> Self {
310 self.is_zero_copy_to_list = is_zctl;
311 self
312 }
313
314 pub fn is_zero_copy_to_list(&self) -> bool {
317 self.is_zero_copy_to_list
318 }
319}
320
321impl Default for ListViewData {
322 fn default() -> Self {
323 Self::new()
324 }
325}
326
327fn fill_referenced_mask<O: IntegerPType, S: IntegerPType>(
333 buf: &mut BitBufferMut,
334 offsets: &[O],
335 sizes: &[S],
336) {
337 let len = offsets.len();
338
339 assert_eq!(
340 len,
341 sizes.len(),
342 "offsets and sizes must be the same length"
343 );
344
345 for i in 0..len {
346 let start: usize = offsets[i].as_();
347 let size: usize = sizes[i].as_();
348 buf.fill_range(start, start + size, true);
349 }
350}
351
352pub trait ListViewArrayExt: ListViewArraySlotsExt {
353 fn nullability(&self) -> crate::dtype::Nullability {
354 match self.as_ref().dtype() {
355 DType::List(_, nullability) => *nullability,
356 _ => unreachable!("ListViewArrayExt requires a list dtype"),
357 }
358 }
359
360 fn listview_validity(&self) -> Validity {
361 child_to_validity(
362 self.as_ref().slots()[ListViewSlots::VALIDITY].as_ref(),
363 self.nullability(),
364 )
365 }
366
367 #[allow(clippy::disallowed_methods)]
368 fn offset_at(&self, index: usize) -> usize {
369 assert!(
370 index < self.as_ref().len(),
371 "Index {index} out of bounds 0..{}",
372 self.as_ref().len()
373 );
374 self.offsets()
375 .as_opt::<Primitive>()
376 .map(|p| match_each_integer_ptype!(p.ptype(), |P| { p.as_slice::<P>()[index].as_() }))
377 .unwrap_or_else(|| {
378 self.offsets()
379 .execute_scalar(index, &mut legacy_session().create_execution_ctx())
380 .vortex_expect("offsets must support execute_scalar")
381 .as_primitive()
382 .as_::<usize>()
383 .vortex_expect("offset must fit in usize")
384 })
385 }
386
387 #[allow(clippy::disallowed_methods)]
388 fn size_at(&self, index: usize) -> usize {
389 assert!(
390 index < self.as_ref().len(),
391 "Index {} out of bounds 0..{}",
392 index,
393 self.as_ref().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, &mut legacy_session().create_execution_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) -> VortexResult<ArrayRef> {
409 let offset = self.offset_at(index);
410 let size = self.size_at(index);
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.as_ref().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);
528 let end = self.offset_at(n_lists - 1) + self.size_at(n_lists - 1);
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(elements: ArrayRef, offsets: ArrayRef, sizes: ArrayRef, validity: Validity) -> Self {
566 let dtype = DType::List(Arc::new(elements.dtype().clone()), validity.nullability());
567 let len = offsets.len();
568 let slots = ListViewData::make_slots(&elements, &offsets, &sizes, &validity, len);
569 ListViewData::validate(&elements, &offsets, &sizes, &validity)
570 .vortex_expect("`ListViewArray` construction failed");
571 let data = ListViewData::new();
572 unsafe {
573 Array::from_parts_unchecked(
574 ArrayParts::new(ListView, dtype, len, data).with_slots(slots),
575 )
576 }
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(
593 ArrayParts::new(ListView, dtype, len, data).with_slots(slots),
594 )
595 })
596 }
597
598 pub unsafe fn new_unchecked(
604 elements: ArrayRef,
605 offsets: ArrayRef,
606 sizes: ArrayRef,
607 validity: Validity,
608 ) -> Self {
609 let dtype = DType::List(Arc::new(elements.dtype().clone()), validity.nullability());
610 let len = offsets.len();
611 let slots = ListViewData::make_slots(&elements, &offsets, &sizes, &validity, len);
612 let data = unsafe { ListViewData::new_unchecked() };
613 unsafe {
614 Array::from_parts_unchecked(
615 ArrayParts::new(ListView, dtype, len, data).with_slots(slots),
616 )
617 }
618 }
619
620 pub unsafe fn with_zero_copy_to_list(self, is_zctl: bool) -> Self {
626 if cfg!(debug_assertions) && is_zctl {
627 #[allow(clippy::disallowed_methods)]
628 let mut ctx = legacy_session().create_execution_ctx();
629 let offsets_primitive = self
630 .offsets()
631 .clone()
632 .execute::<PrimitiveArray>(&mut ctx)
633 .vortex_expect("offsets must canonicalize to primitive");
634 let sizes_primitive = self
635 .sizes()
636 .clone()
637 .execute::<PrimitiveArray>(&mut ctx)
638 .vortex_expect("sizes must canonicalize to primitive");
639 validate_zctl(self.elements(), offsets_primitive, sizes_primitive)
640 .vortex_expect("Failed to validate zero-copy to list flag");
641 }
642 let dtype = self.dtype().clone();
643 let len = self.len();
644 let slots: ArraySlots = self.slots().iter().cloned().collect();
645 let data = unsafe { self.into_data().with_zero_copy_to_list(is_zctl) };
646 unsafe {
647 Array::from_parts_unchecked(
648 ArrayParts::new(ListView, dtype, len, data).with_slots(slots),
649 )
650 }
651 }
652
653 pub fn into_data_parts(self) -> ListViewDataParts {
654 let elements = self.slots()[ListViewSlots::ELEMENTS]
655 .clone()
656 .vortex_expect("ListViewArray elements slot");
657 let offsets = self.slots()[ListViewSlots::OFFSETS]
658 .clone()
659 .vortex_expect("ListViewArray offsets slot");
660 let sizes = self.slots()[ListViewSlots::SIZES]
661 .clone()
662 .vortex_expect("ListViewArray sizes slot");
663 let validity = self.listview_validity();
664 ListViewDataParts {
665 elements_dtype: Arc::new(elements.dtype().clone()),
666 elements,
667 offsets,
668 sizes,
669 validity,
670 }
671 }
672}
673
674fn validate_offsets_and_sizes<O, S>(
676 offsets_slice: &[O],
677 sizes_slice: &[S],
678 elements_len: u64,
679) -> VortexResult<()>
680where
681 O: IntegerPType,
682 S: IntegerPType,
683{
684 debug_assert_eq!(offsets_slice.len(), sizes_slice.len());
685
686 #[allow(clippy::absurd_extreme_comparisons, unused_comparisons)]
687 for i in 0..offsets_slice.len() {
688 let offset = offsets_slice[i];
689 let size = sizes_slice[i];
690
691 vortex_ensure!(offset >= O::zero(), "cannot have negative offsets");
692 vortex_ensure!(size >= S::zero(), "cannot have negative size");
693
694 let offset_u64 = offset
695 .to_u64()
696 .ok_or_else(|| vortex_err!("offset[{i}] = {offset:?} cannot be converted to u64"))?;
697
698 let size_u64 = size
699 .to_u64()
700 .ok_or_else(|| vortex_err!("size[{i}] = {size:?} cannot be converted to u64"))?;
701
702 let end = offset_u64.checked_add(size_u64).ok_or_else(|| {
704 vortex_err!("offset[{i}] ({offset_u64}) + size[{i}] ({size_u64}) would overflow u64")
705 })?;
706
707 if offset_u64 == elements_len {
708 vortex_ensure!(
709 size_u64 == 0,
710 "views to the end of the elements array (length {elements_len}) must have size 0 \
711 (had size {size_u64})"
712 );
713 }
714
715 vortex_ensure!(
716 end <= elements_len,
717 "offset[{i}] + size[{i}] = {offset_u64} + {size_u64} = {end} \
718 exceeds elements length {elements_len}",
719 );
720 }
721
722 Ok(())
723}
724
725#[allow(clippy::disallowed_methods)]
728fn validate_zctl(
729 elements: &ArrayRef,
730 offsets_primitive: PrimitiveArray,
731 sizes_primitive: PrimitiveArray,
732) -> VortexResult<()> {
733 let mut ctx = legacy_session().create_execution_ctx();
736 if let Some(is_sorted) = offsets_primitive.statistics().compute_is_sorted(&mut ctx) {
737 vortex_ensure!(is_sorted, "offsets must be sorted");
738 } else {
739 vortex_bail!("offsets must report is_sorted statistic");
740 }
741
742 fn validate_monotonic_ends<O: IntegerPType, S: IntegerPType>(
745 offsets_slice: &[O],
746 sizes_slice: &[S],
747 len: usize,
748 ) -> VortexResult<()> {
749 let mut max_end = 0usize;
750
751 for i in 0..len {
752 let offset = offsets_slice[i].to_usize().unwrap_or(usize::MAX);
753 let size = sizes_slice[i].to_usize().unwrap_or(usize::MAX);
754
755 vortex_ensure!(
757 offset >= max_end,
758 "Zero-copy-to-list requires views to be non-overlapping and ordered: \
759 view[{}] starts at {} but previous views extend to {}",
760 i,
761 offset,
762 max_end
763 );
764
765 let end = offset.saturating_add(size);
767 max_end = max_end.max(end);
768 }
769
770 Ok(())
771 }
772
773 let offsets_dtype = offsets_primitive.dtype();
774 let sizes_dtype = sizes_primitive.dtype();
775 let len = offsets_primitive.len();
776
777 let offsets_unsigned =
779 offsets_primitive.reinterpret_cast(offsets_dtype.as_ptype().to_unsigned());
780 let sizes_unsigned = sizes_primitive.reinterpret_cast(sizes_dtype.as_ptype().to_unsigned());
781
782 match_each_unsigned_integer_ptype!(offsets_unsigned.ptype(), |O| {
784 match_each_unsigned_integer_ptype!(sizes_unsigned.ptype(), |S| {
785 let offsets_slice = offsets_unsigned.as_slice::<O>();
786 let sizes_slice = sizes_unsigned.as_slice::<S>();
787
788 validate_monotonic_ends(offsets_slice, sizes_slice, len)?;
789 })
790 });
791
792 let mut element_references = vec![0u8; elements.len()];
797
798 fn count_references<O: IntegerPType, S: IntegerPType>(
799 element_references: &mut [u8],
800 offsets_primitive: PrimitiveArray,
801 sizes_primitive: PrimitiveArray,
802 ) {
803 let offsets_slice = offsets_primitive.as_slice::<O>();
804 let sizes_slice = sizes_primitive.as_slice::<S>();
805
806 for i in 0..offsets_slice.len() {
809 let offset: usize = offsets_slice[i].as_();
810 let size: usize = sizes_slice[i].as_();
811 for j in offset..offset + size {
812 element_references[j] = element_references[j].saturating_add(1);
813 }
814 }
815 }
816
817 match_each_unsigned_integer_ptype!(offsets_unsigned.ptype(), |O| {
818 match_each_unsigned_integer_ptype!(sizes_unsigned.ptype(), |S| {
819 count_references::<O, S>(&mut element_references, offsets_unsigned, sizes_unsigned);
820 })
821 });
822
823 let leftmost_used = element_references
825 .iter()
826 .position(|&references| references != 0);
827 let rightmost_used = element_references
828 .iter()
829 .rposition(|&references| references != 0);
830
831 if let (Some(first_ref), Some(last_ref)) = (leftmost_used, rightmost_used) {
832 vortex_ensure!(
833 element_references[first_ref..=last_ref]
834 .iter()
835 .all(|&references| references != 0),
836 "found gap in elements array between first and last referenced elements"
837 );
838 }
839
840 vortex_ensure!(element_references.iter().all(|&references| references <= 1));
841
842 Ok(())
843}