1use std::any::Any;
5use std::sync::Arc;
6
7use num_traits::AsPrimitive;
8use vortex_buffer::BufferAllocatorRef;
9use vortex_error::VortexExpect;
10use vortex_error::VortexResult;
11use vortex_error::vortex_bail;
12use vortex_error::vortex_ensure_eq;
13use vortex_error::vortex_panic;
14
15use crate::ArrayRef;
16use crate::Canonical;
17use crate::ExecutionCtx;
18use crate::IntoArray;
19use crate::array::ArrayView;
20use crate::arrays::List;
21use crate::arrays::ListArray;
22use crate::arrays::ListView;
23use crate::arrays::ListViewArray;
24use crate::arrays::PrimitiveArray;
25use crate::arrays::list::ListArraySlotsExt;
26use crate::arrays::listview::ListViewArraySlotsExt;
27use crate::arrays::listview::ListViewRebuildMode;
28use crate::builders::ArrayBuilder;
29use crate::builders::ChildBuilder;
30use crate::builders::DEFAULT_BUILDER_CAPACITY;
31use crate::builders::PrimitiveBuilder;
32use crate::builders::ValidityBuilder;
33use crate::dtype::DType;
34use crate::dtype::IntegerPType;
35use crate::dtype::Nullability;
36use crate::dtype::Nullability::NonNullable;
37use crate::dtype::OffsetBuilderPType;
38use crate::match_each_integer_ptype;
39use crate::scalar::ListScalar;
40use crate::scalar::Scalar;
41
42pub struct ListBuilder<O: OffsetBuilderPType> {
45 dtype: DType,
47
48 elements_builder: ChildBuilder,
50
51 offsets_builder: PrimitiveBuilder<O>,
53
54 nulls: ValidityBuilder,
56}
57
58impl<O: OffsetBuilderPType> ListBuilder<O> {
59 #[deprecated(note = "use `new_in` with an explicit allocator")]
61 pub fn new(value_dtype: Arc<DType>, nullability: Nullability) -> Self {
62 Self::new_in(value_dtype, nullability, BufferAllocatorRef::static_ref())
63 }
64
65 pub fn new_in(
67 value_dtype: Arc<DType>,
68 nullability: Nullability,
69 allocator: &BufferAllocatorRef,
70 ) -> Self {
71 Self::with_capacity_in(
72 value_dtype,
73 nullability,
74 DEFAULT_BUILDER_CAPACITY * 2,
77 DEFAULT_BUILDER_CAPACITY,
78 allocator,
79 )
80 }
81
82 #[deprecated(note = "use `with_capacity_in` with an explicit allocator")]
90 pub fn with_capacity(
91 value_dtype: Arc<DType>,
92 nullability: Nullability,
93 elements_capacity: usize,
94 capacity: usize,
95 ) -> Self {
96 Self::with_capacity_in(
97 value_dtype,
98 nullability,
99 elements_capacity,
100 capacity,
101 BufferAllocatorRef::static_ref(),
102 )
103 }
104
105 pub fn with_capacity_in(
107 value_dtype: Arc<DType>,
108 nullability: Nullability,
109 elements_capacity: usize,
110 capacity: usize,
111 allocator: &BufferAllocatorRef,
112 ) -> Self {
113 let elements_builder =
114 ChildBuilder::with_capacity(value_dtype.as_ref(), elements_capacity, allocator);
115 let mut offsets_builder =
116 PrimitiveBuilder::<O>::with_capacity_in(NonNullable, capacity + 1, allocator);
117
118 offsets_builder.append_zero();
120
121 Self {
122 elements_builder,
123 offsets_builder,
124 nulls: ValidityBuilder::new(capacity, allocator),
125 dtype: DType::List(value_dtype, nullability),
126 }
127 }
128
129 pub fn append_array_as_list(
136 &mut self,
137 array: &ArrayRef,
138 ctx: &mut ExecutionCtx,
139 ) -> VortexResult<()> {
140 vortex_ensure_eq!(array.dtype(), self.element_dtype());
141
142 self.elements_builder.append_array(array, ctx)?;
143 self.nulls.append_non_null();
144 self.offsets_builder.append_value(
145 O::from_usize(self.elements_builder.len())
146 .vortex_expect("Failed to convert from usize to O"),
147 );
148
149 Ok(())
150 }
151
152 pub fn append_value(&mut self, value: ListScalar) -> VortexResult<()> {
154 match value.elements() {
155 None => {
156 if self.dtype.nullability() == NonNullable {
157 vortex_bail!("Cannot append null value to non-nullable list");
158 }
159 self.append_null();
160 }
161 Some(elements) => {
162 for scalar in elements {
163 self.elements_builder.append_scalar(&scalar)?;
166 }
167
168 self.nulls.append_non_null();
169 self.offsets_builder.append_value(
170 O::from_usize(self.elements_builder.len())
171 .vortex_expect("Failed to convert from usize to O"),
172 );
173 }
174 }
175
176 Ok(())
177 }
178
179 pub fn finish_into_list(&mut self) -> ListArray {
181 assert_eq!(
182 self.offsets_builder.len(),
183 self.nulls.len() + 1,
184 "offsets length must be one more than nulls length."
185 );
186
187 unsafe {
190 ListArray::new_unchecked(
191 self.elements_builder.finish(),
192 self.offsets_builder.finish(),
193 self.nulls.finish_with_nullability(self.dtype.nullability()),
194 )
195 }
196 }
197
198 pub fn element_dtype(&self) -> &DType {
201 let DType::List(element_dtype, _) = &self.dtype else {
202 vortex_panic!("`ListBuilder` has an incorrect dtype: {}", self.dtype);
203 };
204
205 element_dtype
206 }
207
208 pub fn append_list_array(
215 &mut self,
216 array: ArrayView<'_, List>,
217 ctx: &mut ExecutionCtx,
218 ) -> VortexResult<()> {
219 if array.is_empty() {
220 return Ok(());
221 }
222
223 self.nulls.append_validity(array.validity()?, array.len());
224
225 let num_lists = array.len();
226 let offsets = array.offsets().clone().execute::<PrimitiveArray>(ctx)?;
227 match_each_integer_ptype!(offsets.ptype(), |OffsetType| {
228 let offsets = offsets.as_slice::<OffsetType>();
229 let first: usize = offsets[0].as_();
230 let last: usize = offsets[num_lists].as_();
231
232 let elements_base = self.elements_builder.len();
235 if last > first {
236 self.elements_builder
237 .append_array(&array.elements().slice(first..last)?, ctx)?;
238 }
239
240 self.offsets_builder.reserve_exact(num_lists);
241 let mut offsets_range = self.offsets_builder.uninit_range(num_lists);
242 for i in 0..num_lists {
243 let end: usize = offsets[i + 1].as_();
244 offsets_range.set_value(
245 i,
246 O::from_usize(end - first + elements_base)
247 .vortex_expect("Failed to convert offset"),
248 );
249 }
250 unsafe { offsets_range.finish() };
253 });
254 Ok(())
255 }
256
257 pub fn append_listview_array(
265 &mut self,
266 array: ArrayView<'_, ListView>,
267 ctx: &mut ExecutionCtx,
268 ) -> VortexResult<()> {
269 if array.is_empty() {
270 return Ok(());
271 }
272
273 self.nulls.append_validity(array.validity()?, array.len());
274
275 let array = array
279 .into_owned()
280 .rebuild(ListViewRebuildMode::MakeZeroCopyToList, ctx)?;
281 debug_assert!(array.is_zero_copy_to_list());
282
283 let elements = array.elements();
285 let offsets = array.offsets().clone().execute::<PrimitiveArray>(ctx)?;
286 let sizes = array.sizes().clone().execute::<PrimitiveArray>(ctx)?;
287
288 match_each_integer_ptype!(offsets.ptype(), |OffsetType| {
289 match_each_integer_ptype!(sizes.ptype(), |SizeType| {
290 extend_from_listview(
291 self,
292 elements,
293 offsets.as_slice::<OffsetType>(),
294 sizes.as_slice::<SizeType>(),
295 ctx,
296 )?
297 })
298 });
299 Ok(())
300 }
301}
302
303fn extend_from_listview<O, OffsetType, SizeType>(
311 builder: &mut ListBuilder<O>,
312 new_elements: &ArrayRef,
313 new_offsets: &[OffsetType],
314 new_sizes: &[SizeType],
315 ctx: &mut ExecutionCtx,
316) -> VortexResult<()>
317where
318 O: OffsetBuilderPType,
319 OffsetType: IntegerPType,
320 SizeType: IntegerPType,
321{
322 let num_lists = new_offsets.len();
323 debug_assert_eq!(num_lists, new_sizes.len());
324
325 let first: usize = new_offsets[0].as_();
328 let last: usize = new_offsets[num_lists - 1].as_() + new_sizes[num_lists - 1].as_();
329
330 let elements_base = builder.elements_builder.len();
331 if last > first {
332 builder
333 .elements_builder
334 .append_array(&new_elements.slice(first..last)?, ctx)?;
335 }
336
337 builder.offsets_builder.reserve_exact(num_lists);
338 let mut offsets_range = builder.offsets_builder.uninit_range(num_lists);
339 for i in 0..num_lists {
340 let end: usize = new_offsets[i].as_() + new_sizes[i].as_();
341 offsets_range.set_value(
342 i,
343 O::from_usize(end - first + elements_base).vortex_expect("Failed to convert offset"),
344 );
345 }
346
347 unsafe { offsets_range.finish() };
350 Ok(())
351}
352
353impl<O: OffsetBuilderPType> ArrayBuilder for ListBuilder<O> {
354 fn as_any(&self) -> &dyn Any {
355 self
356 }
357
358 fn as_any_mut(&mut self) -> &mut dyn Any {
359 self
360 }
361
362 fn dtype(&self) -> &DType {
363 &self.dtype
364 }
365
366 fn len(&self) -> usize {
367 self.nulls.len()
368 }
369
370 fn append_zeros(&mut self, n: usize) {
371 let curr_len = self.elements_builder.len();
372 for _ in 0..n {
373 self.offsets_builder.append_value(
374 O::from_usize(curr_len).vortex_expect("Failed to convert from usize to <O>"),
375 )
376 }
377 self.nulls.append_n_non_nulls(n);
378 }
379
380 unsafe fn append_nulls_unchecked(&mut self, n: usize) {
381 let curr_len = self.elements_builder.len();
382 for _ in 0..n {
383 self.offsets_builder.append_value(
386 O::from_usize(curr_len).vortex_expect("Failed to convert from usize to <O>"),
387 )
388 }
389 self.nulls.append_n_nulls(n);
390 }
391
392 fn append_scalar(&mut self, scalar: &Scalar) -> VortexResult<()> {
393 vortex_ensure_eq!(
394 scalar.dtype(),
395 self.dtype(),
396 "ListBuilder received a scalar with the wrong dtype"
397 );
398
399 self.append_value(scalar.as_list())
400 }
401
402 fn reserve_exact(&mut self, additional: usize) {
403 self.elements_builder.reserve_exact(additional);
404 self.offsets_builder.reserve_exact(additional);
405 self.nulls.reserve_exact(additional);
406 }
407
408 fn finish(&mut self) -> ArrayRef {
409 self.finish_into_list().into_array()
410 }
411
412 fn finish_into_canonical(&mut self, ctx: &mut ExecutionCtx) -> Canonical {
413 let listview = self
414 .finish()
415 .execute::<ListViewArray>(ctx)
416 .vortex_expect("list builder should canonicalize to listview");
417 Canonical::List(listview)
418 }
419}
420
421#[cfg(test)]
422mod tests {
423 use std::sync::Arc;
424
425 use Nullability::NonNullable;
426 use Nullability::Nullable;
427 use vortex_buffer::BufferAllocatorRef;
428 use vortex_buffer::buffer;
429 use vortex_error::VortexExpect;
430 use vortex_error::VortexResult;
431
432 use crate::IntoArray;
433 use crate::array_session;
434 use crate::arrays::ChunkedArray;
435 use crate::arrays::ListViewArray;
436 use crate::arrays::PrimitiveArray;
437 use crate::arrays::list::ListArraySlotsExt;
438 use crate::arrays::listview::ListViewArrayExt;
439 use crate::arrays::listview::ListViewArraySlotsExt;
440 use crate::assert_arrays_eq;
441 use crate::builders::ArrayBuilder;
442 use crate::builders::ListViewBuilder;
443 use crate::builders::builder_with_capacity_in;
444 use crate::builders::list::ListArray;
445 use crate::builders::list::ListBuilder;
446 use crate::dtype::DType;
447 use crate::dtype::Nullability;
448 use crate::dtype::OffsetBuilderPType;
449 use crate::dtype::PType::I32;
450 use crate::executor::VortexSessionExecute;
451 use crate::scalar::Scalar;
452 use crate::validity::Validity;
453
454 #[test]
455 fn test_empty() {
456 let mut builder = ListBuilder::<u32>::with_capacity_in(
457 Arc::new(I32.into()),
458 NonNullable,
459 0,
460 0,
461 BufferAllocatorRef::static_ref(),
462 );
463
464 let list = builder.finish();
465 assert_eq!(list.len(), 0);
466 assert_eq!(
467 list.dtype(),
468 &DType::List(Arc::new(I32.into()), NonNullable)
469 );
470 }
471
472 #[test]
473 fn test_values() {
474 let dtype: Arc<DType> = Arc::new(I32.into());
475 let mut builder = ListBuilder::<u32>::with_capacity_in(
476 Arc::clone(&dtype),
477 NonNullable,
478 0,
479 0,
480 BufferAllocatorRef::static_ref(),
481 );
482
483 builder
484 .append_value(
485 Scalar::list(
486 Arc::clone(&dtype),
487 vec![1i32.into(), 2i32.into(), 3i32.into()],
488 NonNullable,
489 )
490 .as_list(),
491 )
492 .unwrap();
493
494 builder
495 .append_value(
496 Scalar::list(
497 dtype,
498 vec![4i32.into(), 5i32.into(), 6i32.into()],
499 NonNullable,
500 )
501 .as_list(),
502 )
503 .unwrap();
504
505 let list = builder.finish();
506 assert_eq!(list.len(), 2);
507
508 let mut ctx = array_session().create_execution_ctx();
509 let list_array = list.execute::<ListViewArray>(&mut ctx).unwrap();
510
511 assert_eq!(list_array.list_elements_at(0, &mut ctx).unwrap().len(), 3);
512 assert_eq!(list_array.list_elements_at(1, &mut ctx).unwrap().len(), 3);
513 let expected = ListArray::new(
514 buffer![1i32, 2, 3, 4, 5, 6].into_array(),
515 buffer![0u32, 3, 6].into_array(),
516 Validity::NonNullable,
517 &mut ctx,
518 );
519 assert_arrays_eq!(list_array, expected, &mut ctx);
520 }
521
522 #[test]
523 fn test_append_empty_list() {
524 let mut ctx = array_session().create_execution_ctx();
525 let dtype: Arc<DType> = Arc::new(I32.into());
526 let mut builder = ListBuilder::<u32>::with_capacity_in(
527 Arc::clone(&dtype),
528 NonNullable,
529 0,
530 0,
531 BufferAllocatorRef::static_ref(),
532 );
533
534 assert!(
535 builder
536 .append_value(Scalar::list_empty(dtype, NonNullable).as_list())
537 .is_ok()
538 );
539 let expected = ListArray::new(
540 buffer![0i32; 0].into_array(),
541 buffer![0u32, 0].into_array(),
542 Validity::NonNullable,
543 &mut ctx,
544 );
545 assert_arrays_eq!(
546 builder.finish(),
547 expected,
548 &mut array_session().create_execution_ctx()
549 );
550 }
551
552 #[test]
553 fn test_nullable_values() {
554 let dtype: Arc<DType> = Arc::new(I32.into());
555 let mut builder = ListBuilder::<u32>::with_capacity_in(
556 Arc::clone(&dtype),
557 Nullable,
558 0,
559 0,
560 BufferAllocatorRef::static_ref(),
561 );
562
563 builder
564 .append_value(
565 Scalar::list(
566 Arc::clone(&dtype),
567 vec![1i32.into(), 2i32.into(), 3i32.into()],
568 NonNullable,
569 )
570 .as_list(),
571 )
572 .unwrap();
573
574 builder
575 .append_value(Scalar::list_empty(Arc::clone(&dtype), NonNullable).as_list())
576 .unwrap();
577
578 builder
579 .append_value(
580 Scalar::list(
581 dtype,
582 vec![4i32.into(), 5i32.into(), 6i32.into()],
583 NonNullable,
584 )
585 .as_list(),
586 )
587 .unwrap();
588
589 let list = builder.finish();
590 assert_eq!(list.len(), 3);
591
592 let mut ctx = array_session().create_execution_ctx();
593 let list_array = list.execute::<ListViewArray>(&mut ctx).unwrap();
594
595 assert_eq!(list_array.list_elements_at(0, &mut ctx).unwrap().len(), 3);
596 assert_eq!(list_array.list_elements_at(1, &mut ctx).unwrap().len(), 0);
597 assert_eq!(list_array.list_elements_at(2, &mut ctx).unwrap().len(), 3);
598 let expected = ListArray::new(
599 buffer![1i32, 2, 3, 4, 5, 6].into_array(),
600 buffer![0u32, 3, 3, 6].into_array(),
601 Validity::AllValid,
602 &mut ctx,
603 );
604 assert_arrays_eq!(list_array, expected, &mut ctx);
605 }
606
607 fn test_extend_builder_gen<O: OffsetBuilderPType>() {
608 let list = ListArray::from_iter_opt_slow::<O, _, _>(
609 [Some(vec![0, 1, 2]), None, Some(vec![4, 5])],
610 Arc::new(I32.into()),
611 )
612 .unwrap()
613 .into_array();
614 assert_eq!(list.len(), 3);
615
616 let mut ctx = array_session().create_execution_ctx();
617
618 let mut builder = ListBuilder::<O>::with_capacity_in(
619 Arc::new(I32.into()),
620 Nullable,
621 18,
622 9,
623 BufferAllocatorRef::static_ref(),
624 );
625 list.append_to_builder(&mut builder, &mut ctx).unwrap();
626 list.append_to_builder(&mut builder, &mut ctx).unwrap();
627 list.slice(0..0)
628 .unwrap()
629 .append_to_builder(&mut builder, &mut ctx)
630 .unwrap();
631 list.slice(1..3)
632 .unwrap()
633 .append_to_builder(&mut builder, &mut ctx)
634 .unwrap();
635
636 let expected = ListArray::from_iter_opt_slow::<O, _, _>(
637 [
638 Some(vec![0, 1, 2]),
639 None,
640 Some(vec![4, 5]),
641 Some(vec![0, 1, 2]),
642 None,
643 Some(vec![4, 5]),
644 None,
645 Some(vec![4, 5]),
646 ],
647 Arc::new(DType::Primitive(I32, NonNullable)),
648 )
649 .unwrap()
650 .into_array()
651 .execute::<ListViewArray>(&mut ctx)
652 .unwrap();
653
654 let actual = builder.finish_into_canonical(&mut ctx).into_listview();
655
656 assert_arrays_eq!(actual.elements(), expected.elements(), &mut ctx);
657
658 assert_arrays_eq!(actual.offsets(), expected.offsets(), &mut ctx);
659
660 assert!(
661 actual
662 .validity()
663 .vortex_expect("list validity should be derivable")
664 .mask_eq(
665 &expected
666 .validity()
667 .vortex_expect("list validity should be derivable"),
668 actual.len(),
669 &mut ctx,
670 )
671 .unwrap(),
672 );
673 let independent = ListArray::new(
674 buffer![0i32, 1, 2, 4, 5, 0, 1, 2, 4, 5, 4, 5].into_array(),
675 buffer![0u32, 3, 3, 5, 8, 8, 10, 10, 12].into_array(),
676 Validity::from_iter([true, false, true, true, false, true, false, true]),
677 &mut ctx,
678 );
679 assert_arrays_eq!(actual, independent, &mut ctx);
680 }
681
682 #[test]
687 fn test_append_to_builder_any_list_builder() -> VortexResult<()> {
688 let mut ctx = array_session().create_execution_ctx();
689
690 let list = ListArray::from_iter_opt_slow::<u64, _, _>(
691 [Some(vec![0, 1, 2]), None, Some(vec![4, 5])],
692 Arc::new(I32.into()),
693 )?
694 .into_array();
695 let listview = list
696 .clone()
697 .execute::<ListViewArray>(&mut ctx)?
698 .into_array();
699 let elem_dtype = || Arc::new(I32.into());
700
701 let mut listview_builder =
704 builder_with_capacity_in(list.dtype(), list.len(), BufferAllocatorRef::static_ref());
705 list.append_to_builder(listview_builder.as_mut(), &mut ctx)?;
706 assert_arrays_eq!(listview_builder.finish(), list, &mut ctx);
707
708 let mut lv_u64_u32 = ListViewBuilder::<u64, u32>::with_capacity_in(
711 elem_dtype(),
712 Nullable,
713 8,
714 4,
715 BufferAllocatorRef::static_ref(),
716 );
717 list.append_to_builder(&mut lv_u64_u32, &mut ctx)?;
718 assert_arrays_eq!(lv_u64_u32.finish(), list, &mut ctx);
719
720 let mut lv_i64_i32 = ListViewBuilder::<i64, i32>::with_capacity_in(
721 elem_dtype(),
722 Nullable,
723 8,
724 4,
725 BufferAllocatorRef::static_ref(),
726 );
727 list.append_to_builder(&mut lv_i64_i32, &mut ctx)?;
728 assert_arrays_eq!(lv_i64_i32.finish(), list, &mut ctx);
729
730 let mut lv_u32_u32 = ListViewBuilder::<u32, u32>::with_capacity_in(
731 elem_dtype(),
732 Nullable,
733 8,
734 4,
735 BufferAllocatorRef::static_ref(),
736 );
737 listview.append_to_builder(&mut lv_u32_u32, &mut ctx)?;
738 assert_arrays_eq!(lv_u32_u32.finish(), list, &mut ctx);
739
740 let mut list_builder = ListBuilder::<u32>::with_capacity_in(
743 elem_dtype(),
744 Nullable,
745 8,
746 4,
747 BufferAllocatorRef::static_ref(),
748 );
749 list.append_to_builder(&mut list_builder, &mut ctx)?;
750 assert_arrays_eq!(list_builder.finish(), list, &mut ctx);
751
752 let mut list_builder_i32 = ListBuilder::<i32>::with_capacity_in(
753 elem_dtype(),
754 Nullable,
755 8,
756 4,
757 BufferAllocatorRef::static_ref(),
758 );
759 listview.append_to_builder(&mut list_builder_i32, &mut ctx)?;
760 assert_arrays_eq!(list_builder_i32.finish(), list, &mut ctx);
761
762 Ok(())
763 }
764
765 #[test]
766 fn test_append_list_arrays_grow_builder() -> VortexResult<()> {
767 let mut ctx = array_session().create_execution_ctx();
768 let dtype: Arc<DType> = Arc::new(I32.into());
769
770 let lists: Vec<Option<Vec<i32>>> =
773 (0..100).map(|i| (i % 10 != 0).then(|| vec![i])).collect();
774 let source = ListArray::from_iter_opt_slow::<u32, _, _>(lists.clone(), Arc::clone(&dtype))?;
775 let expected = ListArray::from_iter_opt_slow::<u32, _, _>(
776 lists.iter().cloned().chain(lists.iter().cloned()),
777 Arc::clone(&dtype),
778 )?;
779
780 let mut builder = ListBuilder::<u32>::with_capacity_in(
782 Arc::clone(&dtype),
783 Nullable,
784 0,
785 0,
786 BufferAllocatorRef::static_ref(),
787 );
788 builder.append_list_array(source.as_view(), &mut ctx)?;
789 builder.append_list_array(source.as_view(), &mut ctx)?;
790 assert_arrays_eq!(builder.finish(), expected, &mut ctx);
791
792 let source_listview = source.into_array().execute::<ListViewArray>(&mut ctx)?;
793 let mut builder = ListBuilder::<u32>::with_capacity_in(
794 dtype,
795 Nullable,
796 0,
797 0,
798 BufferAllocatorRef::static_ref(),
799 );
800 builder.append_listview_array(source_listview.as_view(), &mut ctx)?;
801 builder.append_listview_array(source_listview.as_view(), &mut ctx)?;
802 assert_arrays_eq!(builder.finish(), expected, &mut ctx);
803
804 Ok(())
805 }
806
807 #[test]
811 fn test_append_listview_array_flattens_overlaps_and_skips_leading_elements() -> VortexResult<()>
812 {
813 let mut ctx = array_session().create_execution_ctx();
814 let dtype: Arc<DType> = Arc::new(I32.into());
815
816 let overlapping = unsafe {
821 ListViewArray::new_unchecked(
822 buffer![10i32, 20, 30].into_array(),
823 buffer![0u32, 1, 0].into_array(),
824 buffer![2u8, 2, 1].into_array(),
825 Validity::from_iter([true, false, true]),
826 )
827 };
828 assert!(!overlapping.is_zero_copy_to_list());
829
830 let sliced = unsafe {
832 ListViewArray::new_unchecked(
833 buffer![40i32, 50, 60, 70].into_array(),
834 buffer![0u32, 2].into_array(),
835 buffer![2u32, 2].into_array(),
836 Validity::AllValid,
837 )
838 .with_zero_copy_to_list(true)
839 }
840 .into_array()
841 .slice(1..2)?
842 .execute::<ListViewArray>(&mut ctx)?;
843
844 let mut builder = ListBuilder::<u32>::with_capacity_in(
845 dtype,
846 Nullable,
847 0,
848 0,
849 BufferAllocatorRef::static_ref(),
850 );
851 builder.append_listview_array(overlapping.as_view(), &mut ctx)?;
852 builder.append_listview_array(sliced.as_view(), &mut ctx)?;
853
854 let list = builder.finish_into_list();
855 assert_arrays_eq!(
856 list.elements(),
857 PrimitiveArray::from_iter([10i32, 20, 10, 60, 70]),
858 &mut ctx
859 );
860 assert_arrays_eq!(
861 list.offsets(),
862 PrimitiveArray::from_iter([0u32, 2, 2, 3, 5]),
863 &mut ctx
864 );
865
866 let expected = ListArray::new(
867 buffer![10i32, 20, 10, 60, 70].into_array(),
868 buffer![0u32, 2, 2, 3, 5].into_array(),
869 Validity::from_iter([true, false, true, true]),
870 &mut ctx,
871 );
872 assert_arrays_eq!(list, expected, &mut ctx);
873 Ok(())
874 }
875
876 #[test]
877 fn test_extend_builder() {
878 test_extend_builder_gen::<i32>();
879 test_extend_builder_gen::<i64>();
880
881 test_extend_builder_gen::<u32>();
882 test_extend_builder_gen::<u64>();
883 }
884
885 #[test]
886 pub fn test_array_with_gap() {
887 let mut ctx = array_session().create_execution_ctx();
888 let one_trailing_unused_element = ListArray::try_new(
889 buffer![1, 2, 3, 4].into_array(),
890 buffer![0, 3].into_array(),
891 Validity::NonNullable,
892 &mut ctx,
893 )
894 .unwrap();
895
896 let second_array = ListArray::try_new(
897 buffer![5, 6].into_array(),
898 buffer![0, 2].into_array(),
899 Validity::NonNullable,
900 &mut ctx,
901 )
902 .unwrap();
903
904 let chunked_list = ChunkedArray::try_new(
905 vec![
906 one_trailing_unused_element.clone().into_array(),
907 second_array.clone().into_array(),
908 ],
909 DType::List(Arc::new(DType::Primitive(I32, NonNullable)), NonNullable),
910 );
911
912 let mut ctx = array_session().create_execution_ctx();
913 let canon_values = chunked_list
914 .unwrap()
915 .as_array()
916 .clone()
917 .execute::<ListViewArray>(&mut ctx)
918 .unwrap();
919
920 assert_eq!(
921 one_trailing_unused_element
922 .execute_scalar(0, &mut array_session().create_execution_ctx())
923 .unwrap(),
924 canon_values
925 .execute_scalar(0, &mut array_session().create_execution_ctx())
926 .unwrap()
927 );
928 assert_eq!(
929 second_array
930 .execute_scalar(0, &mut array_session().create_execution_ctx())
931 .unwrap(),
932 canon_values
933 .execute_scalar(1, &mut array_session().create_execution_ctx())
934 .unwrap()
935 );
936 }
937
938 #[test]
939 fn test_append_scalar() {
940 let dtype: Arc<DType> = Arc::new(I32.into());
941 let mut builder = ListBuilder::<u64>::with_capacity_in(
942 Arc::clone(&dtype),
943 Nullable,
944 20,
945 10,
946 BufferAllocatorRef::static_ref(),
947 );
948
949 let list_scalar1 =
951 Scalar::list(Arc::clone(&dtype), vec![1i32.into(), 2i32.into()], Nullable);
952 builder.append_scalar(&list_scalar1).unwrap();
953
954 let list_scalar2 = Scalar::list(
956 Arc::clone(&dtype),
957 vec![3i32.into(), 4i32.into(), 5i32.into()],
958 Nullable,
959 );
960 builder.append_scalar(&list_scalar2).unwrap();
961
962 let null_scalar = Scalar::null(DType::List(Arc::clone(&dtype), Nullable));
964 builder.append_scalar(&null_scalar).unwrap();
965
966 let array = builder.finish_into_list();
967 assert_eq!(array.len(), 3);
968
969 let mut ctx = array_session().create_execution_ctx();
970
971 let scalar0 = array.execute_scalar(0, &mut ctx).unwrap();
974 let list0 = scalar0.as_list();
975 assert_eq!(list0.len(), 2);
976 if let Some(list0_items) = list0.elements() {
977 assert_eq!(list0_items[0].as_primitive().typed_value::<i32>(), Some(1));
978 assert_eq!(list0_items[1].as_primitive().typed_value::<i32>(), Some(2));
979 }
980
981 let scalar1 = array.execute_scalar(1, &mut ctx).unwrap();
982 let list1 = scalar1.as_list();
983 assert_eq!(list1.len(), 3);
984 if let Some(list1_items) = list1.elements() {
985 assert_eq!(list1_items[0].as_primitive().typed_value::<i32>(), Some(3));
986 assert_eq!(list1_items[1].as_primitive().typed_value::<i32>(), Some(4));
987 assert_eq!(list1_items[2].as_primitive().typed_value::<i32>(), Some(5));
988 }
989
990 let scalar2 = array.execute_scalar(2, &mut ctx).unwrap();
991 let list2 = scalar2.as_list();
992 assert!(list2.is_null()); assert!(
996 array
997 .validity()
998 .vortex_expect("list validity should be derivable")
999 .execute_is_valid(0, &mut ctx)
1000 .unwrap()
1001 );
1002 assert!(
1003 array
1004 .validity()
1005 .vortex_expect("list validity should be derivable")
1006 .execute_is_valid(1, &mut ctx)
1007 .unwrap()
1008 );
1009 assert!(
1010 !array
1011 .validity()
1012 .vortex_expect("list validity should be derivable")
1013 .execute_is_valid(2, &mut ctx)
1014 .unwrap()
1015 );
1016
1017 let mut builder = ListBuilder::<u64>::with_capacity_in(
1019 dtype,
1020 NonNullable,
1021 20,
1022 10,
1023 BufferAllocatorRef::static_ref(),
1024 );
1025 let wrong_scalar = Scalar::from(42i32);
1026 assert!(builder.append_scalar(&wrong_scalar).is_err());
1027 }
1028
1029 #[test]
1030 fn test_append_array_as_list() {
1031 let dtype: Arc<DType> = Arc::new(I32.into());
1032 let mut ctx = array_session().create_execution_ctx();
1033 let mut builder = ListBuilder::<u32>::with_capacity_in(
1034 Arc::clone(&dtype),
1035 NonNullable,
1036 20,
1037 10,
1038 BufferAllocatorRef::static_ref(),
1039 );
1040
1041 let arr1 = buffer![1i32, 2, 3].into_array();
1043 builder.append_array_as_list(&arr1, &mut ctx).unwrap();
1044
1045 builder
1047 .append_value(
1048 Scalar::list(
1049 Arc::clone(&dtype),
1050 vec![10i32.into(), 11i32.into()],
1051 NonNullable,
1052 )
1053 .as_list(),
1054 )
1055 .unwrap();
1056
1057 let arr2 = buffer![4i32, 5].into_array();
1059 builder.append_array_as_list(&arr2, &mut ctx).unwrap();
1060
1061 let arr3 = buffer![0i32; 0].into_array();
1063 builder.append_array_as_list(&arr3, &mut ctx).unwrap();
1064
1065 builder
1067 .append_value(Scalar::list_empty(Arc::clone(&dtype), NonNullable).as_list())
1068 .unwrap();
1069
1070 let list = builder.finish_into_list();
1071 assert_eq!(list.len(), 5);
1072
1073 assert_arrays_eq!(
1075 list.elements(),
1076 PrimitiveArray::from_iter([1i32, 2, 3, 10, 11, 4, 5]),
1077 &mut ctx
1078 );
1079
1080 assert_arrays_eq!(
1082 list.offsets(),
1083 PrimitiveArray::from_iter([0u32, 3, 5, 7, 7, 7]),
1084 &mut ctx
1085 );
1086
1087 let mut builder = ListBuilder::<u32>::with_capacity_in(
1089 dtype,
1090 NonNullable,
1091 20,
1092 10,
1093 BufferAllocatorRef::static_ref(),
1094 );
1095 let wrong_dtype_arr = buffer![1i64, 2, 3].into_array();
1096 assert!(
1097 builder
1098 .append_array_as_list(&wrong_dtype_arr, &mut ctx)
1099 .is_err()
1100 );
1101 }
1102}