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