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::ListArraySlotsExt;
26use crate::arrays::listview::ListViewArraySlotsExt;
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::ListArraySlotsExt;
395    use crate::arrays::listview::ListViewArrayExt;
396    use crate::arrays::listview::ListViewArraySlotsExt;
397    use crate::assert_arrays_eq;
398    use crate::builders::ArrayBuilder;
399    use crate::builders::ListViewBuilder;
400    use crate::builders::builder_with_capacity;
401    use crate::builders::list::ListArray;
402    use crate::builders::list::ListBuilder;
403    use crate::dtype::DType;
404    use crate::dtype::IntegerPType;
405    use crate::dtype::Nullability;
406    use crate::dtype::PType::I32;
407    use crate::executor::VortexSessionExecute;
408    use crate::scalar::Scalar;
409    use crate::validity::Validity;
410
411    #[test]
412    fn test_empty() {
413        let mut builder =
414            ListBuilder::<u32>::with_capacity(Arc::new(I32.into()), NonNullable, 0, 0);
415
416        let list = builder.finish();
417        assert_eq!(list.len(), 0);
418    }
419
420    #[test]
421    fn test_values() {
422        let dtype: Arc<DType> = Arc::new(I32.into());
423        let mut builder = ListBuilder::<u32>::with_capacity(Arc::clone(&dtype), NonNullable, 0, 0);
424
425        builder
426            .append_value(
427                Scalar::list(
428                    Arc::clone(&dtype),
429                    vec![1i32.into(), 2i32.into(), 3i32.into()],
430                    NonNullable,
431                )
432                .as_list(),
433            )
434            .unwrap();
435
436        builder
437            .append_value(
438                Scalar::list(
439                    dtype,
440                    vec![4i32.into(), 5i32.into(), 6i32.into()],
441                    NonNullable,
442                )
443                .as_list(),
444            )
445            .unwrap();
446
447        let list = builder.finish();
448        assert_eq!(list.len(), 2);
449
450        let mut ctx = array_session().create_execution_ctx();
451        let list_array = list.execute::<ListViewArray>(&mut ctx).unwrap();
452
453        assert_eq!(list_array.list_elements_at(0).unwrap().len(), 3);
454        assert_eq!(list_array.list_elements_at(1).unwrap().len(), 3);
455    }
456
457    #[test]
458    fn test_append_empty_list() {
459        let dtype: Arc<DType> = Arc::new(I32.into());
460        let mut builder = ListBuilder::<u32>::with_capacity(Arc::clone(&dtype), NonNullable, 0, 0);
461
462        assert!(
463            builder
464                .append_value(Scalar::list_empty(dtype, NonNullable).as_list())
465                .is_ok()
466        )
467    }
468
469    #[test]
470    fn test_nullable_values() {
471        let dtype: Arc<DType> = Arc::new(I32.into());
472        let mut builder = ListBuilder::<u32>::with_capacity(Arc::clone(&dtype), Nullable, 0, 0);
473
474        builder
475            .append_value(
476                Scalar::list(
477                    Arc::clone(&dtype),
478                    vec![1i32.into(), 2i32.into(), 3i32.into()],
479                    NonNullable,
480                )
481                .as_list(),
482            )
483            .unwrap();
484
485        builder
486            .append_value(Scalar::list_empty(Arc::clone(&dtype), NonNullable).as_list())
487            .unwrap();
488
489        builder
490            .append_value(
491                Scalar::list(
492                    dtype,
493                    vec![4i32.into(), 5i32.into(), 6i32.into()],
494                    NonNullable,
495                )
496                .as_list(),
497            )
498            .unwrap();
499
500        let list = builder.finish();
501        assert_eq!(list.len(), 3);
502
503        let mut ctx = array_session().create_execution_ctx();
504        let list_array = list.execute::<ListViewArray>(&mut ctx).unwrap();
505
506        assert_eq!(list_array.list_elements_at(0).unwrap().len(), 3);
507        assert_eq!(list_array.list_elements_at(1).unwrap().len(), 0);
508        assert_eq!(list_array.list_elements_at(2).unwrap().len(), 3);
509    }
510
511    fn test_extend_builder_gen<O: IntegerPType>() {
512        let list = ListArray::from_iter_opt_slow::<O, _, _>(
513            [Some(vec![0, 1, 2]), None, Some(vec![4, 5])],
514            Arc::new(I32.into()),
515        )
516        .unwrap()
517        .into_array();
518        assert_eq!(list.len(), 3);
519
520        let mut ctx = array_session().create_execution_ctx();
521
522        let mut builder = ListBuilder::<O>::with_capacity(Arc::new(I32.into()), Nullable, 18, 9);
523        list.append_to_builder(&mut builder, &mut ctx).unwrap();
524        list.append_to_builder(&mut builder, &mut ctx).unwrap();
525        list.slice(0..0)
526            .unwrap()
527            .append_to_builder(&mut builder, &mut ctx)
528            .unwrap();
529        list.slice(1..3)
530            .unwrap()
531            .append_to_builder(&mut builder, &mut ctx)
532            .unwrap();
533
534        let expected = ListArray::from_iter_opt_slow::<O, _, _>(
535            [
536                Some(vec![0, 1, 2]),
537                None,
538                Some(vec![4, 5]),
539                Some(vec![0, 1, 2]),
540                None,
541                Some(vec![4, 5]),
542                None,
543                Some(vec![4, 5]),
544            ],
545            Arc::new(DType::Primitive(I32, NonNullable)),
546        )
547        .unwrap()
548        .into_array()
549        .execute::<ListViewArray>(&mut ctx)
550        .unwrap();
551
552        let actual = builder.finish_into_canonical(&mut ctx).into_listview();
553
554        assert_arrays_eq!(actual.elements(), expected.elements(), &mut ctx);
555
556        assert_arrays_eq!(actual.offsets(), expected.offsets(), &mut ctx);
557
558        assert!(
559            actual
560                .validity()
561                .vortex_expect("list validity should be derivable")
562                .mask_eq(
563                    &expected
564                        .validity()
565                        .vortex_expect("list validity should be derivable"),
566                    actual.len(),
567                    &mut ctx,
568                )
569                .unwrap(),
570        );
571    }
572
573    /// `append_to_builder` must handle any list builder kind without assuming the offset/size
574    /// integer types produced by `builder_with_capacity`. It appends a `List`-encoded array and a
575    /// `ListView`-encoded array into `ListViewBuilder`s and `ListBuilder`s with assorted (and
576    /// non-`u64`) offset/size types.
577    #[test]
578    fn test_append_to_builder_any_list_builder() -> VortexResult<()> {
579        let mut ctx = array_session().create_execution_ctx();
580
581        let list = ListArray::from_iter_opt_slow::<u64, _, _>(
582            [Some(vec![0, 1, 2]), None, Some(vec![4, 5])],
583            Arc::new(I32.into()),
584        )?
585        .into_array();
586        let listview = list
587            .clone()
588            .execute::<ListViewArray>(&mut ctx)?
589            .into_array();
590        let elem_dtype = || Arc::new(I32.into());
591
592        // `builder_with_capacity` produces a `ListViewBuilder` for `DType::List`; appending the
593        // `List`-encoded array must dispatch into it instead of bailing.
594        let mut listview_builder = builder_with_capacity(list.dtype(), list.len());
595        list.append_to_builder(listview_builder.as_mut(), &mut ctx)?;
596        assert_arrays_eq!(listview_builder.finish(), list, &mut ctx);
597
598        // A `ListViewBuilder` with non-`u64` (including signed) offset and size types must work
599        // for both source encodings.
600        let mut lv_u32_u8 = ListViewBuilder::<u32, u8>::with_capacity(elem_dtype(), Nullable, 8, 4);
601        list.append_to_builder(&mut lv_u32_u8, &mut ctx)?;
602        assert_arrays_eq!(lv_u32_u8.finish(), list, &mut ctx);
603
604        let mut lv_i32_i16 =
605            ListViewBuilder::<i32, i16>::with_capacity(elem_dtype(), Nullable, 8, 4);
606        list.append_to_builder(&mut lv_i32_i16, &mut ctx)?;
607        assert_arrays_eq!(lv_i32_i16.finish(), list, &mut ctx);
608
609        let mut lv_u16_u16 =
610            ListViewBuilder::<u16, u16>::with_capacity(elem_dtype(), Nullable, 8, 4);
611        listview.append_to_builder(&mut lv_u16_u16, &mut ctx)?;
612        assert_arrays_eq!(lv_u16_u16.finish(), list, &mut ctx);
613
614        // Both source encodings appended into `ListBuilder`s with non-`u64` (including signed)
615        // offset types.
616        let mut list_builder = ListBuilder::<u32>::with_capacity(elem_dtype(), Nullable, 8, 4);
617        list.append_to_builder(&mut list_builder, &mut ctx)?;
618        assert_arrays_eq!(list_builder.finish(), list, &mut ctx);
619
620        let mut list_builder_i16 = ListBuilder::<i16>::with_capacity(elem_dtype(), Nullable, 8, 4);
621        listview.append_to_builder(&mut list_builder_i16, &mut ctx)?;
622        assert_arrays_eq!(list_builder_i16.finish(), list, &mut ctx);
623
624        Ok(())
625    }
626
627    #[test]
628    fn test_append_list_arrays_grow_builder() -> VortexResult<()> {
629        let mut ctx = array_session().create_execution_ctx();
630        let dtype: Arc<DType> = Arc::new(I32.into());
631
632        // Enough lists to exceed the offsets capacity of a zero-capacity builder, so appending
633        // must grow the builder rather than panic in `uninit_range`.
634        let lists: Vec<Option<Vec<i32>>> =
635            (0..100).map(|i| (i % 10 != 0).then(|| vec![i])).collect();
636        let source = ListArray::from_iter_opt_slow::<u32, _, _>(lists.clone(), Arc::clone(&dtype))?;
637        let expected = ListArray::from_iter_opt_slow::<u32, _, _>(
638            lists.iter().cloned().chain(lists.iter().cloned()),
639            Arc::clone(&dtype),
640        )?;
641
642        // Appending twice checks growth from a non-empty builder and offset rebasing.
643        let mut builder = ListBuilder::<u32>::with_capacity(Arc::clone(&dtype), Nullable, 0, 0);
644        builder.append_list_array(source.as_view(), &mut ctx)?;
645        builder.append_list_array(source.as_view(), &mut ctx)?;
646        assert_arrays_eq!(builder.finish(), expected, &mut ctx);
647
648        let source_listview = source.into_array().execute::<ListViewArray>(&mut ctx)?;
649        let mut builder = ListBuilder::<u32>::with_capacity(dtype, Nullable, 0, 0);
650        builder.append_listview_array(source_listview.as_view(), &mut ctx)?;
651        builder.append_listview_array(source_listview.as_view(), &mut ctx)?;
652        assert_arrays_eq!(builder.finish(), expected, &mut ctx);
653
654        Ok(())
655    }
656
657    #[test]
658    fn test_extend_builder() {
659        test_extend_builder_gen::<i8>();
660        test_extend_builder_gen::<i16>();
661        test_extend_builder_gen::<i32>();
662        test_extend_builder_gen::<i64>();
663
664        test_extend_builder_gen::<u8>();
665        test_extend_builder_gen::<u16>();
666        test_extend_builder_gen::<u32>();
667        test_extend_builder_gen::<u64>();
668    }
669
670    #[test]
671    pub fn test_array_with_gap() {
672        let one_trailing_unused_element = ListArray::try_new(
673            buffer![1, 2, 3, 4].into_array(),
674            buffer![0, 3].into_array(),
675            Validity::NonNullable,
676        )
677        .unwrap();
678
679        let second_array = ListArray::try_new(
680            buffer![5, 6].into_array(),
681            buffer![0, 2].into_array(),
682            Validity::NonNullable,
683        )
684        .unwrap();
685
686        let chunked_list = ChunkedArray::try_new(
687            vec![
688                one_trailing_unused_element.clone().into_array(),
689                second_array.clone().into_array(),
690            ],
691            DType::List(Arc::new(DType::Primitive(I32, NonNullable)), NonNullable),
692        );
693
694        let mut ctx = array_session().create_execution_ctx();
695        let canon_values = chunked_list
696            .unwrap()
697            .as_array()
698            .clone()
699            .execute::<ListViewArray>(&mut ctx)
700            .unwrap();
701
702        assert_eq!(
703            one_trailing_unused_element
704                .execute_scalar(0, &mut array_session().create_execution_ctx())
705                .unwrap(),
706            canon_values
707                .execute_scalar(0, &mut array_session().create_execution_ctx())
708                .unwrap()
709        );
710        assert_eq!(
711            second_array
712                .execute_scalar(0, &mut array_session().create_execution_ctx())
713                .unwrap(),
714            canon_values
715                .execute_scalar(1, &mut array_session().create_execution_ctx())
716                .unwrap()
717        );
718    }
719
720    #[test]
721    fn test_append_scalar() {
722        let dtype: Arc<DType> = Arc::new(I32.into());
723        let mut builder = ListBuilder::<u64>::with_capacity(Arc::clone(&dtype), Nullable, 20, 10);
724
725        // Test appending a valid list.
726        let list_scalar1 =
727            Scalar::list(Arc::clone(&dtype), vec![1i32.into(), 2i32.into()], Nullable);
728        builder.append_scalar(&list_scalar1).unwrap();
729
730        // Test appending another list.
731        let list_scalar2 = Scalar::list(
732            Arc::clone(&dtype),
733            vec![3i32.into(), 4i32.into(), 5i32.into()],
734            Nullable,
735        );
736        builder.append_scalar(&list_scalar2).unwrap();
737
738        // Test appending null value.
739        let null_scalar = Scalar::null(DType::List(Arc::clone(&dtype), Nullable));
740        builder.append_scalar(&null_scalar).unwrap();
741
742        let array = builder.finish_into_list();
743        assert_eq!(array.len(), 3);
744
745        let mut ctx = array_session().create_execution_ctx();
746
747        // Check actual values using scalar_at.
748
749        let scalar0 = array.execute_scalar(0, &mut ctx).unwrap();
750        let list0 = scalar0.as_list();
751        assert_eq!(list0.len(), 2);
752        if let Some(list0_items) = list0.elements() {
753            assert_eq!(list0_items[0].as_primitive().typed_value::<i32>(), Some(1));
754            assert_eq!(list0_items[1].as_primitive().typed_value::<i32>(), Some(2));
755        }
756
757        let scalar1 = array.execute_scalar(1, &mut ctx).unwrap();
758        let list1 = scalar1.as_list();
759        assert_eq!(list1.len(), 3);
760        if let Some(list1_items) = list1.elements() {
761            assert_eq!(list1_items[0].as_primitive().typed_value::<i32>(), Some(3));
762            assert_eq!(list1_items[1].as_primitive().typed_value::<i32>(), Some(4));
763            assert_eq!(list1_items[2].as_primitive().typed_value::<i32>(), Some(5));
764        }
765
766        let scalar2 = array.execute_scalar(2, &mut ctx).unwrap();
767        let list2 = scalar2.as_list();
768        assert!(list2.is_null()); // This should be null.
769
770        // Check validity.
771        assert!(
772            array
773                .validity()
774                .vortex_expect("list validity should be derivable")
775                .execute_is_valid(0, &mut ctx)
776                .unwrap()
777        );
778        assert!(
779            array
780                .validity()
781                .vortex_expect("list validity should be derivable")
782                .execute_is_valid(1, &mut ctx)
783                .unwrap()
784        );
785        assert!(
786            !array
787                .validity()
788                .vortex_expect("list validity should be derivable")
789                .execute_is_valid(2, &mut ctx)
790                .unwrap()
791        );
792
793        // Test wrong dtype error.
794        let mut builder = ListBuilder::<u64>::with_capacity(dtype, NonNullable, 20, 10);
795        let wrong_scalar = Scalar::from(42i32);
796        assert!(builder.append_scalar(&wrong_scalar).is_err());
797    }
798
799    #[test]
800    fn test_append_array_as_list() {
801        let dtype: Arc<DType> = Arc::new(I32.into());
802        let mut ctx = array_session().create_execution_ctx();
803        let mut builder =
804            ListBuilder::<u32>::with_capacity(Arc::clone(&dtype), NonNullable, 20, 10);
805
806        // Append a primitive array as a single list entry.
807        let arr1 = buffer![1i32, 2, 3].into_array();
808        builder.append_array_as_list(&arr1, &mut ctx).unwrap();
809
810        // Interleave with a list scalar.
811        builder
812            .append_value(
813                Scalar::list(
814                    Arc::clone(&dtype),
815                    vec![10i32.into(), 11i32.into()],
816                    NonNullable,
817                )
818                .as_list(),
819            )
820            .unwrap();
821
822        // Append another primitive array as a single list entry.
823        let arr2 = buffer![4i32, 5].into_array();
824        builder.append_array_as_list(&arr2, &mut ctx).unwrap();
825
826        // Append an empty array as a single list entry (empty list).
827        let arr3 = buffer![0i32; 0].into_array();
828        builder.append_array_as_list(&arr3, &mut ctx).unwrap();
829
830        // Interleave with another list scalar (empty list).
831        builder
832            .append_value(Scalar::list_empty(Arc::clone(&dtype), NonNullable).as_list())
833            .unwrap();
834
835        let list = builder.finish_into_list();
836        assert_eq!(list.len(), 5);
837
838        // Verify elements array: [1, 2, 3, 10, 11, 4, 5].
839        assert_arrays_eq!(
840            list.elements(),
841            PrimitiveArray::from_iter([1i32, 2, 3, 10, 11, 4, 5]),
842            &mut ctx
843        );
844
845        // Verify offsets array.
846        assert_arrays_eq!(
847            list.offsets(),
848            PrimitiveArray::from_iter([0u32, 3, 5, 7, 7, 7]),
849            &mut ctx
850        );
851
852        // Test dtype mismatch error.
853        let mut builder = ListBuilder::<u32>::with_capacity(dtype, NonNullable, 20, 10);
854        let wrong_dtype_arr = buffer![1i64, 2, 3].into_array();
855        assert!(
856            builder
857                .append_array_as_list(&wrong_dtype_arr, &mut ctx)
858                .is_err()
859        );
860    }
861}