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