1use std::fmt::Display;
5use std::fmt::Formatter;
6use std::sync::Arc;
7
8use num_traits::AsPrimitive;
9use smallvec::smallvec;
10use vortex_buffer::BitBufferMut;
11use vortex_error::VortexExpect;
12use vortex_error::VortexResult;
13use vortex_error::vortex_bail;
14use vortex_error::vortex_ensure;
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::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
45pub(super) const ELEMENTS_SLOT: usize = 0;
48pub(super) const OFFSETS_SLOT: usize = 1;
53pub(super) const SIZES_SLOT: usize = 2;
58pub(super) const VALIDITY_SLOT: usize = 3;
60pub(super) const NUM_SLOTS: usize = 4;
61pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["elements", "offsets", "sizes", "validity"];
62
63#[derive(Clone, Debug)]
127pub struct ListViewData {
128 is_zero_copy_to_list: bool,
137}
138
139impl Display for ListViewData {
140 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
141 write!(f, "is_zero_copy_to_list: {}", self.is_zero_copy_to_list)
142 }
143}
144
145pub struct ListViewDataParts {
146 pub elements_dtype: Arc<DType>,
147
148 pub elements: ArrayRef,
150
151 pub offsets: ArrayRef,
153
154 pub sizes: ArrayRef,
156
157 pub validity: Validity,
159}
160
161impl ListViewData {
162 pub(crate) fn make_slots(
163 elements: &ArrayRef,
164 offsets: &ArrayRef,
165 sizes: &ArrayRef,
166 validity: &Validity,
167 len: usize,
168 ) -> ArraySlots {
169 smallvec![
170 Some(elements.clone()),
171 Some(offsets.clone()),
172 Some(sizes.clone()),
173 validity_to_child(validity, len),
174 ]
175 }
176
177 pub fn new() -> Self {
184 Self {
185 is_zero_copy_to_list: false,
186 }
187 }
188
189 pub fn try_new() -> VortexResult<Self> {
196 Ok(Self::new())
197 }
198
199 pub unsafe fn new_unchecked() -> Self {
219 Self::new()
220 }
221
222 pub fn validate(
224 elements: &ArrayRef,
225 offsets: &ArrayRef,
226 sizes: &ArrayRef,
227 validity: &Validity,
228 ) -> VortexResult<()> {
229 vortex_ensure!(
231 offsets.dtype().is_int() && !offsets.dtype().is_nullable(),
232 "offsets must be non-nullable integer array, got {}",
233 offsets.dtype()
234 );
235 vortex_ensure!(
236 sizes.dtype().is_int() && !sizes.dtype().is_nullable(),
237 "sizes must be non-nullable integer array, got {}",
238 sizes.dtype()
239 );
240
241 vortex_ensure!(
243 offsets.len() == sizes.len(),
244 "offsets and sizes must have the same length, got {} and {}",
245 offsets.len(),
246 sizes.len()
247 );
248
249 if let Some(validity_len) = validity.maybe_len() {
251 vortex_ensure!(
252 validity_len == offsets.len(),
253 "validity with size {validity_len} does not match array size {}",
254 offsets.len()
255 );
256 }
257
258 if offsets.is_host() && sizes.is_host() {
260 #[allow(clippy::disallowed_methods)]
261 let mut ctx = legacy_session().create_execution_ctx();
262 let offsets_primitive = offsets.clone().execute::<PrimitiveArray>(&mut ctx)?;
263 let sizes_primitive = sizes.clone().execute::<PrimitiveArray>(&mut ctx)?;
264 let offsets_primitive =
267 offsets_primitive.reinterpret_cast(offsets_primitive.ptype().to_unsigned());
268 let sizes_primitive =
269 sizes_primitive.reinterpret_cast(sizes_primitive.ptype().to_unsigned());
270
271 match_each_unsigned_integer_ptype!(offsets_primitive.ptype(), |O| {
273 match_each_unsigned_integer_ptype!(sizes_primitive.ptype(), |S| {
274 let offsets_slice = offsets_primitive.as_slice::<O>();
275 let sizes_slice = sizes_primitive.as_slice::<S>();
276
277 validate_offsets_and_sizes::<O, S>(
278 offsets_slice,
279 sizes_slice,
280 elements.len() as u64,
281 )?;
282 })
283 });
284 }
285
286 Ok(())
287 }
288
289 pub unsafe fn with_zero_copy_to_list(mut self, is_zctl: bool) -> Self {
308 self.is_zero_copy_to_list = is_zctl;
309 self
310 }
311
312 pub fn is_zero_copy_to_list(&self) -> bool {
315 self.is_zero_copy_to_list
316 }
317}
318
319impl Default for ListViewData {
320 fn default() -> Self {
321 Self::new()
322 }
323}
324
325fn fill_referenced_mask<O: IntegerPType, S: IntegerPType>(
331 buf: &mut BitBufferMut,
332 offsets: &[O],
333 sizes: &[S],
334) {
335 let len = offsets.len();
336
337 assert_eq!(
338 len,
339 sizes.len(),
340 "offsets and sizes must be the same length"
341 );
342
343 for i in 0..len {
344 let start: usize = offsets[i].as_();
345 let size: usize = sizes[i].as_();
346 buf.fill_range(start, start + size, true);
347 }
348}
349
350pub trait ListViewArrayExt: TypedArrayRef<ListView> {
351 fn nullability(&self) -> crate::dtype::Nullability {
352 match self.as_ref().dtype() {
353 DType::List(_, nullability) => *nullability,
354 _ => unreachable!("ListViewArrayExt requires a list dtype"),
355 }
356 }
357
358 fn elements(&self) -> &ArrayRef {
359 self.as_ref().slots()[ELEMENTS_SLOT]
360 .as_ref()
361 .vortex_expect("ListViewArray elements slot")
362 }
363
364 fn offsets(&self) -> &ArrayRef {
365 self.as_ref().slots()[OFFSETS_SLOT]
366 .as_ref()
367 .vortex_expect("ListViewArray offsets slot")
368 }
369
370 fn sizes(&self) -> &ArrayRef {
371 self.as_ref().slots()[SIZES_SLOT]
372 .as_ref()
373 .vortex_expect("ListViewArray sizes slot")
374 }
375
376 fn listview_validity(&self) -> Validity {
377 child_to_validity(
378 self.as_ref().slots()[VALIDITY_SLOT].as_ref(),
379 self.nullability(),
380 )
381 }
382
383 #[allow(clippy::disallowed_methods)]
384 fn offset_at(&self, index: usize) -> usize {
385 assert!(
386 index < self.as_ref().len(),
387 "Index {index} out of bounds 0..{}",
388 self.as_ref().len()
389 );
390 self.offsets()
391 .as_opt::<Primitive>()
392 .map(|p| match_each_integer_ptype!(p.ptype(), |P| { p.as_slice::<P>()[index].as_() }))
393 .unwrap_or_else(|| {
394 self.offsets()
395 .execute_scalar(index, &mut legacy_session().create_execution_ctx())
396 .vortex_expect("offsets must support execute_scalar")
397 .as_primitive()
398 .as_::<usize>()
399 .vortex_expect("offset must fit in usize")
400 })
401 }
402
403 #[allow(clippy::disallowed_methods)]
404 fn size_at(&self, index: usize) -> usize {
405 assert!(
406 index < self.as_ref().len(),
407 "Index {} out of bounds 0..{}",
408 index,
409 self.as_ref().len()
410 );
411 self.sizes()
412 .as_opt::<Primitive>()
413 .map(|p| match_each_integer_ptype!(p.ptype(), |P| { p.as_slice::<P>()[index].as_() }))
414 .unwrap_or_else(|| {
415 self.sizes()
416 .execute_scalar(index, &mut legacy_session().create_execution_ctx())
417 .vortex_expect("sizes must support execute_scalar")
418 .as_primitive()
419 .as_::<usize>()
420 .vortex_expect("size must fit in usize")
421 })
422 }
423
424 fn list_elements_at(&self, index: usize) -> VortexResult<ArrayRef> {
425 let offset = self.offset_at(index);
426 let size = self.size_at(index);
427 self.elements().slice(offset..offset + size)
428 }
429
430 fn compute_referenced_elements_mask(&self, ctx: &mut ExecutionCtx) -> VortexResult<Mask> {
441 assert!(!self.elements().is_empty());
442 let len = self.elements().len();
443
444 let offsets_primitive = self.offsets().clone().execute::<PrimitiveArray>(ctx)?;
445 let sizes_primitive = self.sizes().clone().execute::<PrimitiveArray>(ctx)?;
446
447 let mut buf = BitBufferMut::new_unset(len);
448
449 let offsets_primitive =
451 offsets_primitive.reinterpret_cast(offsets_primitive.ptype().to_unsigned());
452 let sizes_primitive =
453 sizes_primitive.reinterpret_cast(sizes_primitive.ptype().to_unsigned());
454 match_each_unsigned_integer_ptype!(offsets_primitive.ptype(), |O| {
455 match_each_unsigned_integer_ptype!(sizes_primitive.ptype(), |S| {
456 fill_referenced_mask::<O, S>(
457 &mut buf,
458 offsets_primitive.as_slice::<O>(),
459 sizes_primitive.as_slice::<S>(),
460 );
461 })
462 });
463
464 Ok(Mask::from_buffer(buf.freeze()))
465 }
466
467 fn compute_density(&self, ctx: &mut ExecutionCtx) -> VortexResult<f32> {
471 if self.elements().is_empty() {
472 return Ok(1.0);
473 }
474
475 if self.sizes().is_empty() {
476 return Ok(0.0);
477 }
478
479 let density = match self.compute_referenced_elements_mask(ctx)? {
480 Mask::AllTrue(_) => 1.0,
481 Mask::AllFalse(_) => 0.0,
482 Mask::Values(values) => values.true_count() as f32 / self.elements().len() as f32,
483 };
484
485 Ok(density)
486 }
487
488 fn upper_bound_density(&self, ctx: &mut ExecutionCtx) -> VortexResult<f32> {
495 let n_elts = self.elements().len();
496 if n_elts == 0 {
497 return Ok(1.0);
498 }
499
500 let sizes = self.sizes();
501 if sizes.is_empty() {
502 return Ok(0.0);
503 }
504
505 let sizes_sum = sizes
507 .statistics()
508 .compute_stat(Stat::Sum, ctx)?
509 .vortex_expect("sizes array has integer ptype elements")
510 .as_primitive()
511 .as_::<u64>()
512 .vortex_expect("integer ptypes can be upcast to u64");
513
514 let estimate = (sizes_sum as f32 / n_elts as f32).min(1.0);
517
518 debug_assert!(estimate >= 0.0);
519
520 Ok(estimate)
521 }
522
523 fn referenced_element_bounds(&self, ctx: &mut ExecutionCtx) -> VortexResult<(usize, usize)> {
536 let n_lists = self.as_ref().len();
537 vortex_ensure!(
538 n_lists > 0,
539 "referenced_element_bounds requires a non-empty array"
540 );
541
542 if self.is_zero_copy_to_list() {
543 let start = self.offset_at(0);
544 let end = self.offset_at(n_lists - 1) + self.size_at(n_lists - 1);
545 return Ok((start, end));
546 }
547
548 let start = self
549 .offsets()
550 .statistics()
551 .compute_min::<usize>(ctx)
552 .vortex_expect("offsets must report a usize min statistic");
553
554 let wide_dtype = DType::from(if self.offsets().dtype().as_ptype().is_unsigned_int() {
557 PType::U64
558 } else {
559 PType::I64
560 });
561 let offsets = self.offsets().cast(wide_dtype.clone())?;
562 let sizes = self.sizes().cast(wide_dtype)?;
563 let end = min_max(
564 &offsets.binary(sizes, Operator::Add)?,
565 ctx,
566 NumericalAggregateOpts::default(),
567 )?
568 .vortex_expect("non-empty array must report a min/max")
569 .max
570 .as_primitive()
571 .as_::<usize>()
572 .vortex_expect("max `offset + size` must fit in a usize");
573
574 Ok((start, end))
575 }
576}
577impl<T: TypedArrayRef<ListView>> ListViewArrayExt for T {}
578
579impl Array<ListView> {
580 pub fn new(elements: ArrayRef, offsets: ArrayRef, sizes: ArrayRef, validity: Validity) -> Self {
582 let dtype = DType::List(Arc::new(elements.dtype().clone()), validity.nullability());
583 let len = offsets.len();
584 let slots = ListViewData::make_slots(&elements, &offsets, &sizes, &validity, len);
585 ListViewData::validate(&elements, &offsets, &sizes, &validity)
586 .vortex_expect("`ListViewArray` construction failed");
587 let data = ListViewData::new();
588 unsafe {
589 Array::from_parts_unchecked(
590 ArrayParts::new(ListView, dtype, len, data).with_slots(slots),
591 )
592 }
593 }
594
595 pub fn try_new(
597 elements: ArrayRef,
598 offsets: ArrayRef,
599 sizes: ArrayRef,
600 validity: Validity,
601 ) -> VortexResult<Self> {
602 let dtype = DType::List(Arc::new(elements.dtype().clone()), validity.nullability());
603 let len = offsets.len();
604 let slots = ListViewData::make_slots(&elements, &offsets, &sizes, &validity, len);
605 ListViewData::validate(&elements, &offsets, &sizes, &validity)?;
606 let data = ListViewData::try_new()?;
607 Ok(unsafe {
608 Array::from_parts_unchecked(
609 ArrayParts::new(ListView, dtype, len, data).with_slots(slots),
610 )
611 })
612 }
613
614 pub unsafe fn new_unchecked(
620 elements: ArrayRef,
621 offsets: ArrayRef,
622 sizes: ArrayRef,
623 validity: Validity,
624 ) -> Self {
625 let dtype = DType::List(Arc::new(elements.dtype().clone()), validity.nullability());
626 let len = offsets.len();
627 let slots = ListViewData::make_slots(&elements, &offsets, &sizes, &validity, len);
628 let data = unsafe { ListViewData::new_unchecked() };
629 unsafe {
630 Array::from_parts_unchecked(
631 ArrayParts::new(ListView, dtype, len, data).with_slots(slots),
632 )
633 }
634 }
635
636 pub unsafe fn with_zero_copy_to_list(self, is_zctl: bool) -> Self {
642 if cfg!(debug_assertions) && is_zctl {
643 #[allow(clippy::disallowed_methods)]
644 let mut ctx = legacy_session().create_execution_ctx();
645 let offsets_primitive = self
646 .offsets()
647 .clone()
648 .execute::<PrimitiveArray>(&mut ctx)
649 .vortex_expect("offsets must canonicalize to primitive");
650 let sizes_primitive = self
651 .sizes()
652 .clone()
653 .execute::<PrimitiveArray>(&mut ctx)
654 .vortex_expect("sizes must canonicalize to primitive");
655 validate_zctl(self.elements(), offsets_primitive, sizes_primitive)
656 .vortex_expect("Failed to validate zero-copy to list flag");
657 }
658 let dtype = self.dtype().clone();
659 let len = self.len();
660 let slots: ArraySlots = self.slots().iter().cloned().collect();
661 let data = unsafe { self.into_data().with_zero_copy_to_list(is_zctl) };
662 unsafe {
663 Array::from_parts_unchecked(
664 ArrayParts::new(ListView, dtype, len, data).with_slots(slots),
665 )
666 }
667 }
668
669 pub fn into_data_parts(self) -> ListViewDataParts {
670 let elements = self.slots()[ELEMENTS_SLOT]
671 .clone()
672 .vortex_expect("ListViewArray elements slot");
673 let offsets = self.slots()[OFFSETS_SLOT]
674 .clone()
675 .vortex_expect("ListViewArray offsets slot");
676 let sizes = self.slots()[SIZES_SLOT]
677 .clone()
678 .vortex_expect("ListViewArray sizes slot");
679 let validity = self.listview_validity();
680 ListViewDataParts {
681 elements_dtype: Arc::new(elements.dtype().clone()),
682 elements,
683 offsets,
684 sizes,
685 validity,
686 }
687 }
688}
689
690fn validate_offsets_and_sizes<O, S>(
692 offsets_slice: &[O],
693 sizes_slice: &[S],
694 elements_len: u64,
695) -> VortexResult<()>
696where
697 O: IntegerPType,
698 S: IntegerPType,
699{
700 debug_assert_eq!(offsets_slice.len(), sizes_slice.len());
701
702 #[allow(clippy::absurd_extreme_comparisons, unused_comparisons)]
703 for i in 0..offsets_slice.len() {
704 let offset = offsets_slice[i];
705 let size = sizes_slice[i];
706
707 vortex_ensure!(offset >= O::zero(), "cannot have negative offsets");
708 vortex_ensure!(size >= S::zero(), "cannot have negative size");
709
710 let offset_u64 = offset
711 .to_u64()
712 .ok_or_else(|| vortex_err!("offset[{i}] = {offset:?} cannot be converted to u64"))?;
713
714 let size_u64 = size
715 .to_u64()
716 .ok_or_else(|| vortex_err!("size[{i}] = {size:?} cannot be converted to u64"))?;
717
718 let end = offset_u64.checked_add(size_u64).ok_or_else(|| {
720 vortex_err!("offset[{i}] ({offset_u64}) + size[{i}] ({size_u64}) would overflow u64")
721 })?;
722
723 if offset_u64 == elements_len {
724 vortex_ensure!(
725 size_u64 == 0,
726 "views to the end of the elements array (length {elements_len}) must have size 0 \
727 (had size {size_u64})"
728 );
729 }
730
731 vortex_ensure!(
732 end <= elements_len,
733 "offset[{i}] + size[{i}] = {offset_u64} + {size_u64} = {end} \
734 exceeds elements length {elements_len}",
735 );
736 }
737
738 Ok(())
739}
740
741#[allow(clippy::disallowed_methods)]
744fn validate_zctl(
745 elements: &ArrayRef,
746 offsets_primitive: PrimitiveArray,
747 sizes_primitive: PrimitiveArray,
748) -> VortexResult<()> {
749 let mut ctx = legacy_session().create_execution_ctx();
752 if let Some(is_sorted) = offsets_primitive.statistics().compute_is_sorted(&mut ctx) {
753 vortex_ensure!(is_sorted, "offsets must be sorted");
754 } else {
755 vortex_bail!("offsets must report is_sorted statistic");
756 }
757
758 fn validate_monotonic_ends<O: IntegerPType, S: IntegerPType>(
761 offsets_slice: &[O],
762 sizes_slice: &[S],
763 len: usize,
764 ) -> VortexResult<()> {
765 let mut max_end = 0usize;
766
767 for i in 0..len {
768 let offset = offsets_slice[i].to_usize().unwrap_or(usize::MAX);
769 let size = sizes_slice[i].to_usize().unwrap_or(usize::MAX);
770
771 vortex_ensure!(
773 offset >= max_end,
774 "Zero-copy-to-list requires views to be non-overlapping and ordered: \
775 view[{}] starts at {} but previous views extend to {}",
776 i,
777 offset,
778 max_end
779 );
780
781 let end = offset.saturating_add(size);
783 max_end = max_end.max(end);
784 }
785
786 Ok(())
787 }
788
789 let offsets_dtype = offsets_primitive.dtype();
790 let sizes_dtype = sizes_primitive.dtype();
791 let len = offsets_primitive.len();
792
793 let offsets_unsigned =
795 offsets_primitive.reinterpret_cast(offsets_dtype.as_ptype().to_unsigned());
796 let sizes_unsigned = sizes_primitive.reinterpret_cast(sizes_dtype.as_ptype().to_unsigned());
797
798 match_each_unsigned_integer_ptype!(offsets_unsigned.ptype(), |O| {
800 match_each_unsigned_integer_ptype!(sizes_unsigned.ptype(), |S| {
801 let offsets_slice = offsets_unsigned.as_slice::<O>();
802 let sizes_slice = sizes_unsigned.as_slice::<S>();
803
804 validate_monotonic_ends(offsets_slice, sizes_slice, len)?;
805 })
806 });
807
808 let mut element_references = vec![0u8; elements.len()];
813
814 fn count_references<O: IntegerPType, S: IntegerPType>(
815 element_references: &mut [u8],
816 offsets_primitive: PrimitiveArray,
817 sizes_primitive: PrimitiveArray,
818 ) {
819 let offsets_slice = offsets_primitive.as_slice::<O>();
820 let sizes_slice = sizes_primitive.as_slice::<S>();
821
822 for i in 0..offsets_slice.len() {
825 let offset: usize = offsets_slice[i].as_();
826 let size: usize = sizes_slice[i].as_();
827 for j in offset..offset + size {
828 element_references[j] = element_references[j].saturating_add(1);
829 }
830 }
831 }
832
833 match_each_unsigned_integer_ptype!(offsets_unsigned.ptype(), |O| {
834 match_each_unsigned_integer_ptype!(sizes_unsigned.ptype(), |S| {
835 count_references::<O, S>(&mut element_references, offsets_unsigned, sizes_unsigned);
836 })
837 });
838
839 let leftmost_used = element_references
841 .iter()
842 .position(|&references| references != 0);
843 let rightmost_used = element_references
844 .iter()
845 .rposition(|&references| references != 0);
846
847 if let (Some(first_ref), Some(last_ref)) = (leftmost_used, rightmost_used) {
848 vortex_ensure!(
849 element_references[first_ref..=last_ref]
850 .iter()
851 .all(|&references| references != 0),
852 "found gap in elements array between first and last referenced elements"
853 );
854 }
855
856 vortex_ensure!(element_references.iter().all(|&references| references <= 1));
857
858 Ok(())
859}