Skip to main content

vortex_array/arrays/listview/
rebuild.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use num_traits::FromPrimitive;
5use vortex_buffer::BufferMut;
6use vortex_error::VortexExpect;
7use vortex_error::VortexResult;
8
9use crate::Canonical;
10use crate::ExecutionCtx;
11use crate::IntoArray;
12use crate::arrays::ConstantArray;
13use crate::arrays::ListViewArray;
14use crate::arrays::PrimitiveArray;
15use crate::arrays::listview::ListViewArrayExt;
16use crate::arrays::primitive::PrimitiveArrayExt;
17use crate::builders::builder_with_capacity;
18use crate::builtins::ArrayBuiltins;
19use crate::dtype::IntegerPType;
20use crate::dtype::Nullability;
21use crate::dtype::PType;
22use crate::match_each_integer_ptype;
23use crate::match_each_unsigned_integer_ptype;
24use crate::scalar::Scalar;
25use crate::scalar_fn::fns::operators::Operator;
26use crate::validity::Validity;
27
28/// The widened offset type used for a rebuilt, zero-copy-to-list array.
29///
30/// Offsets are widened to at least 32 bits (Arrow only permits 32/64-bit list offsets) while
31/// preserving signedness, since a signed result keeps the array zero-copyable to Arrow's
32/// `ListArray`.
33fn rebuilt_offset_ptype(offsets_ptype: PType) -> PType {
34    match offsets_ptype {
35        PType::U8 | PType::U16 | PType::U32 => PType::U32,
36        PType::U64 => PType::U64,
37        PType::I8 | PType::I16 | PType::I32 => PType::I32,
38        PType::I64 => PType::I64,
39        _ => unreachable!("invalid offsets PType"),
40    }
41}
42
43/// Density threshold to decide whether to rebuild a sparse `ListViewArray`.
44///
45/// A `ListViewArray` can accumulate unreferenced bytes in its `elements` buffer after
46/// metadata-only operations like `take` and `filter`. When density (referenced fraction of `elements`)
47/// falls below this threshold, the benefits of a rebuild may outweigh its cost.
48///
49/// This is a somewhat arbitrary rule-of-thumb and may be suboptimal depending on different use cases and
50/// list element dtypes.
51pub const DEFAULT_REBUILD_DENSITY_THRESHOLD: f32 = 0.1;
52
53/// Waste threshold to decide whether to trim a zero-copy-to-list `ListViewArray`.
54///
55/// A zero-copy-to-list array has no overlaps and no interior gaps, so its only unreferenced bytes
56/// are leading and trailing elements. Trimming those is much cheaper than a full rebuild (a lazy
57/// `elements` slice plus an `O(num_lists)` offset adjustment, with no element data copy), so we use
58/// a more aggressive threshold than [`DEFAULT_REBUILD_DENSITY_THRESHOLD`].
59///
60/// When the unreferenced (leading + trailing) fraction of `elements` exceeds this threshold, we trim.
61pub const DEFAULT_TRIM_ELEMENTS_THRESHOLD: f32 = 0.05;
62
63/// Modes for rebuilding a [`ListViewArray`].
64pub enum ListViewRebuildMode {
65    /// Removes all unused data and flattens out all list data, such that the array is zero-copyable
66    /// to a [`ListArray`].
67    ///
68    /// This mode will deduplicate all overlapping list views, such that the [`ListViewArray`] looks
69    /// like a [`ListArray`] but with an additional `sizes` array.
70    ///
71    /// [`ListArray`]: crate::arrays::ListArray
72    MakeZeroCopyToList,
73
74    /// Removes any leading or trailing elements that are unused / not referenced by any views in
75    /// the [`ListViewArray`].
76    ///
77    /// If the referenced `[start, end)` bounds are already known, prefer calling
78    /// [`trim_elements`](ListViewArray::trim_elements) directly to avoid recomputing them.
79    TrimElements,
80
81    /// Equivalent to `MakeZeroCopyToList` plus `TrimElements`.
82    ///
83    /// This is useful when concatenating multiple [`ListViewArray`]s together to create a new
84    /// [`ListViewArray`] that is also zero-copy to a [`ListArray`].
85    ///
86    /// [`ListArray`]: crate::arrays::ListArray
87    MakeExact,
88
89    // TODO(connor)[ListView]: Implement some version of this.
90    /// Finds the shortest packing / overlapping of list elements.
91    ///
92    /// This problem is known to be NP-hard, so maybe when someone proves that P=NP we can implement
93    /// this algorithm (but in all seriousness there are many approximate algorithms that could
94    /// work well here).
95    OverlapCompression,
96}
97
98impl ListViewArray {
99    /// Rebuilds the [`ListViewArray`] according to the specified mode.
100    pub fn rebuild(
101        &self,
102        mode: ListViewRebuildMode,
103        ctx: &mut ExecutionCtx,
104    ) -> VortexResult<ListViewArray> {
105        if self.is_empty() {
106            // SAFETY: An empty array is trivially zero-copyable to a `ListArray`.
107            return Ok(unsafe { self.clone().with_zero_copy_to_list(true) });
108        }
109
110        match mode {
111            ListViewRebuildMode::MakeZeroCopyToList => self.rebuild_zero_copy_to_list(ctx),
112            ListViewRebuildMode::TrimElements => self.rebuild_trim_elements(ctx),
113            ListViewRebuildMode::MakeExact => self.rebuild_make_exact(ctx),
114            ListViewRebuildMode::OverlapCompression => unimplemented!("Does P=NP?"),
115        }
116    }
117
118    /// Rebuilds a [`ListViewArray`], removing all data overlaps and creating a flattened layout.
119    ///
120    /// This is useful when the `elements` child array of the [`ListViewArray`] might have
121    /// overlapping, duplicate, and garbage data, and we want to have fully sequential data like
122    /// a [`ListArray`].
123    ///
124    /// [`ListArray`]: crate::arrays::ListArray
125    fn rebuild_zero_copy_to_list(&self, ctx: &mut ExecutionCtx) -> VortexResult<ListViewArray> {
126        if self.is_zero_copy_to_list() {
127            // Note that since everything in `ListViewArray` is `Arc`ed, this is quite cheap.
128            return Ok(self.clone());
129        }
130
131        let offsets_ptype = self.offsets().dtype().as_ptype();
132        let sizes_ptype = self.sizes().dtype().as_ptype();
133
134        // One of the main purposes behind adding this "zero-copyable to `ListArray`" optimization
135        // is that we want to pass data to systems that expect Arrow data.
136        // The arrow specification only allows for `i32` and `i64` offset and sizes types, so in
137        // order to also make `ListView` zero-copyable to **Arrow**'s `ListArray` (not just Vortex's
138        // `ListArray`), we rebuild the offsets as 32-bit or 64-bit integer types.
139        // TODO(connor)[ListView]: This is true for `sizes` as well, we should do this conversion
140        // for sizes as well.
141        match_each_unsigned_integer_ptype!(sizes_ptype.to_unsigned(), |S| {
142            match offsets_ptype.to_unsigned() {
143                PType::U8 => self.naive_rebuild::<u8, u32, S>(ctx),
144                PType::U16 => self.naive_rebuild::<u16, u32, S>(ctx),
145                PType::U32 => self.naive_rebuild::<u32, u32, S>(ctx),
146                PType::U64 => self.naive_rebuild::<u64, u64, S>(ctx),
147                _ => unreachable!("invalid offsets PType"),
148            }
149        })
150    }
151
152    /// Picks between [`rebuild_with_take`](Self::rebuild_with_take) and
153    /// [`rebuild_list_by_list`](Self::rebuild_list_by_list) based on element dtype and average
154    /// list size.
155    fn naive_rebuild<O: IntegerPType, NewOffset: IntegerPType, S: IntegerPType>(
156        &self,
157        ctx: &mut ExecutionCtx,
158    ) -> VortexResult<ListViewArray> {
159        let sizes_canonical = self.sizes().clone().execute::<PrimitiveArray>(ctx)?;
160        let sizes_canonical =
161            sizes_canonical.reinterpret_cast(sizes_canonical.ptype().to_unsigned());
162        let total: u64 = sizes_canonical
163            .as_slice::<S>()
164            .iter()
165            .map(|s| (*s).as_() as u64)
166            .sum();
167        if Self::should_use_take(total, self.len()) {
168            self.rebuild_with_take::<O, NewOffset, S>(ctx)
169        } else {
170            self.rebuild_list_by_list::<O, NewOffset, S>(ctx)
171        }
172    }
173
174    /// Returns `true` when we are confident that `rebuild_with_take` will
175    /// outperform `rebuild_list_by_list`.
176    ///
177    /// Take is dramatically faster for small lists (often 10-100×) because it
178    /// avoids per-list builder overhead. LBL is the safer default for larger
179    /// lists since its sequential memcpy scales well. We only choose take when
180    /// the average list size is small enough that take clearly dominates.
181    fn should_use_take(total_output_elements: u64, num_lists: usize) -> bool {
182        if num_lists == 0 {
183            return true;
184        }
185        let avg = total_output_elements / num_lists as u64;
186        avg < 128
187    }
188
189    /// Rebuilds elements using a single bulk `take`: collect all element indices into a flat
190    /// `BufferMut<u64>`, perform a single `take`.
191    fn rebuild_with_take<O: IntegerPType, NewOffset: IntegerPType, S: IntegerPType>(
192        &self,
193        ctx: &mut ExecutionCtx,
194    ) -> VortexResult<ListViewArray> {
195        let new_offset_ptype = rebuilt_offset_ptype(self.offsets().dtype().as_ptype());
196        let size_ptype = self.sizes().dtype().as_ptype();
197
198        let offsets_canonical = self.offsets().clone().execute::<PrimitiveArray>(ctx)?;
199        let offsets_canonical =
200            offsets_canonical.reinterpret_cast(offsets_canonical.ptype().to_unsigned());
201        let offsets_slice = offsets_canonical.as_slice::<O>();
202        let sizes_canonical = self.sizes().clone().execute::<PrimitiveArray>(ctx)?;
203        let sizes_canonical =
204            sizes_canonical.reinterpret_cast(sizes_canonical.ptype().to_unsigned());
205        let sizes_slice = sizes_canonical.as_slice::<S>();
206
207        let len = offsets_slice.len();
208
209        let mut new_offsets = BufferMut::<NewOffset>::with_capacity(len);
210        let mut new_sizes = BufferMut::<S>::with_capacity(len);
211        let mut take_indices = BufferMut::<u64>::with_capacity(self.elements().len());
212
213        // Resolve validity to a mask once instead of probing it per row: `execute_is_valid`
214        // executes a scalar on every call for array-backed validity, which is O(len) work repeated
215        // `len` times.
216        let validity = self.validity()?.execute_mask(len, ctx)?;
217
218        let mut n_elements = NewOffset::zero();
219        for index in 0..len {
220            if !validity.value(index) {
221                new_offsets.push(n_elements);
222                new_sizes.push(S::zero());
223                continue;
224            }
225
226            let offset = offsets_slice[index];
227            let size = sizes_slice[index];
228            let start = offset.as_();
229            let stop = start + size.as_();
230
231            new_offsets.push(n_elements);
232            new_sizes.push(size);
233            take_indices.extend(start as u64..stop as u64);
234            n_elements += num_traits::cast(size).vortex_expect("Cast failed");
235        }
236
237        let elements = self.elements().take(take_indices.into_array())?;
238        // Built unsigned; reinterpret back to the signed-preserving result types.
239        let offsets = PrimitiveArray::new(new_offsets.freeze(), Validity::NonNullable)
240            .reinterpret_cast(new_offset_ptype)
241            .into_array();
242        let sizes = PrimitiveArray::new(new_sizes.freeze(), Validity::NonNullable)
243            .reinterpret_cast(size_ptype)
244            .into_array();
245
246        // SAFETY: same invariants as `rebuild_list_by_list` — offsets are sequential and
247        // non-overlapping, all (offset, size) pairs reference valid elements, and the validity
248        // array is preserved from the original.
249        Ok(unsafe {
250            ListViewArray::new_unchecked(elements, offsets, sizes, self.validity()?)
251                .with_zero_copy_to_list(true)
252        })
253    }
254
255    /// Rebuilds elements list-by-list: canonicalize elements upfront, then for each list `slice`
256    /// the relevant range and `append_to_builder` into a typed builder.
257    fn rebuild_list_by_list<O: IntegerPType, NewOffset: IntegerPType, S: IntegerPType>(
258        &self,
259        ctx: &mut ExecutionCtx,
260    ) -> VortexResult<ListViewArray> {
261        let element_dtype = self
262            .dtype()
263            .as_list_element_opt()
264            .vortex_expect("somehow had a canonical list that was not a list");
265
266        let new_offset_ptype = rebuilt_offset_ptype(self.offsets().dtype().as_ptype());
267        let size_ptype = self.sizes().dtype().as_ptype();
268
269        let offsets_canonical = self.offsets().clone().execute::<PrimitiveArray>(ctx)?;
270        let offsets_canonical =
271            offsets_canonical.reinterpret_cast(offsets_canonical.ptype().to_unsigned());
272        let offsets_slice = offsets_canonical.as_slice::<O>();
273        let sizes_canonical = self.sizes().clone().execute::<PrimitiveArray>(ctx)?;
274        let sizes_canonical =
275            sizes_canonical.reinterpret_cast(sizes_canonical.ptype().to_unsigned());
276        let sizes_slice = sizes_canonical.as_slice::<S>();
277
278        let len = offsets_slice.len();
279
280        let mut new_offsets = BufferMut::<NewOffset>::with_capacity(len);
281        // TODO(connor)[ListView]: Do we really need to do this?
282        // The only reason we need to rebuild the sizes here is that the validity may indicate that
283        // a list is null even though it has a non-zero size. This rebuild will set the size of all
284        // null lists to 0.
285        let mut new_sizes = BufferMut::<S>::with_capacity(len);
286
287        // Canonicalize the elements up front as we will be slicing the elements quite a lot.
288        let elements_canonical = self
289            .elements()
290            .clone()
291            .execute::<Canonical>(ctx)?
292            .into_array();
293
294        // Note that we do not know what the exact capacity should be of the new elements since
295        // there could be overlaps in the existing `ListViewArray`.
296        let mut new_elements_builder =
297            builder_with_capacity(element_dtype.as_ref(), self.elements().len());
298
299        // Resolve validity to a mask once instead of probing it per row (see `rebuild_with_take`).
300        let validity = self.validity()?.execute_mask(len, ctx)?;
301
302        let mut n_elements = NewOffset::zero();
303        for index in 0..len {
304            if !validity.value(index) {
305                // For NULL lists, place them after the previous item's data to maintain the
306                // no-overlap invariant for zero-copy to `ListArray` arrays.
307                new_offsets.push(n_elements);
308                new_sizes.push(S::zero());
309                continue;
310            }
311
312            let offset = offsets_slice[index];
313            let size = sizes_slice[index];
314
315            let start = offset.as_();
316            let stop = start + size.as_();
317
318            new_offsets.push(n_elements);
319            new_sizes.push(size);
320            elements_canonical
321                .slice(start..stop)?
322                .append_to_builder(new_elements_builder.as_mut(), ctx)?;
323
324            n_elements += num_traits::cast(size).vortex_expect("Cast failed");
325        }
326
327        // Built unsigned; reinterpret back to the signed-preserving result types.
328        let offsets = PrimitiveArray::new(new_offsets.freeze(), Validity::NonNullable)
329            .reinterpret_cast(new_offset_ptype)
330            .into_array();
331        let sizes = PrimitiveArray::new(new_sizes.freeze(), Validity::NonNullable)
332            .reinterpret_cast(size_ptype)
333            .into_array();
334        let elements = new_elements_builder.finish();
335
336        debug_assert_eq!(
337            n_elements.as_(),
338            elements.len(),
339            "The accumulated elements somehow had the wrong length"
340        );
341
342        // SAFETY:
343        // - All offsets are sequential and non-overlapping (`n_elements` tracks running total).
344        // - Each `offset[i] + size[i]` equals `offset[i+1]` for all valid indices (including null
345        //   lists).
346        // - All elements referenced by (offset, size) pairs exist within the new `elements` array.
347        // - The validity array is preserved from the original array unchanged
348        // - The array satisfies the zero-copy-to-list property by having sorted offsets, no gaps,
349        //   and no overlaps.
350        Ok(unsafe {
351            ListViewArray::new_unchecked(elements, offsets, sizes, self.validity()?)
352                .with_zero_copy_to_list(true)
353        })
354    }
355
356    /// Rebuilds a [`ListViewArray`] by trimming any unused / unreferenced leading and trailing
357    /// elements, which is defined as a contiguous run of values in the `elements` array that are
358    /// not referenced by any views in the corresponding [`ListViewArray`].
359    fn rebuild_trim_elements(&self, ctx: &mut ExecutionCtx) -> VortexResult<ListViewArray> {
360        let (start, end) = self.referenced_element_bounds(ctx)?;
361
362        // SAFETY: we calculated valid start and end bounds
363        unsafe { self.trim_elements(start, end) }
364    }
365
366    /// Unsafely trims `elements` to the referenced half-open range `[start, end)`, adjusting every offset
367    /// down by `start`. The result preserves the original `is_zero_copy_to_list` flag.
368    ///
369    /// # SAFETY
370    ///
371    /// `start` must be the minimum value of `offsets`, and end should be the maximum value of `offsets[i] + size[i]`
372    /// over all indices `i` of offsets. Otherwise, `offsets` and `sizes` may hold references to elements that no longer exist
373    /// and the array will be corrupted.
374    pub unsafe fn trim_elements(&self, start: usize, end: usize) -> VortexResult<ListViewArray> {
375        let adjusted_offsets = match_each_integer_ptype!(self.offsets().dtype().as_ptype(), |O| {
376            let offset = <O as FromPrimitive>::from_usize(start)
377                .vortex_expect("unable to convert the min offset `start` into a `usize`");
378            let scalar = Scalar::primitive(offset, Nullability::NonNullable);
379
380            self.offsets()
381                .clone()
382                .binary(
383                    ConstantArray::new(scalar, self.offsets().len()).into_array(),
384                    Operator::Sub,
385                )
386                .vortex_expect("was somehow unable to adjust offsets down by their minimum")
387        });
388
389        let sliced_elements = self.elements().slice(start..end)?;
390
391        // SAFETY: The only thing we changed was the elements (which we verify through mins and
392        // maxes that all adjusted offsets + sizes are within the correct bounds), so the parameters
393        // are valid. And if the original array was zero-copyable to list, trimming elements doesn't
394        // change that property.
395        Ok(unsafe {
396            ListViewArray::new_unchecked(
397                sliced_elements,
398                adjusted_offsets,
399                self.sizes().clone(),
400                self.validity()?,
401            )
402            .with_zero_copy_to_list(self.is_zero_copy_to_list())
403        })
404    }
405
406    fn rebuild_make_exact(&self, ctx: &mut ExecutionCtx) -> VortexResult<ListViewArray> {
407        if self.is_zero_copy_to_list() {
408            self.rebuild_trim_elements(ctx)
409        } else {
410            // When we completely rebuild the `ListViewArray`, we get the benefit that we also trim
411            // any leading and trailing garbage data.
412            self.rebuild_zero_copy_to_list(ctx)
413        }
414    }
415}
416
417#[cfg(test)]
418mod tests {
419    use vortex_buffer::BitBuffer;
420    use vortex_error::VortexResult;
421
422    use super::super::tests::common::SESSION;
423    use super::ListViewRebuildMode;
424    use crate::IntoArray;
425    use crate::VortexSessionExecute;
426    use crate::arrays::ListViewArray;
427    use crate::arrays::PrimitiveArray;
428    use crate::arrays::listview::ListViewArrayExt;
429    use crate::assert_arrays_eq;
430    use crate::dtype::Nullability;
431    use crate::validity::Validity;
432
433    #[test]
434    fn test_rebuild_flatten_removes_overlaps() -> VortexResult<()> {
435        // Create a list view with overlapping lists: [A, B, C]
436        // List 0: offset=0, size=3 -> [A, B, C]
437        // List 1: offset=1, size=2 -> [B, C] (overlaps with List 0)
438        let elements = PrimitiveArray::from_iter(vec![1i32, 2, 3]).into_array();
439        let offsets = PrimitiveArray::from_iter(vec![0u32, 1]).into_array();
440        let sizes = PrimitiveArray::from_iter(vec![3u32, 2]).into_array();
441
442        let listview = ListViewArray::new(elements, offsets, sizes, Validity::NonNullable);
443
444        let mut ctx = SESSION.create_execution_ctx();
445        let flattened = listview.rebuild(ListViewRebuildMode::MakeZeroCopyToList, &mut ctx)?;
446
447        // After flatten: elements should be [A, B, C, B, C] = [1, 2, 3, 2, 3]
448        // Lists should be sequential with no overlaps
449        assert_eq!(flattened.elements().len(), 5);
450
451        // Offsets should be sequential
452        assert_eq!(flattened.offset_at(0), 0);
453        assert_eq!(flattened.size_at(0), 3);
454        assert_eq!(flattened.offset_at(1), 3);
455        assert_eq!(flattened.size_at(1), 2);
456
457        // Verify the data is correct
458        assert_arrays_eq!(
459            flattened.list_elements_at(0)?,
460            PrimitiveArray::from_iter([1i32, 2, 3]),
461            &mut ctx
462        );
463
464        assert_arrays_eq!(
465            flattened.list_elements_at(1)?,
466            PrimitiveArray::from_iter([2i32, 3]),
467            &mut ctx
468        );
469        Ok(())
470    }
471
472    #[test]
473    fn test_rebuild_flatten_with_nullable() -> VortexResult<()> {
474        use crate::arrays::BoolArray;
475
476        // Create a nullable list view with a null list
477        let elements = PrimitiveArray::from_iter(vec![1i32, 2, 3]).into_array();
478        let offsets = PrimitiveArray::from_iter(vec![0u32, 1, 2]).into_array();
479        let sizes = PrimitiveArray::from_iter(vec![2u32, 1, 1]).into_array();
480        let validity = Validity::Array(
481            BoolArray::new(
482                BitBuffer::from(vec![true, false, true]),
483                Validity::NonNullable,
484            )
485            .into_array(),
486        );
487
488        let listview = ListViewArray::new(elements, offsets, sizes, validity);
489
490        let mut ctx = SESSION.create_execution_ctx();
491        let flattened = listview.rebuild(ListViewRebuildMode::MakeZeroCopyToList, &mut ctx)?;
492
493        // Verify nullability is preserved
494        assert_eq!(flattened.dtype().nullability(), Nullability::Nullable);
495        assert!(flattened.validity()?.execute_is_valid(0, &mut ctx)?);
496        assert!(!flattened.validity()?.execute_is_valid(1, &mut ctx)?);
497        assert!(flattened.validity()?.execute_is_valid(2, &mut ctx)?);
498
499        // Verify valid lists contain correct data
500        assert_arrays_eq!(
501            flattened.list_elements_at(0)?,
502            PrimitiveArray::from_iter([1i32, 2]),
503            &mut ctx
504        );
505
506        assert_arrays_eq!(
507            flattened.list_elements_at(2)?,
508            PrimitiveArray::from_iter([3i32]),
509            &mut ctx
510        );
511        Ok(())
512    }
513
514    #[test]
515    fn test_rebuild_trim_elements_basic() -> VortexResult<()> {
516        // Test trimming both leading and trailing unused elements while preserving gaps in the
517        // middle.
518        // Elements: [_, _, A, B, _, C, D, _, _]
519        //            0  1  2  3  4  5  6  7  8
520        // List 0: offset=2, size=2 -> [A, B]
521        // List 1: offset=5, size=2 -> [C, D]
522        // Should trim to: [A, B, _, C, D] with adjusted offsets.
523        let elements =
524            PrimitiveArray::from_iter(vec![99i32, 98, 1, 2, 97, 3, 4, 96, 95]).into_array();
525        let offsets = PrimitiveArray::from_iter(vec![2u32, 5]).into_array();
526        let sizes = PrimitiveArray::from_iter(vec![2u32, 2]).into_array();
527
528        let listview = ListViewArray::new(elements, offsets, sizes, Validity::NonNullable);
529
530        let mut ctx = SESSION.create_execution_ctx();
531        let trimmed = listview.rebuild(ListViewRebuildMode::TrimElements, &mut ctx)?;
532
533        // After trimming: elements should be [A, B, _, C, D] = [1, 2, 97, 3, 4].
534        assert_eq!(trimmed.elements().len(), 5);
535
536        // Offsets should be adjusted: old offset 2 -> new offset 0, old offset 5 -> new offset 3.
537        assert_eq!(trimmed.offset_at(0), 0);
538        assert_eq!(trimmed.size_at(0), 2);
539        assert_eq!(trimmed.offset_at(1), 3);
540        assert_eq!(trimmed.size_at(1), 2);
541
542        // Verify the data is correct.
543        assert_arrays_eq!(
544            trimmed.list_elements_at(0)?,
545            PrimitiveArray::from_iter([1i32, 2]),
546            &mut ctx
547        );
548
549        assert_arrays_eq!(
550            trimmed.list_elements_at(1)?,
551            PrimitiveArray::from_iter([3i32, 4]),
552            &mut ctx
553        );
554
555        // Note that element at index 2 (97) is preserved as a gap.
556        let all_elements = trimmed
557            .elements()
558            .clone()
559            .execute::<PrimitiveArray>(&mut ctx)?;
560        assert_eq!(all_elements.execute_scalar(2, &mut ctx)?, 97i32.into());
561        Ok(())
562    }
563
564    #[test]
565    fn test_rebuild_with_trailing_nulls_regression() -> VortexResult<()> {
566        // Regression test for issue #5412
567        // Tests that zero-copy-to-list arrays with trailing NULLs correctly calculate
568        // offsets for NULL items to maintain no-overlap invariant
569
570        // Create a ListViewArray with trailing NULLs
571        let elements = PrimitiveArray::from_iter(vec![1i32, 2, 3, 4]).into_array();
572        let offsets = PrimitiveArray::from_iter(vec![0u32, 2, 0, 0]).into_array();
573        let sizes = PrimitiveArray::from_iter(vec![2u32, 2, 0, 0]).into_array();
574        let validity = Validity::from_iter(vec![true, true, false, false]);
575
576        let listview = ListViewArray::new(elements, offsets, sizes, validity);
577
578        // First rebuild to make it zero-copy-to-list
579        let mut ctx = SESSION.create_execution_ctx();
580        let rebuilt = listview.rebuild(ListViewRebuildMode::MakeZeroCopyToList, &mut ctx)?;
581        assert!(rebuilt.is_zero_copy_to_list());
582
583        // Verify NULL items have correct offsets (should not reuse previous offsets)
584        // After rebuild: offsets should be [0, 2, 4, 4] for zero-copy-to-list
585        assert_eq!(rebuilt.offset_at(0), 0);
586        assert_eq!(rebuilt.offset_at(1), 2);
587        assert_eq!(rebuilt.offset_at(2), 4); // NULL should be at position 4
588        assert_eq!(rebuilt.offset_at(3), 4); // Second NULL also at position 4
589
590        // All sizes should be correct
591        assert_eq!(rebuilt.size_at(0), 2);
592        assert_eq!(rebuilt.size_at(1), 2);
593        assert_eq!(rebuilt.size_at(2), 0); // NULL has size 0
594        assert_eq!(rebuilt.size_at(3), 0); // NULL has size 0
595
596        // Now rebuild with MakeExact (which calls naive_rebuild then trim_elements)
597        // This should not panic (issue #5412)
598        let exact = rebuilt.rebuild(ListViewRebuildMode::MakeExact, &mut ctx)?;
599
600        // Verify the result is still valid
601        assert!(exact.is_valid(0, &mut ctx)?);
602        assert!(exact.is_valid(1, &mut ctx)?);
603        assert!(!exact.is_valid(2, &mut ctx)?);
604        assert!(!exact.is_valid(3, &mut ctx)?);
605
606        // Verify data is preserved
607        assert_arrays_eq!(
608            exact.list_elements_at(0)?,
609            PrimitiveArray::from_iter([1i32, 2]),
610            &mut ctx
611        );
612
613        assert_arrays_eq!(
614            exact.list_elements_at(1)?,
615            PrimitiveArray::from_iter([3i32, 4]),
616            &mut ctx
617        );
618        Ok(())
619    }
620
621    /// Regression test for <https://github.com/vortex-data/vortex/issues/6773>.
622    /// u32 offsets exceed u16::MAX, so u16 sizes are widened to u32 for the add.
623    #[test]
624    fn test_rebuild_trim_elements_offsets_wider_than_sizes() -> VortexResult<()> {
625        let mut elems = vec![0i32; 70_005];
626        elems[70_000] = 10;
627        elems[70_001] = 20;
628        elems[70_002] = 30;
629        elems[70_003] = 40;
630        let elements = PrimitiveArray::from_iter(elems).into_array();
631        let offsets = PrimitiveArray::from_iter(vec![70_000u32, 70_002]).into_array();
632        let sizes = PrimitiveArray::from_iter(vec![2u16, 2]).into_array();
633
634        let listview = ListViewArray::new(elements, offsets, sizes, Validity::NonNullable);
635        let mut ctx = SESSION.create_execution_ctx();
636        let trimmed = listview.rebuild(ListViewRebuildMode::TrimElements, &mut ctx)?;
637        assert_arrays_eq!(
638            trimmed.list_elements_at(1)?,
639            PrimitiveArray::from_iter([30i32, 40]),
640            &mut ctx
641        );
642        Ok(())
643    }
644
645    /// Regression test for <https://github.com/vortex-data/vortex/issues/6773>.
646    /// u32 sizes exceed u16::MAX, so u16 offsets are widened to u32 for the add.
647    #[test]
648    fn test_rebuild_trim_elements_sizes_wider_than_offsets() -> VortexResult<()> {
649        let mut elems = vec![0i32; 70_001];
650        elems[3] = 30;
651        elems[4] = 40;
652        let elements = PrimitiveArray::from_iter(elems).into_array();
653        let offsets = PrimitiveArray::from_iter(vec![1u16, 3]).into_array();
654        let sizes = PrimitiveArray::from_iter(vec![70_000u32, 2]).into_array();
655
656        let listview = ListViewArray::new(elements, offsets, sizes, Validity::NonNullable);
657        let mut ctx = SESSION.create_execution_ctx();
658        let trimmed = listview.rebuild(ListViewRebuildMode::TrimElements, &mut ctx)?;
659        assert_arrays_eq!(
660            trimmed.list_elements_at(1)?,
661            PrimitiveArray::from_iter([30i32, 40]),
662            &mut ctx
663        );
664        Ok(())
665    }
666
667    /// Rebuild with signed offsets/sizes: exercises the unsigned-reinterpret read path and asserts
668    /// the result offset/size dtypes preserve signedness (widened to >=32-bit for offsets).
669    #[test]
670    fn test_rebuild_preserves_signed_offset_and_size_types() -> VortexResult<()> {
671        use crate::dtype::PType;
672
673        // Overlapping lists force an actual rebuild rather than the zero-copy fast path.
674        let elements = PrimitiveArray::from_iter(vec![1i32, 2, 3]).into_array();
675        let offsets = PrimitiveArray::from_iter(vec![0i32, 1]).into_array();
676        let sizes = PrimitiveArray::from_iter(vec![3i16, 2]).into_array();
677
678        let listview = ListViewArray::new(elements, offsets, sizes, Validity::NonNullable);
679        let mut ctx = SESSION.create_execution_ctx();
680        let rebuilt = listview.rebuild(ListViewRebuildMode::MakeZeroCopyToList, &mut ctx)?;
681
682        // Values: [1,2,3] and [2,3].
683        assert_arrays_eq!(
684            rebuilt.list_elements_at(0)?,
685            PrimitiveArray::from_iter([1i32, 2, 3]),
686            &mut ctx
687        );
688        assert_arrays_eq!(
689            rebuilt.list_elements_at(1)?,
690            PrimitiveArray::from_iter([2i32, 3]),
691            &mut ctx
692        );
693
694        // Signed input -> signed result (offsets widened to i32, sizes kept i16).
695        assert_eq!(rebuilt.offsets().dtype().as_ptype(), PType::I32);
696        assert_eq!(rebuilt.sizes().dtype().as_ptype(), PType::I16);
697        Ok(())
698    }
699
700    // ── should_use_take heuristic tests ────────────────────────────────────
701
702    #[test]
703    fn heuristic_zero_lists_uses_take() {
704        assert!(ListViewArray::should_use_take(0, 0));
705    }
706
707    #[test]
708    fn heuristic_small_lists_use_take() {
709        // avg = 127 → take
710        assert!(ListViewArray::should_use_take(127_000, 1_000));
711        // avg = 128 → LBL
712        assert!(!ListViewArray::should_use_take(128_000, 1_000));
713    }
714
715    /// Regression test for <https://github.com/vortex-data/vortex/issues/6973>.
716    /// Both offsets and sizes are u8, and offset + size exceeds u8::MAX.
717    #[test]
718    fn test_rebuild_trim_elements_sum_overflows_type() -> VortexResult<()> {
719        let elements = PrimitiveArray::from_iter(vec![0i32; 261]).into_array();
720        let offsets = PrimitiveArray::from_iter(vec![215u8, 0]).into_array();
721        let sizes = PrimitiveArray::from_iter(vec![46u8, 10]).into_array();
722
723        let listview = ListViewArray::new(elements, offsets, sizes, Validity::NonNullable);
724        let mut ctx = SESSION.create_execution_ctx();
725        let trimmed = listview.rebuild(ListViewRebuildMode::TrimElements, &mut ctx)?;
726
727        // min(offsets) = 0, so nothing to trim; output should equal input.
728        assert_arrays_eq!(trimmed, listview, &mut ctx);
729        Ok(())
730    }
731}