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 ListArray::try_new(
188 self.elements_builder.finish(),
189 self.offsets_builder.finish(),
190 self.nulls.finish_with_nullability(self.dtype.nullability()),
191 )
192 .vortex_expect("Buffer, offsets, and validity must have same length.")
193 }
194
195 pub fn element_dtype(&self) -> &DType {
198 let DType::List(element_dtype, _) = &self.dtype else {
199 vortex_panic!("`ListBuilder` has an incorrect dtype: {}", self.dtype);
200 };
201
202 element_dtype
203 }
204
205 pub fn append_list_array(
212 &mut self,
213 array: ArrayView<'_, List>,
214 ctx: &mut ExecutionCtx,
215 ) -> VortexResult<()> {
216 if array.is_empty() {
217 return Ok(());
218 }
219
220 self.nulls.append_validity(array.validity()?, array.len());
221
222 let num_lists = array.len();
223 let offsets = array.offsets().clone().execute::<PrimitiveArray>(ctx)?;
224 match_each_integer_ptype!(offsets.ptype(), |OffsetType| {
225 let offsets = offsets.as_slice::<OffsetType>();
226 let first: usize = offsets[0].as_();
227 let last: usize = offsets[num_lists].as_();
228
229 let elements_base = self.elements_builder.len();
232 if last > first {
233 self.elements_builder
234 .append_array(&array.elements().slice(first..last)?, ctx)?;
235 }
236
237 self.offsets_builder.reserve_exact(num_lists);
238 let mut offsets_range = self.offsets_builder.uninit_range(num_lists);
239 for i in 0..num_lists {
240 let end: usize = offsets[i + 1].as_();
241 offsets_range.set_value(
242 i,
243 O::from_usize(end - first + elements_base)
244 .vortex_expect("Failed to convert offset"),
245 );
246 }
247 unsafe { offsets_range.finish() };
250 });
251 Ok(())
252 }
253
254 pub fn append_listview_array(
262 &mut self,
263 array: ArrayView<'_, ListView>,
264 ctx: &mut ExecutionCtx,
265 ) -> VortexResult<()> {
266 if array.is_empty() {
267 return Ok(());
268 }
269
270 self.nulls.append_validity(array.validity()?, array.len());
271
272 let array = array
276 .into_owned()
277 .rebuild(ListViewRebuildMode::MakeZeroCopyToList, ctx)?;
278 debug_assert!(array.is_zero_copy_to_list());
279
280 let elements = array.elements();
282 let offsets = array.offsets().clone().execute::<PrimitiveArray>(ctx)?;
283 let sizes = array.sizes().clone().execute::<PrimitiveArray>(ctx)?;
284
285 match_each_integer_ptype!(offsets.ptype(), |OffsetType| {
286 match_each_integer_ptype!(sizes.ptype(), |SizeType| {
287 extend_from_listview(
288 self,
289 elements,
290 offsets.as_slice::<OffsetType>(),
291 sizes.as_slice::<SizeType>(),
292 ctx,
293 )?
294 })
295 });
296 Ok(())
297 }
298}
299
300fn extend_from_listview<O, OffsetType, SizeType>(
308 builder: &mut ListBuilder<O>,
309 new_elements: &ArrayRef,
310 new_offsets: &[OffsetType],
311 new_sizes: &[SizeType],
312 ctx: &mut ExecutionCtx,
313) -> VortexResult<()>
314where
315 O: OffsetBuilderPType,
316 OffsetType: IntegerPType,
317 SizeType: IntegerPType,
318{
319 let num_lists = new_offsets.len();
320 debug_assert_eq!(num_lists, new_sizes.len());
321
322 let first: usize = new_offsets[0].as_();
325 let last: usize = new_offsets[num_lists - 1].as_() + new_sizes[num_lists - 1].as_();
326
327 let elements_base = builder.elements_builder.len();
328 if last > first {
329 builder
330 .elements_builder
331 .append_array(&new_elements.slice(first..last)?, ctx)?;
332 }
333
334 builder.offsets_builder.reserve_exact(num_lists);
335 let mut offsets_range = builder.offsets_builder.uninit_range(num_lists);
336 for i in 0..num_lists {
337 let end: usize = new_offsets[i].as_() + new_sizes[i].as_();
338 offsets_range.set_value(
339 i,
340 O::from_usize(end - first + elements_base).vortex_expect("Failed to convert offset"),
341 );
342 }
343
344 unsafe { offsets_range.finish() };
347 Ok(())
348}
349
350impl<O: OffsetBuilderPType> ArrayBuilder for ListBuilder<O> {
351 fn as_any(&self) -> &dyn Any {
352 self
353 }
354
355 fn as_any_mut(&mut self) -> &mut dyn Any {
356 self
357 }
358
359 fn dtype(&self) -> &DType {
360 &self.dtype
361 }
362
363 fn len(&self) -> usize {
364 self.nulls.len()
365 }
366
367 fn append_zeros(&mut self, n: usize) {
368 let curr_len = self.elements_builder.len();
369 for _ in 0..n {
370 self.offsets_builder.append_value(
371 O::from_usize(curr_len).vortex_expect("Failed to convert from usize to <O>"),
372 )
373 }
374 self.nulls.append_n_non_nulls(n);
375 }
376
377 unsafe fn append_nulls_unchecked(&mut self, n: usize) {
378 let curr_len = self.elements_builder.len();
379 for _ in 0..n {
380 self.offsets_builder.append_value(
383 O::from_usize(curr_len).vortex_expect("Failed to convert from usize to <O>"),
384 )
385 }
386 self.nulls.append_n_nulls(n);
387 }
388
389 fn append_scalar(&mut self, scalar: &Scalar) -> VortexResult<()> {
390 vortex_ensure_eq!(
391 scalar.dtype(),
392 self.dtype(),
393 "ListBuilder received a scalar with the wrong dtype"
394 );
395
396 self.append_value(scalar.as_list())
397 }
398
399 fn reserve_exact(&mut self, additional: usize) {
400 self.elements_builder.reserve_exact(additional);
401 self.offsets_builder.reserve_exact(additional);
402 self.nulls.reserve_exact(additional);
403 }
404
405 fn finish(&mut self) -> ArrayRef {
406 self.finish_into_list().into_array()
407 }
408
409 fn finish_into_canonical(&mut self, ctx: &mut ExecutionCtx) -> Canonical {
410 let listview = self
411 .finish()
412 .execute::<ListViewArray>(ctx)
413 .vortex_expect("list builder should canonicalize to listview");
414 Canonical::List(listview)
415 }
416}
417
418#[cfg(test)]
419mod tests {
420 use std::sync::Arc;
421
422 use Nullability::NonNullable;
423 use Nullability::Nullable;
424 use vortex_buffer::BufferAllocatorRef;
425 use vortex_buffer::buffer;
426 use vortex_error::VortexExpect;
427 use vortex_error::VortexResult;
428
429 use crate::IntoArray;
430 use crate::array_session;
431 use crate::arrays::ChunkedArray;
432 use crate::arrays::ListViewArray;
433 use crate::arrays::PrimitiveArray;
434 use crate::arrays::list::ListArraySlotsExt;
435 use crate::arrays::listview::ListViewArrayExt;
436 use crate::arrays::listview::ListViewArraySlotsExt;
437 use crate::assert_arrays_eq;
438 use crate::builders::ArrayBuilder;
439 use crate::builders::ListViewBuilder;
440 use crate::builders::builder_with_capacity_in;
441 use crate::builders::list::ListArray;
442 use crate::builders::list::ListBuilder;
443 use crate::dtype::DType;
444 use crate::dtype::Nullability;
445 use crate::dtype::OffsetBuilderPType;
446 use crate::dtype::PType::I32;
447 use crate::executor::VortexSessionExecute;
448 use crate::scalar::Scalar;
449 use crate::validity::Validity;
450
451 #[test]
452 fn test_empty() {
453 let mut builder = ListBuilder::<u32>::with_capacity_in(
454 Arc::new(I32.into()),
455 NonNullable,
456 0,
457 0,
458 BufferAllocatorRef::static_ref(),
459 );
460
461 let list = builder.finish();
462 assert_eq!(list.len(), 0);
463 }
464
465 #[test]
466 fn test_values() {
467 let dtype: Arc<DType> = Arc::new(I32.into());
468 let mut builder = ListBuilder::<u32>::with_capacity_in(
469 Arc::clone(&dtype),
470 NonNullable,
471 0,
472 0,
473 BufferAllocatorRef::static_ref(),
474 );
475
476 builder
477 .append_value(
478 Scalar::list(
479 Arc::clone(&dtype),
480 vec![1i32.into(), 2i32.into(), 3i32.into()],
481 NonNullable,
482 )
483 .as_list(),
484 )
485 .unwrap();
486
487 builder
488 .append_value(
489 Scalar::list(
490 dtype,
491 vec![4i32.into(), 5i32.into(), 6i32.into()],
492 NonNullable,
493 )
494 .as_list(),
495 )
496 .unwrap();
497
498 let list = builder.finish();
499 assert_eq!(list.len(), 2);
500
501 let mut ctx = array_session().create_execution_ctx();
502 let list_array = list.execute::<ListViewArray>(&mut ctx).unwrap();
503
504 assert_eq!(list_array.list_elements_at(0).unwrap().len(), 3);
505 assert_eq!(list_array.list_elements_at(1).unwrap().len(), 3);
506 }
507
508 #[test]
509 fn test_append_empty_list() {
510 let dtype: Arc<DType> = Arc::new(I32.into());
511 let mut builder = ListBuilder::<u32>::with_capacity_in(
512 Arc::clone(&dtype),
513 NonNullable,
514 0,
515 0,
516 BufferAllocatorRef::static_ref(),
517 );
518
519 assert!(
520 builder
521 .append_value(Scalar::list_empty(dtype, NonNullable).as_list())
522 .is_ok()
523 )
524 }
525
526 #[test]
527 fn test_nullable_values() {
528 let dtype: Arc<DType> = Arc::new(I32.into());
529 let mut builder = ListBuilder::<u32>::with_capacity_in(
530 Arc::clone(&dtype),
531 Nullable,
532 0,
533 0,
534 BufferAllocatorRef::static_ref(),
535 );
536
537 builder
538 .append_value(
539 Scalar::list(
540 Arc::clone(&dtype),
541 vec![1i32.into(), 2i32.into(), 3i32.into()],
542 NonNullable,
543 )
544 .as_list(),
545 )
546 .unwrap();
547
548 builder
549 .append_value(Scalar::list_empty(Arc::clone(&dtype), NonNullable).as_list())
550 .unwrap();
551
552 builder
553 .append_value(
554 Scalar::list(
555 dtype,
556 vec![4i32.into(), 5i32.into(), 6i32.into()],
557 NonNullable,
558 )
559 .as_list(),
560 )
561 .unwrap();
562
563 let list = builder.finish();
564 assert_eq!(list.len(), 3);
565
566 let mut ctx = array_session().create_execution_ctx();
567 let list_array = list.execute::<ListViewArray>(&mut ctx).unwrap();
568
569 assert_eq!(list_array.list_elements_at(0).unwrap().len(), 3);
570 assert_eq!(list_array.list_elements_at(1).unwrap().len(), 0);
571 assert_eq!(list_array.list_elements_at(2).unwrap().len(), 3);
572 }
573
574 fn test_extend_builder_gen<O: OffsetBuilderPType>() {
575 let list = ListArray::from_iter_opt_slow::<O, _, _>(
576 [Some(vec![0, 1, 2]), None, Some(vec![4, 5])],
577 Arc::new(I32.into()),
578 )
579 .unwrap()
580 .into_array();
581 assert_eq!(list.len(), 3);
582
583 let mut ctx = array_session().create_execution_ctx();
584
585 let mut builder = ListBuilder::<O>::with_capacity_in(
586 Arc::new(I32.into()),
587 Nullable,
588 18,
589 9,
590 BufferAllocatorRef::static_ref(),
591 );
592 list.append_to_builder(&mut builder, &mut ctx).unwrap();
593 list.append_to_builder(&mut builder, &mut ctx).unwrap();
594 list.slice(0..0)
595 .unwrap()
596 .append_to_builder(&mut builder, &mut ctx)
597 .unwrap();
598 list.slice(1..3)
599 .unwrap()
600 .append_to_builder(&mut builder, &mut ctx)
601 .unwrap();
602
603 let expected = ListArray::from_iter_opt_slow::<O, _, _>(
604 [
605 Some(vec![0, 1, 2]),
606 None,
607 Some(vec![4, 5]),
608 Some(vec![0, 1, 2]),
609 None,
610 Some(vec![4, 5]),
611 None,
612 Some(vec![4, 5]),
613 ],
614 Arc::new(DType::Primitive(I32, NonNullable)),
615 )
616 .unwrap()
617 .into_array()
618 .execute::<ListViewArray>(&mut ctx)
619 .unwrap();
620
621 let actual = builder.finish_into_canonical(&mut ctx).into_listview();
622
623 assert_arrays_eq!(actual.elements(), expected.elements(), &mut ctx);
624
625 assert_arrays_eq!(actual.offsets(), expected.offsets(), &mut ctx);
626
627 assert!(
628 actual
629 .validity()
630 .vortex_expect("list validity should be derivable")
631 .mask_eq(
632 &expected
633 .validity()
634 .vortex_expect("list validity should be derivable"),
635 actual.len(),
636 &mut ctx,
637 )
638 .unwrap(),
639 );
640 }
641
642 #[test]
647 fn test_append_to_builder_any_list_builder() -> VortexResult<()> {
648 let mut ctx = array_session().create_execution_ctx();
649
650 let list = ListArray::from_iter_opt_slow::<u64, _, _>(
651 [Some(vec![0, 1, 2]), None, Some(vec![4, 5])],
652 Arc::new(I32.into()),
653 )?
654 .into_array();
655 let listview = list
656 .clone()
657 .execute::<ListViewArray>(&mut ctx)?
658 .into_array();
659 let elem_dtype = || Arc::new(I32.into());
660
661 let mut listview_builder =
664 builder_with_capacity_in(list.dtype(), list.len(), BufferAllocatorRef::static_ref());
665 list.append_to_builder(listview_builder.as_mut(), &mut ctx)?;
666 assert_arrays_eq!(listview_builder.finish(), list, &mut ctx);
667
668 let mut lv_u64_u32 = ListViewBuilder::<u64, u32>::with_capacity_in(
671 elem_dtype(),
672 Nullable,
673 8,
674 4,
675 BufferAllocatorRef::static_ref(),
676 );
677 list.append_to_builder(&mut lv_u64_u32, &mut ctx)?;
678 assert_arrays_eq!(lv_u64_u32.finish(), list, &mut ctx);
679
680 let mut lv_i64_i32 = ListViewBuilder::<i64, i32>::with_capacity_in(
681 elem_dtype(),
682 Nullable,
683 8,
684 4,
685 BufferAllocatorRef::static_ref(),
686 );
687 list.append_to_builder(&mut lv_i64_i32, &mut ctx)?;
688 assert_arrays_eq!(lv_i64_i32.finish(), list, &mut ctx);
689
690 let mut lv_u32_u32 = ListViewBuilder::<u32, u32>::with_capacity_in(
691 elem_dtype(),
692 Nullable,
693 8,
694 4,
695 BufferAllocatorRef::static_ref(),
696 );
697 listview.append_to_builder(&mut lv_u32_u32, &mut ctx)?;
698 assert_arrays_eq!(lv_u32_u32.finish(), list, &mut ctx);
699
700 let mut list_builder = ListBuilder::<u32>::with_capacity_in(
703 elem_dtype(),
704 Nullable,
705 8,
706 4,
707 BufferAllocatorRef::static_ref(),
708 );
709 list.append_to_builder(&mut list_builder, &mut ctx)?;
710 assert_arrays_eq!(list_builder.finish(), list, &mut ctx);
711
712 let mut list_builder_i32 = ListBuilder::<i32>::with_capacity_in(
713 elem_dtype(),
714 Nullable,
715 8,
716 4,
717 BufferAllocatorRef::static_ref(),
718 );
719 listview.append_to_builder(&mut list_builder_i32, &mut ctx)?;
720 assert_arrays_eq!(list_builder_i32.finish(), list, &mut ctx);
721
722 Ok(())
723 }
724
725 #[test]
726 fn test_append_list_arrays_grow_builder() -> VortexResult<()> {
727 let mut ctx = array_session().create_execution_ctx();
728 let dtype: Arc<DType> = Arc::new(I32.into());
729
730 let lists: Vec<Option<Vec<i32>>> =
733 (0..100).map(|i| (i % 10 != 0).then(|| vec![i])).collect();
734 let source = ListArray::from_iter_opt_slow::<u32, _, _>(lists.clone(), Arc::clone(&dtype))?;
735 let expected = ListArray::from_iter_opt_slow::<u32, _, _>(
736 lists.iter().cloned().chain(lists.iter().cloned()),
737 Arc::clone(&dtype),
738 )?;
739
740 let mut builder = ListBuilder::<u32>::with_capacity_in(
742 Arc::clone(&dtype),
743 Nullable,
744 0,
745 0,
746 BufferAllocatorRef::static_ref(),
747 );
748 builder.append_list_array(source.as_view(), &mut ctx)?;
749 builder.append_list_array(source.as_view(), &mut ctx)?;
750 assert_arrays_eq!(builder.finish(), expected, &mut ctx);
751
752 let source_listview = source.into_array().execute::<ListViewArray>(&mut ctx)?;
753 let mut builder = ListBuilder::<u32>::with_capacity_in(
754 dtype,
755 Nullable,
756 0,
757 0,
758 BufferAllocatorRef::static_ref(),
759 );
760 builder.append_listview_array(source_listview.as_view(), &mut ctx)?;
761 builder.append_listview_array(source_listview.as_view(), &mut ctx)?;
762 assert_arrays_eq!(builder.finish(), expected, &mut ctx);
763
764 Ok(())
765 }
766
767 #[test]
771 fn test_append_listview_array_flattens_overlaps_and_skips_leading_elements() -> VortexResult<()>
772 {
773 let mut ctx = array_session().create_execution_ctx();
774 let dtype: Arc<DType> = Arc::new(I32.into());
775
776 let overlapping = unsafe {
781 ListViewArray::new_unchecked(
782 buffer![10i32, 20, 30].into_array(),
783 buffer![0u32, 1, 0].into_array(),
784 buffer![2u8, 2, 1].into_array(),
785 Validity::from_iter([true, false, true]),
786 )
787 };
788 assert!(!overlapping.is_zero_copy_to_list());
789
790 let sliced = unsafe {
792 ListViewArray::new_unchecked(
793 buffer![40i32, 50, 60, 70].into_array(),
794 buffer![0u32, 2].into_array(),
795 buffer![2u32, 2].into_array(),
796 Validity::AllValid,
797 )
798 .with_zero_copy_to_list(true)
799 }
800 .into_array()
801 .slice(1..2)?
802 .execute::<ListViewArray>(&mut ctx)?;
803
804 let mut builder = ListBuilder::<u32>::with_capacity_in(
805 dtype,
806 Nullable,
807 0,
808 0,
809 BufferAllocatorRef::static_ref(),
810 );
811 builder.append_listview_array(overlapping.as_view(), &mut ctx)?;
812 builder.append_listview_array(sliced.as_view(), &mut ctx)?;
813
814 let list = builder.finish_into_list();
815 assert_arrays_eq!(
816 list.elements(),
817 PrimitiveArray::from_iter([10i32, 20, 10, 60, 70]),
818 &mut ctx
819 );
820 assert_arrays_eq!(
821 list.offsets(),
822 PrimitiveArray::from_iter([0u32, 2, 2, 3, 5]),
823 &mut ctx
824 );
825
826 Ok(())
827 }
828
829 #[test]
830 fn test_extend_builder() {
831 test_extend_builder_gen::<i32>();
832 test_extend_builder_gen::<i64>();
833
834 test_extend_builder_gen::<u32>();
835 test_extend_builder_gen::<u64>();
836 }
837
838 #[test]
839 pub fn test_array_with_gap() {
840 let one_trailing_unused_element = ListArray::try_new(
841 buffer![1, 2, 3, 4].into_array(),
842 buffer![0, 3].into_array(),
843 Validity::NonNullable,
844 )
845 .unwrap();
846
847 let second_array = ListArray::try_new(
848 buffer![5, 6].into_array(),
849 buffer![0, 2].into_array(),
850 Validity::NonNullable,
851 )
852 .unwrap();
853
854 let chunked_list = ChunkedArray::try_new(
855 vec![
856 one_trailing_unused_element.clone().into_array(),
857 second_array.clone().into_array(),
858 ],
859 DType::List(Arc::new(DType::Primitive(I32, NonNullable)), NonNullable),
860 );
861
862 let mut ctx = array_session().create_execution_ctx();
863 let canon_values = chunked_list
864 .unwrap()
865 .as_array()
866 .clone()
867 .execute::<ListViewArray>(&mut ctx)
868 .unwrap();
869
870 assert_eq!(
871 one_trailing_unused_element
872 .execute_scalar(0, &mut array_session().create_execution_ctx())
873 .unwrap(),
874 canon_values
875 .execute_scalar(0, &mut array_session().create_execution_ctx())
876 .unwrap()
877 );
878 assert_eq!(
879 second_array
880 .execute_scalar(0, &mut array_session().create_execution_ctx())
881 .unwrap(),
882 canon_values
883 .execute_scalar(1, &mut array_session().create_execution_ctx())
884 .unwrap()
885 );
886 }
887
888 #[test]
889 fn test_append_scalar() {
890 let dtype: Arc<DType> = Arc::new(I32.into());
891 let mut builder = ListBuilder::<u64>::with_capacity_in(
892 Arc::clone(&dtype),
893 Nullable,
894 20,
895 10,
896 BufferAllocatorRef::static_ref(),
897 );
898
899 let list_scalar1 =
901 Scalar::list(Arc::clone(&dtype), vec![1i32.into(), 2i32.into()], Nullable);
902 builder.append_scalar(&list_scalar1).unwrap();
903
904 let list_scalar2 = Scalar::list(
906 Arc::clone(&dtype),
907 vec![3i32.into(), 4i32.into(), 5i32.into()],
908 Nullable,
909 );
910 builder.append_scalar(&list_scalar2).unwrap();
911
912 let null_scalar = Scalar::null(DType::List(Arc::clone(&dtype), Nullable));
914 builder.append_scalar(&null_scalar).unwrap();
915
916 let array = builder.finish_into_list();
917 assert_eq!(array.len(), 3);
918
919 let mut ctx = array_session().create_execution_ctx();
920
921 let scalar0 = array.execute_scalar(0, &mut ctx).unwrap();
924 let list0 = scalar0.as_list();
925 assert_eq!(list0.len(), 2);
926 if let Some(list0_items) = list0.elements() {
927 assert_eq!(list0_items[0].as_primitive().typed_value::<i32>(), Some(1));
928 assert_eq!(list0_items[1].as_primitive().typed_value::<i32>(), Some(2));
929 }
930
931 let scalar1 = array.execute_scalar(1, &mut ctx).unwrap();
932 let list1 = scalar1.as_list();
933 assert_eq!(list1.len(), 3);
934 if let Some(list1_items) = list1.elements() {
935 assert_eq!(list1_items[0].as_primitive().typed_value::<i32>(), Some(3));
936 assert_eq!(list1_items[1].as_primitive().typed_value::<i32>(), Some(4));
937 assert_eq!(list1_items[2].as_primitive().typed_value::<i32>(), Some(5));
938 }
939
940 let scalar2 = array.execute_scalar(2, &mut ctx).unwrap();
941 let list2 = scalar2.as_list();
942 assert!(list2.is_null()); assert!(
946 array
947 .validity()
948 .vortex_expect("list validity should be derivable")
949 .execute_is_valid(0, &mut ctx)
950 .unwrap()
951 );
952 assert!(
953 array
954 .validity()
955 .vortex_expect("list validity should be derivable")
956 .execute_is_valid(1, &mut ctx)
957 .unwrap()
958 );
959 assert!(
960 !array
961 .validity()
962 .vortex_expect("list validity should be derivable")
963 .execute_is_valid(2, &mut ctx)
964 .unwrap()
965 );
966
967 let mut builder = ListBuilder::<u64>::with_capacity_in(
969 dtype,
970 NonNullable,
971 20,
972 10,
973 BufferAllocatorRef::static_ref(),
974 );
975 let wrong_scalar = Scalar::from(42i32);
976 assert!(builder.append_scalar(&wrong_scalar).is_err());
977 }
978
979 #[test]
980 fn test_append_array_as_list() {
981 let dtype: Arc<DType> = Arc::new(I32.into());
982 let mut ctx = array_session().create_execution_ctx();
983 let mut builder = ListBuilder::<u32>::with_capacity_in(
984 Arc::clone(&dtype),
985 NonNullable,
986 20,
987 10,
988 BufferAllocatorRef::static_ref(),
989 );
990
991 let arr1 = buffer![1i32, 2, 3].into_array();
993 builder.append_array_as_list(&arr1, &mut ctx).unwrap();
994
995 builder
997 .append_value(
998 Scalar::list(
999 Arc::clone(&dtype),
1000 vec![10i32.into(), 11i32.into()],
1001 NonNullable,
1002 )
1003 .as_list(),
1004 )
1005 .unwrap();
1006
1007 let arr2 = buffer![4i32, 5].into_array();
1009 builder.append_array_as_list(&arr2, &mut ctx).unwrap();
1010
1011 let arr3 = buffer![0i32; 0].into_array();
1013 builder.append_array_as_list(&arr3, &mut ctx).unwrap();
1014
1015 builder
1017 .append_value(Scalar::list_empty(Arc::clone(&dtype), NonNullable).as_list())
1018 .unwrap();
1019
1020 let list = builder.finish_into_list();
1021 assert_eq!(list.len(), 5);
1022
1023 assert_arrays_eq!(
1025 list.elements(),
1026 PrimitiveArray::from_iter([1i32, 2, 3, 10, 11, 4, 5]),
1027 &mut ctx
1028 );
1029
1030 assert_arrays_eq!(
1032 list.offsets(),
1033 PrimitiveArray::from_iter([0u32, 3, 5, 7, 7, 7]),
1034 &mut ctx
1035 );
1036
1037 let mut builder = ListBuilder::<u32>::with_capacity_in(
1039 dtype,
1040 NonNullable,
1041 20,
1042 10,
1043 BufferAllocatorRef::static_ref(),
1044 );
1045 let wrong_dtype_arr = buffer![1i64, 2, 3].into_array();
1046 assert!(
1047 builder
1048 .append_array_as_list(&wrong_dtype_arr, &mut ctx)
1049 .is_err()
1050 );
1051 }
1052}