1use 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
28fn 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
43pub const DEFAULT_REBUILD_DENSITY_THRESHOLD: f32 = 0.1;
52
53pub const DEFAULT_TRIM_ELEMENTS_THRESHOLD: f32 = 0.05;
62
63pub enum ListViewRebuildMode {
65 MakeZeroCopyToList,
73
74 TrimElements,
80
81 MakeExact,
88
89 OverlapCompression,
96}
97
98impl ListViewArray {
99 pub fn rebuild(
101 &self,
102 mode: ListViewRebuildMode,
103 ctx: &mut ExecutionCtx,
104 ) -> VortexResult<ListViewArray> {
105 if self.is_empty() {
106 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 fn rebuild_zero_copy_to_list(&self, ctx: &mut ExecutionCtx) -> VortexResult<ListViewArray> {
126 if self.is_zero_copy_to_list() {
127 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 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 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 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 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 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 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 Ok(unsafe {
250 ListViewArray::new_unchecked(elements, offsets, sizes, self.validity()?)
251 .with_zero_copy_to_list(true)
252 })
253 }
254
255 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 let mut new_sizes = BufferMut::<S>::with_capacity(len);
286
287 let elements_canonical = self
289 .elements()
290 .clone()
291 .execute::<Canonical>(ctx)?
292 .into_array();
293
294 let mut new_elements_builder =
297 builder_with_capacity(element_dtype.as_ref(), self.elements().len());
298
299 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 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 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 Ok(unsafe {
351 ListViewArray::new_unchecked(elements, offsets, sizes, self.validity()?)
352 .with_zero_copy_to_list(true)
353 })
354 }
355
356 fn rebuild_trim_elements(&self, ctx: &mut ExecutionCtx) -> VortexResult<ListViewArray> {
360 let (start, end) = self.referenced_element_bounds(ctx)?;
361
362 unsafe { self.trim_elements(start, end) }
364 }
365
366 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 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 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 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 assert_eq!(flattened.elements().len(), 5);
450
451 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 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 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 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 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 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 assert_eq!(trimmed.elements().len(), 5);
535
536 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 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 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 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 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 assert_eq!(rebuilt.offset_at(0), 0);
586 assert_eq!(rebuilt.offset_at(1), 2);
587 assert_eq!(rebuilt.offset_at(2), 4); assert_eq!(rebuilt.offset_at(3), 4); assert_eq!(rebuilt.size_at(0), 2);
592 assert_eq!(rebuilt.size_at(1), 2);
593 assert_eq!(rebuilt.size_at(2), 0); assert_eq!(rebuilt.size_at(3), 0); let exact = rebuilt.rebuild(ListViewRebuildMode::MakeExact, &mut ctx)?;
599
600 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 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 #[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 #[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 #[test]
670 fn test_rebuild_preserves_signed_offset_and_size_types() -> VortexResult<()> {
671 use crate::dtype::PType;
672
673 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 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 assert_eq!(rebuilt.offsets().dtype().as_ptype(), PType::I32);
696 assert_eq!(rebuilt.sizes().dtype().as_ptype(), PType::I16);
697 Ok(())
698 }
699
700 #[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 assert!(ListViewArray::should_use_take(127_000, 1_000));
711 assert!(!ListViewArray::should_use_take(128_000, 1_000));
713 }
714
715 #[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 assert_arrays_eq!(trimmed, listview, &mut ctx);
729 Ok(())
730 }
731}