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