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 =
265 recursive_list_from_list_view(ext_array.storage_array().clone())?;
266
267 if !Arc::ptr_eq(&converted_storage, ext_array.storage_array()) {
269 ExtensionArray::new(ext_array.ext_dtype().clone(), converted_storage).into_array()
270 } else {
271 ext_array.into_array()
272 }
273 }
274 _ => unreachable!(),
275 })
276}
277
278#[cfg(test)]
279mod tests {
280 use std::sync::Arc;
281
282 use vortex_buffer::buffer;
283 use vortex_error::VortexResult;
284
285 use super::super::tests::common::create_basic_listview;
286 use super::super::tests::common::create_empty_lists_listview;
287 use super::super::tests::common::create_nullable_listview;
288 use super::super::tests::common::create_overlapping_listview;
289 use super::recursive_list_from_list_view;
290 use crate::IntoArray;
291 use crate::LEGACY_SESSION;
292 use crate::VortexSessionExecute;
293 use crate::arrays::BoolArray;
294 use crate::arrays::FixedSizeListArray;
295 use crate::arrays::ListArray;
296 use crate::arrays::ListViewArray;
297 use crate::arrays::PrimitiveArray;
298 use crate::arrays::StructArray;
299 use crate::arrays::listview::list_from_list_view;
300 use crate::arrays::listview::list_view_from_list;
301 use crate::assert_arrays_eq;
302 use crate::dtype::FieldNames;
303 use crate::validity::Validity;
304 use crate::vtable::ValidityHelper;
305
306 #[test]
307 fn test_list_to_listview_basic() -> VortexResult<()> {
308 let elements = buffer![0i32, 1, 2, 3, 4, 5, 6, 7, 8, 9].into_array();
310 let offsets = buffer![0u32, 3, 5, 7, 10].into_array();
311 let list_array =
312 ListArray::try_new(elements.clone(), offsets.clone(), Validity::NonNullable)?;
313
314 let mut ctx = LEGACY_SESSION.create_execution_ctx();
315 let list_view = list_view_from_list(list_array.clone(), &mut ctx)?;
316
317 assert_eq!(list_view.len(), 4);
319 assert_arrays_eq!(elements, list_view.elements().clone());
320
321 let expected_offsets = buffer![0u32, 3, 5, 7].into_array();
323 assert_arrays_eq!(expected_offsets, list_view.offsets().clone());
324
325 let expected_sizes = buffer![3u32, 2, 2, 3].into_array();
327 assert_arrays_eq!(expected_sizes, list_view.sizes().clone());
328
329 assert_arrays_eq!(list_array, list_view);
331 Ok(())
332 }
333
334 #[test]
335 fn test_listview_to_list_zero_copy() -> VortexResult<()> {
336 let list_view = create_basic_listview();
337 let list_array = list_from_list_view(list_view.clone())?;
338
339 assert_arrays_eq!(list_view.elements().clone(), list_array.elements().clone());
341
342 let list_array_offsets_without_last = list_array.offsets().slice(0..list_view.len())?;
345 assert_arrays_eq!(list_view.offsets().clone(), list_array_offsets_without_last);
346
347 assert_arrays_eq!(list_view, list_array);
349 Ok(())
350 }
351
352 #[test]
353 fn test_empty_array_conversions() -> VortexResult<()> {
354 let empty_elements = PrimitiveArray::from_iter::<[i32; 0]>([]).into_array();
356 let empty_offsets = buffer![0u32].into_array();
357 let empty_list =
358 ListArray::try_new(empty_elements.clone(), empty_offsets, Validity::NonNullable)?;
359
360 let mut ctx = LEGACY_SESSION.create_execution_ctx();
363 let empty_list_view = list_view_from_list(empty_list.clone(), &mut ctx)?;
364 assert_eq!(empty_list_view.len(), 0);
365
366 let converted_back = list_from_list_view(empty_list_view)?;
368 assert_eq!(converted_back.len(), 0);
369 assert_eq!(empty_list.len(), converted_back.len());
372 Ok(())
373 }
374
375 #[test]
376 fn test_nullable_conversions() -> VortexResult<()> {
377 let elements = buffer![10i32, 20, 30, 40, 50].into_array();
379 let offsets = buffer![0u32, 2, 4, 5].into_array();
380 let validity = Validity::Array(BoolArray::from_iter(vec![true, false, true]).into_array());
381 let nullable_list =
382 ListArray::try_new(elements.clone(), offsets.clone(), validity.clone())?;
383
384 let mut ctx = LEGACY_SESSION.create_execution_ctx();
385 let nullable_list_view = list_view_from_list(nullable_list.clone(), &mut ctx)?;
386
387 assert_eq!(nullable_list_view.validity(), &validity);
389 assert_eq!(nullable_list_view.len(), 3);
390
391 let converted_back = list_from_list_view(nullable_list_view)?;
393 assert_arrays_eq!(nullable_list, converted_back);
394 Ok(())
395 }
396
397 #[test]
398 fn test_non_zero_copy_listview_to_list() -> VortexResult<()> {
399 let list_view = create_overlapping_listview();
401 let list_array = list_from_list_view(list_view.clone())?;
402
403 for i in 0..list_array.len() {
405 let start = list_array.offset_at(i)?;
406 let end = list_array.offset_at(i + 1)?;
407 assert!(end >= start, "Offsets should be monotonic after conversion");
408 }
409
410 assert_arrays_eq!(list_view, list_array);
412 Ok(())
413 }
414
415 #[test]
416 fn test_empty_sublists() -> VortexResult<()> {
417 let empty_lists_view = create_empty_lists_listview();
418
419 let list_array = list_from_list_view(empty_lists_view.clone())?;
421 assert_eq!(list_array.len(), 4);
422
423 for i in 0..list_array.len() {
425 assert_eq!(list_array.list_elements_at(i)?.len(), 0);
426 }
427
428 let mut ctx = LEGACY_SESSION.create_execution_ctx();
430 let converted_back = list_view_from_list(list_array, &mut ctx)?;
431 assert_arrays_eq!(empty_lists_view, converted_back);
432 Ok(())
433 }
434
435 #[test]
436 fn test_different_offset_types() -> VortexResult<()> {
437 let elements = buffer![1i32, 2, 3, 4, 5].into_array();
439 let i32_offsets = buffer![0i32, 2, 5].into_array();
440 let list_i32 =
441 ListArray::try_new(elements.clone(), i32_offsets.clone(), Validity::NonNullable)?;
442
443 let mut ctx = LEGACY_SESSION.create_execution_ctx();
444 let list_view_i32 = list_view_from_list(list_i32.clone(), &mut ctx)?;
445 assert_eq!(list_view_i32.offsets().dtype(), i32_offsets.dtype());
446 assert_eq!(list_view_i32.sizes().dtype(), i32_offsets.dtype());
447
448 let i64_offsets = buffer![0i64, 2, 5].into_array();
450 let list_i64 =
451 ListArray::try_new(elements.clone(), i64_offsets.clone(), Validity::NonNullable)?;
452
453 let list_view_i64 = list_view_from_list(list_i64.clone(), &mut ctx)?;
454 assert_eq!(list_view_i64.offsets().dtype(), i64_offsets.dtype());
455 assert_eq!(list_view_i64.sizes().dtype(), i64_offsets.dtype());
456
457 assert_arrays_eq!(list_i32, list_view_i32);
459 assert_arrays_eq!(list_i64, list_view_i64);
460 Ok(())
461 }
462
463 #[test]
464 fn test_round_trip_conversions() -> VortexResult<()> {
465 let mut ctx = LEGACY_SESSION.create_execution_ctx();
466
467 let original = create_basic_listview();
469 let to_list = list_from_list_view(original.clone())?;
470 let back_to_view = list_view_from_list(to_list, &mut ctx)?;
471 assert_arrays_eq!(original, back_to_view);
472
473 let nullable = create_nullable_listview();
475 let nullable_to_list = list_from_list_view(nullable.clone())?;
476 let nullable_back = list_view_from_list(nullable_to_list, &mut ctx)?;
477 assert_arrays_eq!(nullable, nullable_back);
478
479 let overlapping = create_overlapping_listview();
481
482 let overlapping_to_list = list_from_list_view(overlapping.clone())?;
483 let overlapping_back = list_view_from_list(overlapping_to_list, &mut ctx)?;
484 assert_arrays_eq!(overlapping, overlapping_back);
485 Ok(())
486 }
487
488 #[test]
489 fn test_single_element_lists() -> VortexResult<()> {
490 let elements = buffer![100i32, 200, 300].into_array();
492 let offsets = buffer![0u32, 1, 2, 3].into_array();
493 let single_elem_list =
494 ListArray::try_new(elements.clone(), offsets, Validity::NonNullable)?;
495
496 let mut ctx = LEGACY_SESSION.create_execution_ctx();
497 let list_view = list_view_from_list(single_elem_list.clone(), &mut ctx)?;
498 assert_eq!(list_view.len(), 3);
499
500 let expected_sizes = buffer![1u32, 1, 1].into_array();
502 assert_arrays_eq!(expected_sizes, list_view.sizes().clone());
503
504 let converted_back = list_from_list_view(list_view)?;
506 assert_arrays_eq!(single_elem_list, converted_back);
507 Ok(())
508 }
509
510 #[test]
511 fn test_mixed_empty_and_non_empty_lists() -> VortexResult<()> {
512 let elements = buffer![1i32, 2, 3, 4, 5, 6].into_array();
514 let offsets = buffer![0u32, 2, 2, 3, 3, 6].into_array();
515 let mixed_list =
516 ListArray::try_new(elements.clone(), offsets.clone(), Validity::NonNullable)?;
517
518 let mut ctx = LEGACY_SESSION.create_execution_ctx();
519 let list_view = list_view_from_list(mixed_list.clone(), &mut ctx)?;
520 assert_eq!(list_view.len(), 5);
521
522 let expected_sizes = buffer![2u32, 0, 1, 0, 3].into_array();
524 assert_arrays_eq!(expected_sizes, list_view.sizes().clone());
525
526 let converted_back = list_from_list_view(list_view)?;
528 assert_arrays_eq!(mixed_list, converted_back);
529 Ok(())
530 }
531
532 #[test]
533 fn test_recursive_simple_listview() -> VortexResult<()> {
534 let list_view = create_basic_listview();
535 let result = recursive_list_from_list_view(list_view.clone().into_array())?;
536
537 assert_eq!(result.len(), list_view.len());
538 assert_arrays_eq!(list_view.into_array(), result);
539 Ok(())
540 }
541
542 #[test]
543 fn test_recursive_nested_listview() -> VortexResult<()> {
544 let inner_elements = buffer![1i32, 2, 3].into_array();
545 let inner_offsets = buffer![0u32, 2].into_array();
546 let inner_sizes = buffer![2u32, 1].into_array();
547 let inner_listview = unsafe {
548 ListViewArray::new_unchecked(
549 inner_elements,
550 inner_offsets,
551 inner_sizes,
552 Validity::NonNullable,
553 )
554 .with_zero_copy_to_list(true)
555 };
556
557 let outer_offsets = buffer![0u32, 1].into_array();
558 let outer_sizes = buffer![1u32, 1].into_array();
559 let outer_listview = unsafe {
560 ListViewArray::new_unchecked(
561 inner_listview.into_array(),
562 outer_offsets,
563 outer_sizes,
564 Validity::NonNullable,
565 )
566 .with_zero_copy_to_list(true)
567 };
568
569 let result = recursive_list_from_list_view(outer_listview.clone().into_array())?;
570
571 assert_eq!(result.len(), 2);
572 assert_arrays_eq!(outer_listview.into_array(), result);
573 Ok(())
574 }
575
576 #[test]
577 fn test_recursive_struct_with_listview_fields() -> VortexResult<()> {
578 let listview_field = create_basic_listview().into_array();
579 let primitive_field = buffer![10i32, 20, 30, 40].into_array();
580
581 let struct_array = StructArray::try_new(
582 FieldNames::from(["lists", "values"]),
583 vec![listview_field, primitive_field],
584 4,
585 Validity::NonNullable,
586 )?;
587
588 let result = recursive_list_from_list_view(struct_array.clone().into_array())?;
589
590 assert_eq!(result.len(), 4);
591 assert_arrays_eq!(struct_array.into_array(), result);
592 Ok(())
593 }
594
595 #[test]
596 fn test_recursive_fixed_size_list_with_listview_elements() -> VortexResult<()> {
597 let lv1_elements = buffer![1i32, 2].into_array();
598 let lv1_offsets = buffer![0u32].into_array();
599 let lv1_sizes = buffer![2u32].into_array();
600 let lv1 = unsafe {
601 ListViewArray::new_unchecked(
602 lv1_elements,
603 lv1_offsets,
604 lv1_sizes,
605 Validity::NonNullable,
606 )
607 .with_zero_copy_to_list(true)
608 };
609
610 let lv2_elements = buffer![3i32, 4].into_array();
611 let lv2_offsets = buffer![0u32].into_array();
612 let lv2_sizes = buffer![2u32].into_array();
613 let lv2 = unsafe {
614 ListViewArray::new_unchecked(
615 lv2_elements,
616 lv2_offsets,
617 lv2_sizes,
618 Validity::NonNullable,
619 )
620 .with_zero_copy_to_list(true)
621 };
622
623 let dtype = lv1.dtype().clone();
624 let chunked_listviews =
625 crate::arrays::ChunkedArray::try_new(vec![lv1.into_array(), lv2.into_array()], dtype)?;
626
627 let fixed_list =
628 FixedSizeListArray::new(chunked_listviews.into_array(), 1, Validity::NonNullable, 2);
629
630 let result = recursive_list_from_list_view(fixed_list.clone().into_array())?;
631
632 assert_eq!(result.len(), 2);
633 assert_arrays_eq!(fixed_list.into_array(), result);
634 Ok(())
635 }
636
637 #[test]
638 fn test_recursive_deep_nesting() -> VortexResult<()> {
639 let innermost_elements = buffer![1i32, 2, 3].into_array();
640 let innermost_offsets = buffer![0u32, 2].into_array();
641 let innermost_sizes = buffer![2u32, 1].into_array();
642 let innermost_listview = unsafe {
643 ListViewArray::new_unchecked(
644 innermost_elements,
645 innermost_offsets,
646 innermost_sizes,
647 Validity::NonNullable,
648 )
649 .with_zero_copy_to_list(true)
650 };
651
652 let struct_array = StructArray::try_new(
653 FieldNames::from(["inner_lists"]),
654 vec![innermost_listview.into_array()],
655 2,
656 Validity::NonNullable,
657 )?;
658
659 let outer_offsets = buffer![0u32, 1].into_array();
660 let outer_sizes = buffer![1u32, 1].into_array();
661 let outer_listview = unsafe {
662 ListViewArray::new_unchecked(
663 struct_array.into_array(),
664 outer_offsets,
665 outer_sizes,
666 Validity::NonNullable,
667 )
668 .with_zero_copy_to_list(true)
669 };
670
671 let result = recursive_list_from_list_view(outer_listview.clone().into_array())?;
672
673 assert_eq!(result.len(), 2);
674 assert_arrays_eq!(outer_listview.into_array(), result);
675 Ok(())
676 }
677
678 #[test]
679 fn test_recursive_primitive_unchanged() -> VortexResult<()> {
680 let prim = buffer![1i32, 2, 3].into_array();
681 let prim_clone = prim.clone();
682 let result = recursive_list_from_list_view(prim)?;
683
684 assert!(Arc::ptr_eq(&result, &prim_clone));
685 Ok(())
686 }
687
688 #[test]
689 fn test_recursive_mixed_listview_and_list() -> VortexResult<()> {
690 let listview = create_basic_listview();
691 let list = list_from_list_view(listview.clone())?;
692
693 let struct_array = StructArray::try_new(
694 FieldNames::from(["listview_field", "list_field"]),
695 vec![listview.into_array(), list.into_array()],
696 4,
697 Validity::NonNullable,
698 )?;
699
700 let result = recursive_list_from_list_view(struct_array.clone().into_array())?;
701
702 assert_eq!(result.len(), 4);
703 assert_arrays_eq!(struct_array.into_array(), result);
704 Ok(())
705 }
706}