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