Skip to main content

vortex_array/builders/
listview.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! ListView Builder Implementation.
5//!
6//! A builder for [`ListViewArray`] that tracks both offsets and sizes.
7//!
8//! Unlike [`ListArray`] which only tracks offsets, [`ListViewArray`] stores both offsets and sizes
9//! in separate arrays for better compression.
10//!
11//! [`ListArray`]: crate::arrays::ListArray
12
13use std::sync::Arc;
14
15use num_traits::ToPrimitive;
16use vortex_buffer::BufferAllocatorRef;
17use vortex_error::VortexExpect;
18use vortex_error::VortexResult;
19use vortex_error::vortex_ensure;
20use vortex_error::vortex_ensure_eq;
21use vortex_error::vortex_panic;
22
23use crate::ArrayRef;
24use crate::Canonical;
25use crate::ExecutionCtx;
26use crate::array::ArrayView;
27use crate::array::IntoArray;
28use crate::arrays::List;
29use crate::arrays::ListView;
30use crate::arrays::ListViewArray;
31use crate::arrays::PrimitiveArray;
32use crate::arrays::list::ListArraySlotsExt;
33use crate::arrays::listview::ListViewArrayExt;
34use crate::arrays::listview::ListViewArraySlotsExt;
35use crate::builders::ArrayBuilder;
36use crate::builders::ChildBuilder;
37use crate::builders::DEFAULT_BUILDER_CAPACITY;
38use crate::builders::PrimitiveBuilder;
39use crate::builders::UninitRange;
40use crate::builders::ValidityBuilder;
41use crate::dtype::DType;
42use crate::dtype::IntegerPType;
43use crate::dtype::Nullability;
44use crate::dtype::OffsetBuilderPType;
45use crate::match_each_integer_ptype;
46use crate::scalar::ListScalar;
47use crate::scalar::Scalar;
48
49/// A builder for creating [`ListViewArray`] instances, parameterized by the [`OffsetBuilderPType`]
50/// of the `offsets` and the `sizes` builders.
51///
52/// This builder tracks both offsets and sizes using potentially different integer types for memory
53/// efficiency. For example, you might use `u64` for offsets but only `u8` for sizes if your lists
54/// are small.
55///
56/// Any combination of [`OffsetBuilderPType`] types is valid, as long as the type of `sizes` can fit
57/// into the type of `offsets`.
58pub struct ListViewBuilder<O: OffsetBuilderPType, S: OffsetBuilderPType> {
59    /// The [`DType`] of the [`ListViewArray`]. This **must** be a [`DType::List`].
60    dtype: DType,
61
62    /// The builder for the underlying elements of the [`ListArray`](crate::arrays::ListArray).
63    elements_builder: ChildBuilder,
64
65    /// The builder for the `offsets` into the `elements` array.
66    offsets_builder: PrimitiveBuilder<O>,
67
68    /// The builder for the `sizes` of each list view.
69    sizes_builder: PrimitiveBuilder<S>,
70
71    /// The null map builder of the [`ListViewArray`].
72    nulls: ValidityBuilder,
73
74    /// Whether the appends so far leave the result zero-copyable to a [`ListArray`].
75    ///
76    /// Only [`append_listview_array`](Self::append_listview_array) can clear this; every other
77    /// append writes its lists back to back.
78    ///
79    /// [`ListArray`]: crate::arrays::ListArray
80    zero_copy_to_list: bool,
81}
82
83impl<O: OffsetBuilderPType, S: OffsetBuilderPType> ListViewBuilder<O, S> {
84    /// Creates a new `ListViewBuilder` with a capacity of [`DEFAULT_BUILDER_CAPACITY`].
85    #[deprecated(note = "use `new_in` with an explicit allocator")]
86    pub fn new(element_dtype: Arc<DType>, nullability: Nullability) -> Self {
87        Self::new_in(element_dtype, nullability, BufferAllocatorRef::static_ref())
88    }
89
90    /// Creates a new `ListViewBuilder` with the default capacity using `allocator`.
91    pub fn new_in(
92        element_dtype: Arc<DType>,
93        nullability: Nullability,
94        allocator: &BufferAllocatorRef,
95    ) -> Self {
96        Self::with_capacity_in(
97            element_dtype,
98            nullability,
99            // We arbitrarily choose 2 times the number of list scalars for the capacity of the
100            // elements builder since we cannot know this ahead of time.
101            DEFAULT_BUILDER_CAPACITY * 2,
102            DEFAULT_BUILDER_CAPACITY,
103            allocator,
104        )
105    }
106
107    /// Create a new [`ListViewArray`] builder with a with the given `capacity`, as well as an
108    /// initial capacity for the `elements` builder (since we cannot know that ahead of time solely
109    /// based on the outer array `capacity`).
110    ///
111    /// # Panics
112    ///
113    /// Panics if the size type `S` cannot fit within the offset type `O`.
114    #[deprecated(note = "use `with_capacity_in` with an explicit allocator")]
115    pub fn with_capacity(
116        element_dtype: Arc<DType>,
117        nullability: Nullability,
118        elements_capacity: usize,
119        capacity: usize,
120    ) -> Self {
121        Self::with_capacity_in(
122            element_dtype,
123            nullability,
124            elements_capacity,
125            capacity,
126            BufferAllocatorRef::static_ref(),
127        )
128    }
129
130    /// Creates a list-view builder with the given capacities using `allocator`.
131    ///
132    /// # Panics
133    ///
134    /// Panics if the size type `S` cannot fit within the offset type `O`.
135    pub fn with_capacity_in(
136        element_dtype: Arc<DType>,
137        nullability: Nullability,
138        elements_capacity: usize,
139        capacity: usize,
140        allocator: &BufferAllocatorRef,
141    ) -> Self {
142        let elements_builder =
143            ChildBuilder::with_capacity(&element_dtype, elements_capacity, allocator);
144
145        let offsets_builder =
146            PrimitiveBuilder::<O>::with_capacity_in(Nullability::NonNullable, capacity, allocator);
147        let sizes_builder =
148            PrimitiveBuilder::<S>::with_capacity_in(Nullability::NonNullable, capacity, allocator);
149
150        let nulls = ValidityBuilder::new(capacity, allocator);
151
152        Self {
153            dtype: DType::List(element_dtype, nullability),
154            elements_builder,
155            offsets_builder,
156            sizes_builder,
157            nulls,
158            zero_copy_to_list: true,
159        }
160    }
161
162    /// Appends an array as a single non-null list entry to the builder.
163    ///
164    /// The input `array` must have the same dtype as the element dtype of this list builder.
165    ///
166    /// Note that the list entry will be non-null but the elements themselves are allowed to be null
167    /// (only if the elements [`DType`] is nullable, of course).
168    pub fn append_array_as_list(
169        &mut self,
170        array: &ArrayRef,
171        ctx: &mut ExecutionCtx,
172    ) -> VortexResult<()> {
173        vortex_ensure_eq!(array.dtype(), self.element_dtype());
174
175        let curr_offset = self.elements_builder.len();
176        let num_elements = array.len();
177
178        // We must assert this even in release mode to ensure that the safety comment in
179        // `finish_into_listview` is correct.
180        assert!(
181            ((curr_offset + num_elements) as u64) < O::max_value_as_u64(),
182            "appending this list would cause an offset overflow"
183        );
184
185        self.elements_builder.append_array(array, ctx)?;
186        self.nulls.append_non_null();
187
188        self.offsets_builder.append_value(
189            O::from_usize(curr_offset).vortex_expect("Failed to convert from usize to `O`"),
190        );
191        self.sizes_builder.append_value(
192            S::from_usize(num_elements).vortex_expect("Failed to convert from usize to `S`"),
193        );
194
195        Ok(())
196    }
197
198    /// Append a list of values to the builder.
199    ///
200    /// This method extends the value builder with the provided values and records
201    /// the offset and size of the new list.
202    pub fn append_value(&mut self, value: ListScalar) -> VortexResult<()> {
203        let Some(elements) = value.elements() else {
204            // If `elements` is `None`, then the `value` is a null value.
205            vortex_ensure!(
206                self.dtype.is_nullable(),
207                "Cannot append null value to non-nullable list builder"
208            );
209            self.append_null();
210            return Ok(());
211        };
212
213        let curr_offset = self.elements_builder.len();
214        let num_elements = elements.len();
215
216        // We must assert this even in release mode to ensure that the safety comment in
217        // `finish_into_listview` is correct.
218        assert!(
219            ((curr_offset + num_elements) as u64) < O::max_value_as_u64(),
220            "appending this list would cause an offset overflow"
221        );
222
223        for scalar in elements {
224            self.elements_builder.append_scalar(&scalar)?;
225        }
226        self.nulls.append_non_null();
227
228        self.offsets_builder.append_value(
229            O::from_usize(curr_offset).vortex_expect("Failed to convert from usize to `O`"),
230        );
231        self.sizes_builder.append_value(
232            S::from_usize(num_elements).vortex_expect("Failed to convert from usize to `S`"),
233        );
234
235        Ok(())
236    }
237
238    /// Appends `array` as `n` identical non-null lists, storing its elements once.
239    ///
240    /// A `ListViewArray` can point many views at one range of elements, so a repeated list costs
241    /// its elements once however many rows it covers. The elements go in as one appended array, so
242    /// a caller that hands over the same `array` on every call - a sparse array filling the gaps
243    /// between its patches, say - stores those elements once for the whole result.
244    ///
245    /// The views share their elements, so the result is no longer zero-copyable to a
246    /// [`ListArray`](crate::arrays::ListArray) unless it covers a single row or empty lists.
247    pub fn append_array_as_repeated_list(
248        &mut self,
249        array: &ArrayRef,
250        n: usize,
251        ctx: &mut ExecutionCtx,
252    ) -> VortexResult<()> {
253        vortex_ensure_eq!(array.dtype(), self.element_dtype());
254
255        if n == 0 {
256            return Ok(());
257        }
258
259        let curr_offset = self.elements_builder.len();
260        let num_elements = array.len();
261
262        // We must assert this even in release mode to ensure that the safety comment in
263        // `finish_into_listview` is correct.
264        assert!(
265            ((curr_offset + num_elements) as u64) < O::max_value_as_u64(),
266            "appending this list would cause an offset overflow"
267        );
268
269        self.elements_builder.append_array(array, ctx)?;
270
271        let offset =
272            O::from_usize(curr_offset).vortex_expect("Failed to convert from usize to `O`");
273        let size = S::from_usize(num_elements).vortex_expect("Failed to convert from usize to `S`");
274        self.offsets_builder.append_n_values(offset, n);
275        self.sizes_builder.append_n_values(size, n);
276        self.nulls.append_n_non_nulls(n);
277
278        if n > 1 && num_elements > 0 {
279            self.zero_copy_to_list = false;
280        }
281
282        Ok(())
283    }
284
285    /// Finishes the builder directly into a [`ListViewArray`].
286    pub fn finish_into_listview(&mut self) -> ListViewArray {
287        debug_assert_eq!(self.offsets_builder.len(), self.sizes_builder.len());
288        debug_assert_eq!(self.offsets_builder.len(), self.nulls.len());
289
290        let elements = self.elements_builder.finish();
291        let offsets = self.offsets_builder.finish();
292        let sizes = self.sizes_builder.finish();
293        let validity = self.nulls.finish_with_nullability(self.dtype.nullability());
294
295        let zero_copy_to_list = std::mem::replace(&mut self.zero_copy_to_list, true);
296
297        // SAFETY:
298        // - Both the offsets and the sizes are non-nullable.
299        // - The offsets, sizes, and validity have the same length since we always appended the same
300        //   amount.
301        // - We checked on construction that the sizes type fits into the offsets.
302        // - In every method that adds values to this builder (`append_value`, `append_scalar`,
303        //   `append_list_array`, and `append_listview_array`), we checked that `offset + size`
304        //   does not overflow. `append_listview_array` rebases the offsets it was handed onto
305        //   exactly the elements it appended, so the source's bound carries over.
306        // - Every append writes its lists back to back, so the result is zero-copyable to a
307        //   `ListArray` unless `zero_copy_to_list` recorded an appended layout we left alone.
308        unsafe {
309            ListViewArray::new_unchecked(elements, offsets, sizes, validity)
310                .with_zero_copy_to_list(zero_copy_to_list)
311        }
312    }
313
314    /// The [`DType`] of the inner elements. Note that this is **not** the same as the [`DType`] of
315    /// the outer `FixedSizeList`.
316    pub fn element_dtype(&self) -> &DType {
317        let DType::List(element_dtype, ..) = &self.dtype else {
318            vortex_panic!("`ListViewBuilder` has an incorrect dtype: {}", self.dtype);
319        };
320
321        element_dtype
322    }
323
324    /// Appends the values of a [`List`]-encoded `array` to this builder.
325    ///
326    /// List encodings dispatch here through
327    /// [`match_each_list_builder!`](crate::match_each_list_builder) because the concrete list
328    /// builders are generic over their offset/size integer types, which cannot be named through a
329    /// `dyn ArrayBuilder`.
330    pub fn append_list_array(
331        &mut self,
332        array: ArrayView<'_, List>,
333        ctx: &mut ExecutionCtx,
334    ) -> VortexResult<()> {
335        if array.is_empty() {
336            return Ok(());
337        }
338
339        self.nulls.append_validity(array.validity()?, array.len());
340
341        let offsets = array.offsets().clone().execute::<PrimitiveArray>(ctx)?;
342        match_each_integer_ptype!(offsets.ptype(), |OffsetType| {
343            extend_from_list(
344                self,
345                array.elements(),
346                offsets.as_slice::<OffsetType>(),
347                ctx,
348            )?
349        });
350        Ok(())
351    }
352
353    /// Appends the values of a [`ListView`]-encoded `array` to this builder.
354    ///
355    /// See [`append_list_array`](Self::append_list_array); this is the same hook for the canonical
356    /// [`ListViewArray`] encoding.
357    ///
358    /// The views keep the layout they arrived in, so overlapping sources keep sharing their
359    /// elements and the finished array reports [`is_zero_copy_to_list`] as `false`. Callers that
360    /// need an exact layout should [`rebuild`](ListViewArray::rebuild) it.
361    ///
362    /// [`is_zero_copy_to_list`]: crate::arrays::listview::ListViewData::is_zero_copy_to_list
363    pub fn append_listview_array(
364        &mut self,
365        array: ArrayView<'_, ListView>,
366        ctx: &mut ExecutionCtx,
367    ) -> VortexResult<()> {
368        if array.is_empty() {
369            return Ok(());
370        }
371
372        let len = array.len();
373
374        // Materialize the metadata once and do the trimming and the rebase by hand. Going through
375        // `rebuild(ListViewRebuildMode::TrimElements)` would subtract the window start from every
376        // offset with a compute kernel, only for the rebase below to add this builder's elements
377        // base straight back on - two passes, one of them through the compute stack, for one
378        // addition per offset. Casting the sizes to the builder's type is another kernel for what
379        // is a copy.
380        let offsets = array.offsets().clone().execute::<PrimitiveArray>(ctx)?;
381        let sizes = array.sizes().clone().execute::<PrimitiveArray>(ctx)?;
382
383        // The window of `elements` that the views actually reference. Everything outside it is
384        // unreachable and must not be appended. An exact source covers its window back to back and
385        // in order, so the first and last view bound it; any other layout has to be searched.
386        let (start, end) = if array.is_zero_copy_to_list() {
387            let last = len - 1;
388            (
389                metadata_at(&offsets, 0),
390                metadata_at(&offsets, last) + metadata_at(&sizes, last),
391            )
392        } else {
393            array.into_owned().referenced_element_bounds(ctx)?
394        };
395
396        // An exact source references every element it carries, back to back, so it lands flush
397        // against the elements already in the builder. Any other layout does not.
398        self.zero_copy_to_list &= array.is_zero_copy_to_list();
399
400        self.nulls.append_validity(array.validity()?, len);
401
402        // Bulk append the referenced elements; the offsets are rebased onto them below.
403        let elements_base = self.elements_builder.len();
404
405        // We must assert this even in release mode to ensure that the safety comment in
406        // `finish_into_listview` is correct.
407        assert!(
408            ((elements_base + (end - start)) as u64) < O::max_value_as_u64(),
409            "appending this list would cause an offset overflow"
410        );
411
412        if end > start {
413            self.elements_builder
414                .append_array(&array.elements().slice(start..end)?, ctx)?;
415        }
416
417        // Every view lies inside `start..end`, so rebasing it onto `elements_base` keeps it inside
418        // the elements this builder now holds - which is what lets `finish_into_listview` build the
419        // array unchecked.
420        assert_eq!(
421            self.elements_builder.len(),
422            elements_base + (end - start),
423            "appending the referenced elements did not extend the child by the window's length"
424        );
425
426        self.offsets_builder.reserve_exact(len);
427        let offsets_range = self.offsets_builder.uninit_range(len);
428        match_each_integer_ptype!(offsets.ptype(), |A| {
429            extend_rebased_offsets::<O, A>(
430                offsets_range,
431                offsets.as_slice::<A>(),
432                start,
433                elements_base,
434            );
435        });
436
437        self.sizes_builder.reserve_exact(len);
438        let mut sizes_range = self.sizes_builder.uninit_range(len);
439        if sizes.ptype() == S::PTYPE {
440            // The sizes already have the builder's type, so there is nothing to convert.
441            sizes_range.copy_from_slice(0, sizes.as_slice::<S>());
442            // SAFETY: `copy_from_slice` initialized all `len` values, and the sizes builder is
443            // non-nullable.
444            unsafe { sizes_range.finish() };
445        } else {
446            match_each_integer_ptype!(sizes.ptype(), |A| {
447                extend_converted_sizes::<S, A>(sizes_range, sizes.as_slice::<A>());
448            });
449        }
450
451        Ok(())
452    }
453}
454
455impl<O: OffsetBuilderPType, S: OffsetBuilderPType> ArrayBuilder for ListViewBuilder<O, S> {
456    fn as_any(&self) -> &dyn std::any::Any {
457        self
458    }
459
460    fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
461        self
462    }
463
464    fn dtype(&self) -> &DType {
465        &self.dtype
466    }
467
468    fn len(&self) -> usize {
469        self.offsets_builder.len()
470    }
471
472    fn append_zeros(&mut self, n: usize) {
473        debug_assert_eq!(self.offsets_builder.len(), self.sizes_builder.len());
474        debug_assert_eq!(self.offsets_builder.len(), self.nulls.len());
475
476        // Get the current position in the elements array.
477        let curr_offset = self.elements_builder.len();
478
479        // Since we consider the "zero" element of a list an empty list, we simply update the
480        // `offsets` and `sizes` metadata to add an empty list.
481        for _ in 0..n {
482            self.offsets_builder.append_value(
483                O::from_usize(curr_offset).vortex_expect("Failed to convert from usize to `O`"),
484            );
485            self.sizes_builder.append_value(S::zero());
486        }
487
488        self.nulls.append_n_non_nulls(n);
489    }
490
491    unsafe fn append_nulls_unchecked(&mut self, n: usize) {
492        debug_assert_eq!(self.offsets_builder.len(), self.sizes_builder.len());
493        debug_assert_eq!(self.offsets_builder.len(), self.nulls.len());
494
495        // Get the current position in the elements array.
496        let curr_offset = self.elements_builder.len();
497
498        // A null list can have any representation, but we choose to use the zero representation.
499        for _ in 0..n {
500            self.offsets_builder.append_value(
501                O::from_usize(curr_offset).vortex_expect("Failed to convert from usize to `O`"),
502            );
503            self.sizes_builder.append_value(S::zero());
504        }
505
506        // This is the only difference from `append_zeros`.
507        self.nulls.append_n_nulls(n);
508    }
509
510    fn append_scalar(&mut self, scalar: &Scalar) -> VortexResult<()> {
511        vortex_ensure_eq!(
512            scalar.dtype(),
513            self.dtype(),
514            "ListViewBuilder received a scalar with the wrong dtype"
515        );
516
517        let list_scalar = scalar.as_list();
518        self.append_value(list_scalar)
519    }
520
521    fn reserve_exact(&mut self, capacity: usize) {
522        self.elements_builder.reserve_exact(capacity * 2);
523        self.offsets_builder.reserve_exact(capacity);
524        self.sizes_builder.reserve_exact(capacity);
525        self.nulls.reserve_exact(capacity);
526    }
527
528    fn finish(&mut self) -> ArrayRef {
529        self.finish_into_listview().into_array()
530    }
531
532    fn finish_into_canonical(&mut self, _ctx: &mut ExecutionCtx) -> Canonical {
533        Canonical::List(self.finish_into_listview())
534    }
535}
536
537/// Appends `ListArray`-layout lists (`n + 1` cumulative offsets) into a [`ListViewBuilder`].
538///
539/// Lists in a `ListArray` are contiguous, so the referenced elements can be appended in bulk,
540/// with the offsets rebased onto the builder's elements and the sizes taken from consecutive
541/// offset differences.
542fn extend_from_list<O, S, OffsetType>(
543    builder: &mut ListViewBuilder<O, S>,
544    elements: &ArrayRef,
545    offsets: &[OffsetType],
546    ctx: &mut ExecutionCtx,
547) -> VortexResult<()>
548where
549    O: OffsetBuilderPType,
550    S: OffsetBuilderPType,
551    OffsetType: IntegerPType,
552{
553    let num_lists = offsets.len() - 1;
554    let first: usize = offsets[0].as_();
555    let last: usize = offsets[num_lists].as_();
556
557    let elements_base = builder.elements_builder.len();
558
559    // We must assert this even in release mode to ensure that the safety comment in
560    // `finish_into_listview` is correct.
561    assert!(
562        ((elements_base + (last - first)) as u64) < O::max_value_as_u64(),
563        "appending this list would cause an offset overflow"
564    );
565
566    if last > first {
567        builder
568            .elements_builder
569            .append_array(&elements.slice(first..last)?, ctx)?;
570    }
571
572    builder.offsets_builder.reserve_exact(num_lists);
573    builder.sizes_builder.reserve_exact(num_lists);
574    let mut offsets_range = builder.offsets_builder.uninit_range(num_lists);
575    let mut sizes_range = builder.sizes_builder.uninit_range(num_lists);
576    for i in 0..num_lists {
577        let start: usize = offsets[i].as_();
578        let end: usize = offsets[i + 1].as_();
579        offsets_range.set_value(
580            i,
581            O::from_usize(start - first + elements_base)
582                .vortex_expect("Failed to convert from usize to `O`"),
583        );
584        sizes_range.set_value(
585            i,
586            S::from_usize(end - start).vortex_expect("Failed to convert from usize to `S`"),
587        );
588    }
589    // SAFETY: We have initialized all `num_lists` values in both ranges, and both the `offsets`
590    // and the `sizes` builders are non-nullable.
591    unsafe { offsets_range.finish() };
592    unsafe { sizes_range.finish() };
593    Ok(())
594}
595
596/// Reads one non-nullable integer value of list view metadata as a `usize`.
597fn metadata_at(metadata: &PrimitiveArray, index: usize) -> usize {
598    match_each_integer_ptype!(metadata.ptype(), |A| {
599        metadata.as_slice::<A>()[index]
600            .to_usize()
601            .vortex_expect("list view metadata must fit in a usize")
602    })
603}
604
605/// Writes `offsets` into `range`, moved off the source's element window and onto the elements the
606/// builder already holds.
607///
608/// `window_start` is the first element the source's views reference, so every offset is at least
609/// `window_start` and the subtraction cannot underflow.
610fn extend_rebased_offsets<O: OffsetBuilderPType, A: IntegerPType>(
611    mut range: UninitRange<O>,
612    offsets: &[A],
613    window_start: usize,
614    elements_base: usize,
615) {
616    debug_assert_eq!(range.len(), offsets.len());
617
618    for (i, &offset) in offsets.iter().enumerate() {
619        let offset = offset
620            .to_usize()
621            .vortex_expect("offsets must always fit in usize");
622        debug_assert!(
623            offset >= window_start,
624            "offset {offset} precedes the referenced window at {window_start}"
625        );
626        let rebased = O::from_usize(offset - window_start + elements_base)
627            .vortex_expect("rebased offset did not fit into the builder's offset type");
628        range.set_value(i, rebased);
629    }
630
631    // SAFETY: We have set all the values in the range, and since `offsets` are non-nullable, we are
632    // done.
633    unsafe { range.finish() };
634}
635
636/// Writes `sizes` into `range`, converting them to the builder's size type.
637///
638/// Sizes that already have the builder's type are copied in bulk by the caller instead.
639fn extend_converted_sizes<S: OffsetBuilderPType, A: IntegerPType>(
640    mut range: UninitRange<S>,
641    sizes: &[A],
642) {
643    debug_assert_eq!(range.len(), sizes.len());
644
645    for (i, &size) in sizes.iter().enumerate() {
646        let size = S::from_usize(
647            size.to_usize()
648                .vortex_expect("sizes must always fit in usize"),
649        )
650        .vortex_expect("size did not fit into the builder's size type");
651        range.set_value(i, size);
652    }
653
654    // SAFETY: We have set all the values in the range, and since `sizes` are non-nullable, we are
655    // done.
656    unsafe { range.finish() };
657}
658
659#[cfg(test)]
660mod tests {
661    use std::sync::Arc;
662
663    use vortex_buffer::BufferAllocatorRef;
664    use vortex_buffer::buffer;
665    use vortex_error::VortexExpect;
666    use vortex_error::VortexResult;
667
668    use super::ListViewBuilder;
669    use crate::IntoArray;
670    use crate::VortexSessionExecute;
671    use crate::array_session;
672    use crate::arrays::ConstantArray;
673    use crate::arrays::ListArray;
674    use crate::arrays::ListViewArray;
675    use crate::arrays::listview::ListViewArrayExt;
676    use crate::arrays::listview::ListViewArraySlotsExt;
677    use crate::assert_arrays_eq;
678    use crate::builders::ArrayBuilder;
679    use crate::builders::listview::PrimitiveArray;
680    use crate::dtype::DType;
681    use crate::dtype::Nullability::NonNullable;
682    use crate::dtype::Nullability::Nullable;
683    use crate::dtype::PType::I32;
684    use crate::scalar::Scalar;
685    use crate::validity::Validity;
686
687    #[test]
688    fn test_empty() {
689        let mut builder = ListViewBuilder::<u32, u32>::with_capacity_in(
690            Arc::new(I32.into()),
691            NonNullable,
692            0,
693            0,
694            BufferAllocatorRef::static_ref(),
695        );
696
697        let listview = builder.finish();
698        assert_eq!(listview.len(), 0);
699        assert_eq!(
700            listview.dtype(),
701            &DType::List(Arc::new(I32.into()), NonNullable)
702        );
703    }
704
705    #[test]
706    fn test_basic_append_and_nulls() {
707        let mut ctx = array_session().create_execution_ctx();
708        let dtype: Arc<DType> = Arc::new(I32.into());
709        let mut builder = ListViewBuilder::<u32, u32>::with_capacity_in(
710            Arc::clone(&dtype),
711            Nullable,
712            0,
713            0,
714            BufferAllocatorRef::static_ref(),
715        );
716
717        // Append a regular list.
718        builder
719            .append_value(
720                Scalar::list(
721                    Arc::clone(&dtype),
722                    vec![1i32.into(), 2i32.into(), 3i32.into()],
723                    NonNullable,
724                )
725                .as_list(),
726            )
727            .unwrap();
728
729        // Append an empty list.
730        builder
731            .append_value(Scalar::list_empty(Arc::clone(&dtype), NonNullable).as_list())
732            .unwrap();
733
734        // Append a null list.
735        builder.append_null();
736
737        // Append another regular list.
738        builder
739            .append_value(
740                Scalar::list(dtype, vec![4i32.into(), 5i32.into()], NonNullable).as_list(),
741            )
742            .unwrap();
743
744        let listview = builder.finish_into_listview();
745        assert_eq!(listview.len(), 4);
746
747        // Check first list: [1, 2, 3].
748        assert_arrays_eq!(
749            listview.list_elements_at(0).unwrap(),
750            PrimitiveArray::from_iter([1i32, 2, 3]),
751            &mut ctx
752        );
753
754        // Check empty list.
755        assert_eq!(listview.list_elements_at(1).unwrap().len(), 0);
756
757        // Check null list.
758        assert!(
759            !listview
760                .validity()
761                .vortex_expect("listview validity should be derivable")
762                .execute_is_valid(2, &mut ctx)
763                .unwrap()
764        );
765
766        // Check last list: [4, 5].
767        assert_arrays_eq!(
768            listview.list_elements_at(3).unwrap(),
769            PrimitiveArray::from_iter([4i32, 5]),
770            &mut ctx
771        );
772        let expected = ListArray::new(
773            buffer![1i32, 2, 3, 4, 5].into_array(),
774            buffer![0u32, 3, 3, 3, 5].into_array(),
775            Validity::from_iter([true, true, false, true]),
776        );
777        assert_arrays_eq!(listview, expected, &mut ctx);
778    }
779
780    #[test]
781    fn test_different_offset_size_types() {
782        let mut ctx = array_session().create_execution_ctx();
783        // Test u64 offsets with u32 sizes.
784        let dtype: Arc<DType> = Arc::new(I32.into());
785        let mut builder = ListViewBuilder::<u64, u32>::with_capacity_in(
786            Arc::clone(&dtype),
787            NonNullable,
788            0,
789            0,
790            BufferAllocatorRef::static_ref(),
791        );
792
793        builder
794            .append_value(
795                Scalar::list(
796                    Arc::clone(&dtype),
797                    vec![1i32.into(), 2i32.into()],
798                    NonNullable,
799                )
800                .as_list(),
801            )
802            .unwrap();
803
804        builder
805            .append_value(
806                Scalar::list(
807                    dtype,
808                    vec![3i32.into(), 4i32.into(), 5i32.into()],
809                    NonNullable,
810                )
811                .as_list(),
812            )
813            .unwrap();
814
815        let listview = builder.finish_into_listview();
816        assert_eq!(listview.len(), 2);
817
818        // Verify first list: [1, 2].
819        assert_arrays_eq!(
820            listview.list_elements_at(0).unwrap(),
821            PrimitiveArray::from_iter([1i32, 2]),
822            &mut ctx
823        );
824
825        // Verify second list: [3, 4, 5].
826        assert_arrays_eq!(
827            listview.list_elements_at(1).unwrap(),
828            PrimitiveArray::from_iter([3i32, 4, 5]),
829            &mut ctx
830        );
831
832        // Test i64 offsets with i32 sizes.
833        let dtype2: Arc<DType> = Arc::new(I32.into());
834        let mut builder2 = ListViewBuilder::<i64, i32>::with_capacity_in(
835            Arc::clone(&dtype2),
836            NonNullable,
837            0,
838            0,
839            BufferAllocatorRef::static_ref(),
840        );
841
842        for i in 0..5 {
843            builder2
844                .append_value(
845                    Scalar::list(Arc::clone(&dtype2), vec![(i * 10).into()], NonNullable).as_list(),
846                )
847                .unwrap();
848        }
849
850        let listview2 = builder2.finish_into_listview();
851        assert_eq!(listview2.len(), 5);
852
853        // Verify the values: [0], [10], [20], [30], [40].
854        for i in 0..5i32 {
855            assert_arrays_eq!(
856                listview2.list_elements_at(i as usize).unwrap(),
857                PrimitiveArray::from_iter([i * 10]),
858                &mut ctx
859            );
860        }
861    }
862
863    #[test]
864    fn test_builder_trait_methods() {
865        let mut ctx = array_session().create_execution_ctx();
866        let dtype: Arc<DType> = Arc::new(I32.into());
867        let mut builder = ListViewBuilder::<u32, u32>::with_capacity_in(
868            Arc::clone(&dtype),
869            Nullable,
870            0,
871            0,
872            BufferAllocatorRef::static_ref(),
873        );
874
875        // Test append_zeros (creates empty lists).
876        builder.append_zeros(2);
877        assert_eq!(builder.len(), 2);
878
879        // Test append_nulls.
880        unsafe {
881            builder.append_nulls_unchecked(2);
882        }
883        assert_eq!(builder.len(), 4);
884
885        // Test append_scalar.
886        let list_scalar = Scalar::list(dtype, vec![10i32.into(), 20i32.into()], Nullable);
887        builder.append_scalar(&list_scalar).unwrap();
888        assert_eq!(builder.len(), 5);
889
890        let listview = builder.finish_into_listview();
891        assert_eq!(listview.len(), 5);
892
893        // First two are empty lists (from append_zeros).
894        assert_eq!(listview.list_elements_at(0).unwrap().len(), 0);
895        assert_eq!(listview.list_elements_at(1).unwrap().len(), 0);
896
897        // Next two are nulls.
898        assert!(
899            !listview
900                .validity()
901                .vortex_expect("listview validity should be derivable")
902                .execute_is_valid(2, &mut ctx)
903                .unwrap()
904        );
905        assert!(
906            !listview
907                .validity()
908                .vortex_expect("listview validity should be derivable")
909                .execute_is_valid(3, &mut ctx)
910                .unwrap()
911        );
912
913        // Last is the regular list: [10, 20].
914        assert_arrays_eq!(
915            listview.list_elements_at(4).unwrap(),
916            PrimitiveArray::from_iter([10i32, 20]),
917            &mut ctx
918        );
919        let expected = ListArray::new(
920            buffer![10i32, 20].into_array(),
921            buffer![0u32, 0, 0, 0, 0, 2].into_array(),
922            Validity::from_iter([true, true, false, false, true]),
923        );
924        assert_arrays_eq!(listview, expected, &mut ctx);
925    }
926
927    #[test]
928    fn test_extend_from_array() {
929        let mut ctx = array_session().create_execution_ctx();
930        let dtype: Arc<DType> = Arc::new(I32.into());
931
932        // Create a source ListArray.
933        let source = ListArray::from_iter_opt_slow::<u32, _, Vec<i32>>(
934            [Some(vec![1, 2, 3]), None, Some(vec![4, 5])],
935            Arc::new(I32.into()),
936        )
937        .unwrap();
938
939        let mut builder = ListViewBuilder::<u32, u32>::with_capacity_in(
940            Arc::clone(&dtype),
941            Nullable,
942            0,
943            0,
944            BufferAllocatorRef::static_ref(),
945        );
946
947        // Add initial data.
948        builder
949            .append_value(Scalar::list(dtype, vec![0i32.into()], NonNullable).as_list())
950            .unwrap();
951
952        // Extend from the ListArray.
953        let source = source
954            .into_array()
955            .execute::<ListViewArray>(&mut ctx)
956            .unwrap();
957        builder
958            .append_listview_array(source.as_view(), &mut ctx)
959            .unwrap();
960
961        // Extend from empty array (should be no-op).
962        let empty_source = ListArray::from_iter_opt_slow::<u32, _, Vec<i32>>(
963            std::iter::empty::<Option<Vec<i32>>>(),
964            Arc::new(I32.into()),
965        )
966        .unwrap();
967        let empty_source = empty_source
968            .into_array()
969            .execute::<ListViewArray>(&mut ctx)
970            .unwrap();
971        builder
972            .append_listview_array(empty_source.as_view(), &mut ctx)
973            .unwrap();
974
975        let listview = builder.finish_into_listview();
976        assert_eq!(listview.len(), 4);
977
978        // Check the extended data.
979        // First list: [0] (initial data).
980        assert_arrays_eq!(
981            listview.list_elements_at(0).unwrap(),
982            PrimitiveArray::from_iter([0i32]),
983            &mut ctx
984        );
985
986        // Second list: [1, 2, 3] (from source).
987        assert_arrays_eq!(
988            listview.list_elements_at(1).unwrap(),
989            PrimitiveArray::from_iter([1i32, 2, 3]),
990            &mut ctx
991        );
992
993        // Third list: null (from source).
994        assert!(
995            !listview
996                .validity()
997                .vortex_expect("listview validity should be derivable")
998                .execute_is_valid(2, &mut ctx)
999                .unwrap()
1000        );
1001
1002        // Fourth list: [4, 5] (from source).
1003        assert_arrays_eq!(
1004            listview.list_elements_at(3).unwrap(),
1005            PrimitiveArray::from_iter([4i32, 5]),
1006            &mut ctx
1007        );
1008        let expected = ListArray::new(
1009            buffer![0i32, 1, 2, 3, 4, 5].into_array(),
1010            buffer![0u32, 1, 4, 4, 6].into_array(),
1011            Validity::from_iter([true, true, false, true]),
1012        );
1013        assert_arrays_eq!(listview, expected, &mut ctx);
1014    }
1015
1016    #[test]
1017    fn test_append_list_array_grows_builder() -> VortexResult<()> {
1018        let mut ctx = array_session().create_execution_ctx();
1019        let dtype: Arc<DType> = Arc::new(I32.into());
1020
1021        // Enough lists to exceed the offsets/sizes capacity of a zero-capacity builder, so
1022        // appending must grow the builder rather than panic in `uninit_range`.
1023        let lists: Vec<Option<Vec<i32>>> =
1024            (0..100).map(|i| (i % 10 != 0).then(|| vec![i])).collect();
1025        let source = ListArray::from_iter_opt_slow::<u32, _, _>(lists.clone(), Arc::clone(&dtype))?;
1026
1027        let mut builder = ListViewBuilder::<u32, u32>::with_capacity_in(
1028            Arc::clone(&dtype),
1029            Nullable,
1030            0,
1031            0,
1032            BufferAllocatorRef::static_ref(),
1033        );
1034        builder.append_list_array(source.as_view(), &mut ctx)?;
1035        // Append a second time to check growth from a non-empty builder and offset rebasing.
1036        builder.append_list_array(source.as_view(), &mut ctx)?;
1037
1038        let listview = builder.finish_into_listview();
1039        assert!(listview.is_zero_copy_to_list());
1040
1041        let expected = ListArray::from_iter_opt_slow::<u32, _, _>(
1042            lists.iter().cloned().chain(lists.iter().cloned()),
1043            dtype,
1044        )?;
1045        assert_arrays_eq!(listview, expected, &mut ctx);
1046
1047        Ok(())
1048    }
1049
1050    /// A constant list array points every view at a single copy of the value; flattening it in the
1051    /// builder would materialize a copy per row.
1052    #[test]
1053    fn test_constant_list_append_keeps_one_copy_of_the_value() -> VortexResult<()> {
1054        let mut ctx = array_session().create_execution_ctx();
1055        let element_dtype: Arc<DType> = Arc::new(I32.into());
1056
1057        const ROWS: usize = 10_000;
1058        let fill = Scalar::list(
1059            Arc::clone(&element_dtype),
1060            vec![1i32.into(), 2i32.into(), 3i32.into()],
1061            NonNullable,
1062        );
1063        let constant = ConstantArray::new(fill, ROWS).into_array();
1064
1065        let mut builder = ListViewBuilder::<u64, u64>::with_capacity_in(
1066            element_dtype,
1067            NonNullable,
1068            0,
1069            0,
1070            BufferAllocatorRef::static_ref(),
1071        );
1072        constant.append_to_builder(&mut builder, &mut ctx)?;
1073        let listview = builder.finish_into_listview();
1074
1075        assert_eq!(listview.len(), ROWS);
1076        assert_eq!(
1077            listview.elements().len(),
1078            3,
1079            "the fill value should be stored once, not once per row",
1080        );
1081        assert!(!listview.is_zero_copy_to_list());
1082        assert_arrays_eq!(&listview.into_array(), &constant, &mut ctx);
1083
1084        Ok(())
1085    }
1086
1087    /// Only the elements the views reference land in the builder: appending a slice out of the
1088    /// middle of a list array leaves the elements on either side of it behind.
1089    #[test]
1090    fn test_append_listview_array_trims_unreferenced_elements() -> VortexResult<()> {
1091        let mut ctx = array_session().create_execution_ctx();
1092        let dtype: Arc<DType> = Arc::new(I32.into());
1093
1094        // Five lists of two elements each, of which we append only the middle three.
1095        let source = ListArray::from_iter_slow::<u32, _>(
1096            (0..5).map(|i| vec![2 * i, 2 * i + 1]),
1097            Arc::clone(&dtype),
1098        )?
1099        .into_array()
1100        .execute::<ListViewArray>(&mut ctx)?;
1101        let middle = source.slice(1..4)?.execute::<ListViewArray>(&mut ctx)?;
1102
1103        let mut builder = ListViewBuilder::<u32, u32>::with_capacity_in(
1104            Arc::clone(&dtype),
1105            NonNullable,
1106            0,
1107            0,
1108            BufferAllocatorRef::static_ref(),
1109        );
1110        builder.append_listview_array(middle.as_view(), &mut ctx)?;
1111        // A second append has to rebase onto the elements already in the builder.
1112        builder.append_listview_array(middle.as_view(), &mut ctx)?;
1113        let listview = builder.finish_into_listview();
1114
1115        assert_eq!(
1116            listview.elements().len(),
1117            12,
1118            "only the six referenced elements of each append should have landed",
1119        );
1120        assert!(
1121            listview.is_zero_copy_to_list(),
1122            "trimming an exact source keeps the result exact",
1123        );
1124
1125        let expected = ListArray::from_iter_slow::<u32, _>(
1126            (1..4).chain(1..4).map(|i| vec![2 * i, 2 * i + 1]),
1127            dtype,
1128        )?;
1129        assert_arrays_eq!(listview, expected, &mut ctx);
1130
1131        Ok(())
1132    }
1133
1134    #[test]
1135    fn test_extend_from_array_overlapping_listview() {
1136        let mut ctx = array_session().create_execution_ctx();
1137        let dtype: Arc<DType> = Arc::new(I32.into());
1138
1139        // Non-ZCTL source:
1140        // - List 0: [10, 20]
1141        // - List 1: null (size is intentionally non-zero in source metadata)
1142        // - List 2: [10]
1143        let source = unsafe {
1144            ListViewArray::new_unchecked(
1145                buffer![10i32, 20, 30].into_array(),
1146                buffer![0u32, 1, 0].into_array(),
1147                buffer![2u8, 2, 1].into_array(),
1148                Validity::from_iter([true, false, true]),
1149            )
1150        };
1151        assert!(!source.is_zero_copy_to_list());
1152
1153        let mut builder = ListViewBuilder::<u32, u32>::with_capacity_in(
1154            Arc::clone(&dtype),
1155            Nullable,
1156            0,
1157            0,
1158            BufferAllocatorRef::static_ref(),
1159        );
1160        builder
1161            .append_listview_array(source.as_view(), &mut ctx)
1162            .unwrap();
1163
1164        let listview = builder.finish_into_listview();
1165        assert_eq!(listview.len(), 3);
1166        // The builder kept the source's overlapping layout.
1167        assert!(!listview.is_zero_copy_to_list());
1168
1169        assert_arrays_eq!(
1170            listview.list_elements_at(0).unwrap(),
1171            PrimitiveArray::from_iter([10i32, 20]),
1172            &mut ctx
1173        );
1174        assert!(
1175            !listview
1176                .validity()
1177                .vortex_expect("listview validity should be derivable")
1178                .execute_is_valid(1, &mut ctx)
1179                .unwrap()
1180        );
1181        // List 1 is null, so the builder no longer rewrites its size to zero.
1182        assert_eq!(listview.size_at(1), source.size_at(1));
1183        assert_arrays_eq!(
1184            listview.list_elements_at(2).unwrap(),
1185            PrimitiveArray::from_iter([10i32]),
1186            &mut ctx
1187        );
1188        let expected = ListArray::new(
1189            buffer![10i32, 20, 10].into_array(),
1190            buffer![0u32, 2, 2, 3].into_array(),
1191            Validity::from_iter([true, false, true]),
1192        );
1193        assert_arrays_eq!(listview, expected, &mut ctx);
1194    }
1195
1196    #[test]
1197    fn test_error_append_null_to_non_nullable() {
1198        let dtype: Arc<DType> = Arc::new(I32.into());
1199        let mut builder = ListViewBuilder::<u32, u32>::with_capacity_in(
1200            Arc::clone(&dtype),
1201            NonNullable,
1202            0,
1203            0,
1204            BufferAllocatorRef::static_ref(),
1205        );
1206
1207        // Create a null list with nullable type (since Scalar::null requires nullable type).
1208        let null_scalar = Scalar::null(DType::List(dtype, Nullable));
1209        let null_list = null_scalar.as_list();
1210
1211        // This should fail because we're trying to append a null to a non-nullable builder.
1212        let result = builder.append_value(null_list);
1213        assert!(result.is_err());
1214        assert!(
1215            result
1216                .unwrap_err()
1217                .to_string()
1218                .contains("null value to non-nullable")
1219        );
1220    }
1221
1222    #[test]
1223    fn test_append_array_as_list() {
1224        let dtype: Arc<DType> = Arc::new(I32.into());
1225        let mut ctx = array_session().create_execution_ctx();
1226        let mut builder = ListViewBuilder::<u32, u32>::with_capacity_in(
1227            Arc::clone(&dtype),
1228            NonNullable,
1229            20,
1230            10,
1231            BufferAllocatorRef::static_ref(),
1232        );
1233
1234        // Append a primitive array as a single list entry.
1235        let arr1 = buffer![1i32, 2, 3].into_array();
1236        builder.append_array_as_list(&arr1, &mut ctx).unwrap();
1237
1238        // Interleave with a list scalar.
1239        builder
1240            .append_value(
1241                Scalar::list(
1242                    Arc::clone(&dtype),
1243                    vec![10i32.into(), 11i32.into()],
1244                    NonNullable,
1245                )
1246                .as_list(),
1247            )
1248            .unwrap();
1249
1250        // Append another primitive array as a single list entry.
1251        let arr2 = buffer![4i32, 5].into_array();
1252        builder.append_array_as_list(&arr2, &mut ctx).unwrap();
1253
1254        // Append an empty array as a single list entry (empty list).
1255        let arr3 = buffer![0i32; 0].into_array();
1256        builder.append_array_as_list(&arr3, &mut ctx).unwrap();
1257
1258        // Interleave with another list scalar.
1259        builder
1260            .append_value(Scalar::list_empty(Arc::clone(&dtype), NonNullable).as_list())
1261            .unwrap();
1262
1263        let listview = builder.finish_into_listview();
1264        assert_eq!(listview.len(), 5);
1265
1266        // Verify elements array: [1, 2, 3, 10, 11, 4, 5].
1267        assert_arrays_eq!(
1268            listview.elements(),
1269            PrimitiveArray::from_iter([1i32, 2, 3, 10, 11, 4, 5]),
1270            &mut ctx
1271        );
1272
1273        // Verify offsets array.
1274        assert_arrays_eq!(
1275            listview.offsets(),
1276            PrimitiveArray::from_iter([0u32, 3, 5, 7, 7]),
1277            &mut ctx
1278        );
1279
1280        // Verify sizes array.
1281        assert_arrays_eq!(
1282            listview.sizes(),
1283            PrimitiveArray::from_iter([3u32, 2, 2, 0, 0]),
1284            &mut ctx
1285        );
1286
1287        // Test dtype mismatch error.
1288        let mut builder = ListViewBuilder::<u32, u32>::with_capacity_in(
1289            dtype,
1290            NonNullable,
1291            20,
1292            10,
1293            BufferAllocatorRef::static_ref(),
1294        );
1295        let wrong_dtype_arr = buffer![1i64, 2, 3].into_array();
1296        assert!(
1297            builder
1298                .append_array_as_list(&wrong_dtype_arr, &mut ctx)
1299                .is_err()
1300        );
1301    }
1302}