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        assert_eq!(
464            list.dtype(),
465            &DType::List(Arc::new(I32.into()), NonNullable)
466        );
467    }
468
469    #[test]
470    fn test_values() {
471        let dtype: Arc<DType> = Arc::new(I32.into());
472        let mut builder = ListBuilder::<u32>::with_capacity_in(
473            Arc::clone(&dtype),
474            NonNullable,
475            0,
476            0,
477            BufferAllocatorRef::static_ref(),
478        );
479
480        builder
481            .append_value(
482                Scalar::list(
483                    Arc::clone(&dtype),
484                    vec![1i32.into(), 2i32.into(), 3i32.into()],
485                    NonNullable,
486                )
487                .as_list(),
488            )
489            .unwrap();
490
491        builder
492            .append_value(
493                Scalar::list(
494                    dtype,
495                    vec![4i32.into(), 5i32.into(), 6i32.into()],
496                    NonNullable,
497                )
498                .as_list(),
499            )
500            .unwrap();
501
502        let list = builder.finish();
503        assert_eq!(list.len(), 2);
504
505        let mut ctx = array_session().create_execution_ctx();
506        let list_array = list.execute::<ListViewArray>(&mut ctx).unwrap();
507
508        assert_eq!(list_array.list_elements_at(0).unwrap().len(), 3);
509        assert_eq!(list_array.list_elements_at(1).unwrap().len(), 3);
510        let expected = ListArray::new(
511            buffer![1i32, 2, 3, 4, 5, 6].into_array(),
512            buffer![0u32, 3, 6].into_array(),
513            Validity::NonNullable,
514        );
515        assert_arrays_eq!(list_array, expected, &mut ctx);
516    }
517
518    #[test]
519    fn test_append_empty_list() {
520        let dtype: Arc<DType> = Arc::new(I32.into());
521        let mut builder = ListBuilder::<u32>::with_capacity_in(
522            Arc::clone(&dtype),
523            NonNullable,
524            0,
525            0,
526            BufferAllocatorRef::static_ref(),
527        );
528
529        assert!(
530            builder
531                .append_value(Scalar::list_empty(dtype, NonNullable).as_list())
532                .is_ok()
533        );
534        let expected = ListArray::new(
535            buffer![0i32; 0].into_array(),
536            buffer![0u32, 0].into_array(),
537            Validity::NonNullable,
538        );
539        assert_arrays_eq!(
540            builder.finish(),
541            expected,
542            &mut array_session().create_execution_ctx()
543        );
544    }
545
546    #[test]
547    fn test_nullable_values() {
548        let dtype: Arc<DType> = Arc::new(I32.into());
549        let mut builder = ListBuilder::<u32>::with_capacity_in(
550            Arc::clone(&dtype),
551            Nullable,
552            0,
553            0,
554            BufferAllocatorRef::static_ref(),
555        );
556
557        builder
558            .append_value(
559                Scalar::list(
560                    Arc::clone(&dtype),
561                    vec![1i32.into(), 2i32.into(), 3i32.into()],
562                    NonNullable,
563                )
564                .as_list(),
565            )
566            .unwrap();
567
568        builder
569            .append_value(Scalar::list_empty(Arc::clone(&dtype), NonNullable).as_list())
570            .unwrap();
571
572        builder
573            .append_value(
574                Scalar::list(
575                    dtype,
576                    vec![4i32.into(), 5i32.into(), 6i32.into()],
577                    NonNullable,
578                )
579                .as_list(),
580            )
581            .unwrap();
582
583        let list = builder.finish();
584        assert_eq!(list.len(), 3);
585
586        let mut ctx = array_session().create_execution_ctx();
587        let list_array = list.execute::<ListViewArray>(&mut ctx).unwrap();
588
589        assert_eq!(list_array.list_elements_at(0).unwrap().len(), 3);
590        assert_eq!(list_array.list_elements_at(1).unwrap().len(), 0);
591        assert_eq!(list_array.list_elements_at(2).unwrap().len(), 3);
592        let expected = ListArray::new(
593            buffer![1i32, 2, 3, 4, 5, 6].into_array(),
594            buffer![0u32, 3, 3, 6].into_array(),
595            Validity::AllValid,
596        );
597        assert_arrays_eq!(list_array, expected, &mut ctx);
598    }
599
600    fn test_extend_builder_gen<O: OffsetBuilderPType>() {
601        let list = ListArray::from_iter_opt_slow::<O, _, _>(
602            [Some(vec![0, 1, 2]), None, Some(vec![4, 5])],
603            Arc::new(I32.into()),
604        )
605        .unwrap()
606        .into_array();
607        assert_eq!(list.len(), 3);
608
609        let mut ctx = array_session().create_execution_ctx();
610
611        let mut builder = ListBuilder::<O>::with_capacity_in(
612            Arc::new(I32.into()),
613            Nullable,
614            18,
615            9,
616            BufferAllocatorRef::static_ref(),
617        );
618        list.append_to_builder(&mut builder, &mut ctx).unwrap();
619        list.append_to_builder(&mut builder, &mut ctx).unwrap();
620        list.slice(0..0)
621            .unwrap()
622            .append_to_builder(&mut builder, &mut ctx)
623            .unwrap();
624        list.slice(1..3)
625            .unwrap()
626            .append_to_builder(&mut builder, &mut ctx)
627            .unwrap();
628
629        let expected = ListArray::from_iter_opt_slow::<O, _, _>(
630            [
631                Some(vec![0, 1, 2]),
632                None,
633                Some(vec![4, 5]),
634                Some(vec![0, 1, 2]),
635                None,
636                Some(vec![4, 5]),
637                None,
638                Some(vec![4, 5]),
639            ],
640            Arc::new(DType::Primitive(I32, NonNullable)),
641        )
642        .unwrap()
643        .into_array()
644        .execute::<ListViewArray>(&mut ctx)
645        .unwrap();
646
647        let actual = builder.finish_into_canonical(&mut ctx).into_listview();
648
649        assert_arrays_eq!(actual.elements(), expected.elements(), &mut ctx);
650
651        assert_arrays_eq!(actual.offsets(), expected.offsets(), &mut ctx);
652
653        assert!(
654            actual
655                .validity()
656                .vortex_expect("list validity should be derivable")
657                .mask_eq(
658                    &expected
659                        .validity()
660                        .vortex_expect("list validity should be derivable"),
661                    actual.len(),
662                    &mut ctx,
663                )
664                .unwrap(),
665        );
666        let independent = ListArray::new(
667            buffer![0i32, 1, 2, 4, 5, 0, 1, 2, 4, 5, 4, 5].into_array(),
668            buffer![0u32, 3, 3, 5, 8, 8, 10, 10, 12].into_array(),
669            Validity::from_iter([true, false, true, true, false, true, false, true]),
670        );
671        assert_arrays_eq!(actual, independent, &mut ctx);
672    }
673
674    /// `append_to_builder` must handle any list builder kind without assuming the offset/size
675    /// integer types produced by `builder_with_capacity`. It appends a `List`-encoded array and a
676    /// `ListView`-encoded array into `ListViewBuilder`s and `ListBuilder`s with assorted (and
677    /// non-`u64`) offset/size types.
678    #[test]
679    fn test_append_to_builder_any_list_builder() -> VortexResult<()> {
680        let mut ctx = array_session().create_execution_ctx();
681
682        let list = ListArray::from_iter_opt_slow::<u64, _, _>(
683            [Some(vec![0, 1, 2]), None, Some(vec![4, 5])],
684            Arc::new(I32.into()),
685        )?
686        .into_array();
687        let listview = list
688            .clone()
689            .execute::<ListViewArray>(&mut ctx)?
690            .into_array();
691        let elem_dtype = || Arc::new(I32.into());
692
693        // `builder_with_capacity` produces a `ListViewBuilder` for `DType::List`; appending the
694        // `List`-encoded array must dispatch into it instead of bailing.
695        let mut listview_builder =
696            builder_with_capacity_in(list.dtype(), list.len(), BufferAllocatorRef::static_ref());
697        list.append_to_builder(listview_builder.as_mut(), &mut ctx)?;
698        assert_arrays_eq!(listview_builder.finish(), list, &mut ctx);
699
700        // A `ListViewBuilder` with non-`u64` (including signed) offset and size types must work
701        // for both source encodings.
702        let mut lv_u64_u32 = ListViewBuilder::<u64, u32>::with_capacity_in(
703            elem_dtype(),
704            Nullable,
705            8,
706            4,
707            BufferAllocatorRef::static_ref(),
708        );
709        list.append_to_builder(&mut lv_u64_u32, &mut ctx)?;
710        assert_arrays_eq!(lv_u64_u32.finish(), list, &mut ctx);
711
712        let mut lv_i64_i32 = ListViewBuilder::<i64, i32>::with_capacity_in(
713            elem_dtype(),
714            Nullable,
715            8,
716            4,
717            BufferAllocatorRef::static_ref(),
718        );
719        list.append_to_builder(&mut lv_i64_i32, &mut ctx)?;
720        assert_arrays_eq!(lv_i64_i32.finish(), list, &mut ctx);
721
722        let mut lv_u32_u32 = ListViewBuilder::<u32, u32>::with_capacity_in(
723            elem_dtype(),
724            Nullable,
725            8,
726            4,
727            BufferAllocatorRef::static_ref(),
728        );
729        listview.append_to_builder(&mut lv_u32_u32, &mut ctx)?;
730        assert_arrays_eq!(lv_u32_u32.finish(), list, &mut ctx);
731
732        // Both source encodings appended into `ListBuilder`s with non-`u64` (including signed)
733        // offset types.
734        let mut list_builder = ListBuilder::<u32>::with_capacity_in(
735            elem_dtype(),
736            Nullable,
737            8,
738            4,
739            BufferAllocatorRef::static_ref(),
740        );
741        list.append_to_builder(&mut list_builder, &mut ctx)?;
742        assert_arrays_eq!(list_builder.finish(), list, &mut ctx);
743
744        let mut list_builder_i32 = ListBuilder::<i32>::with_capacity_in(
745            elem_dtype(),
746            Nullable,
747            8,
748            4,
749            BufferAllocatorRef::static_ref(),
750        );
751        listview.append_to_builder(&mut list_builder_i32, &mut ctx)?;
752        assert_arrays_eq!(list_builder_i32.finish(), list, &mut ctx);
753
754        Ok(())
755    }
756
757    #[test]
758    fn test_append_list_arrays_grow_builder() -> VortexResult<()> {
759        let mut ctx = array_session().create_execution_ctx();
760        let dtype: Arc<DType> = Arc::new(I32.into());
761
762        // Enough lists to exceed the offsets capacity of a zero-capacity builder, so appending
763        // must grow the builder rather than panic in `uninit_range`.
764        let lists: Vec<Option<Vec<i32>>> =
765            (0..100).map(|i| (i % 10 != 0).then(|| vec![i])).collect();
766        let source = ListArray::from_iter_opt_slow::<u32, _, _>(lists.clone(), Arc::clone(&dtype))?;
767        let expected = ListArray::from_iter_opt_slow::<u32, _, _>(
768            lists.iter().cloned().chain(lists.iter().cloned()),
769            Arc::clone(&dtype),
770        )?;
771
772        // Appending twice checks growth from a non-empty builder and offset rebasing.
773        let mut builder = ListBuilder::<u32>::with_capacity_in(
774            Arc::clone(&dtype),
775            Nullable,
776            0,
777            0,
778            BufferAllocatorRef::static_ref(),
779        );
780        builder.append_list_array(source.as_view(), &mut ctx)?;
781        builder.append_list_array(source.as_view(), &mut ctx)?;
782        assert_arrays_eq!(builder.finish(), expected, &mut ctx);
783
784        let source_listview = source.into_array().execute::<ListViewArray>(&mut ctx)?;
785        let mut builder = ListBuilder::<u32>::with_capacity_in(
786            dtype,
787            Nullable,
788            0,
789            0,
790            BufferAllocatorRef::static_ref(),
791        );
792        builder.append_listview_array(source_listview.as_view(), &mut ctx)?;
793        builder.append_listview_array(source_listview.as_view(), &mut ctx)?;
794        assert_arrays_eq!(builder.finish(), expected, &mut ctx);
795
796        Ok(())
797    }
798
799    /// A `ListArray`'s offsets can only describe contiguous, in-order lists, so an overlapping
800    /// source has to be flattened before its elements can be appended in bulk. A sliced source,
801    /// meanwhile, keeps the layout it has and is appended from wherever its first list starts.
802    #[test]
803    fn test_append_listview_array_flattens_overlaps_and_skips_leading_elements() -> VortexResult<()>
804    {
805        let mut ctx = array_session().create_execution_ctx();
806        let dtype: Arc<DType> = Arc::new(I32.into());
807
808        // Overlapping source, so not zero-copyable to a list:
809        // - List 0: [10, 20]
810        // - List 1: null (size is intentionally non-zero in the source metadata)
811        // - List 2: [10], sharing the elements list 0 already referenced
812        let overlapping = unsafe {
813            ListViewArray::new_unchecked(
814                buffer![10i32, 20, 30].into_array(),
815                buffer![0u32, 1, 0].into_array(),
816                buffer![2u8, 2, 1].into_array(),
817                Validity::from_iter([true, false, true]),
818            )
819        };
820        assert!(!overlapping.is_zero_copy_to_list());
821
822        // Zero-copyable source sliced past its first list, so its elements start at offset 2.
823        let sliced = unsafe {
824            ListViewArray::new_unchecked(
825                buffer![40i32, 50, 60, 70].into_array(),
826                buffer![0u32, 2].into_array(),
827                buffer![2u32, 2].into_array(),
828                Validity::AllValid,
829            )
830            .with_zero_copy_to_list(true)
831        }
832        .into_array()
833        .slice(1..2)?
834        .execute::<ListViewArray>(&mut ctx)?;
835
836        let mut builder = ListBuilder::<u32>::with_capacity_in(
837            dtype,
838            Nullable,
839            0,
840            0,
841            BufferAllocatorRef::static_ref(),
842        );
843        builder.append_listview_array(overlapping.as_view(), &mut ctx)?;
844        builder.append_listview_array(sliced.as_view(), &mut ctx)?;
845
846        let list = builder.finish_into_list();
847        assert_arrays_eq!(
848            list.elements(),
849            PrimitiveArray::from_iter([10i32, 20, 10, 60, 70]),
850            &mut ctx
851        );
852        assert_arrays_eq!(
853            list.offsets(),
854            PrimitiveArray::from_iter([0u32, 2, 2, 3, 5]),
855            &mut ctx
856        );
857
858        let expected = ListArray::new(
859            buffer![10i32, 20, 10, 60, 70].into_array(),
860            buffer![0u32, 2, 2, 3, 5].into_array(),
861            Validity::from_iter([true, false, true, true]),
862        );
863        assert_arrays_eq!(list, expected, &mut ctx);
864        Ok(())
865    }
866
867    #[test]
868    fn test_extend_builder() {
869        test_extend_builder_gen::<i32>();
870        test_extend_builder_gen::<i64>();
871
872        test_extend_builder_gen::<u32>();
873        test_extend_builder_gen::<u64>();
874    }
875
876    #[test]
877    pub fn test_array_with_gap() {
878        let one_trailing_unused_element = ListArray::try_new(
879            buffer![1, 2, 3, 4].into_array(),
880            buffer![0, 3].into_array(),
881            Validity::NonNullable,
882        )
883        .unwrap();
884
885        let second_array = ListArray::try_new(
886            buffer![5, 6].into_array(),
887            buffer![0, 2].into_array(),
888            Validity::NonNullable,
889        )
890        .unwrap();
891
892        let chunked_list = ChunkedArray::try_new(
893            vec![
894                one_trailing_unused_element.clone().into_array(),
895                second_array.clone().into_array(),
896            ],
897            DType::List(Arc::new(DType::Primitive(I32, NonNullable)), NonNullable),
898        );
899
900        let mut ctx = array_session().create_execution_ctx();
901        let canon_values = chunked_list
902            .unwrap()
903            .as_array()
904            .clone()
905            .execute::<ListViewArray>(&mut ctx)
906            .unwrap();
907
908        assert_eq!(
909            one_trailing_unused_element
910                .execute_scalar(0, &mut array_session().create_execution_ctx())
911                .unwrap(),
912            canon_values
913                .execute_scalar(0, &mut array_session().create_execution_ctx())
914                .unwrap()
915        );
916        assert_eq!(
917            second_array
918                .execute_scalar(0, &mut array_session().create_execution_ctx())
919                .unwrap(),
920            canon_values
921                .execute_scalar(1, &mut array_session().create_execution_ctx())
922                .unwrap()
923        );
924    }
925
926    #[test]
927    fn test_append_scalar() {
928        let dtype: Arc<DType> = Arc::new(I32.into());
929        let mut builder = ListBuilder::<u64>::with_capacity_in(
930            Arc::clone(&dtype),
931            Nullable,
932            20,
933            10,
934            BufferAllocatorRef::static_ref(),
935        );
936
937        // Test appending a valid list.
938        let list_scalar1 =
939            Scalar::list(Arc::clone(&dtype), vec![1i32.into(), 2i32.into()], Nullable);
940        builder.append_scalar(&list_scalar1).unwrap();
941
942        // Test appending another list.
943        let list_scalar2 = Scalar::list(
944            Arc::clone(&dtype),
945            vec![3i32.into(), 4i32.into(), 5i32.into()],
946            Nullable,
947        );
948        builder.append_scalar(&list_scalar2).unwrap();
949
950        // Test appending null value.
951        let null_scalar = Scalar::null(DType::List(Arc::clone(&dtype), Nullable));
952        builder.append_scalar(&null_scalar).unwrap();
953
954        let array = builder.finish_into_list();
955        assert_eq!(array.len(), 3);
956
957        let mut ctx = array_session().create_execution_ctx();
958
959        // Check actual values using scalar_at.
960
961        let scalar0 = array.execute_scalar(0, &mut ctx).unwrap();
962        let list0 = scalar0.as_list();
963        assert_eq!(list0.len(), 2);
964        if let Some(list0_items) = list0.elements() {
965            assert_eq!(list0_items[0].as_primitive().typed_value::<i32>(), Some(1));
966            assert_eq!(list0_items[1].as_primitive().typed_value::<i32>(), Some(2));
967        }
968
969        let scalar1 = array.execute_scalar(1, &mut ctx).unwrap();
970        let list1 = scalar1.as_list();
971        assert_eq!(list1.len(), 3);
972        if let Some(list1_items) = list1.elements() {
973            assert_eq!(list1_items[0].as_primitive().typed_value::<i32>(), Some(3));
974            assert_eq!(list1_items[1].as_primitive().typed_value::<i32>(), Some(4));
975            assert_eq!(list1_items[2].as_primitive().typed_value::<i32>(), Some(5));
976        }
977
978        let scalar2 = array.execute_scalar(2, &mut ctx).unwrap();
979        let list2 = scalar2.as_list();
980        assert!(list2.is_null()); // This should be null.
981
982        // Check validity.
983        assert!(
984            array
985                .validity()
986                .vortex_expect("list validity should be derivable")
987                .execute_is_valid(0, &mut ctx)
988                .unwrap()
989        );
990        assert!(
991            array
992                .validity()
993                .vortex_expect("list validity should be derivable")
994                .execute_is_valid(1, &mut ctx)
995                .unwrap()
996        );
997        assert!(
998            !array
999                .validity()
1000                .vortex_expect("list validity should be derivable")
1001                .execute_is_valid(2, &mut ctx)
1002                .unwrap()
1003        );
1004
1005        // Test wrong dtype error.
1006        let mut builder = ListBuilder::<u64>::with_capacity_in(
1007            dtype,
1008            NonNullable,
1009            20,
1010            10,
1011            BufferAllocatorRef::static_ref(),
1012        );
1013        let wrong_scalar = Scalar::from(42i32);
1014        assert!(builder.append_scalar(&wrong_scalar).is_err());
1015    }
1016
1017    #[test]
1018    fn test_append_array_as_list() {
1019        let dtype: Arc<DType> = Arc::new(I32.into());
1020        let mut ctx = array_session().create_execution_ctx();
1021        let mut builder = ListBuilder::<u32>::with_capacity_in(
1022            Arc::clone(&dtype),
1023            NonNullable,
1024            20,
1025            10,
1026            BufferAllocatorRef::static_ref(),
1027        );
1028
1029        // Append a primitive array as a single list entry.
1030        let arr1 = buffer![1i32, 2, 3].into_array();
1031        builder.append_array_as_list(&arr1, &mut ctx).unwrap();
1032
1033        // Interleave with a list scalar.
1034        builder
1035            .append_value(
1036                Scalar::list(
1037                    Arc::clone(&dtype),
1038                    vec![10i32.into(), 11i32.into()],
1039                    NonNullable,
1040                )
1041                .as_list(),
1042            )
1043            .unwrap();
1044
1045        // Append another primitive array as a single list entry.
1046        let arr2 = buffer![4i32, 5].into_array();
1047        builder.append_array_as_list(&arr2, &mut ctx).unwrap();
1048
1049        // Append an empty array as a single list entry (empty list).
1050        let arr3 = buffer![0i32; 0].into_array();
1051        builder.append_array_as_list(&arr3, &mut ctx).unwrap();
1052
1053        // Interleave with another list scalar (empty list).
1054        builder
1055            .append_value(Scalar::list_empty(Arc::clone(&dtype), NonNullable).as_list())
1056            .unwrap();
1057
1058        let list = builder.finish_into_list();
1059        assert_eq!(list.len(), 5);
1060
1061        // Verify elements array: [1, 2, 3, 10, 11, 4, 5].
1062        assert_arrays_eq!(
1063            list.elements(),
1064            PrimitiveArray::from_iter([1i32, 2, 3, 10, 11, 4, 5]),
1065            &mut ctx
1066        );
1067
1068        // Verify offsets array.
1069        assert_arrays_eq!(
1070            list.offsets(),
1071            PrimitiveArray::from_iter([0u32, 3, 5, 7, 7, 7]),
1072            &mut ctx
1073        );
1074
1075        // Test dtype mismatch error.
1076        let mut builder = ListBuilder::<u32>::with_capacity_in(
1077            dtype,
1078            NonNullable,
1079            20,
1080            10,
1081            BufferAllocatorRef::static_ref(),
1082        );
1083        let wrong_dtype_arr = buffer![1i64, 2, 3].into_array();
1084        assert!(
1085            builder
1086                .append_array_as_list(&wrong_dtype_arr, &mut ctx)
1087                .is_err()
1088        );
1089    }
1090}