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