1use std::sync::Arc;
5
6use vortex_error::VortexExpect;
7use vortex_error::VortexResult;
8
9use crate::ArrayRef;
10use crate::Canonical;
11use crate::DynArray;
12use crate::ExecutionCtx;
13use crate::IntoArray;
14use crate::ToCanonical;
15use crate::arrays::ExtensionArray;
16use crate::arrays::FixedSizeListArray;
17use crate::arrays::ListArray;
18use crate::arrays::ListViewArray;
19use crate::arrays::PrimitiveArray;
20use crate::arrays::StructArray;
21use crate::arrays::listview::ListViewRebuildMode;
22use crate::builders::PrimitiveBuilder;
23use crate::dtype::IntegerPType;
24use crate::dtype::Nullability;
25use crate::match_each_integer_ptype;
26use crate::vtable::ValidityHelper;
27
28pub fn list_view_from_list(list: ListArray, ctx: &mut ExecutionCtx) -> VortexResult<ListViewArray> {
33 if list.is_empty() {
36 return Ok(Canonical::empty(list.dtype()).into_listview());
37 }
38
39 let list = list.reset_offsets(false).vortex_expect("This can't fail");
43
44 let list_offsets = list.offsets().clone();
45
46 let sizes = match_each_integer_ptype!(list_offsets.dtype().as_ptype(), |O| {
49 build_sizes_from_offsets::<O>(&list, ctx)?
50 });
51
52 debug_assert_eq!(list_offsets.len(), list.len() + 1);
54 let adjusted_offsets = list_offsets.slice(0..list.len())?;
55
56 Ok(unsafe {
60 ListViewArray::new_unchecked(
61 list.elements().clone(),
62 adjusted_offsets,
63 sizes,
64 list.validity().clone(),
65 )
66 .with_zero_copy_to_list(true)
67 })
68}
69
70fn build_sizes_from_offsets<O: IntegerPType>(
72 list: &ListArray,
73 ctx: &mut ExecutionCtx,
74) -> VortexResult<ArrayRef> {
75 let len = list.len();
76 let mut sizes_builder = PrimitiveBuilder::<O>::with_capacity(Nullability::NonNullable, len);
77
78 let mut sizes_range = sizes_builder.uninit_range(len);
80
81 let offsets = list.offsets().clone().execute::<PrimitiveArray>(ctx)?;
82 let offsets_slice = offsets.as_slice::<O>();
83 debug_assert_eq!(len + 1, offsets_slice.len());
84 debug_assert!(offsets_slice.is_sorted());
85
86 for i in 0..len {
88 let size = offsets_slice[i + 1] - offsets_slice[i];
89 sizes_range.set_value(i, size);
90 }
91
92 unsafe {
94 sizes_range.finish();
95 }
96
97 Ok(sizes_builder.finish_into_primitive().into_array())
98}
99
100pub fn list_from_list_view(list_view: ListViewArray) -> VortexResult<ListArray> {
110 let zctl_array = list_view.rebuild(ListViewRebuildMode::MakeExact)?;
112 debug_assert!(zctl_array.is_zero_copy_to_list());
113
114 let list_offsets = match_each_integer_ptype!(zctl_array.offsets().dtype().as_ptype(), |O| {
115 unsafe { build_list_offsets_from_list_view::<O>(&zctl_array) }
118 });
119
120 Ok(unsafe {
124 ListArray::new_unchecked(
125 zctl_array.elements().clone(),
126 list_offsets,
127 zctl_array.validity().clone(),
128 )
129 })
130}
131
132unsafe fn build_list_offsets_from_list_view<O: IntegerPType>(
142 list_view: &ListViewArray,
143) -> ArrayRef {
144 let len = list_view.len();
145 let mut offsets_builder =
146 PrimitiveBuilder::<O>::with_capacity(Nullability::NonNullable, len + 1);
147
148 let mut offsets_range = offsets_builder.uninit_range(len + 1);
150
151 let offsets = list_view.offsets().to_primitive();
152 let offsets_slice = offsets.as_slice::<O>();
153 debug_assert!(offsets_slice.is_sorted());
154
155 offsets_range.copy_from_slice(0, offsets_slice);
157
158 let final_offset = if len != 0 {
160 let last_offset = offsets_slice[len - 1];
161
162 let last_size = list_view.size_at(len - 1);
163 let last_size =
164 O::from_usize(last_size).vortex_expect("size somehow did not fit into offsets");
165
166 last_offset + last_size
167 } else {
168 O::zero()
169 };
170
171 offsets_range.set_value(len, final_offset);
172
173 unsafe {
175 offsets_range.finish();
176 }
177
178 offsets_builder.finish_into_primitive().into_array()
179}
180
181pub fn recursive_list_from_list_view(array: ArrayRef) -> VortexResult<ArrayRef> {
185 if !array.dtype().is_nested() {
186 return Ok(array);
187 }
188
189 let canonical = array.to_canonical()?;
190
191 Ok(match canonical {
192 Canonical::List(listview) => {
193 let converted_elements = recursive_list_from_list_view(listview.elements().clone())?;
194 debug_assert_eq!(converted_elements.len(), listview.elements().len());
195
196 let listview_with_converted_elements =
198 if !Arc::ptr_eq(&converted_elements, listview.elements()) {
199 unsafe {
202 ListViewArray::new_unchecked(
203 converted_elements,
204 listview.offsets().clone(),
205 listview.sizes().clone(),
206 listview.validity().clone(),
207 )
208 .with_zero_copy_to_list(listview.is_zero_copy_to_list())
209 }
210 } else {
211 listview
212 };
213
214 let list_array = list_from_list_view(listview_with_converted_elements)?;
216 list_array.into_array()
217 }
218 Canonical::FixedSizeList(fixed_size_list) => {
219 let converted_elements =
220 recursive_list_from_list_view(fixed_size_list.elements().clone())?;
221
222 if !Arc::ptr_eq(&converted_elements, fixed_size_list.elements()) {
224 FixedSizeListArray::try_new(
225 converted_elements,
226 fixed_size_list.list_size(),
227 fixed_size_list.validity().clone(),
228 fixed_size_list.len(),
229 )
230 .vortex_expect(
231 "FixedSizeListArray reconstruction should not fail with valid components",
232 )
233 .into_array()
234 } else {
235 fixed_size_list.into_array()
236 }
237 }
238 Canonical::Struct(struct_array) => {
239 let fields = struct_array.unmasked_fields();
240 let mut converted_fields = Vec::with_capacity(fields.len());
241 let mut any_changed = false;
242
243 for field in fields.iter() {
244 let converted_field = recursive_list_from_list_view(field.clone())?;
245 any_changed |= !Arc::ptr_eq(&converted_field, field);
247 converted_fields.push(converted_field);
248 }
249
250 if any_changed {
251 StructArray::try_new(
252 struct_array.names().clone(),
253 converted_fields,
254 struct_array.len(),
255 struct_array.validity().clone(),
256 )
257 .vortex_expect("StructArray reconstruction should not fail with valid components")
258 .into_array()
259 } else {
260 struct_array.into_array()
261 }
262 }
263 Canonical::Extension(ext_array) => {
264 let converted_storage = recursive_list_from_list_view(ext_array.storage().clone())?;
265
266 if !Arc::ptr_eq(&converted_storage, ext_array.storage()) {
268 ExtensionArray::new(ext_array.ext_dtype().clone(), converted_storage).into_array()
269 } else {
270 ext_array.into_array()
271 }
272 }
273 _ => unreachable!(),
274 })
275}
276
277#[cfg(test)]
278mod tests {
279 use std::sync::Arc;
280
281 use vortex_buffer::buffer;
282 use vortex_error::VortexResult;
283
284 use super::super::tests::common::create_basic_listview;
285 use super::super::tests::common::create_empty_lists_listview;
286 use super::super::tests::common::create_nullable_listview;
287 use super::super::tests::common::create_overlapping_listview;
288 use super::recursive_list_from_list_view;
289 use crate::IntoArray;
290 use crate::LEGACY_SESSION;
291 use crate::VortexSessionExecute;
292 use crate::arrays::BoolArray;
293 use crate::arrays::FixedSizeListArray;
294 use crate::arrays::ListArray;
295 use crate::arrays::ListViewArray;
296 use crate::arrays::PrimitiveArray;
297 use crate::arrays::StructArray;
298 use crate::arrays::listview::list_from_list_view;
299 use crate::arrays::listview::list_view_from_list;
300 use crate::assert_arrays_eq;
301 use crate::dtype::FieldNames;
302 use crate::validity::Validity;
303 use crate::vtable::ValidityHelper;
304
305 #[test]
306 fn test_list_to_listview_basic() -> VortexResult<()> {
307 let elements = buffer![0i32, 1, 2, 3, 4, 5, 6, 7, 8, 9].into_array();
309 let offsets = buffer![0u32, 3, 5, 7, 10].into_array();
310 let list_array =
311 ListArray::try_new(elements.clone(), offsets.clone(), Validity::NonNullable)?;
312
313 let mut ctx = LEGACY_SESSION.create_execution_ctx();
314 let list_view = list_view_from_list(list_array.clone(), &mut ctx)?;
315
316 assert_eq!(list_view.len(), 4);
318 assert_arrays_eq!(elements, list_view.elements().clone());
319
320 let expected_offsets = buffer![0u32, 3, 5, 7].into_array();
322 assert_arrays_eq!(expected_offsets, list_view.offsets().clone());
323
324 let expected_sizes = buffer![3u32, 2, 2, 3].into_array();
326 assert_arrays_eq!(expected_sizes, list_view.sizes().clone());
327
328 assert_arrays_eq!(list_array, list_view);
330 Ok(())
331 }
332
333 #[test]
334 fn test_listview_to_list_zero_copy() -> VortexResult<()> {
335 let list_view = create_basic_listview();
336 let list_array = list_from_list_view(list_view.clone())?;
337
338 assert_arrays_eq!(list_view.elements().clone(), list_array.elements().clone());
340
341 let list_array_offsets_without_last = list_array.offsets().slice(0..list_view.len())?;
344 assert_arrays_eq!(list_view.offsets().clone(), list_array_offsets_without_last);
345
346 assert_arrays_eq!(list_view, list_array);
348 Ok(())
349 }
350
351 #[test]
352 fn test_empty_array_conversions() -> VortexResult<()> {
353 let empty_elements = PrimitiveArray::from_iter::<[i32; 0]>([]).into_array();
355 let empty_offsets = buffer![0u32].into_array();
356 let empty_list =
357 ListArray::try_new(empty_elements.clone(), empty_offsets, Validity::NonNullable)?;
358
359 let mut ctx = LEGACY_SESSION.create_execution_ctx();
362 let empty_list_view = list_view_from_list(empty_list.clone(), &mut ctx)?;
363 assert_eq!(empty_list_view.len(), 0);
364
365 let converted_back = list_from_list_view(empty_list_view)?;
367 assert_eq!(converted_back.len(), 0);
368 assert_eq!(empty_list.len(), converted_back.len());
371 Ok(())
372 }
373
374 #[test]
375 fn test_nullable_conversions() -> VortexResult<()> {
376 let elements = buffer![10i32, 20, 30, 40, 50].into_array();
378 let offsets = buffer![0u32, 2, 4, 5].into_array();
379 let validity = Validity::Array(BoolArray::from_iter(vec![true, false, true]).into_array());
380 let nullable_list =
381 ListArray::try_new(elements.clone(), offsets.clone(), validity.clone())?;
382
383 let mut ctx = LEGACY_SESSION.create_execution_ctx();
384 let nullable_list_view = list_view_from_list(nullable_list.clone(), &mut ctx)?;
385
386 assert_eq!(nullable_list_view.validity(), &validity);
388 assert_eq!(nullable_list_view.len(), 3);
389
390 let converted_back = list_from_list_view(nullable_list_view)?;
392 assert_arrays_eq!(nullable_list, converted_back);
393 Ok(())
394 }
395
396 #[test]
397 fn test_non_zero_copy_listview_to_list() -> VortexResult<()> {
398 let list_view = create_overlapping_listview();
400 let list_array = list_from_list_view(list_view.clone())?;
401
402 for i in 0..list_array.len() {
404 let start = list_array.offset_at(i)?;
405 let end = list_array.offset_at(i + 1)?;
406 assert!(end >= start, "Offsets should be monotonic after conversion");
407 }
408
409 assert_arrays_eq!(list_view, list_array);
411 Ok(())
412 }
413
414 #[test]
415 fn test_empty_sublists() -> VortexResult<()> {
416 let empty_lists_view = create_empty_lists_listview();
417
418 let list_array = list_from_list_view(empty_lists_view.clone())?;
420 assert_eq!(list_array.len(), 4);
421
422 for i in 0..list_array.len() {
424 assert_eq!(list_array.list_elements_at(i)?.len(), 0);
425 }
426
427 let mut ctx = LEGACY_SESSION.create_execution_ctx();
429 let converted_back = list_view_from_list(list_array, &mut ctx)?;
430 assert_arrays_eq!(empty_lists_view, converted_back);
431 Ok(())
432 }
433
434 #[test]
435 fn test_different_offset_types() -> VortexResult<()> {
436 let elements = buffer![1i32, 2, 3, 4, 5].into_array();
438 let i32_offsets = buffer![0i32, 2, 5].into_array();
439 let list_i32 =
440 ListArray::try_new(elements.clone(), i32_offsets.clone(), Validity::NonNullable)?;
441
442 let mut ctx = LEGACY_SESSION.create_execution_ctx();
443 let list_view_i32 = list_view_from_list(list_i32.clone(), &mut ctx)?;
444 assert_eq!(list_view_i32.offsets().dtype(), i32_offsets.dtype());
445 assert_eq!(list_view_i32.sizes().dtype(), i32_offsets.dtype());
446
447 let i64_offsets = buffer![0i64, 2, 5].into_array();
449 let list_i64 =
450 ListArray::try_new(elements.clone(), i64_offsets.clone(), Validity::NonNullable)?;
451
452 let list_view_i64 = list_view_from_list(list_i64.clone(), &mut ctx)?;
453 assert_eq!(list_view_i64.offsets().dtype(), i64_offsets.dtype());
454 assert_eq!(list_view_i64.sizes().dtype(), i64_offsets.dtype());
455
456 assert_arrays_eq!(list_i32, list_view_i32);
458 assert_arrays_eq!(list_i64, list_view_i64);
459 Ok(())
460 }
461
462 #[test]
463 fn test_round_trip_conversions() -> VortexResult<()> {
464 let mut ctx = LEGACY_SESSION.create_execution_ctx();
465
466 let original = create_basic_listview();
468 let to_list = list_from_list_view(original.clone())?;
469 let back_to_view = list_view_from_list(to_list, &mut ctx)?;
470 assert_arrays_eq!(original, back_to_view);
471
472 let nullable = create_nullable_listview();
474 let nullable_to_list = list_from_list_view(nullable.clone())?;
475 let nullable_back = list_view_from_list(nullable_to_list, &mut ctx)?;
476 assert_arrays_eq!(nullable, nullable_back);
477
478 let overlapping = create_overlapping_listview();
480
481 let overlapping_to_list = list_from_list_view(overlapping.clone())?;
482 let overlapping_back = list_view_from_list(overlapping_to_list, &mut ctx)?;
483 assert_arrays_eq!(overlapping, overlapping_back);
484 Ok(())
485 }
486
487 #[test]
488 fn test_single_element_lists() -> VortexResult<()> {
489 let elements = buffer![100i32, 200, 300].into_array();
491 let offsets = buffer![0u32, 1, 2, 3].into_array();
492 let single_elem_list =
493 ListArray::try_new(elements.clone(), offsets, Validity::NonNullable)?;
494
495 let mut ctx = LEGACY_SESSION.create_execution_ctx();
496 let list_view = list_view_from_list(single_elem_list.clone(), &mut ctx)?;
497 assert_eq!(list_view.len(), 3);
498
499 let expected_sizes = buffer![1u32, 1, 1].into_array();
501 assert_arrays_eq!(expected_sizes, list_view.sizes().clone());
502
503 let converted_back = list_from_list_view(list_view)?;
505 assert_arrays_eq!(single_elem_list, converted_back);
506 Ok(())
507 }
508
509 #[test]
510 fn test_mixed_empty_and_non_empty_lists() -> VortexResult<()> {
511 let elements = buffer![1i32, 2, 3, 4, 5, 6].into_array();
513 let offsets = buffer![0u32, 2, 2, 3, 3, 6].into_array();
514 let mixed_list =
515 ListArray::try_new(elements.clone(), offsets.clone(), Validity::NonNullable)?;
516
517 let mut ctx = LEGACY_SESSION.create_execution_ctx();
518 let list_view = list_view_from_list(mixed_list.clone(), &mut ctx)?;
519 assert_eq!(list_view.len(), 5);
520
521 let expected_sizes = buffer![2u32, 0, 1, 0, 3].into_array();
523 assert_arrays_eq!(expected_sizes, list_view.sizes().clone());
524
525 let converted_back = list_from_list_view(list_view)?;
527 assert_arrays_eq!(mixed_list, converted_back);
528 Ok(())
529 }
530
531 #[test]
532 fn test_recursive_simple_listview() -> VortexResult<()> {
533 let list_view = create_basic_listview();
534 let result = recursive_list_from_list_view(list_view.clone().into_array())?;
535
536 assert_eq!(result.len(), list_view.len());
537 assert_arrays_eq!(list_view.into_array(), result);
538 Ok(())
539 }
540
541 #[test]
542 fn test_recursive_nested_listview() -> VortexResult<()> {
543 let inner_elements = buffer![1i32, 2, 3].into_array();
544 let inner_offsets = buffer![0u32, 2].into_array();
545 let inner_sizes = buffer![2u32, 1].into_array();
546 let inner_listview = unsafe {
547 ListViewArray::new_unchecked(
548 inner_elements,
549 inner_offsets,
550 inner_sizes,
551 Validity::NonNullable,
552 )
553 .with_zero_copy_to_list(true)
554 };
555
556 let outer_offsets = buffer![0u32, 1].into_array();
557 let outer_sizes = buffer![1u32, 1].into_array();
558 let outer_listview = unsafe {
559 ListViewArray::new_unchecked(
560 inner_listview.into_array(),
561 outer_offsets,
562 outer_sizes,
563 Validity::NonNullable,
564 )
565 .with_zero_copy_to_list(true)
566 };
567
568 let result = recursive_list_from_list_view(outer_listview.clone().into_array())?;
569
570 assert_eq!(result.len(), 2);
571 assert_arrays_eq!(outer_listview.into_array(), result);
572 Ok(())
573 }
574
575 #[test]
576 fn test_recursive_struct_with_listview_fields() -> VortexResult<()> {
577 let listview_field = create_basic_listview().into_array();
578 let primitive_field = buffer![10i32, 20, 30, 40].into_array();
579
580 let struct_array = StructArray::try_new(
581 FieldNames::from(["lists", "values"]),
582 vec![listview_field, primitive_field],
583 4,
584 Validity::NonNullable,
585 )?;
586
587 let result = recursive_list_from_list_view(struct_array.clone().into_array())?;
588
589 assert_eq!(result.len(), 4);
590 assert_arrays_eq!(struct_array.into_array(), result);
591 Ok(())
592 }
593
594 #[test]
595 fn test_recursive_fixed_size_list_with_listview_elements() -> VortexResult<()> {
596 let lv1_elements = buffer![1i32, 2].into_array();
597 let lv1_offsets = buffer![0u32].into_array();
598 let lv1_sizes = buffer![2u32].into_array();
599 let lv1 = unsafe {
600 ListViewArray::new_unchecked(
601 lv1_elements,
602 lv1_offsets,
603 lv1_sizes,
604 Validity::NonNullable,
605 )
606 .with_zero_copy_to_list(true)
607 };
608
609 let lv2_elements = buffer![3i32, 4].into_array();
610 let lv2_offsets = buffer![0u32].into_array();
611 let lv2_sizes = buffer![2u32].into_array();
612 let lv2 = unsafe {
613 ListViewArray::new_unchecked(
614 lv2_elements,
615 lv2_offsets,
616 lv2_sizes,
617 Validity::NonNullable,
618 )
619 .with_zero_copy_to_list(true)
620 };
621
622 let dtype = lv1.dtype().clone();
623 let chunked_listviews =
624 crate::arrays::ChunkedArray::try_new(vec![lv1.into_array(), lv2.into_array()], dtype)?;
625
626 let fixed_list =
627 FixedSizeListArray::new(chunked_listviews.into_array(), 1, Validity::NonNullable, 2);
628
629 let result = recursive_list_from_list_view(fixed_list.clone().into_array())?;
630
631 assert_eq!(result.len(), 2);
632 assert_arrays_eq!(fixed_list.into_array(), result);
633 Ok(())
634 }
635
636 #[test]
637 fn test_recursive_deep_nesting() -> VortexResult<()> {
638 let innermost_elements = buffer![1i32, 2, 3].into_array();
639 let innermost_offsets = buffer![0u32, 2].into_array();
640 let innermost_sizes = buffer![2u32, 1].into_array();
641 let innermost_listview = unsafe {
642 ListViewArray::new_unchecked(
643 innermost_elements,
644 innermost_offsets,
645 innermost_sizes,
646 Validity::NonNullable,
647 )
648 .with_zero_copy_to_list(true)
649 };
650
651 let struct_array = StructArray::try_new(
652 FieldNames::from(["inner_lists"]),
653 vec![innermost_listview.into_array()],
654 2,
655 Validity::NonNullable,
656 )?;
657
658 let outer_offsets = buffer![0u32, 1].into_array();
659 let outer_sizes = buffer![1u32, 1].into_array();
660 let outer_listview = unsafe {
661 ListViewArray::new_unchecked(
662 struct_array.into_array(),
663 outer_offsets,
664 outer_sizes,
665 Validity::NonNullable,
666 )
667 .with_zero_copy_to_list(true)
668 };
669
670 let result = recursive_list_from_list_view(outer_listview.clone().into_array())?;
671
672 assert_eq!(result.len(), 2);
673 assert_arrays_eq!(outer_listview.into_array(), result);
674 Ok(())
675 }
676
677 #[test]
678 fn test_recursive_primitive_unchanged() -> VortexResult<()> {
679 let prim = buffer![1i32, 2, 3].into_array();
680 let prim_clone = prim.clone();
681 let result = recursive_list_from_list_view(prim)?;
682
683 assert!(Arc::ptr_eq(&result, &prim_clone));
684 Ok(())
685 }
686
687 #[test]
688 fn test_recursive_mixed_listview_and_list() -> VortexResult<()> {
689 let listview = create_basic_listview();
690 let list = list_from_list_view(listview.clone())?;
691
692 let struct_array = StructArray::try_new(
693 FieldNames::from(["listview_field", "list_field"]),
694 vec![listview.into_array(), list.into_array()],
695 4,
696 Validity::NonNullable,
697 )?;
698
699 let result = recursive_list_from_list_view(struct_array.clone().into_array())?;
700
701 assert_eq!(result.len(), 4);
702 assert_arrays_eq!(struct_array.into_array(), result);
703 Ok(())
704 }
705}