Skip to main content

vortex_array/builders/
list.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::any::Any;
5use std::sync::Arc;
6
7use num_traits::AsPrimitive;
8use vortex_error::VortexExpect;
9use vortex_error::VortexResult;
10use vortex_error::vortex_bail;
11use vortex_error::vortex_ensure;
12use vortex_error::vortex_panic;
13use vortex_mask::Mask;
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::ListArrayExt;
26use crate::arrays::listview::ListViewArrayExt;
27use crate::builders::ArrayBuilder;
28use crate::builders::DEFAULT_BUILDER_CAPACITY;
29use crate::builders::LazyBitBufferBuilder;
30use crate::builders::PrimitiveBuilder;
31use crate::builders::builder_with_capacity;
32use crate::dtype::DType;
33use crate::dtype::IntegerPType;
34use crate::dtype::Nullability;
35use crate::dtype::Nullability::NonNullable;
36use crate::match_each_integer_ptype;
37use crate::scalar::ListScalar;
38use crate::scalar::Scalar;
39
40/// The builder for building a [`ListArray`], parametrized by the [`IntegerPType`] of the `offsets`
41/// builder.
42pub struct ListBuilder<O: IntegerPType> {
43    /// The [`DType`] of the [`ListArray`]. This **must** be a [`DType::List`].
44    dtype: DType,
45
46    /// The builder for the underlying elements of the [`ListArray`].
47    elements_builder: Box<dyn ArrayBuilder>,
48
49    /// The builder for the `offsets` into the `elements` array.
50    offsets_builder: PrimitiveBuilder<O>,
51
52    /// The null map builder of the [`ListArray`].
53    nulls: LazyBitBufferBuilder,
54}
55
56impl<O: IntegerPType> ListBuilder<O> {
57    /// Creates a new `ListBuilder` with a capacity of [`DEFAULT_BUILDER_CAPACITY`].
58    pub fn new(value_dtype: Arc<DType>, nullability: Nullability) -> Self {
59        Self::with_capacity(
60            value_dtype,
61            nullability,
62            // We arbitrarily choose 2 times the number of list scalars for the capacity of the
63            // elements builder since we cannot know this ahead of time.
64            DEFAULT_BUILDER_CAPACITY * 2,
65            DEFAULT_BUILDER_CAPACITY,
66        )
67    }
68
69    /// Create a new [`ListArray`] builder with a with the given `capacity`, as well as an initial
70    /// capacity for the `elements` builder (since we cannot know that ahead of time solely based on
71    /// the outer array `capacity`).
72    ///
73    /// # Notes
74    ///
75    /// The number of offsets is one more than the length (# of list scalars) in the array.
76    pub fn with_capacity(
77        value_dtype: Arc<DType>,
78        nullability: Nullability,
79        elements_capacity: usize,
80        capacity: usize,
81    ) -> Self {
82        let elements_builder = builder_with_capacity(value_dtype.as_ref(), elements_capacity);
83        let mut offsets_builder = PrimitiveBuilder::<O>::with_capacity(NonNullable, capacity + 1);
84
85        // The first offset is always 0 and represents an empty list.
86        offsets_builder.append_zero();
87
88        Self {
89            elements_builder,
90            offsets_builder,
91            nulls: LazyBitBufferBuilder::new(capacity),
92            dtype: DType::List(value_dtype, nullability),
93        }
94    }
95
96    /// Appends an array as a single non-null list entry to the builder.
97    ///
98    /// The input `array` must have the same dtype as the element dtype of this list builder.
99    ///
100    /// Note that the list entry will be non-null but the elements themselves are allowed to be null
101    /// (only if the elements [`DType`] in nullable, of course).
102    pub fn append_array_as_list(
103        &mut self,
104        array: &ArrayRef,
105        ctx: &mut ExecutionCtx,
106    ) -> VortexResult<()> {
107        vortex_ensure!(
108            array.dtype() == self.element_dtype(),
109            "Array dtype {:?} does not match list element dtype {:?}",
110            array.dtype(),
111            self.element_dtype()
112        );
113
114        self.elements_builder.reserve_exact(array.len());
115        array.append_to_builder(self.elements_builder.as_mut(), ctx)?;
116        self.nulls.append_non_null();
117        self.offsets_builder.append_value(
118            O::from_usize(self.elements_builder.len())
119                .vortex_expect("Failed to convert from usize to O"),
120        );
121
122        Ok(())
123    }
124
125    /// Appends a list `value` to the builder.
126    pub fn append_value(&mut self, value: ListScalar) -> VortexResult<()> {
127        match value.elements() {
128            None => {
129                if self.dtype.nullability() == NonNullable {
130                    vortex_bail!("Cannot append null value to non-nullable list");
131                }
132                self.append_null();
133            }
134            Some(elements) => {
135                for scalar in elements {
136                    // TODO(connor): This is slow, we should be able to append multiple values at
137                    // once, or the list scalar should hold an Array
138                    self.elements_builder.append_scalar(&scalar)?;
139                }
140
141                self.nulls.append_non_null();
142                self.offsets_builder.append_value(
143                    O::from_usize(self.elements_builder.len())
144                        .vortex_expect("Failed to convert from usize to O"),
145                );
146            }
147        }
148
149        Ok(())
150    }
151
152    /// Finishes the builder directly into a [`ListArray`].
153    pub fn finish_into_list(&mut self) -> ListArray {
154        assert_eq!(
155            self.offsets_builder.len(),
156            self.nulls.len() + 1,
157            "offsets length must be one more than nulls length."
158        );
159
160        ListArray::try_new(
161            self.elements_builder.finish(),
162            self.offsets_builder.finish(),
163            self.nulls.finish_with_nullability(self.dtype.nullability()),
164        )
165        .vortex_expect("Buffer, offsets, and validity must have same length.")
166    }
167
168    /// The [`DType`] of the inner elements. Note that this is **not** the same as the [`DType`] of
169    /// the outer `List`.
170    pub fn element_dtype(&self) -> &DType {
171        let DType::List(element_dtype, _) = &self.dtype else {
172            vortex_panic!("`ListBuilder` has an incorrect dtype: {}", self.dtype);
173        };
174
175        element_dtype
176    }
177}
178
179/// Appends `ListViewArray`-layout lists (`n` offsets and sizes) into a [`ListBuilder`], converting
180/// into the `ListArray` (`n + 1` offsets) layout.
181fn extend_from_listview<O, OffsetType, SizeType>(
182    builder: &mut ListBuilder<O>,
183    new_elements: &ArrayRef,
184    new_offsets: &[OffsetType],
185    new_sizes: &[SizeType],
186    ctx: &mut ExecutionCtx,
187) -> VortexResult<()>
188where
189    O: IntegerPType,
190    OffsetType: IntegerPType,
191    SizeType: IntegerPType,
192{
193    let num_lists = new_offsets.len();
194    debug_assert_eq!(num_lists, new_sizes.len());
195
196    let total_elements: usize = new_sizes.iter().map(|size| size.as_()).sum();
197    builder.elements_builder.reserve_exact(total_elements);
198
199    let mut curr_offset = builder.elements_builder.len();
200    builder.offsets_builder.reserve_exact(num_lists);
201    let mut offsets_range = builder.offsets_builder.uninit_range(num_lists);
202
203    // We need to append each list individually, converting from `ListViewArray` format to
204    // the `ListArray` format that `ListBuilder` expects.
205    for i in 0..new_offsets.len() {
206        let offset: usize = new_offsets[i].as_();
207        let size: usize = new_sizes[i].as_();
208
209        if size > 0 {
210            let list_elements = new_elements
211                .slice(offset..offset + size)
212                .vortex_expect("list builder slice");
213            list_elements.append_to_builder(builder.elements_builder.as_mut(), ctx)?;
214            curr_offset += size;
215        }
216
217        let new_offset = O::from_usize(curr_offset).vortex_expect("Failed to convert offset");
218
219        offsets_range.set_value(i, new_offset);
220    }
221
222    // SAFETY: We have initialized all `num_lists` values, and since the `offsets` array is
223    // non-nullable, we are done.
224    unsafe { offsets_range.finish() };
225    Ok(())
226}
227
228impl<O: IntegerPType> ArrayBuilder for ListBuilder<O> {
229    fn as_any(&self) -> &dyn Any {
230        self
231    }
232
233    fn as_any_mut(&mut self) -> &mut dyn Any {
234        self
235    }
236
237    fn dtype(&self) -> &DType {
238        &self.dtype
239    }
240
241    fn len(&self) -> usize {
242        self.nulls.len()
243    }
244
245    fn append_zeros(&mut self, n: usize) {
246        let curr_len = self.elements_builder.len();
247        for _ in 0..n {
248            self.offsets_builder.append_value(
249                O::from_usize(curr_len).vortex_expect("Failed to convert from usize to <O>"),
250            )
251        }
252        self.nulls.append_n_non_nulls(n);
253    }
254
255    unsafe fn append_nulls_unchecked(&mut self, n: usize) {
256        let curr_len = self.elements_builder.len();
257        for _ in 0..n {
258            // A list with a null element is can be a list with a zero-span offset and a validity
259            // bit set
260            self.offsets_builder.append_value(
261                O::from_usize(curr_len).vortex_expect("Failed to convert from usize to <O>"),
262            )
263        }
264        self.nulls.append_n_nulls(n);
265    }
266
267    fn append_scalar(&mut self, scalar: &Scalar) -> VortexResult<()> {
268        vortex_ensure!(
269            scalar.dtype() == self.dtype(),
270            "ListBuilder expected scalar with dtype {}, got {}",
271            self.dtype(),
272            scalar.dtype()
273        );
274
275        self.append_value(scalar.as_list())
276    }
277
278    fn reserve_exact(&mut self, additional: usize) {
279        self.elements_builder.reserve_exact(additional);
280        self.offsets_builder.reserve_exact(additional);
281        self.nulls.reserve_exact(additional);
282    }
283
284    unsafe fn set_validity_unchecked(&mut self, validity: Mask) {
285        self.nulls = LazyBitBufferBuilder::from_validity_mask(validity);
286    }
287
288    fn finish(&mut self) -> ArrayRef {
289        self.finish_into_list().into_array()
290    }
291
292    fn finish_into_canonical(&mut self, ctx: &mut ExecutionCtx) -> Canonical {
293        let listview = self
294            .finish()
295            .execute::<ListViewArray>(ctx)
296            .vortex_expect("list builder should canonicalize to listview");
297        Canonical::List(listview)
298    }
299
300    fn append_list_array(
301        &mut self,
302        array: ArrayView<'_, List>,
303        ctx: &mut ExecutionCtx,
304    ) -> VortexResult<()> {
305        if array.is_empty() {
306            return Ok(());
307        }
308
309        self.nulls
310            .append_validity_mask(&array.validity()?.execute_mask(array.len(), ctx)?);
311
312        let num_lists = array.len();
313        let offsets = array.offsets().clone().execute::<PrimitiveArray>(ctx)?;
314        match_each_integer_ptype!(offsets.ptype(), |OffsetType| {
315            let offsets = offsets.as_slice::<OffsetType>();
316            let first: usize = offsets[0].as_();
317            let last: usize = offsets[num_lists].as_();
318
319            // Lists in a `ListArray` are contiguous, so the referenced elements can be appended
320            // in bulk and the offsets rebased onto this builder's elements.
321            let elements_base = self.elements_builder.len();
322            if last > first {
323                self.elements_builder.reserve_exact(last - first);
324                array
325                    .elements()
326                    .slice(first..last)?
327                    .append_to_builder(self.elements_builder.as_mut(), ctx)?;
328            }
329
330            self.offsets_builder.reserve_exact(num_lists);
331            let mut offsets_range = self.offsets_builder.uninit_range(num_lists);
332            for i in 0..num_lists {
333                let end: usize = offsets[i + 1].as_();
334                offsets_range.set_value(
335                    i,
336                    O::from_usize(end - first + elements_base)
337                        .vortex_expect("Failed to convert offset"),
338                );
339            }
340            // SAFETY: We have initialized all `num_lists` values, and since the `offsets` array is
341            // non-nullable, we are done.
342            unsafe { offsets_range.finish() };
343        });
344        Ok(())
345    }
346
347    fn append_listview_array(
348        &mut self,
349        array: ArrayView<'_, ListView>,
350        ctx: &mut ExecutionCtx,
351    ) -> VortexResult<()> {
352        if array.is_empty() {
353            return Ok(());
354        }
355
356        self.nulls
357            .append_validity_mask(&array.validity()?.execute_mask(array.len(), ctx)?);
358
359        // Note that `ListViewArray` has `n` offsets and sizes, not `n+1` offsets like `ListArray`.
360        let elements = array.elements();
361        let offsets = array.offsets().clone().execute::<PrimitiveArray>(ctx)?;
362        let sizes = array.sizes().clone().execute::<PrimitiveArray>(ctx)?;
363
364        match_each_integer_ptype!(offsets.ptype(), |OffsetType| {
365            match_each_integer_ptype!(sizes.ptype(), |SizeType| {
366                extend_from_listview(
367                    self,
368                    elements,
369                    offsets.as_slice::<OffsetType>(),
370                    sizes.as_slice::<SizeType>(),
371                    ctx,
372                )?
373            })
374        });
375        Ok(())
376    }
377}
378
379#[cfg(test)]
380mod tests {
381    use std::sync::Arc;
382
383    use Nullability::NonNullable;
384    use Nullability::Nullable;
385    use vortex_buffer::buffer;
386    use vortex_error::VortexExpect;
387    use vortex_error::VortexResult;
388
389    use crate::IntoArray;
390    use crate::array_session;
391    use crate::arrays::ChunkedArray;
392    use crate::arrays::ListViewArray;
393    use crate::arrays::PrimitiveArray;
394    use crate::arrays::list::ListArrayExt;
395    use crate::arrays::listview::ListViewArrayExt;
396    use crate::assert_arrays_eq;
397    use crate::builders::ArrayBuilder;
398    use crate::builders::ListViewBuilder;
399    use crate::builders::builder_with_capacity;
400    use crate::builders::list::ListArray;
401    use crate::builders::list::ListBuilder;
402    use crate::dtype::DType;
403    use crate::dtype::IntegerPType;
404    use crate::dtype::Nullability;
405    use crate::dtype::PType::I32;
406    use crate::executor::VortexSessionExecute;
407    use crate::scalar::Scalar;
408    use crate::validity::Validity;
409
410    #[test]
411    fn test_empty() {
412        let mut builder =
413            ListBuilder::<u32>::with_capacity(Arc::new(I32.into()), NonNullable, 0, 0);
414
415        let list = builder.finish();
416        assert_eq!(list.len(), 0);
417    }
418
419    #[test]
420    fn test_values() {
421        let dtype: Arc<DType> = Arc::new(I32.into());
422        let mut builder = ListBuilder::<u32>::with_capacity(Arc::clone(&dtype), NonNullable, 0, 0);
423
424        builder
425            .append_value(
426                Scalar::list(
427                    Arc::clone(&dtype),
428                    vec![1i32.into(), 2i32.into(), 3i32.into()],
429                    NonNullable,
430                )
431                .as_list(),
432            )
433            .unwrap();
434
435        builder
436            .append_value(
437                Scalar::list(
438                    dtype,
439                    vec![4i32.into(), 5i32.into(), 6i32.into()],
440                    NonNullable,
441                )
442                .as_list(),
443            )
444            .unwrap();
445
446        let list = builder.finish();
447        assert_eq!(list.len(), 2);
448
449        let mut ctx = array_session().create_execution_ctx();
450        let list_array = list.execute::<ListViewArray>(&mut ctx).unwrap();
451
452        assert_eq!(list_array.list_elements_at(0).unwrap().len(), 3);
453        assert_eq!(list_array.list_elements_at(1).unwrap().len(), 3);
454    }
455
456    #[test]
457    fn test_append_empty_list() {
458        let dtype: Arc<DType> = Arc::new(I32.into());
459        let mut builder = ListBuilder::<u32>::with_capacity(Arc::clone(&dtype), NonNullable, 0, 0);
460
461        assert!(
462            builder
463                .append_value(Scalar::list_empty(dtype, NonNullable).as_list())
464                .is_ok()
465        )
466    }
467
468    #[test]
469    fn test_nullable_values() {
470        let dtype: Arc<DType> = Arc::new(I32.into());
471        let mut builder = ListBuilder::<u32>::with_capacity(Arc::clone(&dtype), Nullable, 0, 0);
472
473        builder
474            .append_value(
475                Scalar::list(
476                    Arc::clone(&dtype),
477                    vec![1i32.into(), 2i32.into(), 3i32.into()],
478                    NonNullable,
479                )
480                .as_list(),
481            )
482            .unwrap();
483
484        builder
485            .append_value(Scalar::list_empty(Arc::clone(&dtype), NonNullable).as_list())
486            .unwrap();
487
488        builder
489            .append_value(
490                Scalar::list(
491                    dtype,
492                    vec![4i32.into(), 5i32.into(), 6i32.into()],
493                    NonNullable,
494                )
495                .as_list(),
496            )
497            .unwrap();
498
499        let list = builder.finish();
500        assert_eq!(list.len(), 3);
501
502        let mut ctx = array_session().create_execution_ctx();
503        let list_array = list.execute::<ListViewArray>(&mut ctx).unwrap();
504
505        assert_eq!(list_array.list_elements_at(0).unwrap().len(), 3);
506        assert_eq!(list_array.list_elements_at(1).unwrap().len(), 0);
507        assert_eq!(list_array.list_elements_at(2).unwrap().len(), 3);
508    }
509
510    fn test_extend_builder_gen<O: IntegerPType>() {
511        let list = ListArray::from_iter_opt_slow::<O, _, _>(
512            [Some(vec![0, 1, 2]), None, Some(vec![4, 5])],
513            Arc::new(I32.into()),
514        )
515        .unwrap()
516        .into_array();
517        assert_eq!(list.len(), 3);
518
519        let mut ctx = array_session().create_execution_ctx();
520
521        let mut builder = ListBuilder::<O>::with_capacity(Arc::new(I32.into()), Nullable, 18, 9);
522        list.append_to_builder(&mut builder, &mut ctx).unwrap();
523        list.append_to_builder(&mut builder, &mut ctx).unwrap();
524        list.slice(0..0)
525            .unwrap()
526            .append_to_builder(&mut builder, &mut ctx)
527            .unwrap();
528        list.slice(1..3)
529            .unwrap()
530            .append_to_builder(&mut builder, &mut ctx)
531            .unwrap();
532
533        let expected = ListArray::from_iter_opt_slow::<O, _, _>(
534            [
535                Some(vec![0, 1, 2]),
536                None,
537                Some(vec![4, 5]),
538                Some(vec![0, 1, 2]),
539                None,
540                Some(vec![4, 5]),
541                None,
542                Some(vec![4, 5]),
543            ],
544            Arc::new(DType::Primitive(I32, NonNullable)),
545        )
546        .unwrap()
547        .into_array()
548        .execute::<ListViewArray>(&mut ctx)
549        .unwrap();
550
551        let actual = builder.finish_into_canonical(&mut ctx).into_listview();
552
553        assert_arrays_eq!(actual.elements(), expected.elements(), &mut ctx);
554
555        assert_arrays_eq!(actual.offsets(), expected.offsets(), &mut ctx);
556
557        assert!(
558            actual
559                .validity()
560                .vortex_expect("list validity should be derivable")
561                .mask_eq(
562                    &expected
563                        .validity()
564                        .vortex_expect("list validity should be derivable"),
565                    actual.len(),
566                    &mut ctx,
567                )
568                .unwrap(),
569        );
570    }
571
572    /// `append_to_builder` must handle any list builder kind without assuming the offset/size
573    /// integer types produced by `builder_with_capacity`. It appends a `List`-encoded array and a
574    /// `ListView`-encoded array into `ListViewBuilder`s and `ListBuilder`s with assorted (and
575    /// non-`u64`) offset/size types.
576    #[test]
577    fn test_append_to_builder_any_list_builder() -> VortexResult<()> {
578        let mut ctx = array_session().create_execution_ctx();
579
580        let list = ListArray::from_iter_opt_slow::<u64, _, _>(
581            [Some(vec![0, 1, 2]), None, Some(vec![4, 5])],
582            Arc::new(I32.into()),
583        )?
584        .into_array();
585        let listview = list
586            .clone()
587            .execute::<ListViewArray>(&mut ctx)?
588            .into_array();
589        let elem_dtype = || Arc::new(I32.into());
590
591        // `builder_with_capacity` produces a `ListViewBuilder` for `DType::List`; appending the
592        // `List`-encoded array must dispatch into it instead of bailing.
593        let mut listview_builder = builder_with_capacity(list.dtype(), list.len());
594        list.append_to_builder(listview_builder.as_mut(), &mut ctx)?;
595        assert_arrays_eq!(listview_builder.finish(), list, &mut ctx);
596
597        // A `ListViewBuilder` with non-`u64` (including signed) offset and size types must work
598        // for both source encodings.
599        let mut lv_u32_u8 = ListViewBuilder::<u32, u8>::with_capacity(elem_dtype(), Nullable, 8, 4);
600        list.append_to_builder(&mut lv_u32_u8, &mut ctx)?;
601        assert_arrays_eq!(lv_u32_u8.finish(), list, &mut ctx);
602
603        let mut lv_i32_i16 =
604            ListViewBuilder::<i32, i16>::with_capacity(elem_dtype(), Nullable, 8, 4);
605        list.append_to_builder(&mut lv_i32_i16, &mut ctx)?;
606        assert_arrays_eq!(lv_i32_i16.finish(), list, &mut ctx);
607
608        let mut lv_u16_u16 =
609            ListViewBuilder::<u16, u16>::with_capacity(elem_dtype(), Nullable, 8, 4);
610        listview.append_to_builder(&mut lv_u16_u16, &mut ctx)?;
611        assert_arrays_eq!(lv_u16_u16.finish(), list, &mut ctx);
612
613        // Both source encodings appended into `ListBuilder`s with non-`u64` (including signed)
614        // offset types.
615        let mut list_builder = ListBuilder::<u32>::with_capacity(elem_dtype(), Nullable, 8, 4);
616        list.append_to_builder(&mut list_builder, &mut ctx)?;
617        assert_arrays_eq!(list_builder.finish(), list, &mut ctx);
618
619        let mut list_builder_i16 = ListBuilder::<i16>::with_capacity(elem_dtype(), Nullable, 8, 4);
620        listview.append_to_builder(&mut list_builder_i16, &mut ctx)?;
621        assert_arrays_eq!(list_builder_i16.finish(), list, &mut ctx);
622
623        Ok(())
624    }
625
626    #[test]
627    fn test_append_list_arrays_grow_builder() -> VortexResult<()> {
628        let mut ctx = array_session().create_execution_ctx();
629        let dtype: Arc<DType> = Arc::new(I32.into());
630
631        // Enough lists to exceed the offsets capacity of a zero-capacity builder, so appending
632        // must grow the builder rather than panic in `uninit_range`.
633        let lists: Vec<Option<Vec<i32>>> =
634            (0..100).map(|i| (i % 10 != 0).then(|| vec![i])).collect();
635        let source = ListArray::from_iter_opt_slow::<u32, _, _>(lists.clone(), Arc::clone(&dtype))?;
636        let expected = ListArray::from_iter_opt_slow::<u32, _, _>(
637            lists.iter().cloned().chain(lists.iter().cloned()),
638            Arc::clone(&dtype),
639        )?;
640
641        // Appending twice checks growth from a non-empty builder and offset rebasing.
642        let mut builder = ListBuilder::<u32>::with_capacity(Arc::clone(&dtype), Nullable, 0, 0);
643        builder.append_list_array(source.as_view(), &mut ctx)?;
644        builder.append_list_array(source.as_view(), &mut ctx)?;
645        assert_arrays_eq!(builder.finish(), expected, &mut ctx);
646
647        let source_listview = source.into_array().execute::<ListViewArray>(&mut ctx)?;
648        let mut builder = ListBuilder::<u32>::with_capacity(dtype, Nullable, 0, 0);
649        builder.append_listview_array(source_listview.as_view(), &mut ctx)?;
650        builder.append_listview_array(source_listview.as_view(), &mut ctx)?;
651        assert_arrays_eq!(builder.finish(), expected, &mut ctx);
652
653        Ok(())
654    }
655
656    #[test]
657    fn test_extend_builder() {
658        test_extend_builder_gen::<i8>();
659        test_extend_builder_gen::<i16>();
660        test_extend_builder_gen::<i32>();
661        test_extend_builder_gen::<i64>();
662
663        test_extend_builder_gen::<u8>();
664        test_extend_builder_gen::<u16>();
665        test_extend_builder_gen::<u32>();
666        test_extend_builder_gen::<u64>();
667    }
668
669    #[test]
670    pub fn test_array_with_gap() {
671        let one_trailing_unused_element = ListArray::try_new(
672            buffer![1, 2, 3, 4].into_array(),
673            buffer![0, 3].into_array(),
674            Validity::NonNullable,
675        )
676        .unwrap();
677
678        let second_array = ListArray::try_new(
679            buffer![5, 6].into_array(),
680            buffer![0, 2].into_array(),
681            Validity::NonNullable,
682        )
683        .unwrap();
684
685        let chunked_list = ChunkedArray::try_new(
686            vec![
687                one_trailing_unused_element.clone().into_array(),
688                second_array.clone().into_array(),
689            ],
690            DType::List(Arc::new(DType::Primitive(I32, NonNullable)), NonNullable),
691        );
692
693        let mut ctx = array_session().create_execution_ctx();
694        let canon_values = chunked_list
695            .unwrap()
696            .as_array()
697            .clone()
698            .execute::<ListViewArray>(&mut ctx)
699            .unwrap();
700
701        assert_eq!(
702            one_trailing_unused_element
703                .execute_scalar(0, &mut array_session().create_execution_ctx())
704                .unwrap(),
705            canon_values
706                .execute_scalar(0, &mut array_session().create_execution_ctx())
707                .unwrap()
708        );
709        assert_eq!(
710            second_array
711                .execute_scalar(0, &mut array_session().create_execution_ctx())
712                .unwrap(),
713            canon_values
714                .execute_scalar(1, &mut array_session().create_execution_ctx())
715                .unwrap()
716        );
717    }
718
719    #[test]
720    fn test_append_scalar() {
721        let dtype: Arc<DType> = Arc::new(I32.into());
722        let mut builder = ListBuilder::<u64>::with_capacity(Arc::clone(&dtype), Nullable, 20, 10);
723
724        // Test appending a valid list.
725        let list_scalar1 =
726            Scalar::list(Arc::clone(&dtype), vec![1i32.into(), 2i32.into()], Nullable);
727        builder.append_scalar(&list_scalar1).unwrap();
728
729        // Test appending another list.
730        let list_scalar2 = Scalar::list(
731            Arc::clone(&dtype),
732            vec![3i32.into(), 4i32.into(), 5i32.into()],
733            Nullable,
734        );
735        builder.append_scalar(&list_scalar2).unwrap();
736
737        // Test appending null value.
738        let null_scalar = Scalar::null(DType::List(Arc::clone(&dtype), Nullable));
739        builder.append_scalar(&null_scalar).unwrap();
740
741        let array = builder.finish_into_list();
742        assert_eq!(array.len(), 3);
743
744        let mut ctx = array_session().create_execution_ctx();
745
746        // Check actual values using scalar_at.
747
748        let scalar0 = array.execute_scalar(0, &mut ctx).unwrap();
749        let list0 = scalar0.as_list();
750        assert_eq!(list0.len(), 2);
751        if let Some(list0_items) = list0.elements() {
752            assert_eq!(list0_items[0].as_primitive().typed_value::<i32>(), Some(1));
753            assert_eq!(list0_items[1].as_primitive().typed_value::<i32>(), Some(2));
754        }
755
756        let scalar1 = array.execute_scalar(1, &mut ctx).unwrap();
757        let list1 = scalar1.as_list();
758        assert_eq!(list1.len(), 3);
759        if let Some(list1_items) = list1.elements() {
760            assert_eq!(list1_items[0].as_primitive().typed_value::<i32>(), Some(3));
761            assert_eq!(list1_items[1].as_primitive().typed_value::<i32>(), Some(4));
762            assert_eq!(list1_items[2].as_primitive().typed_value::<i32>(), Some(5));
763        }
764
765        let scalar2 = array.execute_scalar(2, &mut ctx).unwrap();
766        let list2 = scalar2.as_list();
767        assert!(list2.is_null()); // This should be null.
768
769        // Check validity.
770        assert!(
771            array
772                .validity()
773                .vortex_expect("list validity should be derivable")
774                .execute_is_valid(0, &mut ctx)
775                .unwrap()
776        );
777        assert!(
778            array
779                .validity()
780                .vortex_expect("list validity should be derivable")
781                .execute_is_valid(1, &mut ctx)
782                .unwrap()
783        );
784        assert!(
785            !array
786                .validity()
787                .vortex_expect("list validity should be derivable")
788                .execute_is_valid(2, &mut ctx)
789                .unwrap()
790        );
791
792        // Test wrong dtype error.
793        let mut builder = ListBuilder::<u64>::with_capacity(dtype, NonNullable, 20, 10);
794        let wrong_scalar = Scalar::from(42i32);
795        assert!(builder.append_scalar(&wrong_scalar).is_err());
796    }
797
798    #[test]
799    fn test_append_array_as_list() {
800        let dtype: Arc<DType> = Arc::new(I32.into());
801        let mut ctx = array_session().create_execution_ctx();
802        let mut builder =
803            ListBuilder::<u32>::with_capacity(Arc::clone(&dtype), NonNullable, 20, 10);
804
805        // Append a primitive array as a single list entry.
806        let arr1 = buffer![1i32, 2, 3].into_array();
807        builder.append_array_as_list(&arr1, &mut ctx).unwrap();
808
809        // Interleave with a list scalar.
810        builder
811            .append_value(
812                Scalar::list(
813                    Arc::clone(&dtype),
814                    vec![10i32.into(), 11i32.into()],
815                    NonNullable,
816                )
817                .as_list(),
818            )
819            .unwrap();
820
821        // Append another primitive array as a single list entry.
822        let arr2 = buffer![4i32, 5].into_array();
823        builder.append_array_as_list(&arr2, &mut ctx).unwrap();
824
825        // Append an empty array as a single list entry (empty list).
826        let arr3 = buffer![0i32; 0].into_array();
827        builder.append_array_as_list(&arr3, &mut ctx).unwrap();
828
829        // Interleave with another list scalar (empty list).
830        builder
831            .append_value(Scalar::list_empty(Arc::clone(&dtype), NonNullable).as_list())
832            .unwrap();
833
834        let list = builder.finish_into_list();
835        assert_eq!(list.len(), 5);
836
837        // Verify elements array: [1, 2, 3, 10, 11, 4, 5].
838        assert_arrays_eq!(
839            list.elements(),
840            PrimitiveArray::from_iter([1i32, 2, 3, 10, 11, 4, 5]),
841            &mut ctx
842        );
843
844        // Verify offsets array.
845        assert_arrays_eq!(
846            list.offsets(),
847            PrimitiveArray::from_iter([0u32, 3, 5, 7, 7, 7]),
848            &mut ctx
849        );
850
851        // Test dtype mismatch error.
852        let mut builder = ListBuilder::<u32>::with_capacity(dtype, NonNullable, 20, 10);
853        let wrong_dtype_arr = buffer![1i64, 2, 3].into_array();
854        assert!(
855            builder
856                .append_array_as_list(&wrong_dtype_arr, &mut ctx)
857                .is_err()
858        );
859    }
860}