1use num_traits::FromPrimitive;
5use vortex_buffer::BufferMut;
6use vortex_error::VortexExpect;
7use vortex_error::VortexResult;
8
9use crate::IntoArray;
10use crate::LEGACY_SESSION;
11use crate::ToCanonical;
12use crate::VortexSessionExecute;
13use crate::aggregate_fn::fns::min_max::min_max;
14use crate::arrays::ConstantArray;
15use crate::arrays::ListViewArray;
16use crate::arrays::listview::ListViewArrayExt;
17use crate::builders::builder_with_capacity;
18use crate::builtins::ArrayBuiltins;
19use crate::dtype::DType;
20use crate::dtype::IntegerPType;
21use crate::dtype::Nullability;
22use crate::dtype::PType;
23use crate::match_each_integer_ptype;
24use crate::scalar::Scalar;
25use crate::scalar_fn::fns::operators::Operator;
26
27pub enum ListViewRebuildMode {
29 MakeZeroCopyToList,
37
38 TrimElements,
41
42 MakeExact,
49
50 OverlapCompression,
57}
58
59impl ListViewArray {
60 pub fn rebuild(&self, mode: ListViewRebuildMode) -> VortexResult<ListViewArray> {
62 if self.is_empty() {
63 return Ok(unsafe { self.clone().with_zero_copy_to_list(true) });
65 }
66
67 match mode {
68 ListViewRebuildMode::MakeZeroCopyToList => self.rebuild_zero_copy_to_list(),
69 ListViewRebuildMode::TrimElements => self.rebuild_trim_elements(),
70 ListViewRebuildMode::MakeExact => self.rebuild_make_exact(),
71 ListViewRebuildMode::OverlapCompression => unimplemented!("Does P=NP?"),
72 }
73 }
74
75 fn rebuild_zero_copy_to_list(&self) -> VortexResult<ListViewArray> {
83 if self.is_zero_copy_to_list() {
84 return Ok(self.clone());
86 }
87
88 let offsets_ptype = self.offsets().dtype().as_ptype();
89 let sizes_ptype = self.sizes().dtype().as_ptype();
90
91 match_each_integer_ptype!(sizes_ptype, |S| {
99 match offsets_ptype {
100 PType::U8 => self.naive_rebuild::<u8, u32, S>(),
101 PType::U16 => self.naive_rebuild::<u16, u32, S>(),
102 PType::U32 => self.naive_rebuild::<u32, u32, S>(),
103 PType::U64 => self.naive_rebuild::<u64, u64, S>(),
104 PType::I8 => self.naive_rebuild::<i8, i32, S>(),
105 PType::I16 => self.naive_rebuild::<i16, i32, S>(),
106 PType::I32 => self.naive_rebuild::<i32, i32, S>(),
107 PType::I64 => self.naive_rebuild::<i64, i64, S>(),
108 _ => unreachable!("invalid offsets PType"),
109 }
110 })
111 }
112
113 fn naive_rebuild<O: IntegerPType, NewOffset: IntegerPType, S: IntegerPType>(
117 &self,
118 ) -> VortexResult<ListViewArray> {
119 let sizes_canonical = self.sizes().to_primitive();
120 let total: u64 = sizes_canonical
121 .as_slice::<S>()
122 .iter()
123 .map(|s| (*s).as_() as u64)
124 .sum();
125 if Self::should_use_take(total, self.len()) {
126 self.rebuild_with_take::<O, NewOffset, S>()
127 } else {
128 self.rebuild_list_by_list::<O, NewOffset, S>()
129 }
130 }
131
132 fn should_use_take(total_output_elements: u64, num_lists: usize) -> bool {
140 if num_lists == 0 {
141 return true;
142 }
143 let avg = total_output_elements / num_lists as u64;
144 avg < 128
145 }
146
147 fn rebuild_with_take<O: IntegerPType, NewOffset: IntegerPType, S: IntegerPType>(
150 &self,
151 ) -> VortexResult<ListViewArray> {
152 let offsets_canonical = self.offsets().to_primitive();
153 let offsets_slice = offsets_canonical.as_slice::<O>();
154 let sizes_canonical = self.sizes().to_primitive();
155 let sizes_slice = sizes_canonical.as_slice::<S>();
156
157 let len = offsets_slice.len();
158
159 let mut new_offsets = BufferMut::<NewOffset>::with_capacity(len);
160 let mut new_sizes = BufferMut::<S>::with_capacity(len);
161 let mut take_indices = BufferMut::<u64>::with_capacity(self.elements().len());
162
163 let mut n_elements = NewOffset::zero();
164 for index in 0..len {
165 if !self.validity()?.is_valid(index)? {
166 new_offsets.push(n_elements);
167 new_sizes.push(S::zero());
168 continue;
169 }
170
171 let offset = offsets_slice[index];
172 let size = sizes_slice[index];
173 let start = offset.as_();
174 let stop = start + size.as_();
175
176 new_offsets.push(n_elements);
177 new_sizes.push(size);
178 take_indices.extend(start as u64..stop as u64);
179 n_elements += num_traits::cast(size).vortex_expect("Cast failed");
180 }
181
182 let elements = self.elements().take(take_indices.into_array())?;
183 let offsets = new_offsets.into_array();
184 let sizes = new_sizes.into_array();
185
186 Ok(unsafe {
190 ListViewArray::new_unchecked(elements, offsets, sizes, self.validity()?)
191 .with_zero_copy_to_list(true)
192 })
193 }
194
195 fn rebuild_list_by_list<O: IntegerPType, NewOffset: IntegerPType, S: IntegerPType>(
198 &self,
199 ) -> VortexResult<ListViewArray> {
200 let element_dtype = self
201 .dtype()
202 .as_list_element_opt()
203 .vortex_expect("somehow had a canonical list that was not a list");
204
205 let offsets_canonical = self.offsets().to_primitive();
206 let offsets_slice = offsets_canonical.as_slice::<O>();
207 let sizes_canonical = self.sizes().to_primitive();
208 let sizes_slice = sizes_canonical.as_slice::<S>();
209
210 let len = offsets_slice.len();
211
212 let mut new_offsets = BufferMut::<NewOffset>::with_capacity(len);
213 let mut new_sizes = BufferMut::<S>::with_capacity(len);
218
219 let elements_canonical = self
221 .elements()
222 .to_canonical()
223 .vortex_expect("canonicalize elements for rebuild")
224 .into_array();
225
226 let mut new_elements_builder =
229 builder_with_capacity(element_dtype.as_ref(), self.elements().len());
230
231 let mut n_elements = NewOffset::zero();
232 for index in 0..len {
233 if !self.validity()?.is_valid(index)? {
234 new_offsets.push(n_elements);
237 new_sizes.push(S::zero());
238 continue;
239 }
240
241 let offset = offsets_slice[index];
242 let size = sizes_slice[index];
243
244 let start = offset.as_();
245 let stop = start + size.as_();
246
247 new_offsets.push(n_elements);
248 new_sizes.push(size);
249 new_elements_builder.extend_from_array(&elements_canonical.slice(start..stop)?);
250
251 n_elements += num_traits::cast(size).vortex_expect("Cast failed");
252 }
253
254 let offsets = new_offsets.into_array();
255 let sizes = new_sizes.into_array();
256 let elements = new_elements_builder.finish();
257
258 debug_assert_eq!(
259 n_elements.as_(),
260 elements.len(),
261 "The accumulated elements somehow had the wrong length"
262 );
263
264 Ok(unsafe {
273 ListViewArray::new_unchecked(elements, offsets, sizes, self.validity()?)
274 .with_zero_copy_to_list(true)
275 })
276 }
277
278 fn rebuild_trim_elements(&self) -> VortexResult<ListViewArray> {
282 let start = if self.is_zero_copy_to_list() {
283 self.offset_at(0)
287 } else {
288 self.offsets().statistics().compute_min().vortex_expect(
289 "[ListViewArray::rebuild]: `offsets` must report min statistic that is a `usize`",
290 )
291 };
292
293 let end = if self.is_zero_copy_to_list() {
294 let last_offset = self.offset_at(self.len() - 1);
297 let last_size = self.size_at(self.len() - 1);
298 last_offset + last_size
299 } else {
300 let wide_dtype = DType::from(if self.offsets().dtype().as_ptype().is_unsigned_int() {
304 PType::U64
305 } else {
306 PType::I64
307 });
308 let offsets = self.offsets().cast(wide_dtype.clone())?;
309 let sizes = self.sizes().cast(wide_dtype)?;
310
311 let mut ctx = LEGACY_SESSION.create_execution_ctx();
312 let min_max = min_max(
313 &offsets
314 .binary(sizes, Operator::Add)
315 .vortex_expect("`offsets + sizes` somehow overflowed"),
316 &mut ctx,
317 )
318 .vortex_expect("Something went wrong while computing min and max")
319 .vortex_expect("We checked that the array was not empty in the top-level `rebuild`");
320
321 min_max
322 .max
323 .as_primitive()
324 .as_::<usize>()
325 .vortex_expect("unable to interpret the max `offset + size` as a `usize`")
326 };
327
328 let adjusted_offsets = match_each_integer_ptype!(self.offsets().dtype().as_ptype(), |O| {
329 let offset = <O as FromPrimitive>::from_usize(start)
330 .vortex_expect("unable to convert the min offset `start` into a `usize`");
331 let scalar = Scalar::primitive(offset, Nullability::NonNullable);
332
333 self.offsets()
334 .clone()
335 .binary(
336 ConstantArray::new(scalar, self.offsets().len()).into_array(),
337 Operator::Sub,
338 )
339 .vortex_expect("was somehow unable to adjust offsets down by their minimum")
340 });
341
342 let sliced_elements = self.elements().slice(start..end)?;
343
344 Ok(unsafe {
349 ListViewArray::new_unchecked(
350 sliced_elements,
351 adjusted_offsets,
352 self.sizes().clone(),
353 self.validity()?,
354 )
355 .with_zero_copy_to_list(self.is_zero_copy_to_list())
356 })
357 }
358
359 fn rebuild_make_exact(&self) -> VortexResult<ListViewArray> {
360 if self.is_zero_copy_to_list() {
361 self.rebuild_trim_elements()
362 } else {
363 self.rebuild_zero_copy_to_list()
366 }
367 }
368}
369
370#[cfg(test)]
371#[allow(clippy::cast_possible_truncation)]
372mod tests {
373 use vortex_buffer::BitBuffer;
374 use vortex_error::VortexResult;
375
376 use super::ListViewRebuildMode;
377 use crate::IntoArray;
378 use crate::ToCanonical;
379 use crate::arrays::ListViewArray;
380 use crate::arrays::PrimitiveArray;
381 use crate::arrays::listview::ListViewArrayExt;
382 use crate::assert_arrays_eq;
383 use crate::dtype::Nullability;
384 use crate::validity::Validity;
385
386 #[test]
387 fn test_rebuild_flatten_removes_overlaps() -> VortexResult<()> {
388 let elements = PrimitiveArray::from_iter(vec![1i32, 2, 3]).into_array();
392 let offsets = PrimitiveArray::from_iter(vec![0u32, 1]).into_array();
393 let sizes = PrimitiveArray::from_iter(vec![3u32, 2]).into_array();
394
395 let listview = ListViewArray::new(elements, offsets, sizes, Validity::NonNullable);
396
397 let flattened = listview.rebuild(ListViewRebuildMode::MakeZeroCopyToList)?;
398
399 assert_eq!(flattened.elements().len(), 5);
402
403 assert_eq!(flattened.offset_at(0), 0);
405 assert_eq!(flattened.size_at(0), 3);
406 assert_eq!(flattened.offset_at(1), 3);
407 assert_eq!(flattened.size_at(1), 2);
408
409 assert_arrays_eq!(
411 flattened.list_elements_at(0).unwrap(),
412 PrimitiveArray::from_iter([1i32, 2, 3])
413 );
414
415 assert_arrays_eq!(
416 flattened.list_elements_at(1).unwrap(),
417 PrimitiveArray::from_iter([2i32, 3])
418 );
419 Ok(())
420 }
421
422 #[test]
423 fn test_rebuild_flatten_with_nullable() -> VortexResult<()> {
424 use crate::arrays::BoolArray;
425
426 let elements = PrimitiveArray::from_iter(vec![1i32, 2, 3]).into_array();
428 let offsets = PrimitiveArray::from_iter(vec![0u32, 1, 2]).into_array();
429 let sizes = PrimitiveArray::from_iter(vec![2u32, 1, 1]).into_array();
430 let validity = Validity::Array(
431 BoolArray::new(
432 BitBuffer::from(vec![true, false, true]),
433 Validity::NonNullable,
434 )
435 .into_array(),
436 );
437
438 let listview = ListViewArray::new(elements, offsets, sizes, validity);
439
440 let flattened = listview.rebuild(ListViewRebuildMode::MakeZeroCopyToList)?;
441
442 assert_eq!(flattened.dtype().nullability(), Nullability::Nullable);
444 assert!(flattened.validity()?.is_valid(0).unwrap());
445 assert!(!flattened.validity()?.is_valid(1).unwrap());
446 assert!(flattened.validity()?.is_valid(2).unwrap());
447
448 assert_arrays_eq!(
450 flattened.list_elements_at(0).unwrap(),
451 PrimitiveArray::from_iter([1i32, 2])
452 );
453
454 assert_arrays_eq!(
455 flattened.list_elements_at(2).unwrap(),
456 PrimitiveArray::from_iter([3i32])
457 );
458 Ok(())
459 }
460
461 #[test]
462 fn test_rebuild_trim_elements_basic() -> VortexResult<()> {
463 let elements =
471 PrimitiveArray::from_iter(vec![99i32, 98, 1, 2, 97, 3, 4, 96, 95]).into_array();
472 let offsets = PrimitiveArray::from_iter(vec![2u32, 5]).into_array();
473 let sizes = PrimitiveArray::from_iter(vec![2u32, 2]).into_array();
474
475 let listview = ListViewArray::new(elements, offsets, sizes, Validity::NonNullable);
476
477 let trimmed = listview.rebuild(ListViewRebuildMode::TrimElements)?;
478
479 assert_eq!(trimmed.elements().len(), 5);
481
482 assert_eq!(trimmed.offset_at(0), 0);
484 assert_eq!(trimmed.size_at(0), 2);
485 assert_eq!(trimmed.offset_at(1), 3);
486 assert_eq!(trimmed.size_at(1), 2);
487
488 assert_arrays_eq!(
490 trimmed.list_elements_at(0).unwrap(),
491 PrimitiveArray::from_iter([1i32, 2])
492 );
493
494 assert_arrays_eq!(
495 trimmed.list_elements_at(1).unwrap(),
496 PrimitiveArray::from_iter([3i32, 4])
497 );
498
499 let all_elements = trimmed.elements().to_primitive();
501 assert_eq!(all_elements.scalar_at(2).unwrap(), 97i32.into());
502 Ok(())
503 }
504
505 #[test]
506 fn test_rebuild_with_trailing_nulls_regression() -> VortexResult<()> {
507 let elements = PrimitiveArray::from_iter(vec![1i32, 2, 3, 4]).into_array();
513 let offsets = PrimitiveArray::from_iter(vec![0u32, 2, 0, 0]).into_array();
514 let sizes = PrimitiveArray::from_iter(vec![2u32, 2, 0, 0]).into_array();
515 let validity = Validity::from_iter(vec![true, true, false, false]);
516
517 let listview = ListViewArray::new(elements, offsets, sizes, validity);
518
519 let rebuilt = listview.rebuild(ListViewRebuildMode::MakeZeroCopyToList)?;
521 assert!(rebuilt.is_zero_copy_to_list());
522
523 assert_eq!(rebuilt.offset_at(0), 0);
526 assert_eq!(rebuilt.offset_at(1), 2);
527 assert_eq!(rebuilt.offset_at(2), 4); assert_eq!(rebuilt.offset_at(3), 4); assert_eq!(rebuilt.size_at(0), 2);
532 assert_eq!(rebuilt.size_at(1), 2);
533 assert_eq!(rebuilt.size_at(2), 0); assert_eq!(rebuilt.size_at(3), 0); let exact = rebuilt.rebuild(ListViewRebuildMode::MakeExact)?;
539
540 assert!(exact.is_valid(0).unwrap());
542 assert!(exact.is_valid(1).unwrap());
543 assert!(!exact.is_valid(2).unwrap());
544 assert!(!exact.is_valid(3).unwrap());
545
546 assert_arrays_eq!(
548 exact.list_elements_at(0).unwrap(),
549 PrimitiveArray::from_iter([1i32, 2])
550 );
551
552 assert_arrays_eq!(
553 exact.list_elements_at(1).unwrap(),
554 PrimitiveArray::from_iter([3i32, 4])
555 );
556 Ok(())
557 }
558
559 #[test]
562 fn test_rebuild_trim_elements_offsets_wider_than_sizes() -> VortexResult<()> {
563 let mut elems = vec![0i32; 70_005];
564 elems[70_000] = 10;
565 elems[70_001] = 20;
566 elems[70_002] = 30;
567 elems[70_003] = 40;
568 let elements = PrimitiveArray::from_iter(elems).into_array();
569 let offsets = PrimitiveArray::from_iter(vec![70_000u32, 70_002]).into_array();
570 let sizes = PrimitiveArray::from_iter(vec![2u16, 2]).into_array();
571
572 let listview = ListViewArray::new(elements, offsets, sizes, Validity::NonNullable);
573 let trimmed = listview.rebuild(ListViewRebuildMode::TrimElements)?;
574 assert_arrays_eq!(
575 trimmed.list_elements_at(1).unwrap(),
576 PrimitiveArray::from_iter([30i32, 40])
577 );
578 Ok(())
579 }
580
581 #[test]
584 fn test_rebuild_trim_elements_sizes_wider_than_offsets() -> VortexResult<()> {
585 let mut elems = vec![0i32; 70_001];
586 elems[3] = 30;
587 elems[4] = 40;
588 let elements = PrimitiveArray::from_iter(elems).into_array();
589 let offsets = PrimitiveArray::from_iter(vec![1u16, 3]).into_array();
590 let sizes = PrimitiveArray::from_iter(vec![70_000u32, 2]).into_array();
591
592 let listview = ListViewArray::new(elements, offsets, sizes, Validity::NonNullable);
593 let trimmed = listview.rebuild(ListViewRebuildMode::TrimElements)?;
594 assert_arrays_eq!(
595 trimmed.list_elements_at(1).unwrap(),
596 PrimitiveArray::from_iter([30i32, 40])
597 );
598 Ok(())
599 }
600
601 #[test]
604 fn heuristic_zero_lists_uses_take() {
605 assert!(ListViewArray::should_use_take(0, 0));
606 }
607
608 #[test]
609 fn heuristic_small_lists_use_take() {
610 assert!(ListViewArray::should_use_take(127_000, 1_000));
612 assert!(!ListViewArray::should_use_take(128_000, 1_000));
614 }
615
616 #[test]
619 fn test_rebuild_trim_elements_sum_overflows_type() -> VortexResult<()> {
620 let elements = PrimitiveArray::from_iter(vec![0i32; 261]).into_array();
621 let offsets = PrimitiveArray::from_iter(vec![215u8, 0]).into_array();
622 let sizes = PrimitiveArray::from_iter(vec![46u8, 10]).into_array();
623
624 let listview = ListViewArray::new(elements, offsets, sizes, Validity::NonNullable);
625 let trimmed = listview.rebuild(ListViewRebuildMode::TrimElements)?;
626
627 assert_arrays_eq!(trimmed, listview);
629 Ok(())
630 }
631}