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 assert_eq!(
464 list.dtype(),
465 &DType::List(Arc::new(I32.into()), NonNullable)
466 );
467 }
468
469 #[test]
470 fn test_values() {
471 let dtype: Arc<DType> = Arc::new(I32.into());
472 let mut builder = ListBuilder::<u32>::with_capacity_in(
473 Arc::clone(&dtype),
474 NonNullable,
475 0,
476 0,
477 BufferAllocatorRef::static_ref(),
478 );
479
480 builder
481 .append_value(
482 Scalar::list(
483 Arc::clone(&dtype),
484 vec![1i32.into(), 2i32.into(), 3i32.into()],
485 NonNullable,
486 )
487 .as_list(),
488 )
489 .unwrap();
490
491 builder
492 .append_value(
493 Scalar::list(
494 dtype,
495 vec![4i32.into(), 5i32.into(), 6i32.into()],
496 NonNullable,
497 )
498 .as_list(),
499 )
500 .unwrap();
501
502 let list = builder.finish();
503 assert_eq!(list.len(), 2);
504
505 let mut ctx = array_session().create_execution_ctx();
506 let list_array = list.execute::<ListViewArray>(&mut ctx).unwrap();
507
508 assert_eq!(list_array.list_elements_at(0).unwrap().len(), 3);
509 assert_eq!(list_array.list_elements_at(1).unwrap().len(), 3);
510 let expected = ListArray::new(
511 buffer![1i32, 2, 3, 4, 5, 6].into_array(),
512 buffer![0u32, 3, 6].into_array(),
513 Validity::NonNullable,
514 );
515 assert_arrays_eq!(list_array, expected, &mut ctx);
516 }
517
518 #[test]
519 fn test_append_empty_list() {
520 let dtype: Arc<DType> = Arc::new(I32.into());
521 let mut builder = ListBuilder::<u32>::with_capacity_in(
522 Arc::clone(&dtype),
523 NonNullable,
524 0,
525 0,
526 BufferAllocatorRef::static_ref(),
527 );
528
529 assert!(
530 builder
531 .append_value(Scalar::list_empty(dtype, NonNullable).as_list())
532 .is_ok()
533 );
534 let expected = ListArray::new(
535 buffer![0i32; 0].into_array(),
536 buffer![0u32, 0].into_array(),
537 Validity::NonNullable,
538 );
539 assert_arrays_eq!(
540 builder.finish(),
541 expected,
542 &mut array_session().create_execution_ctx()
543 );
544 }
545
546 #[test]
547 fn test_nullable_values() {
548 let dtype: Arc<DType> = Arc::new(I32.into());
549 let mut builder = ListBuilder::<u32>::with_capacity_in(
550 Arc::clone(&dtype),
551 Nullable,
552 0,
553 0,
554 BufferAllocatorRef::static_ref(),
555 );
556
557 builder
558 .append_value(
559 Scalar::list(
560 Arc::clone(&dtype),
561 vec![1i32.into(), 2i32.into(), 3i32.into()],
562 NonNullable,
563 )
564 .as_list(),
565 )
566 .unwrap();
567
568 builder
569 .append_value(Scalar::list_empty(Arc::clone(&dtype), NonNullable).as_list())
570 .unwrap();
571
572 builder
573 .append_value(
574 Scalar::list(
575 dtype,
576 vec![4i32.into(), 5i32.into(), 6i32.into()],
577 NonNullable,
578 )
579 .as_list(),
580 )
581 .unwrap();
582
583 let list = builder.finish();
584 assert_eq!(list.len(), 3);
585
586 let mut ctx = array_session().create_execution_ctx();
587 let list_array = list.execute::<ListViewArray>(&mut ctx).unwrap();
588
589 assert_eq!(list_array.list_elements_at(0).unwrap().len(), 3);
590 assert_eq!(list_array.list_elements_at(1).unwrap().len(), 0);
591 assert_eq!(list_array.list_elements_at(2).unwrap().len(), 3);
592 let expected = ListArray::new(
593 buffer![1i32, 2, 3, 4, 5, 6].into_array(),
594 buffer![0u32, 3, 3, 6].into_array(),
595 Validity::AllValid,
596 );
597 assert_arrays_eq!(list_array, expected, &mut ctx);
598 }
599
600 fn test_extend_builder_gen<O: OffsetBuilderPType>() {
601 let list = ListArray::from_iter_opt_slow::<O, _, _>(
602 [Some(vec![0, 1, 2]), None, Some(vec![4, 5])],
603 Arc::new(I32.into()),
604 )
605 .unwrap()
606 .into_array();
607 assert_eq!(list.len(), 3);
608
609 let mut ctx = array_session().create_execution_ctx();
610
611 let mut builder = ListBuilder::<O>::with_capacity_in(
612 Arc::new(I32.into()),
613 Nullable,
614 18,
615 9,
616 BufferAllocatorRef::static_ref(),
617 );
618 list.append_to_builder(&mut builder, &mut ctx).unwrap();
619 list.append_to_builder(&mut builder, &mut ctx).unwrap();
620 list.slice(0..0)
621 .unwrap()
622 .append_to_builder(&mut builder, &mut ctx)
623 .unwrap();
624 list.slice(1..3)
625 .unwrap()
626 .append_to_builder(&mut builder, &mut ctx)
627 .unwrap();
628
629 let expected = ListArray::from_iter_opt_slow::<O, _, _>(
630 [
631 Some(vec![0, 1, 2]),
632 None,
633 Some(vec![4, 5]),
634 Some(vec![0, 1, 2]),
635 None,
636 Some(vec![4, 5]),
637 None,
638 Some(vec![4, 5]),
639 ],
640 Arc::new(DType::Primitive(I32, NonNullable)),
641 )
642 .unwrap()
643 .into_array()
644 .execute::<ListViewArray>(&mut ctx)
645 .unwrap();
646
647 let actual = builder.finish_into_canonical(&mut ctx).into_listview();
648
649 assert_arrays_eq!(actual.elements(), expected.elements(), &mut ctx);
650
651 assert_arrays_eq!(actual.offsets(), expected.offsets(), &mut ctx);
652
653 assert!(
654 actual
655 .validity()
656 .vortex_expect("list validity should be derivable")
657 .mask_eq(
658 &expected
659 .validity()
660 .vortex_expect("list validity should be derivable"),
661 actual.len(),
662 &mut ctx,
663 )
664 .unwrap(),
665 );
666 let independent = ListArray::new(
667 buffer![0i32, 1, 2, 4, 5, 0, 1, 2, 4, 5, 4, 5].into_array(),
668 buffer![0u32, 3, 3, 5, 8, 8, 10, 10, 12].into_array(),
669 Validity::from_iter([true, false, true, true, false, true, false, true]),
670 );
671 assert_arrays_eq!(actual, independent, &mut ctx);
672 }
673
674 #[test]
679 fn test_append_to_builder_any_list_builder() -> VortexResult<()> {
680 let mut ctx = array_session().create_execution_ctx();
681
682 let list = ListArray::from_iter_opt_slow::<u64, _, _>(
683 [Some(vec![0, 1, 2]), None, Some(vec![4, 5])],
684 Arc::new(I32.into()),
685 )?
686 .into_array();
687 let listview = list
688 .clone()
689 .execute::<ListViewArray>(&mut ctx)?
690 .into_array();
691 let elem_dtype = || Arc::new(I32.into());
692
693 let mut listview_builder =
696 builder_with_capacity_in(list.dtype(), list.len(), BufferAllocatorRef::static_ref());
697 list.append_to_builder(listview_builder.as_mut(), &mut ctx)?;
698 assert_arrays_eq!(listview_builder.finish(), list, &mut ctx);
699
700 let mut lv_u64_u32 = ListViewBuilder::<u64, u32>::with_capacity_in(
703 elem_dtype(),
704 Nullable,
705 8,
706 4,
707 BufferAllocatorRef::static_ref(),
708 );
709 list.append_to_builder(&mut lv_u64_u32, &mut ctx)?;
710 assert_arrays_eq!(lv_u64_u32.finish(), list, &mut ctx);
711
712 let mut lv_i64_i32 = ListViewBuilder::<i64, i32>::with_capacity_in(
713 elem_dtype(),
714 Nullable,
715 8,
716 4,
717 BufferAllocatorRef::static_ref(),
718 );
719 list.append_to_builder(&mut lv_i64_i32, &mut ctx)?;
720 assert_arrays_eq!(lv_i64_i32.finish(), list, &mut ctx);
721
722 let mut lv_u32_u32 = ListViewBuilder::<u32, u32>::with_capacity_in(
723 elem_dtype(),
724 Nullable,
725 8,
726 4,
727 BufferAllocatorRef::static_ref(),
728 );
729 listview.append_to_builder(&mut lv_u32_u32, &mut ctx)?;
730 assert_arrays_eq!(lv_u32_u32.finish(), list, &mut ctx);
731
732 let mut list_builder = ListBuilder::<u32>::with_capacity_in(
735 elem_dtype(),
736 Nullable,
737 8,
738 4,
739 BufferAllocatorRef::static_ref(),
740 );
741 list.append_to_builder(&mut list_builder, &mut ctx)?;
742 assert_arrays_eq!(list_builder.finish(), list, &mut ctx);
743
744 let mut list_builder_i32 = ListBuilder::<i32>::with_capacity_in(
745 elem_dtype(),
746 Nullable,
747 8,
748 4,
749 BufferAllocatorRef::static_ref(),
750 );
751 listview.append_to_builder(&mut list_builder_i32, &mut ctx)?;
752 assert_arrays_eq!(list_builder_i32.finish(), list, &mut ctx);
753
754 Ok(())
755 }
756
757 #[test]
758 fn test_append_list_arrays_grow_builder() -> VortexResult<()> {
759 let mut ctx = array_session().create_execution_ctx();
760 let dtype: Arc<DType> = Arc::new(I32.into());
761
762 let lists: Vec<Option<Vec<i32>>> =
765 (0..100).map(|i| (i % 10 != 0).then(|| vec![i])).collect();
766 let source = ListArray::from_iter_opt_slow::<u32, _, _>(lists.clone(), Arc::clone(&dtype))?;
767 let expected = ListArray::from_iter_opt_slow::<u32, _, _>(
768 lists.iter().cloned().chain(lists.iter().cloned()),
769 Arc::clone(&dtype),
770 )?;
771
772 let mut builder = ListBuilder::<u32>::with_capacity_in(
774 Arc::clone(&dtype),
775 Nullable,
776 0,
777 0,
778 BufferAllocatorRef::static_ref(),
779 );
780 builder.append_list_array(source.as_view(), &mut ctx)?;
781 builder.append_list_array(source.as_view(), &mut ctx)?;
782 assert_arrays_eq!(builder.finish(), expected, &mut ctx);
783
784 let source_listview = source.into_array().execute::<ListViewArray>(&mut ctx)?;
785 let mut builder = ListBuilder::<u32>::with_capacity_in(
786 dtype,
787 Nullable,
788 0,
789 0,
790 BufferAllocatorRef::static_ref(),
791 );
792 builder.append_listview_array(source_listview.as_view(), &mut ctx)?;
793 builder.append_listview_array(source_listview.as_view(), &mut ctx)?;
794 assert_arrays_eq!(builder.finish(), expected, &mut ctx);
795
796 Ok(())
797 }
798
799 #[test]
803 fn test_append_listview_array_flattens_overlaps_and_skips_leading_elements() -> VortexResult<()>
804 {
805 let mut ctx = array_session().create_execution_ctx();
806 let dtype: Arc<DType> = Arc::new(I32.into());
807
808 let overlapping = unsafe {
813 ListViewArray::new_unchecked(
814 buffer![10i32, 20, 30].into_array(),
815 buffer![0u32, 1, 0].into_array(),
816 buffer![2u8, 2, 1].into_array(),
817 Validity::from_iter([true, false, true]),
818 )
819 };
820 assert!(!overlapping.is_zero_copy_to_list());
821
822 let sliced = unsafe {
824 ListViewArray::new_unchecked(
825 buffer![40i32, 50, 60, 70].into_array(),
826 buffer![0u32, 2].into_array(),
827 buffer![2u32, 2].into_array(),
828 Validity::AllValid,
829 )
830 .with_zero_copy_to_list(true)
831 }
832 .into_array()
833 .slice(1..2)?
834 .execute::<ListViewArray>(&mut ctx)?;
835
836 let mut builder = ListBuilder::<u32>::with_capacity_in(
837 dtype,
838 Nullable,
839 0,
840 0,
841 BufferAllocatorRef::static_ref(),
842 );
843 builder.append_listview_array(overlapping.as_view(), &mut ctx)?;
844 builder.append_listview_array(sliced.as_view(), &mut ctx)?;
845
846 let list = builder.finish_into_list();
847 assert_arrays_eq!(
848 list.elements(),
849 PrimitiveArray::from_iter([10i32, 20, 10, 60, 70]),
850 &mut ctx
851 );
852 assert_arrays_eq!(
853 list.offsets(),
854 PrimitiveArray::from_iter([0u32, 2, 2, 3, 5]),
855 &mut ctx
856 );
857
858 let expected = ListArray::new(
859 buffer![10i32, 20, 10, 60, 70].into_array(),
860 buffer![0u32, 2, 2, 3, 5].into_array(),
861 Validity::from_iter([true, false, true, true]),
862 );
863 assert_arrays_eq!(list, expected, &mut ctx);
864 Ok(())
865 }
866
867 #[test]
868 fn test_extend_builder() {
869 test_extend_builder_gen::<i32>();
870 test_extend_builder_gen::<i64>();
871
872 test_extend_builder_gen::<u32>();
873 test_extend_builder_gen::<u64>();
874 }
875
876 #[test]
877 pub fn test_array_with_gap() {
878 let one_trailing_unused_element = ListArray::try_new(
879 buffer![1, 2, 3, 4].into_array(),
880 buffer![0, 3].into_array(),
881 Validity::NonNullable,
882 )
883 .unwrap();
884
885 let second_array = ListArray::try_new(
886 buffer![5, 6].into_array(),
887 buffer![0, 2].into_array(),
888 Validity::NonNullable,
889 )
890 .unwrap();
891
892 let chunked_list = ChunkedArray::try_new(
893 vec![
894 one_trailing_unused_element.clone().into_array(),
895 second_array.clone().into_array(),
896 ],
897 DType::List(Arc::new(DType::Primitive(I32, NonNullable)), NonNullable),
898 );
899
900 let mut ctx = array_session().create_execution_ctx();
901 let canon_values = chunked_list
902 .unwrap()
903 .as_array()
904 .clone()
905 .execute::<ListViewArray>(&mut ctx)
906 .unwrap();
907
908 assert_eq!(
909 one_trailing_unused_element
910 .execute_scalar(0, &mut array_session().create_execution_ctx())
911 .unwrap(),
912 canon_values
913 .execute_scalar(0, &mut array_session().create_execution_ctx())
914 .unwrap()
915 );
916 assert_eq!(
917 second_array
918 .execute_scalar(0, &mut array_session().create_execution_ctx())
919 .unwrap(),
920 canon_values
921 .execute_scalar(1, &mut array_session().create_execution_ctx())
922 .unwrap()
923 );
924 }
925
926 #[test]
927 fn test_append_scalar() {
928 let dtype: Arc<DType> = Arc::new(I32.into());
929 let mut builder = ListBuilder::<u64>::with_capacity_in(
930 Arc::clone(&dtype),
931 Nullable,
932 20,
933 10,
934 BufferAllocatorRef::static_ref(),
935 );
936
937 let list_scalar1 =
939 Scalar::list(Arc::clone(&dtype), vec![1i32.into(), 2i32.into()], Nullable);
940 builder.append_scalar(&list_scalar1).unwrap();
941
942 let list_scalar2 = Scalar::list(
944 Arc::clone(&dtype),
945 vec![3i32.into(), 4i32.into(), 5i32.into()],
946 Nullable,
947 );
948 builder.append_scalar(&list_scalar2).unwrap();
949
950 let null_scalar = Scalar::null(DType::List(Arc::clone(&dtype), Nullable));
952 builder.append_scalar(&null_scalar).unwrap();
953
954 let array = builder.finish_into_list();
955 assert_eq!(array.len(), 3);
956
957 let mut ctx = array_session().create_execution_ctx();
958
959 let scalar0 = array.execute_scalar(0, &mut ctx).unwrap();
962 let list0 = scalar0.as_list();
963 assert_eq!(list0.len(), 2);
964 if let Some(list0_items) = list0.elements() {
965 assert_eq!(list0_items[0].as_primitive().typed_value::<i32>(), Some(1));
966 assert_eq!(list0_items[1].as_primitive().typed_value::<i32>(), Some(2));
967 }
968
969 let scalar1 = array.execute_scalar(1, &mut ctx).unwrap();
970 let list1 = scalar1.as_list();
971 assert_eq!(list1.len(), 3);
972 if let Some(list1_items) = list1.elements() {
973 assert_eq!(list1_items[0].as_primitive().typed_value::<i32>(), Some(3));
974 assert_eq!(list1_items[1].as_primitive().typed_value::<i32>(), Some(4));
975 assert_eq!(list1_items[2].as_primitive().typed_value::<i32>(), Some(5));
976 }
977
978 let scalar2 = array.execute_scalar(2, &mut ctx).unwrap();
979 let list2 = scalar2.as_list();
980 assert!(list2.is_null()); assert!(
984 array
985 .validity()
986 .vortex_expect("list validity should be derivable")
987 .execute_is_valid(0, &mut ctx)
988 .unwrap()
989 );
990 assert!(
991 array
992 .validity()
993 .vortex_expect("list validity should be derivable")
994 .execute_is_valid(1, &mut ctx)
995 .unwrap()
996 );
997 assert!(
998 !array
999 .validity()
1000 .vortex_expect("list validity should be derivable")
1001 .execute_is_valid(2, &mut ctx)
1002 .unwrap()
1003 );
1004
1005 let mut builder = ListBuilder::<u64>::with_capacity_in(
1007 dtype,
1008 NonNullable,
1009 20,
1010 10,
1011 BufferAllocatorRef::static_ref(),
1012 );
1013 let wrong_scalar = Scalar::from(42i32);
1014 assert!(builder.append_scalar(&wrong_scalar).is_err());
1015 }
1016
1017 #[test]
1018 fn test_append_array_as_list() {
1019 let dtype: Arc<DType> = Arc::new(I32.into());
1020 let mut ctx = array_session().create_execution_ctx();
1021 let mut builder = ListBuilder::<u32>::with_capacity_in(
1022 Arc::clone(&dtype),
1023 NonNullable,
1024 20,
1025 10,
1026 BufferAllocatorRef::static_ref(),
1027 );
1028
1029 let arr1 = buffer![1i32, 2, 3].into_array();
1031 builder.append_array_as_list(&arr1, &mut ctx).unwrap();
1032
1033 builder
1035 .append_value(
1036 Scalar::list(
1037 Arc::clone(&dtype),
1038 vec![10i32.into(), 11i32.into()],
1039 NonNullable,
1040 )
1041 .as_list(),
1042 )
1043 .unwrap();
1044
1045 let arr2 = buffer![4i32, 5].into_array();
1047 builder.append_array_as_list(&arr2, &mut ctx).unwrap();
1048
1049 let arr3 = buffer![0i32; 0].into_array();
1051 builder.append_array_as_list(&arr3, &mut ctx).unwrap();
1052
1053 builder
1055 .append_value(Scalar::list_empty(Arc::clone(&dtype), NonNullable).as_list())
1056 .unwrap();
1057
1058 let list = builder.finish_into_list();
1059 assert_eq!(list.len(), 5);
1060
1061 assert_arrays_eq!(
1063 list.elements(),
1064 PrimitiveArray::from_iter([1i32, 2, 3, 10, 11, 4, 5]),
1065 &mut ctx
1066 );
1067
1068 assert_arrays_eq!(
1070 list.offsets(),
1071 PrimitiveArray::from_iter([0u32, 3, 5, 7, 7, 7]),
1072 &mut ctx
1073 );
1074
1075 let mut builder = ListBuilder::<u32>::with_capacity_in(
1077 dtype,
1078 NonNullable,
1079 20,
1080 10,
1081 BufferAllocatorRef::static_ref(),
1082 );
1083 let wrong_dtype_arr = buffer![1i64, 2, 3].into_array();
1084 assert!(
1085 builder
1086 .append_array_as_list(&wrong_dtype_arr, &mut ctx)
1087 .is_err()
1088 );
1089 }
1090}